feat(mistral): add agent support

python/instrumentation/openinference-instrumentation-mistralai/examples/agent_completions.py

@@ -0,0 +1,25 @@
+from mistralai import Mistral
+from opentelemetry import trace as trace_api
+from opentelemetry.sdk import trace as trace_sdk
+from opentelemetry.sdk.trace.export import ConsoleSpanExporter, SimpleSpanProcessor
+
+from openinference.instrumentation.mistralai import MistralAIInstrumentor
+
+tracer_provider = trace_sdk.TracerProvider()
+tracer_provider.add_span_processor(SimpleSpanProcessor(ConsoleSpanExporter()))
+trace_api.set_tracer_provider(tracer_provider)
+
+MistralAIInstrumentor().instrument()
+
+if __name__ == "__main__":
+    client = Mistral(api_key="redacted")
+    response = client.agents.complete(
+        agent_id="ag:ad73bfd7:20240912:python-codegen-agent:0375a7cf",
+        messages=[
+            {
+                "role": "user",
+                "content":  "is it too much"
+            },
+        ],
+    )
+    print(response)

python/instrumentation/openinference-instrumentation-mistralai/src/openinference/instrumentation/mistralai/__init__.py

@@ -34,6 +34,10 @@ class MistralAIInstrumentor(BaseInstrumentor):  # type: ignore
         "_original_sync_stream_chat_method",
         "_original_async_chat_method",
         "_original_async_stream_chat_method",
+        "_original_sync_agent_method",
+        "_original_sync_stream_agent_method",
+        "_original_async_agent_method",
+        "_original_async_stream_agent_method",
     )
 
     def instrumentation_dependencies(self) -> Collection[str]:
@@ -54,6 +58,7 @@ def _instrument(self, **kwargs: Any) -> None:
         try:
             import mistralai
             from mistralai.chat import Chat
+            from mistralai.agents import Agents
         except ImportError as err:
             raise Exception(
                 "Could not import mistralai. Please install with `pip install mistralai`."
@@ -63,6 +68,10 @@ def _instrument(self, **kwargs: Any) -> None:
         self._original_sync_stream_chat_method = Chat.stream
         self._original_async_chat_method = Chat.complete_async
         self._original_async_stream_chat_method = Chat.stream_async
+        self._original_sync_agent_method = Agents.complete
+        self._original_sync_stream_agent_method = Agents.stream
+        self._original_async_agent_method = Agents.complete_async
+        self._original_async_stream_agent_method = Agents.stream_async
         wrap_function_wrapper(
             module="mistralai.chat",
             name="Chat.complete",
@@ -87,10 +96,39 @@ def _instrument(self, **kwargs: Any) -> None:
             wrapper=_AsyncStreamChatWrapper(self._tracer, mistralai),
         )
 
+        wrap_function_wrapper(
+            module="mistralai.agents",
+            name="Agents.complete",
+            wrapper=_SyncChatWrapper(self._tracer, mistralai),
+        )
+
+        wrap_function_wrapper(
+            module="mistralai.agents",
+            name="Agents.stream",
+            wrapper=_SyncChatWrapper(self._tracer, mistralai),
+        )
+
+        wrap_function_wrapper(
+            module="mistralai.agents",
+            name="Agents.complete_async",
+            wrapper=_AsyncChatWrapper(self._tracer, mistralai),
+        )
+
+        wrap_function_wrapper(
+            module="mistralai.agents",
+            name="Agents.stream_async",
+            wrapper=_AsyncStreamChatWrapper(self._tracer, mistralai),
+        )
+
     def _uninstrument(self, **kwargs: Any) -> None:
         from mistralai.chat import Chat
+        from mistralai.agent import Agent
 
         Chat.complete = self._original_sync_chat_method  # type: ignore
         Chat.stream = self._original_sync_stream_chat_method  # type: ignore
         Chat.complete_async = self._original_async_chat_method  # type: ignore
         Chat.stream_async = self._original_async_stream_chat_method  # type: ignore
+        Agent.complete = self._original_sync_agent_method  # type: ignore
+        Agent.stream = self._original_sync_stream_agent_method  # type: ignore
+        Agent.complete_async = self._original_async_agent_method  # type: ignore
+        Agent.stream_async = self._original_async_stream_agent_method  # type: ignore

python/instrumentation/openinference-instrumentation-mistralai/src/openinference/instrumentation/mistralai/_chat_wrapper.py

@@ -218,7 +218,7 @@ def __call__(
             return wrapped(*args, **kwargs)
         try:
             request_parameters = self._parse_args(signature(wrapped), instance, *args, **kwargs)
-            span_name = "MistralClient.chat"
+            span_name = "MistralClient.agents"
         except Exception:
             logger.exception("Failed to parse request args")
             return wrapped(*args, **kwargs)
@@ -264,7 +264,7 @@ async def __call__(
             return await wrapped(*args, **kwargs)
         try:
             request_parameters = self._parse_args(signature(wrapped), instance, *args, **kwargs)
-            span_name = "MistralAsyncClient.chat"
+            span_name = "MistralAsyncClient.agents"
         except Exception:
             logger.exception("Failed to parse request args")
             return await wrapped(*args, **kwargs)
@@ -310,7 +310,7 @@ def __call__(
             return wrapped(*args, **kwargs)
         try:
             request_parameters = self._parse_args(signature(wrapped), instance, *args, **kwargs)
-            span_name = "MistralAsyncClient.chat"
+            span_name = "MistralAsyncClient.agents"
         except Exception:
             logger.exception("Failed to parse request args")
             return wrapped(*args, **kwargs)