Merge pull request #296 from Scale3-Labs/obinna/S3EN-2344-instrument-mistral

Obinna/s3 en 2344 instrument mistral

Author: alizenhom

README.md

@@ -246,6 +246,7 @@ Langtrace automatically captures traces from the following vendors:
 | Groq         | LLM             | :x:                | :white_check_mark:              |
 | Perplexity   | LLM             | :white_check_mark: | :white_check_mark:              |
 | Gemini       | LLM             | :x:                | :white_check_mark:              |
+| Mistral      | LLM             | :x:                | :white_check_mark:              |
 | Langchain    | Framework       | :x:                | :white_check_mark:              |
 | LlamaIndex   | Framework       | :white_check_mark: | :white_check_mark:              |
 | Langgraph    | Framework       | :x:                | :white_check_mark:              |

pyproject.toml

@@ -18,7 +18,7 @@ classifiers=[
     "Operating System :: OS Independent",
 ]
 dependencies = [
-  'trace-attributes==7.0.3',
+  'trace-attributes==7.0.4',
   'opentelemetry-api>=1.25.0',
   'opentelemetry-sdk>=1.25.0',
   'opentelemetry-instrumentation>=0.47b0',
@@ -52,7 +52,8 @@ dev = [
     "ollama",
     "groq",
     "google-generativeai",
-    "google-cloud-aiplatform"
+    "google-cloud-aiplatform",
+    "mistralai"
 ]
 
 test = [

src/examples/mistral_example/__init__.py

@@ -0,0 +1,14 @@
+import asyncio
+from examples.mistral_example.complete import chat_complete
+from examples.mistral_example.complete_async import complete_async
+from examples.mistral_example.embeddings import embeddings_create
+from langtrace_python_sdk import langtrace, with_langtrace_root_span
+
+langtrace.init()
+
+class MistralRunner:
+    @with_langtrace_root_span("Mistral")
+    def run(self):
+        chat_complete()
+        asyncio.run(complete_async())
+        embeddings_create()

src/examples/mistral_example/complete.py

@@ -0,0 +1,18 @@
+import os
+from langtrace_python_sdk import with_langtrace_root_span
+from mistralai import Mistral
+
+@with_langtrace_root_span("chat_complete")
+def chat_complete():
+    model = "mistral-large-latest"
+    client = Mistral(api_key=os.environ["MISTRAL_API_KEY"])
+    chat_response = client.chat.complete(
+        model= model,
+        messages = [
+            {
+                "role": "user",
+                "content": "I need 10 cocktail recipes with tequila other than the classics like margarita, tequila"
+            },
+        ]
+    )
+    print(chat_response.choices[0].message.content)

src/examples/mistral_example/complete_async.py

@@ -0,0 +1,16 @@
+import os
+from langtrace_python_sdk import with_langtrace_root_span
+from mistralai import Mistral
+
+@with_langtrace_root_span("chat_complete_async")
+async def complete_async():
+    client = Mistral(api_key=os.environ["MISTRAL_API_KEY"])
+    res = await client.chat.complete_async(model="mistral-small-latest", messages=[
+        {
+            "content": "Which locations should I visit when I travel to New york? Answer in one short sentence.",
+            "role": "user",
+        },
+    ])
+    if res is not None:
+        # handle response
+        print(res.choices[0].message.content)

src/examples/mistral_example/embeddings.py

@@ -0,0 +1,17 @@
+import os
+from langtrace_python_sdk import with_langtrace_root_span
+from mistralai import Mistral
+
+
+@with_langtrace_root_span("create_embeddings")
+def embeddings_create():
+    model = "mistral-embed"
+
+    client = Mistral(api_key=os.environ["MISTRAL_API_KEY"])
+
+    embeddings_batch_response = client.embeddings.create(
+        model=model,
+        inputs=["Embed this sentence.", "As well as this one."],
+    )
+
+    print(embeddings_batch_response.data[0].embedding)

src/langtrace_python_sdk/constants/instrumentation/common.py

@@ -29,6 +29,7 @@
     "OLLAMA": "Ollama",
     "VERTEXAI": "VertexAI",
     "GEMINI": "Gemini",
+    "MISTRAL": "Mistral",
 }
 
 LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY = "langtrace_additional_attributes"

src/langtrace_python_sdk/constants/instrumentation/mistral.py

@@ -0,0 +1,24 @@
+from langtrace.trace_attributes import MistralMethods
+
+APIS = {
+    "CHAT_COMPLETE": {
+        "METHOD": MistralMethods.CHAT_COMPLETE.value,
+        "ENDPOINT": "/v1/chat/completions",
+    },
+    "ASYNC_CHAT_COMPLETE": {
+        "METHOD": MistralMethods.ASYNC_CHAT_COMPLETE.value,
+        "ENDPOINT": "/v1/chat/completions",
+    },
+    "CHAT_STREAM": {
+        "METHOD": MistralMethods.CHAT_STREAM.value,
+        "ENDPOINT": "/v1/chat/completions",
+    },
+    "EMBEDDINGS_CREATE": {
+        "METHOD": MistralMethods.EMBEDDINGS_CREATE.value,
+        "ENDPOINT": "/v1/embeddings",
+    },
+    "ASYNC_EMBEDDINGS_CREATE": {
+        "METHOD": MistralMethods.ASYNC_EMBEDDINGS_CREATE.value,
+        "ENDPOINT": "/v1/embeddings",
+    },
+}

src/langtrace_python_sdk/instrumentation/__init__.py

@@ -16,6 +16,7 @@
 from .dspy import DspyInstrumentation
 from .vertexai import VertexAIInstrumentation
 from .gemini import GeminiInstrumentation
+from .mistral import MistralInstrumentation
 
 __all__ = [
     "AnthropicInstrumentation",
@@ -36,4 +37,5 @@
     "DspyInstrumentation",
     "VertexAIInstrumentation",
     "GeminiInstrumentation",
+    "MistralInstrumentation",
 ]

src/langtrace_python_sdk/instrumentation/mistral/__init__.py

@@ -0,0 +1,3 @@
+from .instrumentation import MistralInstrumentation
+
+__all__ = ["MistralInstrumentation"]