handle attribute setting

Author: obinnascale3

src/langtrace_python_sdk/instrumentation/mistral/instrumentation.py

@@ -20,9 +20,7 @@
 
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 from opentelemetry.trace import get_tracer
-from wrapt import wrap_function_wrapper
-
-from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from wrapt import wrap_function_wrapper as _W
 
 from langtrace_python_sdk.instrumentation.mistral.patch import (
     chat_complete,
@@ -34,41 +32,41 @@
 class MistralInstrumentation(BaseInstrumentor):
 
     def instrumentation_dependencies(self) -> Collection[str]:
-        return ["mistralai >= 1.0.1", "trace-attributes >= 4.0.5"]
+        return ["mistralai >= 1.0.1"]
 
     def _instrument(self, **kwargs):
         tracer_provider = kwargs.get("tracer_provider")
         tracer = get_tracer(__name__, "", tracer_provider)
         version = importlib.metadata.version("mistralai")
 
-        wrap_function_wrapper(
-            "mistralai.chat",
-            "Chat.complete",
-            chat_complete("mistral.chat.complete", version, tracer),
+        _W(
+            module="mistralai.chat",
+            name="Chat.complete",
+            wrapper=chat_complete("mistral.chat.complete", version, tracer),
         )
 
-        wrap_function_wrapper(
-            "mistralai.chat",
-            "Chat.stream",
-            chat_complete("mistral.chat.stream", version, tracer, is_streaming=True),
+        _W(
+            module="mistralai.chat",
+            name="Chat.stream",
+            wrapper=chat_complete("mistral.chat.stream", version, tracer, is_streaming=True),
         )
 
-        wrap_function_wrapper(
-            "mistralai.chat",
-            "Chat.complete_async",
-            chat_complete("mistral.chat.complete_async", version, tracer, is_async=True),
+        _W(
+            module="mistralai.chat",
+            name="Chat.complete_async",
+            wrapper=chat_complete("mistral.chat.complete_async", version, tracer, is_async=True),
         )
 
-        wrap_function_wrapper(
-            "mistralai.embeddings",
-            "Embeddings.create",
-            embeddings_create("mistral.embeddings.create", version, tracer),
+        _W(
+            module="mistralai.embeddings",
+            name="Embeddings.create",
+            wrapper=embeddings_create("mistral.embeddings.create", version, tracer),
         )
 
-        wrap_function_wrapper(
-            "mistralai.embeddings",
-            "Embeddings.create_async",
-            embeddings_create("mistral.embeddings.create_async", version, tracer, is_async=True),
+        _W(
+            module="mistralai.embeddings",
+            name="Embeddings.create_async",
+            wrapper=embeddings_create("mistral.embeddings.create_async", version, tracer, is_async=True),
         )
 
     def _uninstrument(self, **kwargs):

src/langtrace_python_sdk/instrumentation/mistral/patch.py

@@ -32,17 +32,15 @@
 from langtrace_python_sdk.constants.instrumentation.mistral import APIS
 from langtrace_python_sdk.utils.llm import (
     calculate_prompt_tokens,
-    get_base_url,
     get_extra_attributes,
     get_langtrace_attributes,
     get_llm_request_attributes,
     get_llm_url,
     get_span_name,
-    get_tool_calls,
-    is_streaming,
     set_event_completion,
     StreamWrapper,
     set_span_attributes,
+    set_usage_attributes,
 )
 
 from langtrace_python_sdk.instrumentation.openai.patch import extract_content
@@ -80,19 +78,6 @@ def traced_method(wrapped, instance, args, kwargs):
             result = wrapped(*args, **kwargs)
             if is_streaming:
                 prompt_tokens = 0
-                for message in kwargs.get("messages", {}):
-                    prompt_tokens += calculate_prompt_tokens(
-                        json.dumps((str(message))), kwargs.get("model")
-                    )
-
-                if (
-                    kwargs.get("functions") is not None
-                ):
-                    for function in kwargs.get("functions"):
-                        prompt_tokens += calculate_prompt_tokens(
-                            json.dumps(function), kwargs.get("model")
-                        )
-
                 return StreamWrapper(
                     result,
                     span,
@@ -207,20 +192,4 @@ def _set_response_attributes(span, kwargs, result):
 
     # Get the usage
     if hasattr(result, "usage") and result.usage is not None:
-        usage = result.usage
-        if usage is not None:
-            set_span_attribute(
-                span,
-                SpanAttributes.LLM_USAGE_PROMPT_TOKENS,
-                result.usage.prompt_tokens,
-            )
-            set_span_attribute(
-                span,
-                SpanAttributes.LLM_USAGE_COMPLETION_TOKENS,
-                result.usage.completion_tokens,
-            )
-            set_span_attribute(
-                span,
-                SpanAttributes.LLM_USAGE_TOTAL_TOKENS,
-                result.usage.total_tokens,
-            )
+        set_usage_attributes(span, result.usage)

src/langtrace_python_sdk/utils/llm.py

@@ -411,6 +411,12 @@ def set_usage_attributes(self, chunk):
                 self.completion_tokens = chunk["eval_count"]
 
     def process_chunk(self, chunk):
+        # Mistral nests the chunk data under a `data` attribute
+        if (
+            hasattr(chunk, "data") and chunk.data is not None
+            and hasattr(chunk.data, "choices") and chunk.data.choices is not None
+        ):
+            chunk = chunk.data
         self.set_response_model(chunk=chunk)
         self.build_streaming_response(chunk=chunk)
         self.set_usage_attributes(chunk=chunk)