✨(tools) add basic document translate tool

CHANGELOG.md

@@ -13,6 +13,7 @@ and this project adheres to
 - ✨(waffle) hide the waffle if not fr theme
 - ✨(front) allow pasting an attachment from clipboard
 - ✨(array) temporarily adjust array
+- ✨(tools) add basic translate tool
 
 ### Changed
 

src/backend/chat/agents/translate.py

@@ -0,0 +1,27 @@
+"""Build the translation agent."""
+
+import dataclasses
+import logging
+
+from django.conf import settings
+
+from .base import BaseAgent
+
+logger = logging.getLogger(__name__)
+
+
+@dataclasses.dataclass(init=False)
+class TranslationAgent(BaseAgent):
+    """Create a Pydantic AI translation Agent instance with the configured settings"""
+
+    def __init__(self, **kwargs):
+        """Initialize the agent with the configured model."""
+        super().__init__(
+            model_hrid=settings.LLM_DEFAULT_MODEL_HRID,
+            output_type=str,
+            **kwargs,
+        )
+
+    def get_tools(self) -> list:
+        """Translation does not need any tools."""
+        return []

src/backend/chat/clients/pydantic_ai.py

@@ -142,6 +142,7 @@
 from chat.tools.document_generic_search_rag import add_document_rag_search_tool_from_setting
 from chat.tools.document_search_rag import add_document_rag_search_tool
 from chat.tools.document_summarize import document_summarize
+from chat.tools.document_translate import document_translate
 from chat.vercel_ai_sdk.core import events_v4, events_v5
 from chat.vercel_ai_sdk.encoder import CURRENT_EVENT_ENCODER_VERSION, EventEncoder
 
@@ -667,6 +668,16 @@ def summarization_system_prompt() -> str:
                 "You may add a follow-up question after the summary if needed."
             )
 
+        @self.conversation_agent.instructions
+        def translation_system_prompt() -> str:
+            return (
+                "When you receive a result from the translation tool, you MUST return it "
+                "directly to the user without any modification, paraphrasing, or additional "
+                "summarization. "
+                "The tool already produces complete translations that should be presented "
+                "verbatim."
+            )
+
         # Inform the model (system-level) that documents are attached and available
         @self.conversation_agent.instructions
         def attached_documents_note() -> str:
@@ -682,6 +693,12 @@ async def summarize(ctx: RunContext, *args, **kwargs) -> ToolReturn:
             """Wrap the document_summarize tool to provide context and add the tool."""
             return await document_summarize(ctx, *args, **kwargs)
 
+        @self.conversation_agent.tool(name="translate", retries=2)
+        @functools.wraps(document_translate)
+        async def translate(ctx: RunContext, *args, **kwargs) -> ToolReturn:
+            """Wrap the document_translate tool to provide context and add the tool."""
+            return await document_translate(ctx, *args, **kwargs)
+
     async def _handle_input_documents(
         self,
         input_documents: List[BinaryContent | DocumentUrl],

src/backend/chat/tests/conftest.py

@@ -15,6 +15,7 @@
 import chat.views
 import conversations.urls
 from chat.agents.summarize import SummarizationAgent
+from chat.agents.translate import TranslationAgent
 from chat.clients.pydantic_ai import AIAgentService
 
 logger = logging.getLogger(__name__)
@@ -91,6 +92,34 @@ def __init__(self, **kwargs):
     yield _mock_agent
 
 
+@pytest.fixture(name="mock_translation_agent")
+def mock_translation_agent_fixture():
+    """Fixture to mock TranslationAgent with a custom model."""
+
+    @contextmanager
+    def _mock_agent(model):
+        """Context manager to mock TranslationAgent with a custom model."""
+        with ExitStack() as stack:
+
+            class TranslationAgentMock(TranslationAgent):
+                """Mocked TranslationAgent to override the model."""
+
+                def __init__(self, **kwargs):
+                    super().__init__(**kwargs)
+                    logger.info("Overriding TranslationAgent model with %s", model)
+                    self._model = model  # pylint: disable=protected-access
+
+            stack.enter_context(
+                patch("chat.agents.translate.TranslationAgent", new=TranslationAgentMock)
+            )
+            stack.enter_context(
+                patch("chat.tools.document_translate.TranslationAgent", new=TranslationAgentMock)
+            )
+            yield
+
+    yield _mock_agent
+
+
 PIXEL_PNG = (
     b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x06\x00"
     b"\x00\x00\x1f\x15\xc4\x89\x00\x00\x00\nIDATx\x9cc\xf8\xff\xff?\x00\x05\xfe\x02\xfe"