From 99c4b6ba293cda9588c3c2dfcd244508a3db4dcd Mon Sep 17 00:00:00 2001 From: Quentin BEY Date: Thu, 2 Oct 2025 16:18:59 +0200 Subject: [PATCH] =?UTF-8?q?=E2=9C=A8(models)=20add=20mistral=20support=20&?= =?UTF-8?q?=20customization?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This allows to connect to a Mistral model, on a Mistral platform, self-hosted. --- CHANGELOG.md | 1 + src/backend/chat/agents/base.py | 114 +++++++++++++++--- src/backend/chat/clients/pydantic_ai.py | 13 +- src/backend/chat/llm_configuration.py | 8 +- .../chat/tests/agents/test_base_agent.py | 81 +++++++++++++ .../chat/tests/test_llm_configuration.py | 24 ++++ .../test_conversation_with_document_upload.py | 2 + .../configuration/llm/default.json | 3 +- 8 files changed, 224 insertions(+), 22 deletions(-) create mode 100644 src/backend/chat/tests/agents/test_base_agent.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 4ec16e6..a006ba1 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/src/backend/chat/agents/base.py b/src/backend/chat/agents/base.py index 3f12c84..becd465 100644 --- a/src/backend/chat/agents/base.py +++ b/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 = ( [ diff --git a/src/backend/chat/clients/pydantic_ai.py b/src/backend/chat/clients/pydantic_ai.py index 1014d7b..6c660fa 100644 --- a/src/backend/chat/clients/pydantic_ai.py +++ b/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 ( diff --git a/src/backend/chat/llm_configuration.py b/src/backend/chat/llm_configuration.py index 820fd68..a5d9169 100644 --- a/src/backend/chat/llm_configuration.py +++ b/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.""" diff --git a/src/backend/chat/tests/agents/test_base_agent.py b/src/backend/chat/tests/agents/test_base_agent.py new file mode 100644 index 0000000..81ac21c --- /dev/null +++ b/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 diff --git a/src/backend/chat/tests/test_llm_configuration.py b/src/backend/chat/tests/test_llm_configuration.py index f7e7e68..f275a14 100644 --- a/src/backend/chat/tests/test_llm_configuration.py +++ b/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 diff --git a/src/backend/chat/tests/views/chat/conversations/test_conversation_with_document_upload.py b/src/backend/chat/tests/views/chat/conversations/test_conversation_with_document_upload.py index 31f6a84..9298ea0 100644 --- a/src/backend/chat/tests/views/chat/conversations/test_conversation_with_document_upload.py +++ b/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' diff --git a/src/backend/conversations/configuration/llm/default.json b/src/backend/conversations/configuration/llm/default.json index 6591dbc..d77bc4b 100644 --- a/src/backend/conversations/configuration/llm/default.json +++ b/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" } ] }