✨(models) add mistral support & customization

This allows to connect to a Mistral model, on a Mistral
platform, self-hosted.

CHANGELOG.md

@@ -45,6 +45,7 @@ and this project adheres to
 - 📈(langfuse) add light instrumentation #26
 - 🚑️(agent) allow Mistral w/ vLLM & tools #36
 - ✨(web-search) add Brave search tool #47
+- ✨(models) add mistral support & customization #51
 
 
 [unreleased]: https://github.com/numerique-gouv/conversations/compare/HEAD...main

src/backend/chat/agents/base.py

@@ -1,18 +1,104 @@
 """Base module for PydanticAI agents."""
 
 import dataclasses
+import logging
 
 from django.conf import settings
 from django.core.exceptions import ImproperlyConfigured
 
+import httpx
 from pydantic_ai import Agent
-from pydantic_ai.models.openai import OpenAIChatModel
-from pydantic_ai.profiles.openai import OpenAIModelProfile
-from pydantic_ai.providers.openai import OpenAIProvider
+from pydantic_ai.models import get_user_agent
+from pydantic_ai.profiles import ModelProfile
 from pydantic_ai.toolsets import FunctionToolset
 
 from chat.tools import get_pydantic_tools_by_name
 
+logger = logging.getLogger(__name__)
+
+
+def prepare_custom_model(configuration: "chat.llm_configuration.LLModel"):
+    """
+    Prepare a custom model instance based on the provided configuration.
+
+    Only few providers are supported at the moment, according to our needs.
+    We define custom models/providers to be able to keep specific configuration
+    when needed.
+    """
+    # pylint: disable=import-outside-toplevel
+
+    match configuration.provider.kind:
+        case "mistral":
+            import pydantic_ai.models.mistral as mistral_models  # noqa: PLC0415
+            from pydantic_ai.providers.mistral import MistralProvider  # noqa: PLC0415
+
+            # --- Monkey patch for pydantic_ai.models.mistral._map_content ---
+            # pylint: disable=protected-access
+
+            # ⚠ WARNING ⚠ WARNING ⚠ WARNING ⚠ WARNING ⚠ WARNING ⚠ WARNING ⚠ WARNING ⚠ WARNING ⚠
+            # |  This workaround is fragile and only works because we are in streaming mode.  |
+            # ⚠ WARNING ⚠ WARNING ⚠ WARNING ⚠ WARNING ⚠ WARNING ⚠ WARNING ⚠ WARNING ⚠ WARNING ⚠
+
+            # The original _map_content raises exceptions for some when responses
+            # contains citation/reference data, which is the case anytime we use
+            # web search or other RAG tool (https://docs.mistral.ai/capabilities/citations/).
+            # We make the patch idempotent using a sentinel attribute so repeated calls
+            # to prepare_custom_model do not re-wrap and do not cause recursive calls.
+            if not getattr(mistral_models, "__safe_map_patched__", False):
+                _original_map_content = mistral_models._map_content  # noqa: SLF001
+
+                def _safe_map_content(*args, **kwargs):
+                    try:
+                        return _original_map_content(*args, **kwargs)
+                    except AssertionError as exc:
+                        logger.debug("Caught exception in _map_content: %s", exc)
+                        return None, []
+
+                # Replace the original module-level function
+                # mistral_models._map_content = _safe_map_content
+                mistral_models.__safe_map_patched__ = True
+            # pylint: enable=protected-access
+            # --- End monkey patch ---
+
+            return mistral_models.MistralModel(
+                model_name=configuration.model_name,
+                profile=(
+                    ModelProfile(**configuration.profile.dict(exclude_unset=True))
+                    if configuration.profile
+                    else None
+                ),
+                provider=MistralProvider(
+                    api_key=configuration.provider.api_key,
+                    base_url=configuration.provider.base_url,
+                    # Disable the use of cached client
+                    http_client=httpx.AsyncClient(
+                        timeout=httpx.Timeout(timeout=600, connect=5),
+                        headers={"User-Agent": get_user_agent()},
+                    ),
+                ),
+            )
+        case "openai":
+            from pydantic_ai.models.openai import OpenAIChatModel  # noqa: PLC0415
+            from pydantic_ai.profiles.openai import OpenAIModelProfile  # noqa: PLC0415
+            from pydantic_ai.providers.openai import OpenAIProvider  # noqa: PLC0415
+
+            return OpenAIChatModel(
+                model_name=configuration.model_name,
+                profile=(
+                    OpenAIModelProfile(**configuration.profile.dict(exclude_unset=True))
+                    if configuration.profile
+                    else None
+                ),
+                provider=OpenAIProvider(
+                    base_url=configuration.provider.base_url,
+                    api_key=configuration.provider.api_key,
+                ),
+            )
+        case _:
+            raise ImproperlyConfigured(
+                f"Unsupported provider kind '{configuration.provider.kind}' for custom model."
+            )
+
 
 @dataclasses.dataclass(init=False)
 class BaseAgent(Agent):
@@ -35,20 +121,14 @@ class BaseAgent(Agent):
                 f"LLM model configuration '{model_hrid}' not found."
             ) from exc
 
-        _model_instance = OpenAIChatModel(
-            model_name=self.configuration.model_name,
-            profile=(
-                OpenAIModelProfile(**self.configuration.profile.dict(exclude_unset=True))
-                if self.configuration.profile
-                else None
-            ),
-            provider=OpenAIProvider(
-                base_url=self.configuration.provider.base_url,
-                api_key=self.configuration.provider.api_key,
-            )
-            if self.configuration.provider
-            else None,
-        )
+        if self.configuration.is_custom:
+            _model_instance = prepare_custom_model(self.configuration)
+        else:
+            # In this case, we rely on PydanticAI's built-in model registry
+            # and configuration: check pydantic_ai.models.KnownModelName
+            # and pydantic_ai.models.infer_model()
+            _model_instance = self.configuration.model_name
+
         _system_prompt = self.configuration.system_prompt
         _base_toolset = (
             [

src/backend/chat/clients/pydantic_ai.py

@@ -377,6 +377,7 @@ class AIAgentService:  # pylint: disable=too-many-instance-attributes
                     "before answering the user request."
                 )
 
+        _tool_is_streaming = False
         async with AsyncExitStack() as stack:
             # MCP servers (if any) can be initialized here
             mcp_servers = [await stack.enter_async_context(mcp) for mcp in get_mcp_servers()]
@@ -476,6 +477,7 @@ class AIAgentService:  # pylint: disable=too-many-instance-attributes
                                     if isinstance(event.delta, TextPartDelta):
                                         yield events_v4.TextPart(text=event.delta.content_delta)
                                     elif isinstance(event.delta, ToolCallPartDelta):
+                                        _tool_is_streaming = True
                                         yield events_v4.ToolCallDeltaPart(
                                             tool_call_id=event.delta.tool_call_id,
                                             args_text_delta=event.delta.args_delta,
@@ -497,9 +499,14 @@ class AIAgentService:  # pylint: disable=too-many-instance-attributes
                                     dataclasses.asdict(event),
                                 )
                                 if isinstance(event, FunctionToolCallEvent):
-                                    # We are already streaming the tool call events don't yield
-                                    # the tool call again
-                                    pass
+                                    if not _tool_is_streaming:
+                                        yield events_v4.ToolCallPart(
+                                            tool_call_id=event.tool_call_id,
+                                            tool_name=event.part.tool_name,
+                                            args=json.loads(event.part.args)
+                                            if event.part.args
+                                            else {},
+                                        )
                                 elif isinstance(event, FunctionToolResultEvent):
                                     if isinstance(event.result, ToolReturnPart):
                                         if event.result.metadata and (

src/backend/chat/llm_configuration.py

@@ -2,7 +2,7 @@
 
 import os
 from functools import lru_cache
-from typing import Annotated, Any, Optional, Self
+from typing import Annotated, Any, Literal, Optional, Self
 
 from pydantic import (
     AfterValidator,
@@ -54,6 +54,7 @@ class LLMProvider(BaseModel):
     hrid: str
     base_url: SettingEnvValue
     api_key: SettingEnvValue
+    kind: Literal["openai", "mistral"] = "openai"
 
 
 class LLMProfile(BaseModel):
@@ -140,6 +141,11 @@ class LLModel(BaseModel):
             )
         return self
 
+    @property
+    def is_custom(self) -> bool:
+        """Return True if the model is a custom model (i.e., defines a provider)."""
+        return self.provider is not None
+
 
 class LLMConfiguration(BaseModel):
     """Model representing the entire LLM configuration."""

src/backend/chat/tests/agents/test_base_agent.py

@@ -0,0 +1,81 @@
+"""Tests for the BaseAgent class and its model initialization logic."""
+# pylint: disable=protected-access
+
+from pydantic_ai.models.mistral import MistralModel
+from pydantic_ai.models.openai import OpenAIChatModel
+
+from chat.agents.base import BaseAgent
+from chat.llm_configuration import LLModel, LLMProvider
+
+
+def test_not_custom_model(monkeypatch, settings):
+    """Test that a model without a provider relies on Pydantic AI detection."""
+    settings.LLM_CONFIGURATIONS = {
+        "gpt-4": LLModel(
+            hrid="gpt-4",
+            model_name="openai:gpt-4",
+            human_readable_name="GPT-4",
+            is_active=True,
+            system_prompt="direct",
+            tools=[],
+        ),
+    }
+
+    # Required for OpenAI models client initialization
+    monkeypatch.setenv("OPENAI_API_KEY", "hello")
+
+    agent = BaseAgent(model_hrid="gpt-4")
+    assert isinstance(agent._model, OpenAIChatModel)
+
+
+def test_custom_model_openai(settings):
+    """Test that a custom OpenAI model is initialized correctly."""
+    settings.LLM_CONFIGURATIONS = {
+        "openai-compatible-model": LLModel(
+            hrid="custom-gpt-4",
+            model_name="gpt-4",
+            human_readable_name="Custom GPT-4",
+            profile=None,
+            provider=LLMProvider(
+                hrid="openai",
+                kind="openai",
+                base_url="https://test.vllm/v1",
+                api_key="testkey",
+            ),
+            is_active=True,
+            system_prompt="direct",
+            tools=[],
+        ),
+    }
+
+    agent = BaseAgent(model_hrid="openai-compatible-model")
+    assert isinstance(agent._model, OpenAIChatModel)
+
+
+def test_custom_model_mistral(settings):
+    """Test that a custom Mistral model is initialized correctly."""
+    settings.LLM_CONFIGURATIONS = {
+        "mistral-model": LLModel(
+            hrid="mistral-model",
+            model_name="mistral-7b-instruct-v0.1",
+            human_readable_name="Mistral 7B Instruct",
+            profile=None,
+            provider=LLMProvider(
+                hrid="mistral",
+                kind="mistral",
+                base_url="https://api.mistral.ai/v1",
+                api_key="testkey",
+            ),
+            is_active=True,
+            system_prompt="direct",
+            tools=[],
+        ),
+    }
+
+    agent = BaseAgent(model_hrid="mistral-model")
+
+    assert isinstance(agent._model, MistralModel)
+
+    import pydantic_ai.models.mistral as mistral_models  # noqa: PLC0415 # pylint: disable=import-outside-toplevel
+
+    assert mistral_models.__safe_map_patched__ is True  # pylint: disable=protected-access

src/backend/chat/tests/test_llm_configuration.py

@@ -167,3 +167,27 @@ def test_load_llm_configuration(tmp_path, monkeypatch):
     assert "gpt-4" in model_map
     assert model_map["gpt-4"].provider.base_url == "env_value"
     assert model_map["gpt-4"].provider.api_key == "setting_value"
+
+
+def test_llmodel_is_custom_property():
+    """Test the is_custom property of LLModel."""
+    provider = LLMProvider(hrid="custom", base_url="direct", api_key="direct")
+    custom_model = LLModel(
+        hrid="custom-model",
+        model_name="custom-model",
+        human_readable_name="Custom Model",
+        provider=provider,
+        is_active=True,
+        system_prompt="direct",
+        tools=[],
+    )
+    non_custom_model = LLModel(
+        hrid="prefixed-model",
+        model_name="openai:prefixed-model",
+        human_readable_name="Prefixed Model",
+        is_active=True,
+        system_prompt="direct",
+        tools=[],
+    )
+    assert custom_model.is_custom is True
+    assert non_custom_model.is_custom is False

src/backend/chat/tests/views/chat/conversations/test_conversation_with_document_upload.py

@@ -279,6 +279,8 @@ def test_post_conversation_with_document_upload(  # pylint: disable=too-many-arg
         '"args":{"documents":[{"identifier":"sample.pdf"}]}}\n'
         'a:{"toolCallId":"XXX","result":{"state":"done"}}\n'
         'b:{"toolCallId":"pyd_ai_YYY","toolName":"document_search_albert_rag"}\n'
+        '9:{"toolCallId":"pyd_ai_YYY","toolName":"document_search_albert_rag",'
+        '"args":{"query":"What does the document say?"}}\n'
         'h:{"sourceType":"url","id":"XXX","url":"sample.pdf","title":null,"providerMetadata":{}}\n'
         'a:{"toolCallId":"pyd_ai_YYY","result":[{"url":"sample.pdf","content":"This '
         'is the content of the PDF.","score":0.9}]}\n'

src/backend/conversations/configuration/llm/default.json

@@ -36,7 +36,8 @@
     {
       "hrid": "default-provider",
       "base_url": "settings.AI_BASE_URL",
-      "api_key": "settings.AI_API_KEY"
+      "api_key": "settings.AI_API_KEY",
+      "kind": "openai"
     }
   ]
 }