diff --git a/src/langtrace_python_sdk/constants/exporter/langtrace_exporter.py b/src/langtrace_python_sdk/constants/exporter/langtrace_exporter.py index 0261ab6f..3bc4c58c 100644 --- a/src/langtrace_python_sdk/constants/exporter/langtrace_exporter.py +++ b/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" diff --git a/src/langtrace_python_sdk/extensions/langtrace_exporter.py b/src/langtrace_python_sdk/extensions/langtrace_exporter.py index bc26802a..664013db 100644 --- a/src/langtrace_python_sdk/extensions/langtrace_exporter.py +++ b/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: diff --git a/src/langtrace_python_sdk/langtrace.py b/src/langtrace_python_sdk/langtrace.py index 1b2c59c8..51005706 100644 --- a/src/langtrace_python_sdk/langtrace.py +++ b/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: diff --git a/src/langtrace_python_sdk/utils/with_root_span.py b/src/langtrace_python_sdk/utils/with_root_span.py index 67e6f66d..f1a5296f 100644 --- a/src/langtrace_python_sdk/utils/with_root_span.py +++ b/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 diff --git a/src/tests/test_session_id.py b/src/tests/test_session_id.py new file mode 100644 index 00000000..32adc667 --- /dev/null +++ b/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