From 392eeece3ebedd5162ac0b045efb9da13a6d160c Mon Sep 17 00:00:00 2001 From: camilleAND Date: Wed, 15 Oct 2025 10:45:04 +0200 Subject: [PATCH] =?UTF-8?q?=E2=9C=A8(summarize)=20new=20summarize=20tool?= =?UTF-8?q?=20integration?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Improve the existing tool to manage bigger documents. --- CHANGELOG.md | 1 + src/backend/chat/agents/summarize.py | 99 ++++++++++++++----- src/backend/chat/clients/pydantic_ai.py | 22 +++-- src/backend/chat/tools/document_search_rag.py | 12 ++- .../chat/components/ToolInvocationItem.tsx | 2 +- 5 files changed, 100 insertions(+), 36 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a94d134..7e7a556 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/src/backend/chat/agents/summarize.py b/src/backend/chat/agents/summarize.py index 38d92d6..4f1d5f4 100644 --- a/src/backend/chat/agents/summarize.py +++ b/src/backend/chat/agents/summarize.py @@ -2,6 +2,7 @@ import dataclasses import logging +import asyncio from django.conf import settings from django.core.files.storage import default_storage @@ -11,6 +12,7 @@ from pydantic_ai import RunContext from pydantic_ai.messages import ToolReturn from .base import BaseAgent +from ..tools.document_search_rag import add_document_rag_search_tool logger = logging.getLogger(__name__) @@ -35,14 +37,18 @@ def read_document_content(doc): return doc.file_name, f.read().decode("utf-8") -async def hand_off_to_summarization_agent(ctx: RunContext) -> ToolReturn: +async def hand_off_to_summarization_agent( + ctx: RunContext, *, instructions: str | None = None +) -> 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. + Summarize the documents for the user, only when asked for. + Instructions are optional but should reflect the user's request. + Examples : + "RĂ©sume ce doc en 2 paragraphes" -> instructions = "rĂ©sumĂ© en 2 paragraphes" + "RĂ©sume ce doc en anglais" -> instructions = "In English" + "RĂ©sume ce doc" -> instructions = "" (default) + Args: + instructions (str | None): The instructions the user gave to use for the summarization """ summarization_agent = SummarizationAgent() @@ -53,6 +59,8 @@ async def hand_off_to_summarization_agent(ctx: RunContext) -> ToolReturn: "Document contents:\n" "{documents_prompt}\n" ) + + # Collect documents content text_attachment = await sync_to_async(list)( ctx.deps.conversation.attachments.filter( content_type__startswith="text/", @@ -61,25 +69,72 @@ async def hand_off_to_summarization_agent(ctx: RunContext) -> ToolReturn: documents = [await read_document_content(doc) for doc in text_attachment] - documents_prompt = "\n\n".join( - [ - (f"\n\n{name}\n\n\n{content}\n\n") - for name, content in documents + # Instructions: rely on tool argument only; model should extract them upstream + if instructions is not None: + instructions_hint: str = instructions.strip() + else: + instructions_hint = "" + + # Helpers + def chunk_text(text: str, size: int = 10000) -> list[str]: + if size <= 0: + return [text] + return [text[i : i + size] for i in range(0, len(text), size)] + + # 2) Chunk documents and summarize each chunk + full_text = "\n\n".join(doc[1] for doc in documents) + chunks = chunk_text(full_text, size=10000) + logger.info( + "[summarize] chunking: %s parts (size~%s), instructions='%s'", + len(chunks), + 10000, + instructions_hint or "", + ) + + async def summarize_chunk(idx, chunk, total_chunks, summarization_agent, ctx): + sum_prompt = ( + "Tu es un agent spĂ©cialisĂ© en synthĂšses de textes. " + "GĂ©nĂšre un rĂ©sumĂ© clair et concis du passage suivant (partie {idx}/{total}) :\n" + "'''\n{context}\n'''\n\n" + ).format(context=chunk, idx=idx, total=total_chunks) + logger.info("[summarize] CHUNK %s/%s prompt=> %s", idx, total_chunks, sum_prompt[0:100]+'...') + resp = await summarization_agent.run(sum_prompt, usage=ctx.usage) + logger.info("[summarize] CHUNK %s/%s response<= %s", idx, total_chunks, resp.output or "") + return resp.output or "" + + # Parallelize the chunk summarization in batches of 5 using asyncio.gather + chunk_summaries: list[str] = [] + batch_size = 5 + for start_idx in range(0, len(chunks), batch_size): + end_idx = start_idx + batch_size + batch_chunks = chunks[start_idx:end_idx] + summarization_tasks = [ + summarize_chunk(idx, chunk, len(chunks), summarization_agent, ctx) + for idx, chunk in enumerate(batch_chunks, start=start_idx + 1) ] - ) + batch_results = await asyncio.gather(*summarization_tasks) + chunk_summaries.extend(batch_results) - formatted_prompt = prompt.format( - user_prompt=ctx.prompt, - documents_prompt=documents_prompt, - ) + if not instructions_hint: + instructions_hint = "Le rĂ©sumĂ© doit ĂȘtre en Français, contenir 2 ou 3 parties." - logger.debug("Summarize prompt: %s", formatted_prompt) - - response = await summarization_agent.run(formatted_prompt, usage=ctx.usage) - - logger.debug("Summarize response: %s", response) + # 3) Merge chunk summaries into a single concise summary + merged_prompt = ( + "Produit une synthĂšse cohĂ©rente Ă  partir des rĂ©sumĂ©s ci-dessous.\n\n" + "'''\n{context}\n'''\n\n" + "Contraintes :\n" + "- RĂ©sumer sans rĂ©pĂ©ter.\n" + "- Harmoniser le style et la terminologie.\n" + "- Le rĂ©sumĂ© final doit ĂȘtre bien structurĂ© et formatĂ© en markdown. \n" + "- Respecter les consignes : {instructions}\n" + "RĂ©ponds directement avec le rĂ©sumĂ© final." + ).format(context="\n\n".join(chunk_summaries), instructions=instructions_hint or "") + logger.info("[summarize] MERGE prompt=> %s", merged_prompt) + merged_resp = await summarization_agent.run(merged_prompt, usage=ctx.usage) + final_summary = (merged_resp.output or "").strip() + logger.info("[summarize] MERGE response<= %s", final_summary) return ToolReturn( - return_value=response.output, + return_value=final_summary, metadata={"sources": {doc[0] for doc in documents}}, ) diff --git a/src/backend/chat/clients/pydantic_ai.py b/src/backend/chat/clients/pydantic_ai.py index e5d8056..f2a5915 100644 --- a/src/backend/chat/clients/pydantic_ai.py +++ b/src/backend/chat/clients/pydantic_ai.py @@ -483,14 +483,20 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes @self.conversation_agent.system_prompt def summarization_system_prompt() -> str: return ( - "When you receive a result from the summarization tool, you MUST return it " - "directly to the user without any modification, paraphrasing, or additional " - "summarization." - "The tool already produces optimized summaries that should be presented " - "verbatim." - "You may translate the summary if required, but you MUST preserve all the " - "information from the original summary." - "You may add a follow-up question after the summary if needed." + "When the user asks to summarize attached document(s), you MUST call the" + " summarize tool. Pass user's instructions if provided, otherwise pass an" + " empty instructions string once the user confirms (e.g. says 'ok'). Do NOT" + " call web search or document_search_rag to produce summaries; rely only on" + " the attached documents stored in context." + ) + + # 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 diff --git a/src/backend/chat/tools/document_search_rag.py b/src/backend/chat/tools/document_search_rag.py index 121b43f..28cf865 100644 --- a/src/backend/chat/tools/document_search_rag.py +++ b/src/backend/chat/tools/document_search_rag.py @@ -20,8 +20,12 @@ def add_document_rag_search_tool(agent: Agent) -> None: Args: ctx (RunContext): The run context containing the conversation. - query (str): The term to search the internet for. + query (str): The query to search the documents for. """ + # Defensive: ctx.deps or ctx.deps.conversation may be unavailable in some flows (start of conversation) + if not getattr(ctx, "deps", None) or not getattr(ctx.deps, "conversation", None): + return ToolReturn(return_value=[], content="", metadata={"sources": set()}) + document_store_backend = import_string(settings.RAG_DOCUMENT_SEARCH_BACKEND) document_store = document_store_backend(ctx.deps.conversation.collection_id) @@ -43,8 +47,6 @@ def add_document_rag_search_tool(agent: Agent) -> None: def document_rag_instructions() -> str: """Dynamic system prompt function to add RAG instructions if any.""" return ( - "If the user wants specific information from a document, invoke " - "web_search_albert_rag with an appropriate query string." - "Do not ask the user for the document; rely on the tool to locate " - "and return relevant passages." + "Use document_search_rag ONLY to retrieve specific passages from attached documents. " + "Do NOT use it to summarize; for summaries, call the summarize tool instead." ) diff --git a/src/frontend/apps/conversations/src/features/chat/components/ToolInvocationItem.tsx b/src/frontend/apps/conversations/src/features/chat/components/ToolInvocationItem.tsx index de5a712..a6f8305 100644 --- a/src/frontend/apps/conversations/src/features/chat/components/ToolInvocationItem.tsx +++ b/src/frontend/apps/conversations/src/features/chat/components/ToolInvocationItem.tsx @@ -68,7 +68,7 @@ export const ToolInvocationItem: React.FC = ({ > - {t('Search...')} + {toolInvocation.toolName === 'summarize' ? t('Summarizing...') : t('Search...')} )}