⚡️(summarize) use semchunk for better doc chunking

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.
2025-11-06 23:17:41 +01:00
parent 392eeece3e
commit 78a2393383
7 changed files with 166 additions and 99 deletions
@@ -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,
@@ -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
@@ -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":"<mocked_uuid>"}\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",
@@ -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?",
@@ -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)
@@ -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(
@@ -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",