Add OTel Tracing (#280)

Author: amenasria

examples/mistral/agents/async_multi_turn_conversation.py

@@ -0,0 +1,69 @@
+import os
+from mistralai import Mistral
+
+from mistralai.extra.run.context import RunContext
+import logging
+import time
+import asyncio
+
+
+MODEL = "mistral-medium-latest"
+
+USER_MESSAGE = """
+Please make the Secret Santa for me
+To properly do it you need to:
+- Get the friend you were assigned to (using the get_secret_santa_assignment function)
+- Read into his gift wishlist what they would like to receive (using the get_gift_wishlist function)
+- Buy the gift (using the buy_gift function)
+- Find the best website to buy the gift using a web search
+- Send it to them (using the send_gift function)
+"""
+
+
+async def main():
+    api_key = os.environ["MISTRAL_API_KEY"]
+    mistral_agent_id = os.environ["MISTRAL_AGENT_ID"]
+    client = Mistral(
+        api_key=api_key, debug_logger=logging.getLogger("mistralai")
+    )
+
+    async with RunContext(
+        agent_id=mistral_agent_id
+    ) as run_context:
+        run_context.register_func(get_secret_santa_assignment)
+        run_context.register_func(get_gift_wishlist)
+        run_context.register_func(buy_gift)
+        run_context.register_func(send_gift)
+
+        await client.beta.conversations.run_async(
+            run_ctx=run_context,
+            inputs=USER_MESSAGE,
+        )
+
+
+def get_secret_santa_assignment():
+    """Get the friend you were assigned to"""
+    time.sleep(2)
+    return "John Doe"
+
+
+def get_gift_wishlist(friend_name: str):
+    """Get the gift wishlist of the friend you were assigned to"""
+    time.sleep(1.5)
+    return ["Book", "Chocolate", "T-Shirt"]
+
+
+def buy_gift(gift_name: str):
+    """Buy the gift you want to send to your friend"""
+    time.sleep(1.1)
+    return f"Bought {gift_name}"
+
+
+def send_gift(friend_name: str, gift_name: str, website: str):
+    """Send the gift to your friend"""
+    time.sleep(2.2)
+    return f"Sent {gift_name} to {friend_name} bought on {website}"
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

poetry.lock

@@ -1,6 +1,6 @@
 [project]
 name = "mistralai"
-version = "1.9.11"
+version = "1.9.12"
 description = "Python Client SDK for the Mistral AI API."
 authors = [{ name = "Mistral" },]
 readme = "README-PYPI.md"
@@ -13,6 +13,10 @@ dependencies = [
     "typing-inspection >=0.4.0",
     "pyyaml (>=6.0.2,<7.0.0)",
     "invoke (>=2.2.0,<3.0.0)",
+    "opentelemetry-sdk (>=1.33.1,<2.0.0)",
+    "opentelemetry-api (>=1.33.1,<2.0.0)",
+    "opentelemetry-exporter-otlp-proto-http (>=1.37.0,<2.0.0)",
+    "opentelemetry-semantic-conventions (>=0.59b0,<0.60)",
 ]
 
 [tool.poetry]

pyproject.toml

@@ -1,5 +1,6 @@
 from .custom_user_agent import CustomUserAgentHook
 from .deprecation_warning import DeprecationWarningHook
+from .tracing import TracingHook
 from .types import Hooks
 
 # This file is only ever generated once on the first generation and then is free to be modified.
@@ -13,5 +14,9 @@ def init_hooks(hooks: Hooks):
     with an instance of a hook that implements that specific Hook interface
     Hooks are registered per SDK instance, and are valid for the lifetime of the SDK instance
     """
+    tracing_hook = TracingHook()
     hooks.register_before_request_hook(CustomUserAgentHook())
     hooks.register_after_success_hook(DeprecationWarningHook())
+    hooks.register_after_success_hook(tracing_hook)
+    hooks.register_before_request_hook(tracing_hook)
+    hooks.register_after_error_hook(tracing_hook)