✨(summarize): new summarize tool integration

CHANGELOG.md

@@ -13,6 +13,7 @@ and this project adheres to
 - 🐛(front) fix target blank links in chat #103
 - 🚑️(posthog) pass str instead of UUID for user PK #134
 - ⚡️(web-search) keep running when tool call fails #137
+- ✨(summarize): new summarize tool integration #78
 
 ### Removed
 

src/backend/chat/agents/summarize.py

@@ -4,11 +4,6 @@
 import logging
 
 from django.conf import settings
-from django.core.files.storage import default_storage
-
-from asgiref.sync import sync_to_async
-from pydantic_ai import RunContext
-from pydantic_ai.messages import ToolReturn
 
 from .base import BaseAgent
 
@@ -26,60 +21,3 @@ def __init__(self, **kwargs):
             output_type=str,
             **kwargs,
         )
-
-
-@sync_to_async
-def read_document_content(doc):
-    """Read document content asynchronously."""
-    with default_storage.open(doc.key) as f:
-        return doc.file_name, f.read().decode("utf-8")
-
-
-async def hand_off_to_summarization_agent(ctx: RunContext) -> ToolReturn:
-    """
-    Generate a complete, ready-to-use summary of the documents in context
-    (do not request the documents to the user).
-    Return this summary directly to the user WITHOUT any modification,
-    or additional summarization.
-    The summary is already optimized and MUST be presented as-is in the final response
-    or translated preserving the information.
-    """
-    summarization_agent = SummarizationAgent()
-
-    prompt = (
-        "Do not mention the user request in your answer.\n"
-        "User request:\n"
-        "{user_prompt}\n\n"
-        "Document contents:\n"
-        "{documents_prompt}\n"
-    )
-    text_attachment = await sync_to_async(list)(
-        ctx.deps.conversation.attachments.filter(
-            content_type__startswith="text/",
-        )
-    )
-
-    documents = [await read_document_content(doc) for doc in text_attachment]
-
-    documents_prompt = "\n\n".join(
-        [
-            (f"<document>\n<name>\n{name}\n</name>\n<content>\n{content}\n</content>\n</document>")
-            for name, content in documents
-        ]
-    )
-
-    formatted_prompt = prompt.format(
-        user_prompt=ctx.prompt,
-        documents_prompt=documents_prompt,
-    )
-
-    logger.debug("Summarize prompt: %s", formatted_prompt)
-
-    response = await summarization_agent.run(formatted_prompt, usage=ctx.usage)
-
-    logger.debug("Summarize response: %s", response)
-
-    return ToolReturn(
-        return_value=response.output,
-        metadata={"sources": {doc[0] for doc in documents}},
-    )

src/backend/chat/clients/pydantic_ai.py

@@ -7,6 +7,7 @@
 """
 
 import dataclasses
+import functools
 import json
 import logging
 import time
@@ -25,7 +26,7 @@
 
 from asgiref.sync import sync_to_async
 from langfuse import get_client
-from pydantic_ai import Agent
+from pydantic_ai import Agent, RunContext
 from pydantic_ai.messages import (
     BinaryContent,
     DocumentUrl,
@@ -59,7 +60,6 @@
     update_history_local_urls,
     update_local_urls,
 )
-from chat.agents.summarize import hand_off_to_summarization_agent
 from chat.ai_sdk_types import (
     LanguageModelV1Source,
     SourceUIPart,
@@ -73,6 +73,7 @@
 )
 from chat.mcp_servers import get_mcp_servers
 from chat.tools.document_search_rag import add_document_rag_search_tool
+from chat.tools.document_summarize import document_summarize
 from chat.vercel_ai_sdk.core import events_v4, events_v5
 from chat.vercel_ai_sdk.encoder import EventEncoder
 
@@ -493,13 +494,20 @@ def summarization_system_prompt() -> str:
                     "You may add a follow-up question after the summary if needed."
                 )
 
-            @self.conversation_agent.tool
-            async def summarize(ctx) -> ToolReturn:
-                """
-                Summarize the documents for the user, only when asked for,
-                the documents are in my context.
-                """
-                return await hand_off_to_summarization_agent(ctx)
+            # Inform the model (system-level) that documents are attached and available
+            @self.conversation_agent.system_prompt
+            def attached_documents_note() -> str:
+                return (
+                    "[Internal context] User documents are attached to this conversation. "
+                    "Do not request re-upload of documents; consider them already available "
+                    "via the internal store."
+                )
+
+            @self.conversation_agent.tool(name="summarize", retries=2)
+            @functools.wraps(document_summarize)
+            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)
         else:
             conversation_documents = [
                 cd

src/backend/chat/tests/conftest.py

@@ -8,6 +8,7 @@
 
 import pytest
 
+from chat.agents.summarize import SummarizationAgent
 from chat.clients.pydantic_ai import AIAgentService
 
 logger = logging.getLogger(__name__)
@@ -50,6 +51,40 @@ def __init__(self, **kwargs):
     yield _mock_service
 
 
+@pytest.fixture(name="mock_summarization_agent")
+def mock_summarization_agent_fixture():
+    """Fixture to mock SummarizationAgent with a custom model."""
+
+    @contextmanager
+    def _mock_agent(model):
+        """Context manager to mock SummarizationAgent with a custom model."""
+        with ExitStack() as stack:
+
+            class SummarizationAgentMock(SummarizationAgent):
+                """Mocked SummarizationAgent to override the model."""
+
+                def __init__(self, **kwargs):
+                    super().__init__(**kwargs)
+                    # We cannot use stack.enter_context(agent.override(model=model))
+                    # Because the agent is used outside of this context manager.
+                    # So we directly override the protected member.
+                    logger.info("Overriding SummarizationAgent model with %s", model)
+                    self._model = model  # pylint: disable=protected-access
+
+            # Mock the SummarizationAgent in all relevant modules, because first import wins
+            stack.enter_context(
+                patch("chat.agents.summarize.SummarizationAgent", new=SummarizationAgentMock)
+            )
+            stack.enter_context(
+                patch(
+                    "chat.tools.document_summarize.SummarizationAgent", new=SummarizationAgentMock
+                )
+            )
+            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"