langtrace-python-sdk 2.1.29__py3-none-any.whl → 2.2.1__py3-none-any.whl

langtrace_python_sdk/instrumentation/ollama/patch.py

@@ -1,41 +1,35 @@
 from langtrace_python_sdk.constants.instrumentation.ollama import APIS
-from importlib_metadata import version as v
-from langtrace_python_sdk.constants import LANGTRACE_SDK_NAME
 from langtrace_python_sdk.utils import set_span_attribute
-from langtrace_python_sdk.utils.silently_fail import silently_fail
-from langtrace_python_sdk.constants.instrumentation.common import (
-    LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY,
-    SERVICE_PROVIDERS,
+from langtrace_python_sdk.utils.llm import (
+    get_extra_attributes,
+    get_langtrace_attributes,
+    get_llm_request_attributes,
+    get_llm_url,
+    set_event_completion,
 )
-from opentelemetry import baggage
+from langtrace_python_sdk.utils.silently_fail import silently_fail
+from langtrace_python_sdk.constants.instrumentation.common import SERVICE_PROVIDERS
 from langtrace.trace_attributes import LLMSpanAttributes, Event
 from opentelemetry.trace import SpanKind
 import json
 from opentelemetry.trace.status import Status, StatusCode
+from langtrace.trace_attributes import SpanAttributes
 
 
 def generic_patch(operation_name, version, tracer):
     def traced_method(wrapped, instance, args, kwargs):
-        base_url = (
-            str(instance._client._base_url)
-            if hasattr(instance, "_client") and hasattr(instance._client, "_base_url")
-            else ""
-        )
         api = APIS[operation_name]
         service_provider = SERVICE_PROVIDERS["OLLAMA"]
-        extra_attributes = baggage.get_baggage(LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY)
         span_attributes = {
-            "langtrace.sdk.name": "langtrace-python-sdk",
-            "langtrace.service.name": service_provider,
-            "langtrace.service.type": "llm",
-            "langtrace.service.version": version,
-            "langtrace.version": v(LANGTRACE_SDK_NAME),
-            "llm.model": kwargs.get("model"),
-            "llm.stream": kwargs.get("stream"),
-            "url.full": base_url,
-            "llm.api": api["ENDPOINT"],
-            "llm.response_format": kwargs.get("format"),
-            **(extra_attributes if extra_attributes is not None else {}),
+            **get_langtrace_attributes(version, service_provider),
+            **get_llm_request_attributes(
+                kwargs,
+                prompts=kwargs.get("messages", None),
+            ),
+            **get_llm_url(instance),
+            SpanAttributes.LLM_PATH: api["ENDPOINT"],
+            SpanAttributes.LLM_RESPONSE_FORMAT: kwargs.get("format"),
+            **get_extra_attributes(),
         }
 
         attributes = LLMSpanAttributes(**span_attributes)
@@ -77,24 +71,14 @@ def ageneric_patch(operation_name, version, tracer):
     async def traced_method(wrapped, instance, args, kwargs):
         api = APIS[operation_name]
         service_provider = SERVICE_PROVIDERS["OLLAMA"]
-        extra_attributes = baggage.get_baggage(LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY)
         span_attributes = {
-            "langtrace.sdk.name": "langtrace-python-sdk",
-            "langtrace.service.name": service_provider,
-            "url.full": "",
-            "llm.api": "",
-            "langtrace.service.type": "llm",
-            "langtrace.service.version": version,
-            "langtrace.version": v(LANGTRACE_SDK_NAME),
-            "llm.model": kwargs.get("model"),
-            "llm.stream": kwargs.get("stream"),
-            "llm.response_format": kwargs.get("format"),
-            "http.timeout": (
-                kwargs.get("keep_alive") if "keep_alive" in kwargs else None
-            ),
-            **(extra_attributes if extra_attributes is not None else {}),
+            **get_langtrace_attributes(version, service_provider),
+            **get_llm_request_attributes(kwargs),
+            **get_llm_url(instance),
+            SpanAttributes.LLM_PATH: api["ENDPOINT"],
+            SpanAttributes.LLM_RESPONSE_FORMAT: kwargs.get("format"),
+            **get_extra_attributes(),
         }
-
         attributes = LLMSpanAttributes(**span_attributes)
         with tracer.start_as_current_span(api["METHOD"], kind=SpanKind.CLIENT) as span:
             _set_input_attributes(span, kwargs, attributes)
@@ -130,23 +114,25 @@ def _set_response_attributes(span, response):
     input_tokens = response.get("prompt_eval_count") or 0
     output_tokens = response.get("eval_count") or 0
     total_tokens = input_tokens + output_tokens
-    usage_dict = {
-        "input_tokens": input_tokens,
-        "output_tokens": output_tokens,
-        "total_tokens": total_tokens,
-    }
 
     if total_tokens > 0:
-        set_span_attribute(span, "llm.token.counts", json.dumps(usage_dict))
-    set_span_attribute(span, "llm.finish_reason", response.get("done_reason"))
+        set_span_attribute(span, SpanAttributes.LLM_USAGE_PROMPT_TOKENS, input_tokens)
+        set_span_attribute(
+            span, SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, output_tokens
+        )
+        set_span_attribute(span, SpanAttributes.LLM_USAGE_TOTAL_TOKENS, total_tokens)
+
+    set_span_attribute(
+        span,
+        SpanAttributes.LLM_RESPONSE_FINISH_REASON,
+        response.get("done_reason"),
+    )
     if "message" in response:
-        set_span_attribute(span, "llm.responses", json.dumps([response.get("message")]))
+        set_event_completion(span, [response.get("message")])
 
     if "response" in response:
-        set_span_attribute(
-            span,
-            "llm.responses",
-            json.dumps([{"role": "assistant", "content": response.get("response")}]),
+        set_event_completion(
+            span, [{"role": "assistant", "content": response.get("response")}]
         )
 
 
@@ -156,26 +142,23 @@ def _set_input_attributes(span, kwargs, attributes):
 
     for field, value in attributes.model_dump(by_alias=True).items():
         set_span_attribute(span, field, value)
-    if "messages" in kwargs:
+
+    if "options" in kwargs:
         set_span_attribute(
             span,
-            "llm.prompts",
-            json.dumps(kwargs.get("messages", [])),
+            SpanAttributes.LLM_REQUEST_TEMPERATURE,
+            options.get("temperature"),
         )
-    if "prompt" in kwargs:
+        set_span_attribute(span, SpanAttributes.LLM_REQUEST_TOP_P, options.get("top_p"))
         set_span_attribute(
             span,
-            "llm.prompts",
-            json.dumps([{"role": "user", "content": kwargs.get("prompt", "")}]),
+            SpanAttributes.LLM_FREQUENCY_PENALTY,
+            options.get("frequency_penalty"),
         )
-    if "options" in kwargs:
-        set_span_attribute(span, "llm.temperature", options.get("temperature"))
-        set_span_attribute(span, "llm.top_p", options.get("top_p"))
         set_span_attribute(
-            span, "llm.frequency_penalty", options.get("frequency_penalty")
-        )
-        set_span_attribute(
-            span, "llm.presence_penalty", options.get("presence_penalty")
+            span,
+            SpanAttributes.LLM_PRESENCE_PENALTY,
+            options.get("presence_penalty"),
         )
 
 
@@ -194,6 +177,14 @@ def _handle_streaming_response(span, response, api):
             if api == "generate":
                 accumulated_tokens["response"] += chunk["response"]
 
+            span.add_event(
+                Event.STREAM_OUTPUT.value,
+                {
+                    SpanAttributes.LLM_CONTENT_COMPLETION_CHUNK: chunk.get("response")
+                    or chunk.get("message").get("content"),
+                },
+            )
+
         _set_response_attributes(span, chunk | accumulated_tokens)
     finally:
         # Finalize span after processing all chunks
@@ -220,6 +211,12 @@ async def _ahandle_streaming_response(span, response, api):
             if api == "generate":
                 accumulated_tokens["response"] += chunk["response"]
 
+            span.add_event(
+                Event.STREAM_OUTPUT.value,
+                {
+                    SpanAttributes.LLM_CONTENT_COMPLETION_CHUNK: json.dumps(chunk),
+                },
+            )
         _set_response_attributes(span, chunk | accumulated_tokens)
     finally:
         # Finalize span after processing all chunks