refactor: migrate session ID logic to get_headers function

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

Author: devin-ai-integration[bot]

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,
@@ -135,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):
@@ -216,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()
@@ -230,6 +236,7 @@ def init(
         service_name=service_name,
         disable_logging=disable_logging,
         headers=headers,
+        session_id=session_id,
     )
 
     if config.disable_logging: