From 0763aed5f745dfddc7839ad3c308354b5e19588f Mon Sep 17 00:00:00 2001
From: Leonardo Crociani <88234007+leonardocrociani@users.noreply.github.com>
Date: Wed, 10 Sep 2025 18:54:26 +0200
Subject: [PATCH 1/6] mistral ai genie

---
 pyproject.toml                |  2 +
 src/chonkie/genie/__init__.py |  2 +
 src/chonkie/genie/mistral.py  | 96 +++++++++++++++++++++++++++++++++++
 3 files changed, 100 insertions(+)
 create mode 100644 src/chonkie/genie/mistral.py

diff --git a/pyproject.toml b/pyproject.toml
index d9f90cbe1..28e0f1b6e 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -84,6 +84,7 @@ openai = [
     "numpy>=2.0.0, <3.0",
     "pydantic>=2.0.0",
 ]
+mistralai = ["mistralai>=1.9.10"]
 semantic = ["tokenizers>=0.16.0", "model2vec>=0.3.0", "numpy>=2.0.0, <3.0"]
 gemini = ["pydantic>=2.0.0", "google-genai>=1.0.0"]
 azure-openai = [
@@ -121,6 +122,7 @@ all = [
     "sentence-transformers>=3.0.0",
     "numpy>=2.0.0, <3.0",
     "openai>=1.0.0",
+    "mistralai>=1.9.10",
     "model2vec>=0.3.0",
     "cohere>=5.13.0",
     "voyageai>=0.3.2",
diff --git a/src/chonkie/genie/__init__.py b/src/chonkie/genie/__init__.py
index 71be72f7e..694aff54b 100644
--- a/src/chonkie/genie/__init__.py
+++ b/src/chonkie/genie/__init__.py
@@ -3,10 +3,12 @@
 from .base import BaseGenie
 from .gemini import GeminiGenie
 from .openai import OpenAIGenie
+from .mistral import MistralGenie
 
 # Add all genie classes to __all__
 __all__ = [
     "BaseGenie",
     "GeminiGenie",
     "OpenAIGenie",
+    "MistralGenie",
 ]
\ No newline at end of file
diff --git a/src/chonkie/genie/mistral.py b/src/chonkie/genie/mistral.py
new file mode 100644
index 000000000..4d6dcdf89
--- /dev/null
+++ b/src/chonkie/genie/mistral.py
@@ -0,0 +1,96 @@
+"""Implementation of the MistralGenie class."""
+import importlib.util as importutil
+import json
+import os
+from typing import TYPE_CHECKING, Any, Dict, Optional
+
+from .base import BaseGenie
+
+if TYPE_CHECKING:
+    from pydantic import BaseModel
+
+
+class MistralGenie(BaseGenie):
+    """Mistral's Genie."""
+
+    def __init__(
+        self,
+        model: str = "mistral-large-latest",
+        api_key: Optional[str] = None
+    ):
+        """Initialize the MistralGenie class.
+
+        Args:
+            model (str): The model to use.
+            api_key (Optional[str]): The API key to use.
+
+        """
+        super().__init__()
+
+        # Lazily import the dependencies
+        self._import_dependencies()
+
+        # Initialize the API key
+        self.api_key = api_key or os.environ.get("MISTRAL_API_KEY")
+        if not self.api_key:
+            raise ValueError(
+                "MistralGenie requires an API key. Either pass the `api_key` "
+                "parameter or set the `MISTRAL_API_KEY` in your environment."
+            )
+
+        # Initialize the client and model
+        self.client = MistralClient(api_key=self.api_key)  # type: ignore
+        self.model = model
+
+    def generate(self, prompt: str) -> str:
+        """Generate a response based on the given prompt."""
+        messages = [{"role": "user", "content": prompt}]
+        response = self.client.chat(model=self.model, messages=messages)
+        return str(response.choices[0].message.content)
+
+    def generate_json(self, prompt: str, schema: "BaseModel") -> Dict[str, Any]:
+        """Generate a JSON response based on the given prompt and schema."""
+        json_schema = json.dumps(schema.model_json_schema(), indent=2)
+        full_prompt = (
+            f"{prompt}\n\n"
+            "Please provide the output in a JSON format that strictly adheres to the "
+            f"following schema:\n```json\n{json_schema}\n```"
+        )
+        messages = [{"role": "user", "content": full_prompt}]
+
+        response = self.client.chat.complete(
+            model=self.model,
+            messages=messages,
+            response_format={"type": "json_object"},
+        )
+
+        try:
+            content = response.choices[0].message.content
+            return dict(json.loads(content))
+        except Exception as e:
+            raise ValueError(f"Failed to parse JSON response: {e}")
+
+    def _is_available(self) -> bool:
+        """Check if all the dependencies are available in the environment."""
+        if (
+            importutil.find_spec("pydantic") is not None
+            and importutil.find_spec("mistralai") is not None
+        ):
+            return True
+        return False
+
+    def _import_dependencies(self) -> None:
+        """Import all the required dependencies."""
+        if self._is_available():
+            global BaseModel, MistralClient
+            from mistralai import Mistral as MistralClient
+            from pydantic import BaseModel
+        else:
+            raise ImportError(
+                "One or more of the required modules are not available: "
+                "[pydantic, mistralai]"
+            )
+
+    def __repr__(self) -> str:
+        """Return a string representation of the MistralGenie instance."""
+        return f"MistralGenie(model={self.model})"
\ No newline at end of file

From bdfa96d14a8ff6979864d1e463be7a1bedf18419 Mon Sep 17 00:00:00 2001
From: Leonardo Crociani <88234007+leonardocrociani@users.noreply.github.com>
Date: Wed, 10 Sep 2025 19:08:08 +0200
Subject: [PATCH 2/6] dependency fix

---
 pyproject.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index 28e0f1b6e..92dba910d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -109,7 +109,7 @@ pinecone = ["pinecone"]
 mongodb = ["pymongo"]
 
 # Optional dependencies for the genie
-genie = ["pydantic>=2.0.0", "google-genai>=1.0.0"]
+genie = ["pydantic>=2.0.0", "google-genai>=1.0.0", "mistralai>=1.9.10", "numpy>=2.0.0, <3.0"]
 
 # All dependencies
 all = [

From 544b92494136e5ae49b0a4e92d4d3ab5a318d074 Mon Sep 17 00:00:00 2001
From: Leonardo Crociani <88234007+leonardocrociani@users.noreply.github.com>
Date: Thu, 11 Sep 2025 14:42:04 +0200
Subject: [PATCH 3/6] Update src/chonkie/genie/mistral.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
---
 src/chonkie/genie/mistral.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/chonkie/genie/mistral.py b/src/chonkie/genie/mistral.py
index 4d6dcdf89..9dafb96bb 100644
--- a/src/chonkie/genie/mistral.py
+++ b/src/chonkie/genie/mistral.py
@@ -45,7 +45,7 @@ def __init__(
     def generate(self, prompt: str) -> str:
         """Generate a response based on the given prompt."""
         messages = [{"role": "user", "content": prompt}]
-        response = self.client.chat(model=self.model, messages=messages)
+        response = self.client.chat.complete(model=self.model, messages=messages)
         return str(response.choices[0].message.content)
 
     def generate_json(self, prompt: str, schema: "BaseModel") -> Dict[str, Any]:

From dc526268547c491bef4973069f6a4e6507d3aaa4 Mon Sep 17 00:00:00 2001
From: Leonardo Crociani <88234007+leonardocrociani@users.noreply.github.com>
Date: Thu, 11 Sep 2025 14:42:13 +0200
Subject: [PATCH 4/6] Update src/chonkie/genie/mistral.py

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
---
 src/chonkie/genie/mistral.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/chonkie/genie/mistral.py b/src/chonkie/genie/mistral.py
index 9dafb96bb..20fe1c212 100644
--- a/src/chonkie/genie/mistral.py
+++ b/src/chonkie/genie/mistral.py
@@ -58,7 +58,7 @@ def generate_json(self, prompt: str, schema: "BaseModel") -> Dict[str, Any]:
         )
         messages = [{"role": "user", "content": full_prompt}]
 
-        response = self.client.chat.complete(
+        response = self.client.chat(
             model=self.model,
             messages=messages,
             response_format={"type": "json_object"},

From 6391973b08c13d83aa32e1e4ce6b3c56f37ed864 Mon Sep 17 00:00:00 2001
From: Leonardo Crociani <88234007+leonardocrociani@users.noreply.github.com>
Date: Thu, 11 Sep 2025 14:42:27 +0200
Subject: [PATCH 5/6] Update src/chonkie/genie/mistral.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
---
 src/chonkie/genie/mistral.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/chonkie/genie/mistral.py b/src/chonkie/genie/mistral.py
index 20fe1c212..c8cdb6ada 100644
--- a/src/chonkie/genie/mistral.py
+++ b/src/chonkie/genie/mistral.py
@@ -67,7 +67,7 @@ def generate_json(self, prompt: str, schema: "BaseModel") -> Dict[str, Any]:
         try:
             content = response.choices[0].message.content
             return dict(json.loads(content))
-        except Exception as e:
+        except json.JSONDecodeError as e:
             raise ValueError(f"Failed to parse JSON response: {e}")
 
     def _is_available(self) -> bool:

From 2f0fc6ad1dd05501a353a639f748919047ac1c78 Mon Sep 17 00:00:00 2001
From: Leonardo Crociani <88234007+leonardocrociani@users.noreply.github.com>
Date: Tue, 7 Oct 2025 13:02:29 +0200
Subject: [PATCH 6/6] Change chat method to chat.complete in mistral.py

---
 src/chonkie/genie/mistral.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/chonkie/genie/mistral.py b/src/chonkie/genie/mistral.py
index c8cdb6ada..1c9bb3e60 100644
--- a/src/chonkie/genie/mistral.py
+++ b/src/chonkie/genie/mistral.py
@@ -58,7 +58,7 @@ def generate_json(self, prompt: str, schema: "BaseModel") -> Dict[str, Any]:
         )
         messages = [{"role": "user", "content": full_prompt}]
 
-        response = self.client.chat(
+        response = self.client.chat.complete(
             model=self.model,
             messages=messages,
             response_format={"type": "json_object"},
@@ -93,4 +93,4 @@ def _import_dependencies(self) -> None:
 
     def __repr__(self) -> str:
         """Return a string representation of the MistralGenie instance."""
-        return f"MistralGenie(model={self.model})"
\ No newline at end of file
+        return f"MistralGenie(model={self.model})"