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",