feat: add session ID support through HTTP headers

Co-Authored-By: [email protected] <[email protected]>

Author: devin-ai-integration[bot]

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

@@ -98,6 +98,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:

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