diff --git a/CHANGELOG.md b/CHANGELOG.md index 869458f..7b21107 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,11 @@ 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 + +- 🔥(posthog) remove posthog middleware for async mode fix #146 ## [0.0.7] - 2025-10-28 diff --git a/src/backend/chat/agents/summarize.py b/src/backend/chat/agents/summarize.py index 38d92d6..5efcbb8 100644 --- a/src/backend/chat/agents/summarize.py +++ b/src/backend/chat/agents/summarize.py @@ -4,11 +4,6 @@ import dataclasses import logging from django.conf import settings -from django.core.files.storage import default_storage - -from asgiref.sync import sync_to_async -from pydantic_ai import RunContext -from pydantic_ai.messages import ToolReturn from .base import BaseAgent @@ -26,60 +21,3 @@ 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 hand_off_to_summarization_agent(ctx: RunContext) -> ToolReturn: - """ - Generate a complete, ready-to-use summary of the documents in context - (do not request the documents to the user). - Return this summary directly to the user WITHOUT any modification, - or additional summarization. - The summary is already optimized and MUST be presented as-is in the final response - or translated preserving the information. - """ - summarization_agent = SummarizationAgent() - - prompt = ( - "Do not mention the user request in your answer.\n" - "User request:\n" - "{user_prompt}\n\n" - "Document contents:\n" - "{documents_prompt}\n" - ) - text_attachment = await sync_to_async(list)( - ctx.deps.conversation.attachments.filter( - content_type__startswith="text/", - ) - ) - - documents = [await read_document_content(doc) for doc in text_attachment] - - documents_prompt = "\n\n".join( - [ - (f"\n\n{name}\n\n\n{content}\n\n") - for name, content in documents - ] - ) - - formatted_prompt = prompt.format( - user_prompt=ctx.prompt, - documents_prompt=documents_prompt, - ) - - logger.debug("Summarize prompt: %s", formatted_prompt) - - response = await summarization_agent.run(formatted_prompt, usage=ctx.usage) - - logger.debug("Summarize response: %s", response) - - return ToolReturn( - return_value=response.output, - metadata={"sources": {doc[0] for doc in documents}}, - ) diff --git a/src/backend/chat/clients/pydantic_ai.py b/src/backend/chat/clients/pydantic_ai.py index e5d8056..65265f6 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 @@ -493,13 +494,20 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes "You may add a follow-up question after the summary if needed." ) - @self.conversation_agent.tool - async def summarize(ctx) -> ToolReturn: - """ - Summarize the documents for the user, only when asked for, - the documents are in my context. - """ - return await hand_off_to_summarization_agent(ctx) + # Inform the model (system-level) that documents are attached and available + @self.conversation_agent.system_prompt + def attached_documents_note() -> str: + return ( + "[Internal context] User documents are attached to this conversation. " + "Do not request re-upload of documents; consider them already available " + "via the internal store." + ) + + @self.conversation_agent.tool(name="summarize", retries=2) + @functools.wraps(document_summarize) + async def summarize(ctx: RunContext, *args, **kwargs) -> ToolReturn: + """Wrap the document_summarize tool to provide context and add the tool.""" + return await document_summarize(ctx, *args, **kwargs) else: conversation_documents = [ cd 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/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..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 @@ -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/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) diff --git a/src/backend/chat/tools/document_search_rag.py b/src/backend/chat/tools/document_search_rag.py index 121b43f..f59c119 100644 --- a/src/backend/chat/tools/document_search_rag.py +++ b/src/backend/chat/tools/document_search_rag.py @@ -20,7 +20,7 @@ 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. """ document_store_backend = import_string(settings.RAG_DOCUMENT_SEARCH_BACKEND) @@ -43,8 +43,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/backend/chat/tools/document_summarize.py b/src/backend/chat/tools/document_summarize.py new file mode 100644 index 0000000..83faca7 --- /dev/null +++ b/src/backend/chat/tools/document_summarize.py @@ -0,0 +1,189 @@ +"""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.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__) + + +@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] + "..." + ) + + 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: + """ + 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 + """ + try: + 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/", + ) + ) + + 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." + ) + + 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 = [] + 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 diff --git a/src/backend/conversations/settings.py b/src/backend/conversations/settings.py index e51c3b0..759dc72 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", ] @@ -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/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) diff --git a/src/backend/pyproject.toml b/src/backend/pyproject.toml index a88b5de..311a4eb 100644 --- a/src/backend/pyproject.toml +++ b/src/backend/pyproject.toml @@ -27,41 +27,42 @@ 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", + "semchunk==3.2.5", + "sentry-sdk==2.43.0", "trafilatura==2.0.0", "uvicorn==0.38.0", "whitenoise==6.11.0", @@ -81,11 +82,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 +95,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", ] 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...'); + })()} )}