From 4aa9c239a6313594545c748f5f101d878f2f538a Mon Sep 17 00:00:00 2001
From: jayesh <jayeshparmar9829@gmail.com>
Date: Thu, 28 Nov 2024 23:43:52 +0530
Subject: [PATCH] mistral with openai

---
 livekit-agents/livekit/agents/llm/llm.py      |  1 +
 .../livekit/plugins/openai/llm.py             | 46 +++++++++++++++++--
 .../livekit/plugins/openai/models.py          | 11 +++++
 3 files changed, 55 insertions(+), 3 deletions(-)

diff --git a/livekit-agents/livekit/agents/llm/llm.py b/livekit-agents/livekit/agents/llm/llm.py
index 351fcc9b1..82d13ba33 100644
--- a/livekit-agents/livekit/agents/llm/llm.py
+++ b/livekit-agents/livekit/agents/llm/llm.py
@@ -50,6 +50,7 @@ class Choice:
 @dataclass
 class LLMCapabilities:
     supports_choices_on_int: bool = True
+    supports_stream_options: bool = True
 
 
 @dataclass
diff --git a/livekit-plugins/livekit-plugins-openai/livekit/plugins/openai/llm.py b/livekit-plugins/livekit-plugins-openai/livekit/plugins/openai/llm.py
index 9b187206b..7e32dda03 100644
--- a/livekit-plugins/livekit-plugins-openai/livekit/plugins/openai/llm.py
+++ b/livekit-plugins/livekit-plugins-openai/livekit/plugins/openai/llm.py
@@ -46,6 +46,7 @@
     ChatModels,
     DeepSeekChatModels,
     GroqChatModels,
+    MistralChatModels,
     OctoChatModels,
     PerplexityChatModels,
     TelnyxChatModels,
@@ -85,7 +86,10 @@ def __init__(
         ``OPENAI_API_KEY`` environmental variable.
         """
         super().__init__()
-        self._capabilities = llm.LLMCapabilities(supports_choices_on_int=True)
+        self._capabilities = llm.LLMCapabilities(
+            supports_choices_on_int=True,
+            supports_stream_options=True,
+        )
 
         self._opts = LLMOptions(
             model=model,
@@ -196,6 +200,42 @@ def with_cerebras(
             tool_choice=tool_choice,
         )
 
+    @staticmethod
+    def with_mistral(
+        *,
+        model: str | MistralChatModels = "mistral-large-latest",
+        api_key: str | None = None,
+        base_url: str | None = "https://api.mistral.ai/v1",
+        user: str | None = None,
+        temperature: float | None = None,
+        parallel_tool_calls: bool | None = None,
+        tool_choice: Union[ToolChoice, Literal["auto", "required", "none"]] = "auto",
+    ) -> LLM:
+        """
+        Create a new instance of Mistral LLM.
+
+        ``api_key`` must be set to your Mistral API key, either using the argument or by setting
+        the ``MISTRAL_API_KEY`` environmental variable.
+        """
+
+        api_key = api_key or os.environ.get("MISTRAL_API_KEY")
+        if api_key is None:
+            raise ValueError(
+                "Mistral API key is required, either as argument or set MISTRAL_API_KEY environmental variable"
+            )
+
+        mistral_llm = LLM(
+            model=model,
+            api_key=api_key,
+            base_url=base_url,
+            user=user,
+            temperature=temperature,
+            parallel_tool_calls=parallel_tool_calls,
+            tool_choice=tool_choice,
+        )
+        mistral_llm._capabilities = llm.LLMCapabilities(supports_stream_options=False)
+        return mistral_llm
+
     @staticmethod
     def with_vertex(
         *,
@@ -660,7 +700,8 @@ def chat(
                     }
                 else:
                     opts["tool_choice"] = tool_choice
-
+        if self._capabilities.supports_stream_options:
+            opts["stream_options"] = {"include_usage": True}
         user = self._opts.user or openai.NOT_GIVEN
         if temperature is None:
             temperature = self._opts.temperature
@@ -672,7 +713,6 @@ def chat(
             model=self._opts.model,
             n=n,
             temperature=temperature,
-            stream_options={"include_usage": True},
             stream=True,
             user=user,
             **opts,
diff --git a/livekit-plugins/livekit-plugins-openai/livekit/plugins/openai/models.py b/livekit-plugins/livekit-plugins-openai/livekit/plugins/openai/models.py
index c2667665d..5d162ded4 100644
--- a/livekit-plugins/livekit-plugins-openai/livekit/plugins/openai/models.py
+++ b/livekit-plugins/livekit-plugins-openai/livekit/plugins/openai/models.py
@@ -37,6 +37,17 @@
 
 # adapters for OpenAI-compatible LLMs
 
+MistralChatModels = Literal[
+    "mistral-large-latest",
+    "ministral-3b-latest",
+    "ministral-8b-latest",
+    "mistral-small-latest",
+    "mistral-large-2411",
+    "ministral-3b-2410",
+    "ministral-8b-2410",
+    "mistral-small-2409",
+]
+
 TelnyxChatModels = Literal[
     "meta-llama/Meta-Llama-3.1-8B-Instruct",
     "meta-llama/Meta-Llama-3.1-70B-Instruct",