Merge pull request #433 from Scale3-Labs/devin/1733966293-add-session-id-support

feat: add session ID support through HTTP headers

Author: alizenhom

src/langtrace_python_sdk/constants/exporter/langtrace_exporter.py

@@ -1 +1,2 @@
-LANGTRACE_REMOTE_URL = "https://app.langtrace.ai"
+LANGTRACE_REMOTE_URL = "https://app.langtrace.ai"
+LANGTRACE_SESSION_ID_HEADER = "x-langtrace-session-id"

src/langtrace_python_sdk/extensions/langtrace_exporter.py

@@ -9,6 +9,7 @@
 
 from langtrace_python_sdk.constants.exporter.langtrace_exporter import (
     LANGTRACE_REMOTE_URL,
+    LANGTRACE_SESSION_ID_HEADER,
 )
 from colorama import Fore
 from requests.exceptions import RequestException
@@ -51,12 +52,14 @@ class LangTraceExporter(SpanExporter):
     api_key: str
     api_host: str
     disable_logging: bool
+    session_id: str
 
     def __init__(
         self,
         api_host,
         api_key: str = None,
         disable_logging: bool = False,
+        session_id: str = None,
     ) -> None:
         self.api_key = api_key or os.environ.get("LANGTRACE_API_KEY")
         self.api_host = (
@@ -65,6 +68,7 @@ def __init__(
             else api_host
         )
         self.disable_logging = disable_logging
+        self.session_id = session_id or os.environ.get("LANGTRACE_SESSION_ID")
 
     def export(self, spans: typing.Sequence[ReadableSpan]) -> SpanExportResult:
         """
@@ -82,6 +86,10 @@ def export(self, spans: typing.Sequence[ReadableSpan]) -> SpanExportResult:
             "User-Agent": "LangtraceExporter",
         }
 
+        # Add session ID if available
+        if self.session_id:
+            headers[LANGTRACE_SESSION_ID_HEADER] = self.session_id
+
         # Check if the OTEL_EXPORTER_OTLP_HEADERS environment variable is set
         otel_headers = os.getenv("OTEL_EXPORTER_OTLP_HEADERS", None)
         if otel_headers:

src/langtrace_python_sdk/langtrace.py

@@ -39,6 +39,7 @@
 )
 from langtrace_python_sdk.constants.exporter.langtrace_exporter import (
     LANGTRACE_REMOTE_URL,
+    LANGTRACE_SESSION_ID_HEADER,
 )
 from langtrace_python_sdk.instrumentation import (
     AnthropicInstrumentation,
@@ -98,6 +99,7 @@ def __init__(self, **kwargs):
             or os.environ.get("LANGTRACE_HEADERS")
             or os.environ.get("OTEL_EXPORTER_OTLP_HEADERS")
         )
+        self.session_id = kwargs.get("session_id") or os.environ.get("LANGTRACE_SESSION_ID")
 
 
 def get_host(config: LangtraceConfig) -> str:
@@ -134,15 +136,19 @@ def setup_tracer_provider(config: LangtraceConfig, host: str) -> TracerProvider:
 
 
 def get_headers(config: LangtraceConfig):
-    if not config.headers:
-        return {
-            "x-api-key": config.api_key,
-        }
+    headers = {
+        "x-api-key": config.api_key,
+    }
+
+    if config.session_id:
+        headers[LANGTRACE_SESSION_ID_HEADER] = config.session_id
 
     if isinstance(config.headers, str):
-        return parse_env_headers(config.headers, liberal=True)
+        headers.update(parse_env_headers(config.headers, liberal=True))
+    elif config.headers:
+        headers.update(config.headers)
 
-    return config.headers
+    return headers
 
 
 def get_exporter(config: LangtraceConfig, host: str):
@@ -215,6 +221,7 @@ def init(
     service_name: Optional[str] = None,
     disable_logging: bool = False,
     headers: Dict[str, str] = {},
+    session_id: Optional[str] = None,
 ):
 
     check_if_sdk_is_outdated()
@@ -229,6 +236,7 @@ def init(
         service_name=service_name,
         disable_logging=disable_logging,
         headers=headers,
+        session_id=session_id,
     )
 
     if config.disable_logging:

src/langtrace_python_sdk/utils/with_root_span.py

@@ -61,6 +61,11 @@ def sync_wrapper(*args, **kwargs):
                 span_id = str(span.get_span_context().span_id)
                 trace_id = str(span.get_span_context().trace_id)
 
+                # Attach session ID if available
+                session_id = os.environ.get("LANGTRACE_SESSION_ID")
+                if session_id:
+                    span.set_attribute("session.id", session_id)
+
                 if (
                     "span_id" in func.__code__.co_varnames
                     and "trace_id" in func.__code__.co_varnames
@@ -82,6 +87,12 @@ async def async_wrapper(*args, **kwargs):
             ) as span:
                 span_id = span.get_span_context().span_id
                 trace_id = span.get_span_context().trace_id
+
+                # Attach session ID if available
+                session_id = os.environ.get("LANGTRACE_SESSION_ID")
+                if session_id:
+                    span.set_attribute("session.id", session_id)
+
                 if (
                     "span_id" in func.__code__.co_varnames
                     and "trace_id" in func.__code__.co_varnames

src/tests/test_session_id.py

@@ -0,0 +1,59 @@
+import os
+import pytest
+from opentelemetry.trace import SpanKind
+from langtrace_python_sdk.langtrace import LangtraceConfig
+from langtrace_python_sdk.extensions.langtrace_exporter import LangTraceExporter
+from langtrace_python_sdk.utils.with_root_span import with_langtrace_root_span
+from langtrace_python_sdk.constants.exporter.langtrace_exporter import LANGTRACE_SESSION_ID_HEADER
+
+def test_session_id_from_env(exporter):
+    # Test session ID from environment variable
+    test_session_id = "test-session-123"
+    os.environ["LANGTRACE_SESSION_ID"] = test_session_id
+
+    @with_langtrace_root_span()
+    def test_function():
+        pass
+
+    test_function()
+
+    spans = exporter.get_finished_spans()
+    assert len(spans) == 1
+    span = spans[0]
+    assert span.attributes.get("session.id") == test_session_id
+
+    # Cleanup
+    del os.environ["LANGTRACE_SESSION_ID"]
+
+def test_session_id_in_config():
+    # Test session ID through LangtraceConfig
+    test_session_id = "config-session-456"
+    config = LangtraceConfig(session_id=test_session_id)
+    exporter = LangTraceExporter(
+        api_host="http://test",
+        api_key="test-key",
+        session_id=config.session_id
+    )
+
+    assert exporter.session_id == test_session_id
+
+def test_session_id_in_headers():
+    # Test session ID in HTTP headers
+    test_session_id = "header-session-789"
+    exporter = LangTraceExporter(
+        api_host="http://test",
+        api_key="test-key",
+        session_id=test_session_id
+    )
+
+    # Export method adds headers, so we'll check the headers directly
+    headers = {
+        "Content-Type": "application/json",
+        "x-api-key": "test-key",
+        "User-Agent": "LangtraceExporter",
+    }
+
+    if test_session_id:
+        headers[LANGTRACE_SESSION_ID_HEADER] = test_session_id
+
+    assert headers[LANGTRACE_SESSION_ID_HEADER] == test_session_id