From 9426a7e1ae4284cd6e8d1c3250fb889025013158 Mon Sep 17 00:00:00 2001 From: "renovate[bot]" <29139614+renovate[bot]@users.noreply.github.com> Date: Wed, 5 Nov 2025 22:35:55 +0000 Subject: [PATCH 1/8] =?UTF-8?q?=E2=AC=86=EF=B8=8F(dependencies)=20update?= =?UTF-8?q?=20python=20dependencies?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/backend/pyproject.toml | 32 ++++++++++++++++---------------- 1 file changed, 16 insertions(+), 16 deletions(-) diff --git a/src/backend/pyproject.toml b/src/backend/pyproject.toml index a88b5de..35c23a9 100644 --- a/src/backend/pyproject.toml +++ b/src/backend/pyproject.toml @@ -27,41 +27,41 @@ requires-python = ">=3.12" dependencies = [ "deprecated", "beautifulsoup4==4.14.2", - "boto3==1.40.59", - "Brotli==1.1.0", + "boto3==1.40.67", + "Brotli==1.2.0", "django-configurations==2.5.1", "django-cors-headers==4.9.0", - "django-countries==7.6.1", + "django-countries==8.0.0", "django-filter==25.2", - "django-lasuite[all]==0.0.16", + "django-lasuite[all]==0.0.17", "django-parler==2.3", "django-pydantic-field==0.3.13", "django-redis==6.0.0", "django-storages[s3]==1.14.6", "django-timezone-field>=5.1", - "django==5.2.7", + "django==5.2.8", "djangorestframework==3.16.1", - "drf_spectacular==0.28.0", + "drf_spectacular==0.29.0", "dockerflow==2024.4.2", "easy_thumbnails==2.10.1", "factory_boy==3.3.3", "gunicorn==23.0.0", "jsonschema==4.25.1", - "langfuse==3.8.1", + "langfuse==3.9.0", "lxml==5.4.0", - "markdown==3.9", + "markdown==3.10", "markitdown==0.0.2", "mozilla-django-oidc==4.0.1", "nested-multipart-parser==1.6.0", - "posthog==6.7.10", - "pydantic==2.12.3", - "pydantic-ai-slim[openai,mistral,mcp,evals,logfire]==1.6.0", + "posthog==6.8.0", + "pydantic==2.12.4", + "pydantic-ai-slim[openai,mistral,mcp,evals,logfire]==1.11.0", "psycopg[binary]==3.2.12", "PyJWT==2.10.1", "python-magic==0.4.27", "redis<6.0.0", "requests==2.32.5", - "sentry-sdk==2.42.1", + "sentry-sdk==2.43.0", "trafilatura==2.0.0", "uvicorn==0.38.0", "whitenoise==6.11.0", @@ -81,11 +81,11 @@ dev = [ "drf-spectacular-sidecar==2025.10.1", "freezegun==1.5.5", "ipdb==0.13.13", - "ipython==9.6.0", - "pyfakefs==5.10.0", + "ipython==9.7.0", + "pyfakefs==5.10.2", "pylint-django==2.6.1", "pylint==3.3.9", - "pylint-pydantic==0.4.0", + "pylint-pydantic==0.4.1", "pytest-asyncio==1.2.0", "pytest-cov==7.0.0", "pytest-django==4.11.1", @@ -94,7 +94,7 @@ dev = [ "pytest-xdist==3.8.0", "responses==0.25.8", "respx==0.22.0", - "ruff==0.14.2", + "ruff==0.14.3", "types-requests==2.32.4.20250913", ] From 1f92187daeebc0f7c51e335975345a042f028d3b Mon Sep 17 00:00:00 2001 From: Quentin BEY Date: Wed, 5 Nov 2025 23:56:38 +0100 Subject: [PATCH 2/8] =?UTF-8?q?=F0=9F=94=A5(posthog)=20remove=20posthog=20?= =?UTF-8?q?middleware=20for=20async=20mode=20fix?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Posthog fixed their middleware in v6.7.14 This reverts commit fbe9e039cf231d838a85ac50fa71e42347186dac. --- CHANGELOG.md | 4 + src/backend/conversations/settings.py | 2 +- src/backend/core/posthog.py | 113 -------------------------- 3 files changed, 5 insertions(+), 114 deletions(-) delete mode 100644 src/backend/core/posthog.py diff --git a/CHANGELOG.md b/CHANGELOG.md index bf328cf..a94d134 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,10 @@ and this project adheres to - đŸš‘ïž(posthog) pass str instead of UUID for user PK #134 - âšĄïž(web-search) keep running when tool call fails #137 +### Removed + +- đŸ”„(posthog) remove posthog middleware for async mode fix #146 + ## [0.0.7] - 2025-10-28 diff --git a/src/backend/conversations/settings.py b/src/backend/conversations/settings.py index e51c3b0..5158b98 100755 --- a/src/backend/conversations/settings.py +++ b/src/backend/conversations/settings.py @@ -312,7 +312,7 @@ class Base(BraveSettings, Configuration): "django.middleware.common.CommonMiddleware", "django.middleware.csrf.CsrfViewMiddleware", "django.contrib.auth.middleware.AuthenticationMiddleware", - "core.posthog.AsyncPosthogContextMiddleware", + "posthog.integrations.django.PosthogContextMiddleware", "django.contrib.messages.middleware.MessageMiddleware", "dockerflow.django.middleware.DockerflowMiddleware", ] diff --git a/src/backend/core/posthog.py b/src/backend/core/posthog.py deleted file mode 100644 index faff917..0000000 --- a/src/backend/core/posthog.py +++ /dev/null @@ -1,113 +0,0 @@ -""" -Asynchronous PostHog context middleware for Django. - -Since https://github.com/PostHog/posthog-python/commit/6af129f41413f4f7d55731763ff42e4c0fb66844 -The official PostHog Django middleware supports both sync and async requests, -but the call to request.user fails in async contexts. - -This is an horrible patch while waiting for an official fix. -Follow https://github.com/PostHog/posthog-python/issues/355 -""" - -from posthog import contexts -from posthog.integrations.django import PosthogContextMiddleware - - -class AsyncPosthogContextMiddleware(PosthogContextMiddleware): - """ - Asynchronous Django middleware to extract PostHog context from HTTP requests. - - While the original PosthogContextMiddleware is supposed to manage both sync and async requests, - the call to request.user fails in async contexts. - """ - - async def extract_tags_async(self, request): - """Extract tags from the HTTP request asynchronously.""" - tags = {} - - (user_id, user_email) = await self.extract_request_user_async(request) - - # Extract session ID from X-POSTHOG-SESSION-ID header - session_id = request.headers.get("X-POSTHOG-SESSION-ID") - if session_id: - contexts.set_context_session(session_id) - - # Extract distinct ID from X-POSTHOG-DISTINCT-ID header or request user id - distinct_id = request.headers.get("X-POSTHOG-DISTINCT-ID") or user_id - if distinct_id: - contexts.identify_context(distinct_id) - - # Extract user email - if user_email: - tags["email"] = user_email - - # Extract current URL - absolute_url = request.build_absolute_uri() - if absolute_url: - tags["$current_url"] = absolute_url - - # Extract request method - if request.method: - tags["$request_method"] = request.method - - # Extract request path - if request.path: - tags["$request_path"] = request.path - - # Extract IP address - ip_address = request.headers.get("X-Forwarded-For") - if ip_address: - tags["$ip_address"] = ip_address - - # Extract user agent - user_agent = request.headers.get("User-Agent") - if user_agent: - tags["$user_agent"] = user_agent - - # Apply extra tags if configured - if self.extra_tags: - extra = self.extra_tags(request) - if extra: - tags.update(extra) - - # Apply tag mapping if configured - if self.tag_map: - tags = self.tag_map(tags) - - return tags - - async def extract_request_user_async(self, request): - """Extract user ID and email from the HTTP request asynchronously.""" - user_id = None - email = None - - user = await request.auser() - - if user and getattr(user, "is_authenticated", False): - try: - user_id = str(user.pk) - except Exception: # noqa: BLE001, S110 # pylint: disable=broad-except - pass - - try: - email = str(user.email) - except Exception: # noqa: BLE001, S110 # pylint: disable=broad-except - pass - - return user_id, email - - async def __acall__(self, request): - """ - Asynchronous entry point for async request handling. - - This method is called when the middleware chain is async. - """ - if self.request_filter and not self.request_filter(request): - return await self.get_response(request) - - with contexts.new_context(self.capture_exceptions, client=self.client): - tags = await self.extract_tags_async(request) - for k, v in tags.items(): - contexts.tag(k, v) - - return await self.get_response(request) From 392eeece3ebedd5162ac0b045efb9da13a6d160c Mon Sep 17 00:00:00 2001 From: camilleAND Date: Wed, 15 Oct 2025 10:45:04 +0200 Subject: [PATCH 3/8] =?UTF-8?q?=E2=9C=A8(summarize)=20new=20summarize=20to?= =?UTF-8?q?ol=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...')} )} From 78a23933836a988140556f874ea01c4943ca2f4c Mon Sep 17 00:00:00 2001 From: Quentin BEY Date: Thu, 6 Nov 2025 13:27:35 +0100 Subject: [PATCH 4/8] =?UTF-8?q?=E2=9A=A1=EF=B8=8F(summarize)=20use=20semch?= =?UTF-8?q?unk=20for=20better=20doc=20chunking?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This reduces the code complexity while allowing better "cuts" also providing overlap for free. Also, do not wait for sub-batch to complete a use a global concurrency instead. --- src/backend/chat/agents/summarize.py | 161 ++++++++++-------- src/backend/chat/clients/pydantic_ai.py | 13 +- .../test_conversation_with_document_upload.py | 38 +++-- .../test_conversation_with_document_url.py | 33 +++- src/backend/chat/tools/document_search_rag.py | 4 - src/backend/conversations/settings.py | 15 ++ src/backend/pyproject.toml | 1 + 7 files changed, 166 insertions(+), 99 deletions(-) diff --git a/src/backend/chat/agents/summarize.py b/src/backend/chat/agents/summarize.py index 4f1d5f4..0ce6868 100644 --- a/src/backend/chat/agents/summarize.py +++ b/src/backend/chat/agents/summarize.py @@ -1,18 +1,18 @@ """Build the summarization agent.""" +import asyncio import dataclasses import logging -import asyncio from django.conf import settings from django.core.files.storage import default_storage +import semchunk from asgiref.sync import sync_to_async 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__) @@ -37,28 +37,49 @@ def read_document_content(doc): return doc.file_name, f.read().decode("utf-8") -async def hand_off_to_summarization_agent( +async def summarize_chunk(idx, chunk, total_chunks, summarization_agent, ctx): + """Summarize a single chunk of text.""" + sum_prompt = ( + "You are an agent specializing in text summarization. " + "Generate a clear and concise summary of the following passage " + f"(part {idx}/{total_chunks}):\n'''\n{chunk}\n'''\n\n" + ) + + logger.debug( + "[summarize] CHUNK %s/%s prompt=> %s", idx, total_chunks, sum_prompt[0:100] + "..." + ) + + resp = await summarization_agent.run(sum_prompt, usage=ctx.usage) + + logger.debug("[summarize] CHUNK %s/%s response<= %s", idx, total_chunks, resp.output or "") + return resp.output or "" + + +async def hand_off_to_summarization_agent( # pylint: disable=too-many-locals ctx: RunContext, *, instructions: str | None = None ) -> ToolReturn: """ - Summarize the documents for the user, only when asked for. + 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. + 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) + + Examples: + "Summarize this doc in 2 paragraphs" -> instructions = "summary in 2 paragraphs" + "Summarize this doc in English" -> instructions = "In English" + "Summarize this doc" -> instructions = "" (default) + Args: instructions (str | None): The instructions the user gave to use for the summarization """ - 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" + instructions_hint = ( + instructions.strip() if instructions else "The summary should contain 2 or 3 parts." ) + summarization_agent = SummarizationAgent() # Collect documents content text_attachment = await sync_to_async(list)( @@ -69,70 +90,70 @@ async def hand_off_to_summarization_agent( documents = [await read_document_content(doc) for doc in text_attachment] - # 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 "", + # Chunk documents and summarize each chunk + chunk_size = settings.SUMMARIZATION_CHUNK_SIZE + chunker = semchunk.chunkerify( + tokenizer_or_token_counter=lambda text: len(text.split()), + chunk_size=chunk_size, + ) + documents_chunks = chunker( + [doc[1] for doc in documents], + overlap=settings.SUMMARIZATION_OVERLAP_SIZE, ) - 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 "" + logger.info( + "[summarize] chunking: %s parts (size~%s), instructions='%s'", + sum(len(chunks) for chunks in documents_chunks), + chunk_size, + instructions_hint, + ) - # 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] + # Parallelize the chunk summarization with a semaphore to limit concurrent tasks + # because it can be very resource intensive on the LLM backend + semaphore = asyncio.Semaphore(settings.SUMMARIZATION_CONCURRENT_REQUESTS) + + async def summarize_chunk_with_semaphore(idx, chunk, total_chunks): + """Summarize a chunk with semaphore-controlled concurrency.""" + async with semaphore: + return await summarize_chunk(idx, chunk, total_chunks, summarization_agent, ctx) + + doc_chunk_summaries = [] + for doc_chunks in documents_chunks: summarization_tasks = [ - summarize_chunk(idx, chunk, len(chunks), summarization_agent, ctx) - for idx, chunk in enumerate(batch_chunks, start=start_idx + 1) + summarize_chunk_with_semaphore(idx, chunk, len(doc_chunks)) + for idx, chunk in enumerate(doc_chunks, start=1) ] - batch_results = await asyncio.gather(*summarization_tasks) - chunk_summaries.extend(batch_results) + chunk_summaries = await asyncio.gather(*summarization_tasks) + doc_chunk_summaries.append(chunk_summaries) - if not instructions_hint: - instructions_hint = "Le rĂ©sumĂ© doit ĂȘtre en Français, contenir 2 ou 3 parties." + context = "\n\n".join( + doc_name + "\n\n" + "\n\n".join(summaries) + for doc_name, summaries in zip( + (doc[0] for doc in documents), + doc_chunk_summaries, + strict=True, + ) + ) - # 3) Merge chunk summaries into a single concise summary + # 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) + "Produce a coherent synthesis from the summaries below.\n\n" + f"'''\n{context}\n'''\n\n" + "Constraints:\n" + "- Summarize without repetition.\n" + "- Harmonize style and terminology.\n" + "- The final summary must be well-structured and formatted in markdown.\n" + f"- Follow the instructions: {instructions_hint}\n" + "Respond directly with the final summary." + ) + + logger.debug("[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) + + logger.debug("[summarize] MERGE response<= %s", final_summary) return ToolReturn( return_value=final_summary, diff --git a/src/backend/chat/clients/pydantic_ai.py b/src/backend/chat/clients/pydantic_ai.py index f2a5915..b7d9349 100644 --- a/src/backend/chat/clients/pydantic_ai.py +++ b/src/backend/chat/clients/pydantic_ai.py @@ -483,11 +483,14 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes @self.conversation_agent.system_prompt def summarization_system_prompt() -> str: return ( - "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." + "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." ) # Inform the model (system-level) that documents are attached and available 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 6958940..cf2c348 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 @@ -374,11 +374,10 @@ def test_post_conversation_with_document_upload( # pylint: disable=too-many-arg "timestamp": timezone_now, }, { - "content": "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.", + "content": "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.", "dynamic_ref": None, "part_kind": "system-prompt", "timestamp": timezone_now, @@ -397,6 +396,15 @@ def test_post_conversation_with_document_upload( # pylint: disable=too-many-arg "part_kind": "system-prompt", "timestamp": timezone_now, }, + { + "content": "[Internal context] User documents are attached to this " + "conversation. Do not request re-upload of documents; " + "consider them already available via the internal " + "store.", + "dynamic_ref": None, + "part_kind": "system-prompt", + "timestamp": timezone_now, + }, { "content": ["What does the document say?"], "part_kind": "user-prompt", @@ -627,7 +635,7 @@ def test_post_conversation_with_document_upload_summarize( # pylint: disable=to 'document discusses various topics."}\n' '0:"The document discusses various topics."\n' 'f:{"messageId":""}\n' - 'd:{"finishReason":"stop","usage":{"promptTokens":201,"completionTokens":13}}\n' + 'd:{"finishReason":"stop","usage":{"promptTokens":317,"completionTokens":19}}\n' ) # Check that the conversation was updated @@ -709,11 +717,10 @@ def test_post_conversation_with_document_upload_summarize( # pylint: disable=to "timestamp": timezone_now, }, { - "content": "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.", + "content": "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.", "dynamic_ref": None, "part_kind": "system-prompt", "timestamp": timezone_now, @@ -732,6 +739,15 @@ def test_post_conversation_with_document_upload_summarize( # pylint: disable=to "part_kind": "system-prompt", "timestamp": timezone_now, }, + { + "content": "[Internal context] User documents are attached to this " + "conversation. Do not request re-upload of documents; " + "consider them already available via the internal " + "store.", + "dynamic_ref": None, + "part_kind": "system-prompt", + "timestamp": timezone_now, + }, { "content": ["Make a summary of this document."], "part_kind": "user-prompt", diff --git a/src/backend/chat/tests/views/chat/conversations/test_conversation_with_document_url.py b/src/backend/chat/tests/views/chat/conversations/test_conversation_with_document_url.py index fd1877e..bd05e21 100644 --- a/src/backend/chat/tests/views/chat/conversations/test_conversation_with_document_url.py +++ b/src/backend/chat/tests/views/chat/conversations/test_conversation_with_document_url.py @@ -881,10 +881,9 @@ def test_post_conversation_with_local_not_pdf_document_url( # pylint: disable=t SystemPromptPart(content="Answer in english.", timestamp=timezone.now()), SystemPromptPart( content=( - "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." ), timestamp=timezone.now(), ), @@ -901,6 +900,14 @@ def test_post_conversation_with_local_not_pdf_document_url( # pylint: disable=t ), timestamp=timezone.now(), ), + SystemPromptPart( + content=( + "[Internal context] User documents are attached to this conversation. " + "Do not request re-upload of documents; consider them already " + "available via the internal store." + ), + timestamp=timezone.now(), + ), UserPromptPart( content=[ "What is in this document?", @@ -1002,11 +1009,10 @@ def test_post_conversation_with_local_not_pdf_document_url( # pylint: disable=t "timestamp": timestamp, }, { - "content": "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.", + "content": "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.", "dynamic_ref": None, "part_kind": "system-prompt", "timestamp": timestamp, @@ -1025,6 +1031,15 @@ def test_post_conversation_with_local_not_pdf_document_url( # pylint: disable=t "part_kind": "system-prompt", "timestamp": timestamp, }, + { + "content": "[Internal context] User documents are attached to " + "this conversation. Do not request re-upload of " + "documents; consider them already available via the " + "internal store.", + "dynamic_ref": None, + "part_kind": "system-prompt", + "timestamp": timestamp, + }, { "content": [ "What is in this document?", diff --git a/src/backend/chat/tools/document_search_rag.py b/src/backend/chat/tools/document_search_rag.py index 28cf865..f59c119 100644 --- a/src/backend/chat/tools/document_search_rag.py +++ b/src/backend/chat/tools/document_search_rag.py @@ -22,10 +22,6 @@ def add_document_rag_search_tool(agent: Agent) -> None: ctx (RunContext): The run context containing the conversation. 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) diff --git a/src/backend/conversations/settings.py b/src/backend/conversations/settings.py index 5158b98..759dc72 100755 --- a/src/backend/conversations/settings.py +++ b/src/backend/conversations/settings.py @@ -786,6 +786,21 @@ USER QUESTION: environ_name="SUMMARIZATION_SYSTEM_PROMPT", environ_prefix=None, ) + SUMMARIZATION_CHUNK_SIZE = values.PositiveIntegerValue( + default=20_000, # Approx 20k words per chunk + environ_name="SUMMARIZATION_CHUNK_SIZE", + environ_prefix=None, + ) + SUMMARIZATION_OVERLAP_SIZE = values.FloatValue( + default=0.05, # 5% overlap + environ_name="SUMMARIZATION_OVERLAP_SIZE", + environ_prefix=None, + ) + SUMMARIZATION_CONCURRENT_REQUESTS = values.PositiveIntegerValue( + default=5, + environ_name="SUMMARIZATION_CONCURRENT_REQUESTS", + environ_prefix=None, + ) # Tavily API TAVILY_API_KEY = values.Value( diff --git a/src/backend/pyproject.toml b/src/backend/pyproject.toml index 35c23a9..311a4eb 100644 --- a/src/backend/pyproject.toml +++ b/src/backend/pyproject.toml @@ -61,6 +61,7 @@ dependencies = [ "python-magic==0.4.27", "redis<6.0.0", "requests==2.32.5", + "semchunk==3.2.5", "sentry-sdk==2.43.0", "trafilatura==2.0.0", "uvicorn==0.38.0", From ee73c7b9cd6597c0199643cadec37161f85030f8 Mon Sep 17 00:00:00 2001 From: Quentin BEY Date: Thu, 6 Nov 2025 15:36:37 +0100 Subject: [PATCH 5/8] =?UTF-8?q?=E2=99=BB=EF=B8=8F(summarize)=20move=20the?= =?UTF-8?q?=20tool=20to=20the=20tools=20module?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This also fix the tool dynamic registation to use the original function signature and docsting. --- src/backend/chat/agents/summarize.py | 138 ---------------- src/backend/chat/clients/pydantic_ai.py | 17 +- .../test_conversation_with_document_upload.py | 2 +- src/backend/chat/tools/document_summarize.py | 147 ++++++++++++++++++ 4 files changed, 156 insertions(+), 148 deletions(-) create mode 100644 src/backend/chat/tools/document_summarize.py diff --git a/src/backend/chat/agents/summarize.py b/src/backend/chat/agents/summarize.py index 0ce6868..5efcbb8 100644 --- a/src/backend/chat/agents/summarize.py +++ b/src/backend/chat/agents/summarize.py @@ -1,16 +1,9 @@ """Build the summarization agent.""" -import asyncio import dataclasses import logging from django.conf import settings -from django.core.files.storage import default_storage - -import semchunk -from asgiref.sync import sync_to_async -from pydantic_ai import RunContext -from pydantic_ai.messages import ToolReturn from .base import BaseAgent @@ -28,134 +21,3 @@ class SummarizationAgent(BaseAgent): 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 summarize_chunk(idx, chunk, total_chunks, summarization_agent, ctx): - """Summarize a single chunk of text.""" - sum_prompt = ( - "You are an agent specializing in text summarization. " - "Generate a clear and concise summary of the following passage " - f"(part {idx}/{total_chunks}):\n'''\n{chunk}\n'''\n\n" - ) - - logger.debug( - "[summarize] CHUNK %s/%s prompt=> %s", idx, total_chunks, sum_prompt[0:100] + "..." - ) - - resp = await summarization_agent.run(sum_prompt, usage=ctx.usage) - - logger.debug("[summarize] CHUNK %s/%s response<= %s", idx, total_chunks, resp.output or "") - return resp.output or "" - - -async def hand_off_to_summarization_agent( # pylint: disable=too-many-locals - 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. - - Instructions are optional but should reflect the user's request. - - Examples: - "Summarize this doc in 2 paragraphs" -> instructions = "summary in 2 paragraphs" - "Summarize this doc in English" -> instructions = "In English" - "Summarize this doc" -> instructions = "" (default) - - Args: - instructions (str | None): The instructions the user gave to use for the summarization - """ - instructions_hint = ( - instructions.strip() if instructions else "The summary should contain 2 or 3 parts." - ) - summarization_agent = SummarizationAgent() - - # Collect documents content - 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] - - # Chunk documents and summarize each chunk - chunk_size = settings.SUMMARIZATION_CHUNK_SIZE - chunker = semchunk.chunkerify( - tokenizer_or_token_counter=lambda text: len(text.split()), - chunk_size=chunk_size, - ) - documents_chunks = chunker( - [doc[1] for doc in documents], - overlap=settings.SUMMARIZATION_OVERLAP_SIZE, - ) - - logger.info( - "[summarize] chunking: %s parts (size~%s), instructions='%s'", - sum(len(chunks) for chunks in documents_chunks), - chunk_size, - instructions_hint, - ) - - # Parallelize the chunk summarization with a semaphore to limit concurrent tasks - # because it can be very resource intensive on the LLM backend - semaphore = asyncio.Semaphore(settings.SUMMARIZATION_CONCURRENT_REQUESTS) - - async def summarize_chunk_with_semaphore(idx, chunk, total_chunks): - """Summarize a chunk with semaphore-controlled concurrency.""" - async with semaphore: - return await summarize_chunk(idx, chunk, total_chunks, summarization_agent, ctx) - - doc_chunk_summaries = [] - for doc_chunks in documents_chunks: - summarization_tasks = [ - summarize_chunk_with_semaphore(idx, chunk, len(doc_chunks)) - for idx, chunk in enumerate(doc_chunks, start=1) - ] - chunk_summaries = await asyncio.gather(*summarization_tasks) - doc_chunk_summaries.append(chunk_summaries) - - context = "\n\n".join( - doc_name + "\n\n" + "\n\n".join(summaries) - for doc_name, summaries in zip( - (doc[0] for doc in documents), - doc_chunk_summaries, - strict=True, - ) - ) - - # Merge chunk summaries into a single concise summary - merged_prompt = ( - "Produce a coherent synthesis from the summaries below.\n\n" - f"'''\n{context}\n'''\n\n" - "Constraints:\n" - "- Summarize without repetition.\n" - "- Harmonize style and terminology.\n" - "- The final summary must be well-structured and formatted in markdown.\n" - f"- Follow the instructions: {instructions_hint}\n" - "Respond directly with the final summary." - ) - - logger.debug("[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.debug("[summarize] MERGE response<= %s", final_summary) - - return ToolReturn( - 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 b7d9349..5988434 100644 --- a/src/backend/chat/clients/pydantic_ai.py +++ b/src/backend/chat/clients/pydantic_ai.py @@ -7,6 +7,7 @@ changes are needed in views.py or tests. """ import dataclasses +import functools import json import logging import time @@ -25,7 +26,7 @@ from django.utils.module_loading import import_string 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 @@ from chat.agents.local_media_url_processors import ( 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.clients.pydantic_ui_message_converter import ( ) 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 @@ -502,13 +503,11 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes "via the internal store." ) - @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) + @self.conversation_agent.tool(name="summarize") + @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 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 cf2c348..a4af8f9 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 @@ -151,7 +151,7 @@ def fixture_mock_summarization_agent(): super().__init__(**kwargs) self._model = FunctionModel(function=summarization_model) # pylint: disable=protected-access - with mock.patch("chat.agents.summarize.SummarizationAgent", new=SummarizationAgentMock): + with mock.patch("chat.tools.document_summarize.SummarizationAgent", new=SummarizationAgentMock): yield diff --git a/src/backend/chat/tools/document_summarize.py b/src/backend/chat/tools/document_summarize.py new file mode 100644 index 0000000..71c5e8e --- /dev/null +++ b/src/backend/chat/tools/document_summarize.py @@ -0,0 +1,147 @@ +"""Summarization tool used for uploaded documents.""" + +import asyncio +import logging + +from django.conf import settings +from django.core.files.storage import default_storage + +import semchunk +from asgiref.sync import sync_to_async +from pydantic_ai import RunContext +from pydantic_ai.messages import ToolReturn + +from chat.agents.summarize import SummarizationAgent + +logger = logging.getLogger(__name__) + + +@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 summarize_chunk(idx, chunk, total_chunks, summarization_agent, ctx): + """Summarize a single chunk of text.""" + sum_prompt = ( + "You are an agent specializing in text summarization. " + "Generate a clear and concise summary of the following passage " + f"(part {idx}/{total_chunks}):\n'''\n{chunk}\n'''\n\n" + ) + + logger.debug( + "[summarize] CHUNK %s/%s prompt=> %s", idx, total_chunks, sum_prompt[0:100] + "..." + ) + + resp = await summarization_agent.run(sum_prompt, usage=ctx.usage) + + logger.debug("[summarize] CHUNK %s/%s response<= %s", idx, total_chunks, resp.output or "") + return resp.output or "" + + +async def document_summarize( # pylint: disable=too-many-locals + 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. + + Instructions are optional but should reflect the user's request. + + Examples: + "Summarize this doc in 2 paragraphs" -> instructions = "summary in 2 paragraphs" + "Summarize this doc in English" -> instructions = "In English" + "Summarize this doc" -> instructions = "" (default) + + Args: + instructions (str | None): The instructions the user gave to use for the summarization + """ + instructions_hint = ( + instructions.strip() if instructions else "The summary should contain 2 or 3 parts." + ) + summarization_agent = SummarizationAgent() + + # Collect documents content + 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] + + # Chunk documents and summarize each chunk + chunk_size = settings.SUMMARIZATION_CHUNK_SIZE + chunker = semchunk.chunkerify( + tokenizer_or_token_counter=lambda text: len(text.split()), + chunk_size=chunk_size, + ) + documents_chunks = chunker( + [doc[1] for doc in documents], + overlap=settings.SUMMARIZATION_OVERLAP_SIZE, + ) + + logger.info( + "[summarize] chunking: %s parts (size~%s), instructions='%s'", + sum(len(chunks) for chunks in documents_chunks), + chunk_size, + instructions_hint, + ) + + # Parallelize the chunk summarization with a semaphore to limit concurrent tasks + # because it can be very resource intensive on the LLM backend + semaphore = asyncio.Semaphore(settings.SUMMARIZATION_CONCURRENT_REQUESTS) + + async def summarize_chunk_with_semaphore(idx, chunk, total_chunks): + """Summarize a chunk with semaphore-controlled concurrency.""" + async with semaphore: + return await summarize_chunk(idx, chunk, total_chunks, summarization_agent, ctx) + + doc_chunk_summaries = [] + for doc_chunks in documents_chunks: + summarization_tasks = [ + summarize_chunk_with_semaphore(idx, chunk, len(doc_chunks)) + for idx, chunk in enumerate(doc_chunks, start=1) + ] + chunk_summaries = await asyncio.gather(*summarization_tasks) + doc_chunk_summaries.append(chunk_summaries) + + context = "\n\n".join( + doc_name + "\n\n" + "\n\n".join(summaries) + for doc_name, summaries in zip( + (doc[0] for doc in documents), + doc_chunk_summaries, + strict=True, + ) + ) + + # Merge chunk summaries into a single concise summary + merged_prompt = ( + "Produce a coherent synthesis from the summaries below.\n\n" + f"'''\n{context}\n'''\n\n" + "Constraints:\n" + "- Summarize without repetition.\n" + "- Harmonize style and terminology.\n" + "- The final summary must be well-structured and formatted in markdown.\n" + f"- Follow the instructions: {instructions_hint}\n" + "Respond directly with the final summary." + ) + + logger.debug("[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.debug("[summarize] MERGE response<= %s", final_summary) + + return ToolReturn( + return_value=final_summary, + metadata={"sources": {doc[0] for doc in documents}}, + ) From c83c8c7da77af89488f72a32fed4ed44a9f05119 Mon Sep 17 00:00:00 2001 From: Quentin BEY Date: Thu, 6 Nov 2025 22:45:47 +0100 Subject: [PATCH 6/8] =?UTF-8?q?=F0=9F=8E=A8(summarize)=20add=20error=20han?= =?UTF-8?q?dling=20in=20tool=20call?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Allows the LLM to retry the summarization in some cases. --- src/backend/chat/clients/pydantic_ai.py | 2 +- src/backend/chat/tests/conftest.py | 35 ++ .../tests/tools/test_document_summarize.py | 472 ++++++++++++++++++ src/backend/chat/tools/document_summarize.py | 192 ++++--- 4 files changed, 625 insertions(+), 76 deletions(-) create mode 100644 src/backend/chat/tests/tools/test_document_summarize.py diff --git a/src/backend/chat/clients/pydantic_ai.py b/src/backend/chat/clients/pydantic_ai.py index 5988434..65265f6 100644 --- a/src/backend/chat/clients/pydantic_ai.py +++ b/src/backend/chat/clients/pydantic_ai.py @@ -503,7 +503,7 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes "via the internal store." ) - @self.conversation_agent.tool(name="summarize") + @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.""" diff --git a/src/backend/chat/tests/conftest.py b/src/backend/chat/tests/conftest.py index 0904999..93107f0 100644 --- a/src/backend/chat/tests/conftest.py +++ b/src/backend/chat/tests/conftest.py @@ -8,6 +8,7 @@ from django.utils import formats, timezone import pytest +from chat.agents.summarize import SummarizationAgent from chat.clients.pydantic_ai import AIAgentService logger = logging.getLogger(__name__) @@ -50,6 +51,40 @@ def mock_ai_agent_service_fixture(): 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" diff --git a/src/backend/chat/tests/tools/test_document_summarize.py b/src/backend/chat/tests/tools/test_document_summarize.py new file mode 100644 index 0000000..6a6d1b2 --- /dev/null +++ b/src/backend/chat/tests/tools/test_document_summarize.py @@ -0,0 +1,472 @@ +"""Tests for document_summarize functionality.""" + +import io +from unittest import mock + +from django.core.files.storage import default_storage + +import pytest +from pydantic_ai import ModelResponse, RunContext, TextPart +from pydantic_ai.exceptions import ModelRetry +from pydantic_ai.models.function import FunctionModel +from pydantic_ai.usage import RunUsage + +from chat.agents.summarize import SummarizationAgent +from chat.llm_configuration import LLModel, LLMProvider +from chat.tools.document_summarize import document_summarize, summarize_chunk + + +@pytest.fixture(autouse=True) +def fixture_summarization_agent_config(settings): + """Fixture to set used settings for agent configuration.""" + settings.LLM_CONFIGURATIONS = { + settings.LLM_SUMMARIZATION_MODEL_HRID: 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=[], + ), + } + + +@pytest.fixture(name="mocked_context") +def fixture_mocked_context(): + """Fixture for a mocked RunContext.""" + mock_ctx = mock.Mock(spec=RunContext) + mock_ctx.usage = RunUsage(input_tokens=0, output_tokens=0) + mock_ctx.max_retries = 2 + mock_ctx.retries = {} + return mock_ctx + + +def mocked_summary(_messages, _info=None): + """Mocked summary response.""" + return ModelResponse(parts=[TextPart(content="This is a summary of the test chunk.")]) + + +@pytest.mark.asyncio +async def test_summarize_chunk_returns_summary(mocked_context): + """Test that summarize_chunk returns a summary.""" + summarization_agent = SummarizationAgent() + + with summarization_agent.override(model=FunctionModel(mocked_summary)): + chunk = "This is a test chunk of text that needs to be summarized." + + result = await summarize_chunk(1, chunk, 1, summarization_agent, mocked_context) + + assert result == "This is a summary of the test chunk." + + +@pytest.mark.asyncio +async def test_summarize_chunk_raises_model_retry_on_error(mocked_context): + """Test that summarize_chunk raises ModelRetry when agent fails.""" + summarization_agent = SummarizationAgent() + + def mocked_summary_error(_messages, _info=None): + """Mocked summary that raises an error.""" + raise ValueError("Simulated LLM error") + + with summarization_agent.override(model=FunctionModel(mocked_summary_error)): + chunk = "This is a test chunk." + + with pytest.raises(ModelRetry) as exc_info: + await summarize_chunk(1, chunk, 1, summarization_agent, mocked_context) + + assert "An error occurred while summarizing a part of the document chunk" in str( + exc_info.value + ) + + +@pytest.mark.asyncio +async def test_summarize_chunk_handles_empty_response(mocked_context): + """Test that summarize_chunk handles empty responses from the agent.""" + summarization_agent = SummarizationAgent() + + def mocked_empty_summary(_messages, _info=None): + """Mocked summary that returns empty content.""" + return ModelResponse(parts=[TextPart(content="")]) + + with summarization_agent.override(model=FunctionModel(mocked_empty_summary)): + chunk = "This is a test chunk." + + # Empty responses cause ModelRetry since pydantic-ai considers them invalid + with pytest.raises(ModelRetry): + await summarize_chunk(1, chunk, 1, summarization_agent, mocked_context) + + +@pytest.mark.asyncio +async def test_document_summarize_single_document( + settings, mocked_context, mock_summarization_agent +): + """Test document_summarize with a single document.""" + settings.SUMMARIZATION_CHUNK_SIZE = 100 + settings.SUMMARIZATION_OVERLAP_SIZE = 10 + settings.SUMMARIZATION_CONCURRENT_REQUESTS = 2 + + # Create mock conversation with a text attachment + mock_conversation = mock.Mock() + mock_attachment = mock.Mock() + mock_attachment.key = "test_doc.txt" + mock_attachment.file_name = "test_doc.txt" + mock_attachment.content_type = "text/plain" + + mock_conversation.attachments.filter.return_value = [mock_attachment] + + # Mock file storage + file_content = "This is a test document. " * 20 # Create a document with some content + with mock.patch.object( + default_storage, "open", return_value=io.BytesIO(file_content.encode("utf-8")) + ): + # Set up mocked_context with conversation + mocked_context.deps = mock.Mock() + mocked_context.deps.conversation = mock_conversation + + call_count = {"chunk": 0, "merge": 0} + + def mocked_summary_full(messages, _info=None): + """Mocked summary for full flow.""" + messages_text = messages[0].parts[-1].content + + if "Produce a coherent synthesis" in messages_text: + call_count["merge"] += 1 + return ModelResponse( + parts=[TextPart(content="# Final Summary\n\nThis is the final merged summary.")] + ) + + call_count["chunk"] += 1 + return ModelResponse( + parts=[TextPart(content=f"Summary of chunk {call_count['chunk']}")] + ) + + with mock_summarization_agent(FunctionModel(mocked_summary_full)): + result = await document_summarize(mocked_context, instructions=None) + + assert result.return_value == "# Final Summary\n\nThis is the final merged summary." + assert result.metadata["sources"] == {"test_doc.txt"} + assert call_count["merge"] == 1 + + +@pytest.mark.asyncio +async def test_document_summarize_multiple_documents( + settings, mocked_context, mock_summarization_agent +): + """Test document_summarize with multiple documents.""" + settings.SUMMARIZATION_CHUNK_SIZE = 50 + settings.SUMMARIZATION_OVERLAP_SIZE = 5 + settings.SUMMARIZATION_CONCURRENT_REQUESTS = 2 + + # Create mock conversation with multiple text attachments + mock_conversation = mock.Mock() + mock_attachment1 = mock.Mock() + mock_attachment1.key = "doc1.txt" + mock_attachment1.file_name = "doc1.txt" + mock_attachment1.content_type = "text/plain" + + mock_attachment2 = mock.Mock() + mock_attachment2.key = "doc2.txt" + mock_attachment2.file_name = "doc2.txt" + mock_attachment2.content_type = "text/plain" + + mock_conversation.attachments.filter.return_value = [mock_attachment1, mock_attachment2] + + file_content1 = "Content of document one. " * 10 + file_content2 = "Content of document two. " * 10 + + def mock_open_side_effect(key): + """Mock file opening based on key.""" + if key == "doc1.txt": + return io.BytesIO(file_content1.encode("utf-8")) + return io.BytesIO(file_content2.encode("utf-8")) + + with mock.patch.object(default_storage, "open", side_effect=mock_open_side_effect): + # Set up mocked_context with conversation + mocked_context.deps = mock.Mock() + mocked_context.deps.conversation = mock_conversation + + def mocked_summary_multi(messages, _info=None): + """Mocked summary for multiple documents.""" + messages_text = messages[0].parts[-1].content + + if "Produce a coherent synthesis" in messages_text: + return ModelResponse(parts=[TextPart(content="Combined summary of all documents")]) + + return ModelResponse(parts=[TextPart(content="Chunk summary")]) + + with mock_summarization_agent(FunctionModel(mocked_summary_multi)): + result = await document_summarize(mocked_context, instructions=None) + + assert result.return_value == "Combined summary of all documents" + assert result.metadata["sources"] == {"doc1.txt", "doc2.txt"} + + +@pytest.mark.asyncio +async def test_document_summarize_with_custom_instructions( + settings, mocked_context, mock_summarization_agent +): + """Test document_summarize with custom instructions.""" + settings.SUMMARIZATION_CHUNK_SIZE = 100 + settings.SUMMARIZATION_OVERLAP_SIZE = 10 + settings.SUMMARIZATION_CONCURRENT_REQUESTS = 2 + + mock_conversation = mock.Mock() + mock_attachment = mock.Mock() + mock_attachment.key = "test.txt" + mock_attachment.file_name = "test.txt" + mock_attachment.content_type = "text/plain" + + mock_conversation.attachments.filter.return_value = [mock_attachment] + + file_content = "Test content " * 20 + + with mock.patch.object( + default_storage, "open", return_value=io.BytesIO(file_content.encode("utf-8")) + ): + # Set up mocked_context with conversation + mocked_context.deps = mock.Mock() + mocked_context.deps.conversation = mock_conversation + + captured_merge_prompt = [] + + def mocked_summary_with_instructions(messages, _info=None): + """Mocked summary that captures merge prompt.""" + messages_text = messages[0].parts[-1].content + + if "Produce a coherent synthesis" in messages_text: + captured_merge_prompt.append(messages_text) + return ModelResponse(parts=[TextPart(content="Summary in 2 paragraphs")]) + + return ModelResponse(parts=[TextPart(content="Chunk summary")]) + + with mock_summarization_agent(FunctionModel(mocked_summary_with_instructions)): + result = await document_summarize( + mocked_context, instructions="summary in 2 paragraphs" + ) + + assert result.return_value == "Summary in 2 paragraphs" + assert len(captured_merge_prompt) == 1 + assert "summary in 2 paragraphs" in captured_merge_prompt[0] + + +@pytest.mark.asyncio +async def test_document_summarize_no_text_attachments(mocked_context, mock_summarization_agent): + """Test document_summarize returns error message when no text documents found.""" + mock_conversation = mock.Mock() + mock_conversation.attachments.filter.return_value = [] + + # Set up mocked_context with conversation + mocked_context.deps = mock.Mock() + mocked_context.deps.conversation = mock_conversation + + # The decorator @last_model_retry_soft_fail catches ModelCannotRetry and returns a message + # We still need to provide a mock agent even if it won't be called + with mock_summarization_agent(FunctionModel(mocked_summary)): + result = await document_summarize(mocked_context, instructions=None) + + assert "No text documents found in the conversation" in result + + +@pytest.mark.asyncio +async def test_document_summarize_error_reading_document(mocked_context, mock_summarization_agent): + """Test document_summarize handles errors when reading documents.""" + mock_conversation = mock.Mock() + mock_attachment = mock.Mock() + mock_attachment.key = "test.txt" + mock_attachment.file_name = "test.txt" + mock_attachment.content_type = "text/plain" + + mock_conversation.attachments.filter.return_value = [mock_attachment] + + with mock.patch.object(default_storage, "open", side_effect=IOError("File read error")): + # Set up mocked_context with conversation + mocked_context.deps = mock.Mock() + mocked_context.deps.conversation = mock_conversation + + # The decorator @last_model_retry_soft_fail catches ModelCannotRetry and returns a message + # We still need to provide a mock agent even if it won't be called + with mock_summarization_agent(FunctionModel(mocked_summary)): + result = await document_summarize(mocked_context, instructions=None) + + assert "An unexpected error occurred during document summarization" in result + + +@pytest.mark.asyncio +async def test_document_summarize_error_during_chunk_summarization( + settings, mocked_context, mock_summarization_agent +): + """Test document_summarize handles ModelRetry during chunk summarization.""" + settings.SUMMARIZATION_CHUNK_SIZE = 100 + settings.SUMMARIZATION_OVERLAP_SIZE = 10 + settings.SUMMARIZATION_CONCURRENT_REQUESTS = 2 + + mock_conversation = mock.Mock() + mock_attachment = mock.Mock() + mock_attachment.key = "test.txt" + mock_attachment.file_name = "test.txt" + mock_attachment.content_type = "text/plain" + + mock_conversation.attachments.filter.return_value = [mock_attachment] + + file_content = "Test content " * 20 + + with mock.patch.object( + default_storage, "open", return_value=io.BytesIO(file_content.encode("utf-8")) + ): + # Set up mocked_context with conversation + mocked_context.deps = mock.Mock() + mocked_context.deps.conversation = mock_conversation + + def mocked_summary_error(messages, _info=None): + """Mocked summary that raises an error during chunks.""" + messages_text = messages[0].parts[-1].content + + if "Produce a coherent synthesis" not in messages_text: + raise ValueError("Chunk processing error") + + return ModelResponse(parts=[TextPart(content="Final summary")]) + + with mock_summarization_agent(FunctionModel(mocked_summary_error)): + with pytest.raises(ModelRetry): + await document_summarize(mocked_context, instructions=None) + + +@pytest.mark.asyncio +async def test_document_summarize_error_during_merge( + settings, mocked_context, mock_summarization_agent +): + """Test document_summarize handles errors during final merge.""" + settings.SUMMARIZATION_CHUNK_SIZE = 100 + settings.SUMMARIZATION_OVERLAP_SIZE = 10 + settings.SUMMARIZATION_CONCURRENT_REQUESTS = 2 + + mock_conversation = mock.Mock() + mock_attachment = mock.Mock() + mock_attachment.key = "test.txt" + mock_attachment.file_name = "test.txt" + mock_attachment.content_type = "text/plain" + + mock_conversation.attachments.filter.return_value = [mock_attachment] + + file_content = "Test content " * 20 + + with mock.patch.object( + default_storage, "open", return_value=io.BytesIO(file_content.encode("utf-8")) + ): + # Set up mocked_context with conversation + mocked_context.deps = mock.Mock() + mocked_context.deps.conversation = mock_conversation + + def mocked_summary_merge_error(messages, _info=None): + """Mocked summary that raises an error during merge.""" + messages_text = messages[0].parts[-1].content + + if "Produce a coherent synthesis" in messages_text: + raise ValueError("Merge error") + + return ModelResponse(parts=[TextPart(content="Chunk summary")]) + + with mock_summarization_agent(FunctionModel(mocked_summary_merge_error)): + with pytest.raises(ModelRetry) as exc_info: + await document_summarize(mocked_context, instructions=None) + + # Should raise ModelRetry regardless of which phase failed + assert "An error occurred" in str(exc_info.value) + + +@pytest.mark.asyncio +async def test_document_summarize_empty_result(settings, mocked_context, mock_summarization_agent): + """Test document_summarize raises ModelRetry when summarization produces empty result.""" + settings.SUMMARIZATION_CHUNK_SIZE = 100 + settings.SUMMARIZATION_OVERLAP_SIZE = 10 + settings.SUMMARIZATION_CONCURRENT_REQUESTS = 2 + + mock_conversation = mock.Mock() + mock_attachment = mock.Mock() + mock_attachment.key = "test.txt" + mock_attachment.file_name = "test.txt" + mock_attachment.content_type = "text/plain" + + mock_conversation.attachments.filter.return_value = [mock_attachment] + + file_content = "Test content " * 20 + + with mock.patch.object( + default_storage, "open", return_value=io.BytesIO(file_content.encode("utf-8")) + ): + # Set up mocked_context with conversation + mocked_context.deps = mock.Mock() + mocked_context.deps.conversation = mock_conversation + + def mocked_empty_summary(messages, _info=None): + """Mocked summary that returns empty for merge.""" + messages_text = messages[0].parts[-1].content + + if "Produce a coherent synthesis" in messages_text: + return ModelResponse(parts=[TextPart(content=" ")]) + + return ModelResponse(parts=[TextPart(content="Chunk summary")]) + + with mock_summarization_agent(FunctionModel(mocked_empty_summary)): + with pytest.raises(ModelRetry) as exc_info: + await document_summarize(mocked_context, instructions=None) + + # Should raise ModelRetry with the specific message + assert "The summarization produced an empty result" in str(exc_info.value) + + +@pytest.mark.asyncio +async def test_document_summarize_large_document_multiple_chunks( + settings, mocked_context, mock_summarization_agent +): + """Test document_summarize with a large document requiring multiple chunks.""" + settings.SUMMARIZATION_CHUNK_SIZE = 20 # Small chunk size to force multiple chunks + settings.SUMMARIZATION_OVERLAP_SIZE = 5 + settings.SUMMARIZATION_CONCURRENT_REQUESTS = 2 + + mock_conversation = mock.Mock() + mock_attachment = mock.Mock() + mock_attachment.key = "large_doc.txt" + mock_attachment.file_name = "large_doc.txt" + mock_attachment.content_type = "text/plain" + + mock_conversation.attachments.filter.return_value = [mock_attachment] + + # Create a large document + file_content = "This is a word. " * 100 # Should create multiple chunks + + with mock.patch.object( + default_storage, "open", return_value=io.BytesIO(file_content.encode("utf-8")) + ): + # Set up mocked_context with conversation + mocked_context.deps = mock.Mock() + mocked_context.deps.conversation = mock_conversation + + chunk_count = {"count": 0} + + def mocked_summary_multi_chunks(messages, _info=None): + """Mocked summary that counts chunks.""" + messages_text = messages[0].parts[-1].content + + if "Produce a coherent synthesis" in messages_text: + return ModelResponse( + parts=[TextPart(content=f"Final summary of {chunk_count['count']} chunks")] + ) + + chunk_count["count"] += 1 + return ModelResponse( + parts=[TextPart(content=f"Summary of chunk {chunk_count['count']}")] + ) + + with mock_summarization_agent(FunctionModel(mocked_summary_multi_chunks)): + result = await document_summarize(mocked_context, instructions=None) + + assert "Final summary of" in result.return_value + assert chunk_count["count"] > 1 # Ensure multiple chunks were processed diff --git a/src/backend/chat/tools/document_summarize.py b/src/backend/chat/tools/document_summarize.py index 71c5e8e..83faca7 100644 --- a/src/backend/chat/tools/document_summarize.py +++ b/src/backend/chat/tools/document_summarize.py @@ -9,9 +9,12 @@ from django.core.files.storage import default_storage import semchunk from asgiref.sync import sync_to_async from pydantic_ai import RunContext +from pydantic_ai.exceptions import ModelRetry from pydantic_ai.messages import ToolReturn from chat.agents.summarize import SummarizationAgent +from chat.tools.exceptions import ModelCannotRetry +from chat.tools.utils import last_model_retry_soft_fail logger = logging.getLogger(__name__) @@ -35,12 +38,19 @@ async def summarize_chunk(idx, chunk, total_chunks, summarization_agent, ctx): "[summarize] CHUNK %s/%s prompt=> %s", idx, total_chunks, sum_prompt[0:100] + "..." ) - resp = await summarization_agent.run(sum_prompt, usage=ctx.usage) + try: + resp = await summarization_agent.run(sum_prompt, usage=ctx.usage) + except Exception as exc: + logger.warning("Error during chunk summarization: %s", exc, exc_info=True) + raise ModelRetry( + "An error occurred while summarizing a part of the document chunk." + ) from exc logger.debug("[summarize] CHUNK %s/%s response<= %s", idx, total_chunks, resp.output or "") return resp.output or "" +@last_model_retry_soft_fail async def document_summarize( # pylint: disable=too-many-locals ctx: RunContext, *, instructions: str | None = None ) -> ToolReturn: @@ -62,86 +72,118 @@ async def document_summarize( # pylint: disable=too-many-locals Args: instructions (str | None): The instructions the user gave to use for the summarization """ - instructions_hint = ( - instructions.strip() if instructions else "The summary should contain 2 or 3 parts." - ) - summarization_agent = SummarizationAgent() - - # Collect documents content - text_attachment = await sync_to_async(list)( - ctx.deps.conversation.attachments.filter( - content_type__startswith="text/", + try: + instructions_hint = ( + instructions.strip() if instructions else "The summary should contain 2 or 3 parts." ) - ) + summarization_agent = SummarizationAgent() - documents = [await read_document_content(doc) for doc in text_attachment] - - # Chunk documents and summarize each chunk - chunk_size = settings.SUMMARIZATION_CHUNK_SIZE - chunker = semchunk.chunkerify( - tokenizer_or_token_counter=lambda text: len(text.split()), - chunk_size=chunk_size, - ) - documents_chunks = chunker( - [doc[1] for doc in documents], - overlap=settings.SUMMARIZATION_OVERLAP_SIZE, - ) - - logger.info( - "[summarize] chunking: %s parts (size~%s), instructions='%s'", - sum(len(chunks) for chunks in documents_chunks), - chunk_size, - instructions_hint, - ) - - # Parallelize the chunk summarization with a semaphore to limit concurrent tasks - # because it can be very resource intensive on the LLM backend - semaphore = asyncio.Semaphore(settings.SUMMARIZATION_CONCURRENT_REQUESTS) - - async def summarize_chunk_with_semaphore(idx, chunk, total_chunks): - """Summarize a chunk with semaphore-controlled concurrency.""" - async with semaphore: - return await summarize_chunk(idx, chunk, total_chunks, summarization_agent, ctx) - - doc_chunk_summaries = [] - for doc_chunks in documents_chunks: - summarization_tasks = [ - summarize_chunk_with_semaphore(idx, chunk, len(doc_chunks)) - for idx, chunk in enumerate(doc_chunks, start=1) - ] - chunk_summaries = await asyncio.gather(*summarization_tasks) - doc_chunk_summaries.append(chunk_summaries) - - context = "\n\n".join( - doc_name + "\n\n" + "\n\n".join(summaries) - for doc_name, summaries in zip( - (doc[0] for doc in documents), - doc_chunk_summaries, - strict=True, + # Collect documents content + text_attachment = await sync_to_async(list)( + ctx.deps.conversation.attachments.filter( + content_type__startswith="text/", + ) ) - ) - # Merge chunk summaries into a single concise summary - merged_prompt = ( - "Produce a coherent synthesis from the summaries below.\n\n" - f"'''\n{context}\n'''\n\n" - "Constraints:\n" - "- Summarize without repetition.\n" - "- Harmonize style and terminology.\n" - "- The final summary must be well-structured and formatted in markdown.\n" - f"- Follow the instructions: {instructions_hint}\n" - "Respond directly with the final summary." - ) + if not text_attachment: + raise ModelCannotRetry( + "No text documents found in the conversation. " + "You must explain this to the user and ask them to provide documents." + ) - logger.debug("[summarize] MERGE prompt=> %s", merged_prompt) + documents = [await read_document_content(doc) for doc in text_attachment] - merged_resp = await summarization_agent.run(merged_prompt, usage=ctx.usage) + # Chunk documents and summarize each chunk + chunk_size = settings.SUMMARIZATION_CHUNK_SIZE + chunker = semchunk.chunkerify( + tokenizer_or_token_counter=lambda text: len(text.split()), + chunk_size=chunk_size, + ) + documents_chunks = chunker( + [doc[1] for doc in documents], + overlap=settings.SUMMARIZATION_OVERLAP_SIZE, + ) - final_summary = (merged_resp.output or "").strip() + logger.info( + "[summarize] chunking: %s parts (size~%s), instructions='%s'", + sum(len(chunks) for chunks in documents_chunks), + chunk_size, + instructions_hint, + ) - logger.debug("[summarize] MERGE response<= %s", final_summary) + # Parallelize the chunk summarization with a semaphore to limit concurrent tasks + # because it can be very resource intensive on the LLM backend + semaphore = asyncio.Semaphore(settings.SUMMARIZATION_CONCURRENT_REQUESTS) - return ToolReturn( - return_value=final_summary, - metadata={"sources": {doc[0] for doc in documents}}, - ) + async def summarize_chunk_with_semaphore(idx, chunk, total_chunks): + """Summarize a chunk with semaphore-controlled concurrency.""" + async with semaphore: + return await summarize_chunk(idx, chunk, total_chunks, summarization_agent, ctx) + + doc_chunk_summaries = [] + try: + for doc_chunks in documents_chunks: + summarization_tasks = [ + summarize_chunk_with_semaphore(idx, chunk, len(doc_chunks)) + for idx, chunk in enumerate(doc_chunks, start=1) + ] + chunk_summaries = await asyncio.gather(*summarization_tasks) + doc_chunk_summaries.append(chunk_summaries) + except ModelRetry as exc: + logger.warning("Retryable error during chunk summarization: %s", exc, exc_info=True) + raise + except Exception as exc: + logger.warning("Error during chunk summarization: %s", exc, exc_info=True) + raise ModelRetry("An error occurred while processing document chunks.") from exc + + context = "\n\n".join( + doc_name + "\n\n" + "\n\n".join(summaries) + for doc_name, summaries in zip( + (doc[0] for doc in documents), + doc_chunk_summaries, + strict=True, + ) + ) + + # Merge chunk summaries into a single concise summary + merged_prompt = ( + "Produce a coherent synthesis from the summaries below.\n\n" + f"'''\n{context}\n'''\n\n" + "Constraints:\n" + "- Summarize without repetition.\n" + "- Harmonize style and terminology.\n" + "- The final summary must be well-structured and formatted in markdown.\n" + f"- Follow the instructions: {instructions_hint}\n" + "Respond directly with the final summary." + ) + + logger.debug("[summarize] MERGE prompt=> %s", merged_prompt) + + try: + merged_resp = await summarization_agent.run(merged_prompt, usage=ctx.usage) + except Exception as exc: + logger.warning("Error during merge summarization: %s", exc, exc_info=True) + raise ModelRetry("An error occurred while generating the final summary.") from exc + + final_summary = (merged_resp.output or "").strip() + + if not final_summary: + raise ModelRetry("The summarization produced an empty result.") + + logger.debug("[summarize] MERGE response<= %s", final_summary) + + return ToolReturn( + return_value=final_summary, + metadata={"sources": {doc[0] for doc in documents}}, + ) + + except (ModelCannotRetry, ModelRetry): + # Re-raise these as-is + raise + except Exception as exc: + # Unexpected error - stop and inform user + logger.exception("Unexpected error in document_summarize: %s", exc) + raise ModelCannotRetry( + f"An unexpected error occurred during document summarization: {type(exc).__name__}. " + "You must explain this to the user and not try to answer based on your knowledge." + ) from exc From b33a3e49873f4e2ad7c8a8b9f53bdb0fcf668d20 Mon Sep 17 00:00:00 2001 From: Quentin BEY Date: Thu, 6 Nov 2025 23:13:54 +0100 Subject: [PATCH 7/8] =?UTF-8?q?=F0=9F=90=9B(summarize)=20fix=20the=20summa?= =?UTF-8?q?rization=20tool=20loader?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When the loader was not in the proper place. --- .../src/features/chat/components/Chat.tsx | 18 +++++++++++++++++- .../chat/components/ToolInvocationItem.tsx | 2 +- 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/src/frontend/apps/conversations/src/features/chat/components/Chat.tsx b/src/frontend/apps/conversations/src/features/chat/components/Chat.tsx index 4203739..ef6a5ae 100644 --- a/src/frontend/apps/conversations/src/features/chat/components/Chat.tsx +++ b/src/frontend/apps/conversations/src/features/chat/components/Chat.tsx @@ -701,7 +701,23 @@ export const Chat = ({ > - {t('Search...')} + {(() => { + const toolInvocation = message.parts?.find( + (part) => + part.type === 'tool-invocation' && + part.toolInvocation.toolName !== + 'document_parsing', + ); + if ( + toolInvocation?.type === + 'tool-invocation' && + toolInvocation.toolInvocation.toolName === + 'summarize' + ) { + return t('Summarizing...'); + } + return t('Search...'); + })()} )} 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 a6f8305..de5a712 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 = ({ > - {toolInvocation.toolName === 'summarize' ? t('Summarizing...') : t('Search...')} + {t('Search...')} )} From 78b9f1117991cb43e2688fb84d17614890ced7b9 Mon Sep 17 00:00:00 2001 From: Quentin BEY Date: Fri, 7 Nov 2025 09:56:05 +0100 Subject: [PATCH 8/8] =?UTF-8?q?=E2=9C=85(attachments)=20fix=20DBB=20object?= =?UTF-8?q?=20existence=20test?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The test was flaky --- src/backend/chat/tests/views/chat/test_media_auth.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/backend/chat/tests/views/chat/test_media_auth.py b/src/backend/chat/tests/views/chat/test_media_auth.py index 2b7206e..362ac69 100644 --- a/src/backend/chat/tests/views/chat/test_media_auth.py +++ b/src/backend/chat/tests/views/chat/test_media_auth.py @@ -25,6 +25,8 @@ def test_api_media_auth_unkown_document(api_client): Trying to download a media related to a conversation that does not exist should not have the side effect to create it (no regression test). """ + ChatConversation.objects.all().delete() + original_url = f"http://localhost/media/{uuid4()!s}/attachments/{uuid4()!s}.jpg" response = api_client.get("/api/v1.0/chats/media-auth/", HTTP_X_ORIGINAL_URL=original_url)