feat: add mistral instrumentation (#1100)

* feat: add mistral instrumentation

* fix: remove requirements on mistralai

* fix: add changelog

CHANGELOG.md

@@ -8,6 +8,12 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
 Nothing unreleased.
 
+## [1.1.305] - 2024-06-26
+
+### Added
+
+- Mistral AI instrumentation
+
 ## [1.1.304] - 2024-06-21
 
 ### Fixed

backend/chainlit/__init__.py

@@ -19,6 +19,9 @@
     )
     from chainlit.llama_index.callbacks import LlamaIndexCallbackHandler
     from chainlit.openai import instrument_openai
+    from chainlit.mistralai import instrument_mistralai
+
+from literalai import ChatGeneration, CompletionGeneration, GenerationMessage
 
 import chainlit.input_widget as input_widget
 from chainlit.action import Action
@@ -57,7 +60,6 @@
 from chainlit.user_session import user_session
 from chainlit.utils import make_module_getattr, wrap_user_function
 from chainlit.version import __version__
-from literalai import ChatGeneration, CompletionGeneration, GenerationMessage
 
 if env_found:
     logger.info("Loaded .env file")
@@ -360,6 +362,7 @@ def acall(self):
         "LlamaIndexCallbackHandler": "chainlit.llama_index.callbacks",
         "HaystackAgentCallbackHandler": "chainlit.haystack.callbacks",
         "instrument_openai": "chainlit.openai",
+        "instrument_mistralai": "chainlit.mistralai",
     }
 )
 
@@ -416,6 +419,7 @@ def acall(self):
     "LlamaIndexCallbackHandler",
     "HaystackAgentCallbackHandler",
     "instrument_openai",
+    "instrument_mistralai",
 ]
 
 

backend/chainlit/mistralai/__init__.py

@@ -0,0 +1,53 @@
+import asyncio
+from typing import Union
+
+from literalai import ChatGeneration, CompletionGeneration
+from literalai.helper import timestamp_utc
+
+from chainlit.context import get_context
+from chainlit.step import Step
+from chainlit.utils import check_module_version
+
+
+def instrument_mistralai():
+    from literalai.instrumentation.mistralai import instrument_mistralai
+
+    def on_new_generation(
+        generation: Union["ChatGeneration", "CompletionGeneration"], timing
+    ):
+        context = get_context()
+
+        parent_id = None
+        if context.current_step:
+            parent_id = context.current_step.id
+        elif context.session.root_message:
+            parent_id = context.session.root_message.id
+
+        step = Step(
+            name=generation.model if generation.model else generation.provider,
+            type="llm",
+            parent_id=parent_id,
+        )
+        step.generation = generation
+        # Convert start/end time from seconds to milliseconds
+        step.start = (
+            timestamp_utc(timing.get("start"))
+            if timing.get("start", None) is not None
+            else None
+        )
+        step.end = (
+            timestamp_utc(timing.get("end"))
+            if timing.get("end", None) is not None
+            else None
+        )
+
+        if isinstance(generation, ChatGeneration):
+            step.input = generation.messages
+            step.output = generation.message_completion  # type: ignore
+        else:
+            step.input = generation.prompt
+            step.output = generation.completion
+
+        asyncio.create_task(step.send())
+
+    instrument_mistralai(None, on_new_generation)

backend/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "chainlit"
-version = "1.1.304"
+version = "1.1.305"
 keywords = [
     'LLM',
     'Agents',
@@ -27,7 +27,7 @@ chainlit = 'chainlit.cli:cli'
 [tool.poetry.dependencies]
 python = ">=3.8.1,<4.0.0"
 httpx = ">=0.23.0"
-literalai = "0.0.604"
+literalai = "0.0.607"
 dataclasses_json = "^0.5.7"
 fastapi = "^0.110.1"
 starlette = "^0.37.2"
@@ -51,7 +51,7 @@ pyjwt = "^2.8.0"
 # Fix https://stackoverflow.com/questions/77582651/why-doesnt-numpy-support-pep517-builds
 numpy = [
     { version = "^1.26", python = ">=3.9" },
-    { version = "^1.24.4", python = "<3.9" }
+    { version = "^1.24.4", python = "<3.9" },
 ]
 
 [tool.poetry.group.tests]
@@ -61,7 +61,7 @@ optional = true
 openai = "^1.11.1"
 langchain = "^0.1.5"
 llama-index = "^0.10.45"
-tenacity = "8.3.0" # pin until fixed /~https://github.com/langchain-ai/langchain/issues/22972
+tenacity = "8.3.0"          # pin until fixed /~https://github.com/langchain-ai/langchain/issues/22972
 transformers = "^4.30.1"
 matplotlib = "3.7.1"
 farm-haystack = "^1.18.0"