langtrace-python-sdk 3.3.13__tar.gz → 3.3.15__tar.gz

{langtrace_python_sdk-3.3.13 → langtrace_python_sdk-3.3.15}/PKG-INFO

@@ -1,10 +1,11 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: langtrace-python-sdk
-Version: 3.3.13
+Version: 3.3.15
 Summary: Python SDK for LangTrace
 Project-URL: Homepage, https://github.com/Scale3-Labs/langtrace-python-sdk
 Author-email: Scale3 Labs <engineering@scale3labs.com>
-License: Apache-2.0
+License-Expression: Apache-2.0
+License-File: LICENSE
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python :: 3

{langtrace_python_sdk-3.3.13 → langtrace_python_sdk-3.3.15}/src/examples/cohere_example/__init__.py

@@ -1,8 +1,11 @@
 from examples.cohere_example.chat import chat_comp
+from examples.cohere_example.chatv2 import chat_v2
+from examples.cohere_example.chat_streamv2 import chat_stream_v2
 from examples.cohere_example.chat_stream import chat_stream
 from examples.cohere_example.tools import tool_calling
 from examples.cohere_example.embed import embed
 from examples.cohere_example.rerank import rerank
+from examples.cohere_example.rerankv2 import rerank_v2
 from langtrace_python_sdk import with_langtrace_root_span
 
 
@@ -10,8 +13,11 @@ class CohereRunner:
 
     @with_langtrace_root_span("Cohere")
     def run(self):
+        chat_v2()
+        chat_stream_v2()
         chat_comp()
         chat_stream()
         tool_calling()
         embed()
         rerank()
+        rerank_v2()

langtrace_python_sdk-3.3.15/src/examples/cohere_example/chat_streamv2.py

@@ -0,0 +1,17 @@
+import os
+from langtrace_python_sdk import langtrace
+import cohere
+
+langtrace.init(api_key=os.getenv("LANGTRACE_API_KEY"))
+co = cohere.ClientV2(api_key=os.getenv("COHERE_API_KEY"))
+
+def chat_stream_v2():
+    res = co.chat_stream(
+        model="command-r-plus-08-2024",
+        messages=[{"role": "user", "content": "Write a title for a blog post about API design. Only output the title text"}],
+    )
+
+    for event in res:
+        if event:
+            if event.type == "content-delta":
+                print(event.delta.message.content.text)

langtrace_python_sdk-3.3.15/src/examples/cohere_example/chatv2.py

@@ -0,0 +1,21 @@
+import os
+from langtrace_python_sdk import langtrace
+import cohere
+
+langtrace.init(api_key=os.getenv("LANGTRACE_API_KEY"))
+
+
+def chat_v2():
+    co = cohere.ClientV2(api_key=os.getenv("COHERE_API_KEY"))
+
+    res = co.chat(
+        model="command-r-plus-08-2024",
+        messages=[
+            {
+                "role": "user",
+                "content": "Write a title for a blog post about API design. Only output the title text.",
+            }
+        ],
+    )
+
+    print(res.message.content[0].text)

langtrace_python_sdk-3.3.15/src/examples/cohere_example/rerankv2.py

@@ -0,0 +1,23 @@
+import os
+from langtrace_python_sdk import langtrace
+import cohere
+
+langtrace.init(api_key=os.getenv("LANGTRACE_API_KEY"))
+co = cohere.ClientV2(api_key=os.getenv("COHERE_API_KEY"))
+
+docs = [
+    "Carson City is the capital city of the American state of Nevada.",
+    "The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean. Its capital is Saipan.",
+    "Capitalization or capitalisation in English grammar is the use of a capital letter at the start of a word. English usage varies from capitalization in other languages.",
+    "Washington, D.C. (also known as simply Washington or D.C., and officially as the District of Columbia) is the capital of the United States. It is a federal district.",
+    "Capital punishment has existed in the United States since beforethe United States was a country. As of 2017, capital punishment is legal in 30 of the 50 states.",
+]
+
+def rerank_v2():
+    response = co.rerank(
+        model="rerank-v3.5",
+        query="What is the capital of the United States?",
+        documents=docs,
+        top_n=3,
+    )
+    print(response)

langtrace_python_sdk-3.3.15/src/langtrace_python_sdk/constants/exporter/langtrace_exporter.py

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

langtrace_python_sdk-3.3.15/src/langtrace_python_sdk/constants/instrumentation/cohere.py

@@ -0,0 +1,42 @@
+APIS = {
+    "CHAT_CREATE": {
+        "URL": "https://api.cohere.ai",
+        "METHOD": "cohere.client.chat",
+        "ENDPOINT": "/v1/chat",
+    },
+    "CHAT_CREATE_V2": {
+        "URL": "https://api.cohere.ai",
+        "METHOD": "cohere.client_v2.chat",
+        "ENDPOINT": "/v2/chat",
+    },
+    "EMBED": {
+        "URL": "https://api.cohere.ai",
+        "METHOD": "cohere.client.embed",
+        "ENDPOINT": "/v1/embed",
+    },
+    "EMBED_V2": {
+        "URL": "https://api.cohere.ai",
+        "METHOD": "cohere.client_v2.embed",
+        "ENDPOINT": "/v2/embed",
+    },
+    "CHAT_STREAM": {
+        "URL": "https://api.cohere.ai",
+        "METHOD": "cohere.client.chat_stream",
+        "ENDPOINT": "/v1/chat",
+    },
+    "CHAT_STREAM_V2": {
+        "URL": "https://api.cohere.ai",
+        "METHOD": "cohere.client_v2.chat_stream",
+        "ENDPOINT": "/v2/chat",
+    },
+    "RERANK": {
+        "URL": "https://api.cohere.ai",
+        "METHOD": "cohere.client.rerank",
+        "ENDPOINT": "/v1/rerank",
+    },
+    "RERANK_V2": {
+        "URL": "https://api.cohere.ai",
+        "METHOD": "cohere.client_v2.rerank",
+        "ENDPOINT": "/v2/rerank",
+    },
+}

{langtrace_python_sdk-3.3.13 → langtrace_python_sdk-3.3.15}/src/langtrace_python_sdk/extensions/langtrace_exporter.py

@@ -9,6 +9,7 @@ from opentelemetry.trace.span import format_trace_id
 
 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 @@ class LangTraceExporter(SpanExporter):
             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 @@ class LangTraceExporter(SpanExporter):
             "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:

{langtrace_python_sdk-3.3.13 → langtrace_python_sdk-3.3.15}/src/langtrace_python_sdk/instrumentation/cohere/instrumentation.py

@@ -23,6 +23,7 @@ from wrapt import wrap_function_wrapper
 
 from langtrace_python_sdk.instrumentation.cohere.patch import (
     chat_create,
+    chat_create_v2,
     chat_stream,
     embed,
     rerank,
@@ -48,6 +49,18 @@ class CohereInstrumentation(BaseInstrumentor):
             chat_create("cohere.client.chat", version, tracer),
         )
 
+        wrap_function_wrapper(
+            "cohere.client_v2",
+            "ClientV2.chat",
+            chat_create_v2("cohere.client_v2.chat", version, tracer),
+        )
+
+        wrap_function_wrapper(
+            "cohere.client_v2",
+            "ClientV2.chat_stream",
+            chat_create_v2("cohere.client_v2.chat", version, tracer, stream=True),
+        )
+
         wrap_function_wrapper(
             "cohere.client",
             "Client.chat_stream",
@@ -60,12 +73,24 @@ class CohereInstrumentation(BaseInstrumentor):
             embed("cohere.client.embed", version, tracer),
         )
 
+        wrap_function_wrapper(
+            "cohere.client_v2",
+            "ClientV2.embed",
+            embed("cohere.client.embed", version, tracer, v2=True),
+        )
+
         wrap_function_wrapper(
             "cohere.client",
             "Client.rerank",
             rerank("cohere.client.rerank", version, tracer),
         )
 
+        wrap_function_wrapper(
+            "cohere.client_v2",
+            "ClientV2.rerank",
+            rerank("cohere.client.rerank", version, tracer, v2=True),
+        )
+
     def _instrument_module(self, module_name):
         pass
 

{langtrace_python_sdk-3.3.13 → langtrace_python_sdk-3.3.15}/src/langtrace_python_sdk/instrumentation/cohere/patch.py

@@ -24,6 +24,7 @@ from langtrace_python_sdk.utils.llm import (
     get_span_name,
     set_event_completion,
     set_usage_attributes,
+    StreamWrapper
 )
 from langtrace.trace_attributes import Event, LLMSpanAttributes
 from langtrace_python_sdk.utils import set_span_attribute
@@ -38,7 +39,7 @@ from langtrace_python_sdk.constants.instrumentation.common import (
 from langtrace.trace_attributes import SpanAttributes
 
 
-def rerank(original_method, version, tracer):
+def rerank(original_method, version, tracer, v2=False):
     """Wrap the `rerank` method."""
 
     def traced_method(wrapped, instance, args, kwargs):
@@ -49,8 +50,8 @@ def rerank(original_method, version, tracer):
             **get_llm_request_attributes(kwargs, operation_name="rerank"),
             **get_llm_url(instance),
             SpanAttributes.LLM_REQUEST_MODEL: kwargs.get("model") or "command-r-plus",
-            SpanAttributes.LLM_URL: APIS["RERANK"]["URL"],
-            SpanAttributes.LLM_PATH: APIS["RERANK"]["ENDPOINT"],
+            SpanAttributes.LLM_URL: APIS["RERANK" if not v2 else "RERANK_V2"]["URL"],
+            SpanAttributes.LLM_PATH: APIS["RERANK" if not v2 else "RERANK_V2"]["ENDPOINT"],
             SpanAttributes.LLM_REQUEST_DOCUMENTS: json.dumps(
                 kwargs.get("documents"), cls=datetime_encoder
             ),
@@ -61,7 +62,7 @@ def rerank(original_method, version, tracer):
         attributes = LLMSpanAttributes(**span_attributes)
 
         span = tracer.start_span(
-            name=get_span_name(APIS["RERANK"]["METHOD"]), kind=SpanKind.CLIENT
+            name=get_span_name(APIS["RERANK" if not v2 else "RERANK_V2"]["METHOD"]), kind=SpanKind.CLIENT
         )
         for field, value in attributes.model_dump(by_alias=True).items():
             set_span_attribute(span, field, value)
@@ -119,7 +120,7 @@ def rerank(original_method, version, tracer):
     return traced_method
 
 
-def embed(original_method, version, tracer):
+def embed(original_method, version, tracer, v2=False):
     """Wrap the `embed` method."""
 
     def traced_method(wrapped, instance, args, kwargs):
@@ -129,8 +130,8 @@ def embed(original_method, version, tracer):
             **get_langtrace_attributes(version, service_provider),
             **get_llm_request_attributes(kwargs, operation_name="embed"),
             **get_llm_url(instance),
-            SpanAttributes.LLM_URL: APIS["EMBED"]["URL"],
-            SpanAttributes.LLM_PATH: APIS["EMBED"]["ENDPOINT"],
+            SpanAttributes.LLM_URL: APIS["EMBED" if not v2 else "EMBED_V2"]["URL"],
+            SpanAttributes.LLM_PATH: APIS["EMBED" if not v2 else "EMBED_V2"]["ENDPOINT"],
             SpanAttributes.LLM_REQUEST_EMBEDDING_INPUTS: json.dumps(
                 kwargs.get("texts")
             ),
@@ -143,7 +144,7 @@ def embed(original_method, version, tracer):
         attributes = LLMSpanAttributes(**span_attributes)
 
         span = tracer.start_span(
-            name=get_span_name(APIS["EMBED"]["METHOD"]),
+            name=get_span_name(APIS["EMBED" if not v2 else "EMBED_V2"]["METHOD"]),
             kind=SpanKind.CLIENT,
         )
         for field, value in attributes.model_dump(by_alias=True).items():
@@ -343,6 +344,103 @@ def chat_create(original_method, version, tracer):
     return traced_method
 
 
+def chat_create_v2(original_method, version, tracer, stream=False):
+    """Wrap the `chat_create` method for Cohere API v2."""
+
+    def traced_method(wrapped, instance, args, kwargs):
+        service_provider = SERVICE_PROVIDERS["COHERE"]
+        
+        messages = kwargs.get("messages", [])
+        if kwargs.get("preamble"):
+            messages = [{"role": "system", "content": kwargs["preamble"]}] + messages
+
+        span_attributes = {
+            **get_langtrace_attributes(version, service_provider),
+            **get_llm_request_attributes(kwargs, prompts=messages),
+            **get_llm_url(instance),
+            SpanAttributes.LLM_REQUEST_MODEL: kwargs.get("model") or "command-r-plus",
+            SpanAttributes.LLM_URL: APIS["CHAT_CREATE_V2"]["URL"],
+            SpanAttributes.LLM_PATH: APIS["CHAT_CREATE_V2"]["ENDPOINT"],
+            **get_extra_attributes(),
+        }
+
+        attributes = LLMSpanAttributes(**span_attributes)
+
+        for attr_name in ["max_input_tokens", "conversation_id", "connectors", "tools", "tool_results"]:
+            value = kwargs.get(attr_name)
+            if value is not None:
+                if attr_name == "max_input_tokens":
+                    attributes.llm_max_input_tokens = str(value)
+                elif attr_name == "conversation_id":
+                    attributes.conversation_id = value
+                else:
+                    setattr(attributes, f"llm_{attr_name}", json.dumps(value))
+
+        span = tracer.start_span(
+            name=get_span_name(APIS["CHAT_CREATE_V2"]["METHOD"]), 
+            kind=SpanKind.CLIENT
+        )
+
+        for field, value in attributes.model_dump(by_alias=True).items():
+            set_span_attribute(span, field, value)
+
+        try:
+            result = wrapped(*args, **kwargs)
+
+            if stream:
+                return StreamWrapper(
+                    result,
+                    span,
+                    tool_calls=kwargs.get("tools") is not None,
+                )
+            else:
+                if hasattr(result, "id") and result.id is not None:
+                    span.set_attribute(SpanAttributes.LLM_GENERATION_ID, result.id)
+                    span.set_attribute(SpanAttributes.LLM_RESPONSE_ID, result.id)
+
+                if (hasattr(result, "message") and 
+                    hasattr(result.message, "content") and 
+                    len(result.message.content) > 0 and
+                    hasattr(result.message.content[0], "text") and
+                    result.message.content[0].text is not None and
+                    result.message.content[0].text != ""):
+                    responses = [{
+                        "role": result.message.role,
+                        "content": result.message.content[0].text
+                    }]
+                    set_event_completion(span, responses)
+                if hasattr(result, "tool_calls") and result.tool_calls is not None:
+                    tool_calls = [tool_call.json() for tool_call in result.tool_calls]
+                    span.set_attribute(
+                        SpanAttributes.LLM_TOOL_RESULTS,
+                        json.dumps(tool_calls)
+                    )
+                if hasattr(result, "usage") and result.usage is not None:
+                    if (hasattr(result.usage, "billed_units") and 
+                        result.usage.billed_units is not None):
+                        usage = result.usage.billed_units
+                        for metric, value in {
+                            "input": usage.input_tokens or 0,
+                            "output": usage.output_tokens or 0,
+                            "total": (usage.input_tokens or 0) + (usage.output_tokens or 0),
+                        }.items():
+                            span.set_attribute(
+                                f"gen_ai.usage.{metric}_tokens",
+                                int(value)
+                            )
+                span.set_status(StatusCode.OK)
+                span.end()
+                return result
+
+        except Exception as error:
+            span.record_exception(error)
+            span.set_status(Status(StatusCode.ERROR, str(error)))
+            span.end()
+            raise
+
+    return traced_method
+
+
 def chat_stream(original_method, version, tracer):
     """Wrap the `messages_stream` method."""
 

{langtrace_python_sdk-3.3.13 → langtrace_python_sdk-3.3.15}/src/langtrace_python_sdk/instrumentation/dspy/patch.py

@@ -50,7 +50,7 @@ def patch_bootstrapfewshot_optimizer(operation_name, version, tracer):
                     ),
                 }
                 span_attributes["dspy.optimizer.module.prog"] = json.dumps(prog)
-        if hasattr(instance, "metric"):
+        if hasattr(instance, "metric") and getattr(instance, "metric") is not None:
             span_attributes["dspy.optimizer.metric"] = getattr(
                 instance, "metric"
             ).__name__

{langtrace_python_sdk-3.3.13 → langtrace_python_sdk-3.3.15}/src/langtrace_python_sdk/instrumentation/gemini/patch.py

@@ -137,7 +137,7 @@ def set_response_attributes(
     if hasattr(result, "text"):
         set_event_completion(span, [{"role": "assistant", "content": result.text}])
 
-    if hasattr(result, "usage_metadata"):
+    if hasattr(result, "usage_metadata") and result.usage_metadata is not None:
         usage = result.usage_metadata
         input_tokens = usage.prompt_token_count
         output_tokens = usage.candidates_token_count
@@ -152,7 +152,7 @@ def build_streaming_response(span, response):
         item_to_yield = item
         complete_response += str(item.text)
         yield item_to_yield
-        if hasattr(item, "usage_metadata"):
+        if hasattr(item, "usage_metadata") and item.usage_metadata is not None:
             usage = item.usage_metadata
             input_tokens = usage.prompt_token_count
             output_tokens = usage.candidates_token_count
@@ -171,7 +171,7 @@ async def abuild_streaming_response(span, response):
         item_to_yield = item
         complete_response += str(item.text)
         yield item_to_yield
-        if hasattr(item, "usage_metadata"):
+        if hasattr(item, "usage_metadata") and item.usage_metadata is not None:
             usage = item.usage_metadata
             input_tokens = usage.prompt_token_count
             output_tokens = usage.candidates_token_count

{langtrace_python_sdk-3.3.13 → langtrace_python_sdk-3.3.15}/src/langtrace_python_sdk/instrumentation/langchain_core/patch.py

@@ -98,7 +98,7 @@ def generic_patch(
                 result = wrapped(*args, **kwargs)
                 if trace_output:
                     span.set_attribute("langchain.outputs", to_json_string(result))
-                    if hasattr(result, "usage_metadata"):
+                    if hasattr(result, "usage_metadata") and result.usage_metadata is not None:
                         span.set_attribute(
                             SpanAttributes.LLM_USAGE_PROMPT_TOKENS,
                             result.usage_metadata["input_tokens"],

{langtrace_python_sdk-3.3.13 → langtrace_python_sdk-3.3.15}/src/langtrace_python_sdk/instrumentation/openai/patch.py

@@ -342,8 +342,12 @@ def async_chat_completions_create(version: str, tracer: Tracer) -> Callable:
             service_provider = SERVICE_PROVIDERS["PPLX"]
         elif "azure" in get_base_url(instance):
             service_provider = SERVICE_PROVIDERS["AZURE"]
+        elif "groq" in get_base_url(instance):
+            service_provider = SERVICE_PROVIDERS["GROQ"]
         elif "x.ai" in get_base_url(instance):
             service_provider = SERVICE_PROVIDERS["XAI"]
+        elif "deepseek" in get_base_url(instance):
+            service_provider = SERVICE_PROVIDERS["DEEPSEEK"]
         llm_prompts = []
         for item in kwargs.get("messages", []):
             tools = get_tool_calls(item)
@@ -431,6 +435,18 @@ def embeddings_create(version: str, tracer: Tracer) -> Callable:
         kwargs: EmbeddingsCreateKwargs,
     ) -> Any:
         service_provider = SERVICE_PROVIDERS["OPENAI"]
+        base_url = get_base_url(instance)
+
+        if "perplexity" in base_url:
+            service_provider = SERVICE_PROVIDERS["PPLX"]
+        elif "azure" in base_url:
+            service_provider = SERVICE_PROVIDERS["AZURE"]
+        elif "groq" in base_url:
+            service_provider = SERVICE_PROVIDERS["GROQ"]
+        elif "x.ai" in base_url:
+            service_provider = SERVICE_PROVIDERS["XAI"]
+        elif "deepseek" in base_url:
+            service_provider = SERVICE_PROVIDERS["DEEPSEEK"]
 
         span_attributes = {
             **get_langtrace_attributes(version, service_provider, vendor_type="llm"),
@@ -469,7 +485,6 @@ def embeddings_create(version: str, tracer: Tracer) -> Callable:
             kind=SpanKind.CLIENT,
             context=set_span_in_context(trace.get_current_span()),
         ) as span:
-
             set_span_attributes(span, attributes)
             try:
                 # Attempt to call the original method
@@ -507,17 +522,27 @@ def async_embeddings_create(version: str, tracer: Tracer) -> Callable:
     ) -> Awaitable[Any]:
 
         service_provider = SERVICE_PROVIDERS["OPENAI"]
+        base_url = get_base_url(instance)
+        if "perplexity" in base_url:
+            service_provider = SERVICE_PROVIDERS["PPLX"]
+        elif "azure" in base_url:
+            service_provider = SERVICE_PROVIDERS["AZURE"]
+        elif "groq" in base_url:
+            service_provider = SERVICE_PROVIDERS["GROQ"]
+        elif "x.ai" in base_url:
+            service_provider = SERVICE_PROVIDERS["XAI"]
+        elif "deepseek" in base_url:
+            service_provider = SERVICE_PROVIDERS["DEEPSEEK"]
 
         span_attributes = {
             **get_langtrace_attributes(version, service_provider, vendor_type="llm"),
             **get_llm_request_attributes(kwargs, operation_name="embed"),
+            **get_llm_url(instance),
             SpanAttributes.LLM_PATH: APIS["EMBEDDINGS_CREATE"]["ENDPOINT"],
             SpanAttributes.LLM_REQUEST_DIMENSIONS: kwargs.get("dimensions"),
             **get_extra_attributes(),  # type: ignore
         }
 
-        attributes = LLMSpanAttributes(**filter_valid_attributes(span_attributes))
-
         encoding_format = kwargs.get("encoding_format")
         if encoding_format is not None:
             if not isinstance(encoding_format, list):
@@ -530,17 +555,31 @@ def async_embeddings_create(version: str, tracer: Tracer) -> Callable:
             span_attributes[SpanAttributes.LLM_REQUEST_EMBEDDING_INPUTS] = json.dumps(
                 [kwargs.get("input", "")]
             )
+            span_attributes[SpanAttributes.LLM_PROMPTS] = json.dumps(
+                [
+                    {
+                        "role": "user",
+                        "content": kwargs.get("input"),
+                    }
+                ]
+            )
+
+        attributes = LLMSpanAttributes(**filter_valid_attributes(span_attributes))
 
         with tracer.start_as_current_span(
             name=get_span_name(APIS["EMBEDDINGS_CREATE"]["METHOD"]),
             kind=SpanKind.CLIENT,
             context=set_span_in_context(trace.get_current_span()),
         ) as span:
-
             set_span_attributes(span, attributes)
             try:
                 # Attempt to call the original method
                 result = await wrapped(*args, **kwargs)
+                usage = getattr(result, "usage", None)
+                if usage:
+                    set_usage_attributes(
+                        span, {"prompt_tokens": getattr(usage, "prompt_tokens", 0)}
+                    )
                 span.set_status(StatusCode.OK)
                 return result
             except Exception as err:

{langtrace_python_sdk-3.3.13 → langtrace_python_sdk-3.3.15}/src/langtrace_python_sdk/instrumentation/vertexai/patch.py

@@ -77,7 +77,7 @@ def set_response_attributes(span: Span, result):
     if hasattr(result, "text"):
         set_event_completion(span, [{"role": "assistant", "content": result.text}])
 
-    if hasattr(result, "usage_metadata"):
+    if hasattr(result, "usage_metadata") and result.usage_metadata is not None:
         usage = result.usage_metadata
         input_tokens = usage.prompt_token_count
         output_tokens = usage.candidates_token_count

{langtrace_python_sdk-3.3.13 → langtrace_python_sdk-3.3.15}/src/langtrace_python_sdk/langtrace.py

@@ -39,6 +39,7 @@ from opentelemetry.exporter.otlp.proto.http.trace_exporter import (
 )
 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 @@ class LangtraceConfig:
             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: