From 6d07ae0a9b6ced88749f0383129e9234250f03e7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Hugues=20de=20Saxc=C3=A9?= <hugues@chainlit.io>
Date: Wed, 26 Jun 2024 01:17:01 +0200
Subject: [PATCH] feat: add mistral instrumentation (#1100)

* feat: add mistral instrumentation

* fix: remove requirements on mistralai

* fix: add changelog
---
 CHANGELOG.md                           |  6 +++
 backend/chainlit/__init__.py           |  6 ++-
 backend/chainlit/mistralai/__init__.py | 53 ++++++++++++++++++++++++++
 backend/pyproject.toml                 |  8 ++--
 4 files changed, 68 insertions(+), 5 deletions(-)
 create mode 100644 backend/chainlit/mistralai/__init__.py

diff --git a/CHANGELOG.md b/CHANGELOG.md
index ae79bafd3f..e11473144f 100644
--- a/CHANGELOG.md
+++ b/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
diff --git a/backend/chainlit/__init__.py b/backend/chainlit/__init__.py
index d22abc46a4..f6d14bfe17 100644
--- a/backend/chainlit/__init__.py
+++ b/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",
 ]
 
 
diff --git a/backend/chainlit/mistralai/__init__.py b/backend/chainlit/mistralai/__init__.py
new file mode 100644
index 0000000000..5133c04827
--- /dev/null
+++ b/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)
diff --git a/backend/pyproject.toml b/backend/pyproject.toml
index 4bebe66e18..0293f59f12 100644
--- a/backend/pyproject.toml
+++ b/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"