diff --git a/CHANGELOG.md b/CHANGELOG.md index 8b13cbd..c83a5dc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,6 +18,7 @@ and this project adheres to - 🎉(conversations) bootstrap backend & frontend #1 - ✨(web-search) add RAG capability to do web search #7 +- ✨(chat) add document RAG on document uploaded by user #8 [unreleased]: https://github.com/numerique-gouv/conversations/compare/HEAD...main diff --git a/src/backend/.pylintrc b/src/backend/.pylintrc index 9807de7..5872279 100644 --- a/src/backend/.pylintrc +++ b/src/backend/.pylintrc @@ -23,7 +23,9 @@ jobs=0 # List of plugins (as comma separated values of python modules names) to load, # usually to register additional checkers. -load-plugins=pylint_django,pylint.extensions.no_self_use +load-plugins=pylint_django, + pylint.extensions.no_self_use, + pylint_pydantic, # Pickle collected data for later comparisons. persistent=yes diff --git a/src/backend/chat/agent_rag/constants.py b/src/backend/chat/agent_rag/constants.py index 0013d6f..0fb7d2c 100644 --- a/src/backend/chat/agent_rag/constants.py +++ b/src/backend/chat/agent_rag/constants.py @@ -32,7 +32,7 @@ class RAGWebResults(BaseModel): def to_prompt(self) -> str: """Convert the web results to a prompt string.""" - _format = " - URL: {url}:\n content: {content}\n\n" + _format = " - From: {url}:\n content: {content}\n\n" return ( "\n\n".join( _format.format(url=result.url, content=result.content) for result in self.data diff --git a/src/backend/chat/agent_rag/document_converter/__init__.py b/src/backend/chat/agent_rag/document_converter/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/backend/chat/agent_rag/document_converter/markitdown.py b/src/backend/chat/agent_rag/document_converter/markitdown.py new file mode 100644 index 0000000..f0761d0 --- /dev/null +++ b/src/backend/chat/agent_rag/document_converter/markitdown.py @@ -0,0 +1,42 @@ +"""Document Converter using MarkItDown""" + +from io import BytesIO +from pathlib import Path +from typing import BinaryIO, Union + +from markitdown import MarkItDown + + +class DocumentConverter: + """Simple document converter that uses MarkItDown to convert documents to Markdown format.""" + + def __init__(self): + """Initialize the DocumentConverter with MarkItDown.""" + self.converter = MarkItDown(enable_plugins=False) + + def convert_raw( # pylint: disable=unused-argument + self, + *, + name: str, + content_type: str, + content: BytesIO, + ) -> str: + """ + Convert a document to Markdown format. + The name, content_type, and content parameters comes from the user input + (vercel SDK Attachment, or BinaryContent). + + Args: + name (str): The name of the document. + content_type (str): The MIME type of the document (e.g., "application/pdf"). + content (BytesIO): The content of the document as a BytesIO stream. + """ + return self._convert(content) + + def _convert(self, document: Union[Path, str, BinaryIO]) -> str: + """ + Convert the given document using the underlying DocumentConverter. + """ + conversion = self.converter.convert(document) + document_markdown = conversion.text_content + return document_markdown diff --git a/src/backend/chat/agent_rag/document_search/__init__.py b/src/backend/chat/agent_rag/document_search/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/backend/chat/agent_rag/document_search/albert_api.py b/src/backend/chat/agent_rag/document_search/albert_api.py new file mode 100644 index 0000000..eddfc98 --- /dev/null +++ b/src/backend/chat/agent_rag/document_search/albert_api.py @@ -0,0 +1,206 @@ +"""Implementation of the Albert API for RAG document search.""" + +import json +import logging +from io import BytesIO +from urllib.parse import urljoin + +from django.conf import settings + +import requests + +from chat.agent_rag.albert_api_constants import Searches +from chat.agent_rag.constants import RAGWebResult, RAGWebResults, RAGWebUsage +from chat.agent_rag.document_converter.markitdown import DocumentConverter +from chat.models import ChatConversation + +logger = logging.getLogger(__name__) + + +class AlbertRagDocumentSearch: + """ + This class is a placeholder for the Albert API implementation. + It is designed to be used with the RAG (Retrieval-Augmented Generation) document search system. + + It provides methods to: + - Create a collection for the search operation. + - Parse documents and convert them to Markdown format: + + Handle PDF parsing using the Albert API. + + Use the DocumentConverter (markitdown) for other formats. + - Store parsed documents in the Albert collection. + - Perform a search operation using the Albert API. + """ + + def __init__(self, conversation: ChatConversation): + # Initialize any necessary parameters or configurations here + self._base_url = settings.ALBERT_API_URL + self._headers = { + "Authorization": f"Bearer {settings.ALBERT_API_KEY}", + } + self._collections_endpoint = urljoin(self._base_url, "/v1/collections") + self._documents_endpoint = urljoin(self._base_url, "/v1/documents") + self._pdf_parser_endpoint = urljoin(self._base_url, "/v1/parse-beta") + self._search_endpoint = urljoin(self._base_url, "/v1/search") + + self.conversation = conversation + + @property + def _albert_collection_id(self): + """ + Generate the collection name based on the conversation ID. + This is used to create or retrieve a collection for the search operation. + """ + return f"conversation-{self.conversation.pk}" + + @property + def collection_id(self) -> int: + """ + Get the collection ID for the current conversation. + + Might be created later by self._create_collection() if it does not exist. + """ + return int(self.conversation.collection_id) if self.conversation.collection_id else None + + def _create_collection(self) -> bool: + """ + Create a temporary collection for the search operation. + This method should handle the logic to create or retrieve an existing collection. + """ + response = requests.post( + self._collections_endpoint, + headers=self._headers, + json={ + "name": self._albert_collection_id, + "description": "Temporary collection for RAG document search", + "visibility": "private", + }, + timeout=settings.ALBERT_API_TIMEOUT, + ) + response.raise_for_status() + self.conversation.collection_id = str(response.json()["id"]) + return True + + def _parse_pdf_document(self, name: str, content_type: str, content: BytesIO) -> str: + """ + Parse the PDF document content and return the text content. + This method should handle the logic to convert the PDF into + a format suitable for the Albert API. + """ + response = requests.post( + self._pdf_parser_endpoint, + headers=self._headers, + files={ + "file": ( + name, + content, + content_type, + ), # Use the name as the filename in the request + "output_format": (None, "markdown"), # Specify the output format as Markdown, + }, + timeout=settings.ALBERT_API_PARSE_TIMEOUT, + ) + response.raise_for_status() + + return "\n\n".join( + document_page["content"] for document_page in response.json().get("data", []) + ) + + def parse_document(self, name: str, content_type: str, content: BytesIO): + """ + Parse the document and prepare it for the search operation. + This method should handle the logic to convert the document + into a format suitable for the Albert API. + + Args: + name (str): The name of the document. + content_type (str): The MIME type of the document (e.g., "application/pdf"). + content (BytesIO): The content of the document as a BytesIO stream. + + Returns: + str: The document content in Markdown format. + """ + # Implement the parsing logic here + if content_type == "application/pdf": + # Handle PDF parsing + markdown_content = self._parse_pdf_document( + name=name, content_type=content_type, content=content + ) + else: + markdown_content = DocumentConverter().convert_raw( + name=name, content_type=content_type, content=content + ) + + return markdown_content + + def _store_document(self, name: str, content: str): + """ + Store the document content in the Albert collection. + This method should handle the logic to send the document content to the Albert API. + + Args: + content (str): The content of the document in Markdown format. + """ + if not self.collection_id and not self._create_collection(): + raise RuntimeError("Failed to create or retrieve the collection.") + + response = requests.post( + urljoin(self._base_url, self._documents_endpoint), + headers=self._headers, + files={ + "file": (f"{name}.md", BytesIO(content.encode("utf-8")), "text/markdown"), + "collection": (None, int(self.collection_id)), + "metadata": (None, json.dumps({"document_name": name})), # undocumented API + }, + timeout=settings.ALBERT_API_TIMEOUT, + ) + logger.debug(response.json()) + response.raise_for_status() + + def parse_and_store_document(self, name: str, content_type: str, content: BytesIO): + """ + Parse the document and store it in the Albert collection. + + Args: + name (str): The name of the document. + content_type (str): The MIME type of the document (e.g., "application/pdf"). + content (BytesIO): The content of the document as a BytesIO stream. + """ + document_content = self.parse_document(name, content_type, content) + self._store_document(name, document_content) + return document_content + + def search(self, query): + """ + Perform a search using the Albert API based on the provided query. + + :param query: The search query string. + :return: Search results from the Albert API. + """ + response = requests.post( + urljoin(self._base_url, self._search_endpoint), + headers=self._headers, + json={ + "collections": [self.collection_id], + "prompt": query, + "score_threshold": 0.6, + }, + timeout=settings.ALBERT_API_TIMEOUT, + ) + response.raise_for_status() + + searches = Searches(**response.json()) + + return RAGWebResults( + data=[ + RAGWebResult( + url=result.chunk.metadata["document_name"], + content=result.chunk.content, + score=result.score, + ) + for result in searches.data + ], + usage=RAGWebUsage( + prompt_tokens=searches.usage.prompt_tokens, + completion_tokens=searches.usage.completion_tokens, + ), + ) diff --git a/src/backend/chat/clients/exceptions.py b/src/backend/chat/clients/exceptions.py new file mode 100644 index 0000000..7dbb6b4 --- /dev/null +++ b/src/backend/chat/clients/exceptions.py @@ -0,0 +1,9 @@ +"""Module containing custom exceptions for chat clients.""" + + +class WebSearchEmptyException(Exception): + """Exception raised when a web search returns no results.""" + + def __init__(self, message="Web search returned no results."): + self.message = message + super().__init__(self.message) diff --git a/src/backend/chat/clients/pydantic_ai.py b/src/backend/chat/clients/pydantic_ai.py index 23a47cf..6ef91de 100644 --- a/src/backend/chat/clients/pydantic_ai.py +++ b/src/backend/chat/clients/pydantic_ai.py @@ -11,16 +11,18 @@ import json import logging import uuid from contextlib import AsyncExitStack -from typing import Dict, List +from typing import Dict, List, Tuple from django.conf import settings from django.core.exceptions import ImproperlyConfigured from django.utils.module_loading import import_string +from django.utils.translation import gettext_lazy as _ from asgiref.sync import sync_to_async from pydantic import BaseModel from pydantic_ai import Agent, NativeOutput from pydantic_ai.messages import ( + BinaryContent, FunctionToolCallEvent, FunctionToolResultEvent, ModelMessage, @@ -37,18 +39,19 @@ from pydantic_ai.messages import ( ToolCallPart, ToolCallPartDelta, ToolReturnPart, - UserContent, UserPromptPart, ) from pydantic_ai.models.openai import OpenAIModel, OpenAIResponsesModelSettings from pydantic_ai.providers.openai import OpenAIProvider +from chat.agent_rag.document_search.albert_api import AlbertRagDocumentSearch from chat.ai_sdk_types import ( LanguageModelV1Source, SourceUIPart, UIMessage, ) from chat.clients.async_to_sync import convert_async_generator_to_sync +from chat.clients.exceptions import WebSearchEmptyException from chat.clients.pydantic_ui_message_converter import ( model_message_to_ui_message, ui_message_to_user_content, @@ -88,6 +91,7 @@ class UserIntent(BaseModel): """Model to represent the detected user intent.""" web_search: bool = False + attachment_summary: bool = False class AIAgentService: @@ -127,12 +131,12 @@ class AIAgentService: # Core agent runner # --------------------------------------------------------------------- # - async def _detect_user_intent(self, user_prompt: List[UserContent]) -> UserIntent: + async def _detect_user_intent(self, user_prompt: str) -> UserIntent: """ Detect the user intent by calling a small LLM. Args: - user_prompt (List[UserContent]): The user prompt to analyze. + user_prompt str: The user prompt to analyze. Returns: UserIntent: The detected user intent, indicating if a web search is needed. Raises: @@ -148,9 +152,8 @@ class AIAgentService: ) if not setting # ie if setting is None or setting == "" ]: - raise ImproperlyConfigured( - f"AI routing model configuration not set: {missing_settings}" - ) + logger.error("AI routing model configuration not set: %s", missing_settings) + return UserIntent() agent = Agent( model=OpenAIModel( @@ -166,8 +169,145 @@ class AIAgentService: result = await agent.run(user_prompt) logger.debug("Detected user intent: %s", result) + + # Disable some intent if the project settings do not allow it + if not settings.RAG_WEB_SEARCH_BACKEND: + # If web search is not enabled, we can skip the intent detection + result.output.web_search = False + logger.info("Web search backend is disabled, skipping intent detection.") + return result.output + def parse_input_documents(self, documents: List[BinaryContent]): + """ + Parse and store input documents in the conversation's document store. + """ + document_store_backend = import_string(settings.RAG_DOCUMENT_SEARCH_BACKEND) + document_store = document_store_backend(self.conversation) + for document in documents: + document_store.parse_and_store_document( + name=document.identifier, + content_type=document.media_type, + content=document.data, + ) + + def perform_rag( + self, + user_prompt: str, + intent_web_search: bool = False, + force_web_search: bool = False, + document_search: bool = False, + ) -> Tuple[str, List[SourceUIPart]]: + """ + Perform RAG (Retrieval-Augmented Generation) based on the conversation settings. + + Args: + web_search (bool): Whether to perform a web search. + document_search (bool): Whether to query attachments. + """ + ui_sources = [] + + if intent_web_search or force_web_search: + web_search_backend = import_string(settings.RAG_WEB_SEARCH_BACKEND) + web_search_results = web_search_backend().web_search(user_prompt) + else: + web_search_results = None + + if force_web_search and web_search_results is None: + logger.error("Forced web search was requested but no results were found.") + raise WebSearchEmptyException() + + if document_search: + document_search_backend = AlbertRagDocumentSearch(self.conversation) + document_search_results = document_search_backend.search(user_prompt) + else: + document_search_results = None + + if web_search_results is None and document_search_results is None: + logger.warning("No web search or document search results found, skipping RAG.") + return "", ui_sources + + prompted_results = "\n\n".join( + search_results.to_prompt() + for search_results in [web_search_results, document_search_results] + if search_results is not None + ) + new_prompt = settings.RAG_WEB_SEARCH_PROMPT_UPDATE.format( + search_results=prompted_results, user_prompt=user_prompt + ) + + _unique_sources = set() + for search_results in [web_search_results, document_search_results]: + if search_results is None: + continue + + for result in search_results.data: + logger.debug("Search result: %s", result.model_dump()) + + # Several chunks may come from the same URL, + # so we need to ensure we don't duplicate sources. + if result.url in _unique_sources: + logger.debug("Skipping duplicated source: %s", result.url) + continue + + _unique_sources.add(result.url) + url_source = LanguageModelV1Source( + source_type="url", + id=str(uuid.uuid4()), + url=result.url, + providerMetadata={}, + ) + ui_sources.append(SourceUIPart(type="source", source=url_source)) + + return new_prompt, ui_sources + + def prepare_prompt( + self, message: UIMessage + ) -> Tuple[str, List[BinaryContent], List[BinaryContent]]: + """ + Prepare the user prompt for the agent. + + This method is used to convert a UIMessage into a format suitable for the agent. + It extracts the user content from the message and returns it as a list of UserContent. + """ + user_content = ui_message_to_user_content(message) + + user_prompt = [] + attachment_images = [] + attachment_documents = [] + attachment_audio = [] + attachment_video = [] + for content in user_content: + if isinstance(content, str): + user_prompt.append(content) + elif isinstance(content, BinaryContent): + if content.is_audio: + attachment_audio.append(content) + elif content.is_video: + attachment_video.append(content) + elif content.is_image: + attachment_images.append(content) + else: + attachment_documents.append(content) + else: + # Should never happen, but just in case + raise ValueError(f"Unsupported UserContent type: {type(content)}") + + if any(attachment_audio): + # Should be handled by the frontend, but just in case + raise ValueError("Audio attachments are not supported in the current implementation.") + if any(attachment_video): + # Should be handled by the frontend, but just in case + raise ValueError("Video attachments are not supported in the current implementation.") + + if len(user_prompt) != 1: + raise ValueError( + "User prompt must contain exactly one text part, " + f"but got {len(user_prompt)} parts: {user_prompt}" + ) + + return user_prompt[0], attachment_images, attachment_documents + async def _run_agent( # noqa: PLR0912, PLR0915 self, messages: List[UIMessage], @@ -178,86 +318,101 @@ class AIAgentService: return history = ModelMessagesTypeAdapter.validate_python(self.conversation.openai_messages) - prompt = ui_message_to_user_content(messages[-1]) + user_prompt, input_images, input_documents = self.prepare_prompt(messages[-1]) + usage = {"promptTokens": 0, "completionTokens": 0} - # Check is the user prompt requires web search - if not settings.RAG_WEB_SEARCH_BACKEND: - # If web search is not enabled, we can skip the intent detection - user_intent = UserIntent(web_search=False) - logger.info("Web search backend is disabled, skipping intent detection.") - elif force_web_search: - # While the only intent detection is web search, we can - # skip the detection if the user has explicitly requested a web search. - user_intent = UserIntent(web_search=True) - logger.info("Web search requested by user, skipping intent detection.") - else: - user_intent: UserIntent = await self._detect_user_intent(prompt) + # Detect the user intent + if not force_web_search: + user_intent: UserIntent = await self._detect_user_intent(user_prompt) logger.info("User intent detected: %s", user_intent.model_dump()) + else: + # If the user has requested a web search, we consider it as the no intent + # of document summarization. + user_intent = UserIntent() - logger.debug("User intent %s", user_intent) + if input_documents and user_intent.attachment_summary: + # If the user has provided documents and requested a summary, + # we need to handle that. + logger.warning( + "Attachment summarization is not supported yet, ignoring the user intent." + ) + user_intent.attachment_summary = False + yield {"type": "3", "payload": "attachment_summary_not_supported"} + return + + # If user uploaded documents and did not enforce a web search, we disable the intent + if input_documents and not force_web_search: + user_intent.web_search = False + logger.info("User intent web search disabled due to input documents.") + + conversation_has_documents = bool(self.conversation.collection_id) + if input_documents: + _tool_call_id = str(uuid.uuid4()) + yield { + "type": "9", + "payload": { + "toolCallId": _tool_call_id, + "toolName": "document_parsing", + "args": { + "documents": [ + { + "identifier": doc.identifier, + } + for doc in input_documents + ], + }, + }, + } + self.parse_input_documents(input_documents) + if not conversation_has_documents: + conversation_has_documents = True + await self.conversation.asave(update_fields=["collection_id", "updated_at"]) + + yield { + "type": "a", + "payload": { + "toolCallId": _tool_call_id, + "result": {"state": "done"}, + }, + } + + # Prepare the prompt for the agent + try: + _new_prompt, _ui_sources = self.perform_rag( + user_prompt=user_prompt, + intent_web_search=user_intent.web_search, + force_web_search=force_web_search, + document_search=conversation_has_documents, + ) + except WebSearchEmptyException: + yield { + "type": "h", + "payload": { + "source_type": "error", + "id": str(uuid.uuid4()), + "error": _("No web search results found."), + }, + } + return _user_initial_prompt_str = None - _ui_sources = [] - if user_intent.web_search: # might be forced by force_web_search - search_backend = import_string(settings.RAG_WEB_SEARCH_BACKEND) - search_results = search_backend().web_search( - " ".join(prompt for prompt in prompt if isinstance(prompt, str)) - ) - - if search_results.data: - for idx, prompt_item in enumerate(prompt): - if isinstance(prompt_item, str): - _user_initial_prompt_str = str(prompt_item) - prompt[idx] = settings.RAG_WEB_SEARCH_PROMPT_UPDATE.format( - search_results=search_results.to_prompt(), - user_prompt=prompt_item, - ) - break - - _unique_sources = set() - for result in search_results.data: - logger.debug("Search result: %s", result.model_dump()) - - # Several chunks may come from the same URL, - # so we need to ensure we don't duplicate sources. - if result.url in _unique_sources: - logger.debug("Skipping duplicated source: %s", result.url) - continue - - _unique_sources.add(result.url) - url_source = LanguageModelV1Source( - source_type="url", - id=str(uuid.uuid4()), - url=result.url, - providerMetadata={}, - ) - _ui_sources.append(SourceUIPart(type="source", source=url_source)) - - yield { - "type": "h", - "payload": url_source.model_dump(mode="json"), - } - elif force_web_search: - logger.warning("Web search was forced but no results were found.") + if _new_prompt: + _user_initial_prompt_str = str(user_prompt) # copy the original user prompt + user_prompt = _new_prompt + for _ui_source in _ui_sources: yield { "type": "h", - "payload": { - "source_type": "error", - "id": str(uuid.uuid4()), - "error": "No web search results found.", - }, + "payload": _ui_source.source.model_dump(mode="json"), } - return - else: - logger.warning("No web search results found, continuing without web search.") async with AsyncExitStack() as stack: # MCP servers (if any) can be initialized here mcp_servers = [await stack.enter_async_context(mcp) for mcp in get_mcp_servers()] async with _build_pydantic_agent(mcp_servers).iter( - prompt, message_history=history + [user_prompt] + input_images, + message_history=history, ) as run: async for node in run: if Agent.is_user_prompt_node(node): diff --git a/src/backend/chat/clients/pydantic_ui_message_converter.py b/src/backend/chat/clients/pydantic_ui_message_converter.py index f407db6..5eda9f0 100644 --- a/src/backend/chat/clients/pydantic_ui_message_converter.py +++ b/src/backend/chat/clients/pydantic_ui_message_converter.py @@ -63,7 +63,11 @@ def ui_message_to_user_content(message: UIMessage) -> List[UserContent]: # Handle data URLs raw_data = base64.b64decode(experimental_attachment.url.split(",")[1]) user_contents.append( - BinaryContent(data=raw_data, media_type=experimental_attachment.contentType) + BinaryContent( + data=raw_data, + media_type=experimental_attachment.contentType, + identifier=experimental_attachment.name, + ) ) else: raise ValueError( diff --git a/src/backend/chat/migrations/0004_chatconversation_collection_id.py b/src/backend/chat/migrations/0004_chatconversation_collection_id.py new file mode 100644 index 0000000..67ba8aa --- /dev/null +++ b/src/backend/chat/migrations/0004_chatconversation_collection_id.py @@ -0,0 +1,21 @@ +# Generated by Django 5.2.3 on 2025-08-04 13:17 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [ + ("chat", "0003_alter_chatconversation_messages"), + ] + + operations = [ + migrations.AddField( + model_name="chatconversation", + name="collection_id", + field=models.CharField( + blank=True, + help_text="Collection ID for the conversation, used for RAG document search", + null=True, + ), + ), + ] diff --git a/src/backend/chat/models.py b/src/backend/chat/models.py index 94e2741..fcaf502 100644 --- a/src/backend/chat/models.py +++ b/src/backend/chat/models.py @@ -66,3 +66,9 @@ class ChatConversation(BaseModel): blank=True, help_text="Agent usage for the chat conversation, provided by OpenAI API", ) + + collection_id = models.CharField( + blank=True, + null=True, + help_text="Collection ID for the conversation, used for RAG document search", + ) diff --git a/src/backend/chat/tests/agent_rag/document_converter/__init__.py b/src/backend/chat/tests/agent_rag/document_converter/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/backend/chat/tests/agent_rag/document_converter/test_markitdown.py b/src/backend/chat/tests/agent_rag/document_converter/test_markitdown.py new file mode 100644 index 0000000..e83a2b4 --- /dev/null +++ b/src/backend/chat/tests/agent_rag/document_converter/test_markitdown.py @@ -0,0 +1,29 @@ +""" +Unit tests for the DocumentConverter. + +Only for coverage as the DocumentConverter is a simple wrapper around MarkItDown. +""" + +from io import BytesIO +from unittest.mock import MagicMock, patch + +from chat.agent_rag.document_converter.markitdown import DocumentConverter + + +@patch("chat.agent_rag.document_converter.markitdown.MarkItDown") +def test_document_converter(mock_markitdown: MagicMock): + """Test that the DocumentConverter calls the underlying MarkItDown converter.""" + mock_conversion = MagicMock() + mock_conversion.text_content = "converted text" + mock_markitdown.return_value.convert.return_value = mock_conversion + + converter = DocumentConverter() + + result = converter.convert_raw( + name="test.pdf", + content_type="application/pdf", + content=BytesIO(b"test content"), + ) + + assert result == "converted text" + converter.converter.convert.assert_called_once() # pylint: disable=no-member diff --git a/src/backend/chat/tests/agent_rag/document_search/__init__.py b/src/backend/chat/tests/agent_rag/document_search/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/backend/chat/tests/agent_rag/document_search/test_albert_api.py b/src/backend/chat/tests/agent_rag/document_search/test_albert_api.py new file mode 100644 index 0000000..ccb2cef --- /dev/null +++ b/src/backend/chat/tests/agent_rag/document_search/test_albert_api.py @@ -0,0 +1,243 @@ +"""Tests for the Albert RAG document search API.""" +# pylint: disable=protected-access + +from io import BytesIO +from unittest.mock import patch + +import pytest +import responses +from requests import HTTPError + +from chat.agent_rag.document_search.albert_api import AlbertRagDocumentSearch +from chat.factories import ChatConversationFactory + +pytestmark = pytest.mark.django_db() + + +def test_albert_collection_id_property(): + """Test the _albert_collection_id property.""" + conversation = ChatConversationFactory() + assert ( + AlbertRagDocumentSearch(conversation)._albert_collection_id + == f"conversation-{conversation.pk}" + ) + + +def test_collection_id_property(): + """Test the collection_id property.""" + conversation = ChatConversationFactory() + search = AlbertRagDocumentSearch(conversation) + + # When collection_id is None + conversation.collection_id = None + assert search.collection_id is None + + # When collection_id is set + conversation.collection_id = "123" + assert search.collection_id == 123 + + +@responses.activate +def test_create_collection_success(): + """Test _create_collection successfully creates a collection.""" + conversation = ChatConversationFactory() + search = AlbertRagDocumentSearch(conversation) + responses.add( + responses.POST, + search._collections_endpoint, + json={"id": "456"}, + status=201, + ) + + assert search._create_collection() is True + assert conversation.collection_id == "456" + + +@responses.activate +def test_create_collection_failure(): + """Test _create_collection handles API errors.""" + conversation = ChatConversationFactory() + search = AlbertRagDocumentSearch(conversation) + responses.add( + responses.POST, + search._collections_endpoint, + status=500, + ) + + with pytest.raises(HTTPError): + search._create_collection() + + +@responses.activate +def test_parse_pdf_document_success(): + """Test _parse_pdf_document successfully parses a PDF.""" + conversation = ChatConversationFactory() + search = AlbertRagDocumentSearch(conversation) + responses.add( + responses.POST, + search._pdf_parser_endpoint, + json={"data": [{"content": "Page 1"}, {"content": "Page 2"}]}, + status=200, + ) + + content = search._parse_pdf_document("test.pdf", "application/pdf", BytesIO(b"pdf_content")) + assert content == "Page 1\n\nPage 2" + + +@responses.activate +def test_parse_pdf_document_failure(): + """Test _parse_pdf_document handles API errors.""" + conversation = ChatConversationFactory() + search = AlbertRagDocumentSearch(conversation) + responses.add( + responses.POST, + search._pdf_parser_endpoint, + status=500, + ) + + with pytest.raises(HTTPError): + search._parse_pdf_document("test.pdf", "application/pdf", BytesIO(b"pdf_content")) + + +@patch("chat.agent_rag.document_search.albert_api.AlbertRagDocumentSearch._parse_pdf_document") +def test_parse_document_pdf(mock_parse_pdf): + """Test parse_document for PDF content.""" + conversation = ChatConversationFactory() + search = AlbertRagDocumentSearch(conversation) + mock_parse_pdf.return_value = "Parsed PDF content" + result = search.parse_document("test.pdf", "application/pdf", BytesIO(b"pdf")) + assert result == "Parsed PDF content" + mock_parse_pdf.assert_called_once() + + +@patch("chat.agent_rag.document_search.albert_api.DocumentConverter") +def test_parse_document_other(mock_converter): + """Test parse_document for other content types.""" + conversation = ChatConversationFactory() + search = AlbertRagDocumentSearch(conversation) + mock_converter.return_value.convert_raw.return_value = "Converted content" + result = search.parse_document("test.txt", "text/plain", BytesIO(b"text")) + assert result == "Converted content" + mock_converter.return_value.convert_raw.assert_called_once() + + +@responses.activate +def test_store_document_success(): + """Test _store_document successfully stores a document.""" + conversation = ChatConversationFactory(collection_id="123") + search = AlbertRagDocumentSearch(conversation) + responses.add( + responses.POST, + search._documents_endpoint, + json={"id": "doc1"}, + status=201, + ) + + search._store_document("test_doc", "some content") + assert len(responses.calls) == 1 + assert responses.calls[0].request.url == search._documents_endpoint + + +@responses.activate +@patch("chat.agent_rag.document_search.albert_api.AlbertRagDocumentSearch._create_collection") +def test_store_document_creates_collection(mock_create_collection): + """Test _store_document creates a collection if one doesn't exist.""" + conversation = ChatConversationFactory(collection_id=None) + search = AlbertRagDocumentSearch(conversation) + + def set_collection_id(*args, **kwargs): + conversation.collection_id = "123" + return True + + mock_create_collection.side_effect = set_collection_id + + responses.add( + responses.POST, + search._documents_endpoint, + json={"id": "doc1"}, + status=201, + ) + + search._store_document("test_doc", "some content") + mock_create_collection.assert_called_once() + assert conversation.collection_id == "123" + + +@patch( + "chat.agent_rag.document_search.albert_api.AlbertRagDocumentSearch._create_collection", + return_value=False, +) +def test_store_document_create_collection_fails(mock_create_collection): + """Test _store_document raises error if collection creation fails.""" + conversation = ChatConversationFactory(collection_id=None) + search = AlbertRagDocumentSearch(conversation) + with pytest.raises(RuntimeError, match="Failed to create or retrieve the collection."): + search._store_document("test_doc", "some content") + mock_create_collection.assert_called_once() + + +@patch("chat.agent_rag.document_search.albert_api.AlbertRagDocumentSearch.parse_document") +@patch("chat.agent_rag.document_search.albert_api.AlbertRagDocumentSearch._store_document") +def test_parse_and_store_document(mock_store, mock_parse): + """Test parse_and_store_document orchestrates parsing and storing.""" + conversation = ChatConversationFactory() + search = AlbertRagDocumentSearch(conversation) + mock_parse.return_value = "parsed content" + name = "test.txt" + content_type = "text/plain" + content = BytesIO(b"text") + + result = search.parse_and_store_document(name, content_type, content) + + assert result == "parsed content" + mock_parse.assert_called_once_with(name, content_type, content) + mock_store.assert_called_once_with(name, "parsed content") + + +@responses.activate +def test_search_success(): + """Test search successfully returns results.""" + conversation = ChatConversationFactory(collection_id="123") + search = AlbertRagDocumentSearch(conversation) + mock_response = { + "data": [ + { + "method": "semantic", + "chunk": { + "id": 1, + "content": "Relevant content snippet.", + "metadata": {"document_name": "doc1.txt"}, + }, + "score": 0.9, + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 20}, + } + responses.post( + url=search._search_endpoint, + json=mock_response, + status=200, + ) + + results = search.search("test query") + + assert len(results.data) == 1 + assert results.data[0].content == "Relevant content snippet." + assert results.data[0].url == "doc1.txt" + assert results.data[0].score == 0.9 + assert results.usage.prompt_tokens == 10 + assert results.usage.completion_tokens == 20 + + +@responses.activate +def test_search_failure(): + """Test search handles API errors.""" + conversation = ChatConversationFactory(collection_id="123") + search = AlbertRagDocumentSearch(conversation) + responses.post( + url=search._search_endpoint, + status=500, + ) + + with pytest.raises(HTTPError): + search.search("test query") diff --git a/src/backend/chat/tests/views/chat/conversations/test_conversation.py b/src/backend/chat/tests/views/chat/conversations/test_conversation.py index e8a47f3..a5b441a 100644 --- a/src/backend/chat/tests/views/chat/conversations/test_conversation.py +++ b/src/backend/chat/tests/views/chat/conversations/test_conversation.py @@ -467,7 +467,7 @@ def test_post_conversation_with_image(api_client, mock_openai_stream_image, mock "v7-jQ3Y5AAAADklEQVQI12P4AIX8EAgALgAD_aNpbtEAAAAASUVORK5CYII=" ), "kind": "binary", - "identifier": None, + "identifier": "FELV-cat.jpg", "media_type": "image/png", "vendor_metadata": None, }, diff --git a/src/backend/chat/tests/views/chat/conversations/test_conversation_with_document_search_rag.py b/src/backend/chat/tests/views/chat/conversations/test_conversation_with_document_search_rag.py new file mode 100644 index 0000000..bc59c0c --- /dev/null +++ b/src/backend/chat/tests/views/chat/conversations/test_conversation_with_document_search_rag.py @@ -0,0 +1,457 @@ +"""Unit tests for chat conversation actions with document search RAG functionality.""" + +# pylint: disable=too-many-lines +import base64 +import json +from io import BytesIO + +from django.utils import timezone + +import httpx +import pytest +import responses +import respx +from freezegun import freeze_time +from rest_framework import status + +from chat.ai_sdk_types import ( + Attachment, + LanguageModelV1Source, + SourceUIPart, + TextUIPart, + UIMessage, +) +from chat.factories import ChatConversationFactory + +# enable database transactions for tests: +# transaction=True ensures that the data are available in the database +# in other threads +pytestmark = pytest.mark.django_db(transaction=True) + + +@pytest.fixture(autouse=True) +def ai_settings(settings): + """Fixture to set AI service URLs for testing.""" + settings.AI_BASE_URL = "https://www.external-ai-service.com/" + settings.AI_API_KEY = "test-api-key" + settings.AI_MODEL = "test-model" + + # Enable Albert API for document search + settings.RAG_DOCUMENT_SEARCH_BACKEND = ( + "chat.agent_rag.document_search.albert_api.AlbertRagDocumentSearch" + ) + settings.ALBERT_API_URL = "https://albert.api.etalab.gouv.fr" + settings.ALBERT_API_KEY = "albert-api-key" + settings.RAG_WEB_SEARCH_PROMPT_UPDATE = ( + "Based on the following document contents:\n\n{search_results}\n\n" + "Please answer the user's question: {user_prompt}" + ) + + # Set up AI routing model settings for intent detection + settings.AI_ROUTING_MODEL = "mini-model" + settings.AI_ROUTING_MODEL_BASE_URL = "https://www.mini-ai-service.com/" + settings.AI_ROUTING_MODEL_API_KEY = "test-routing-api-key" + settings.AI_ROUTING_SYSTEM_PROMPT = ( + "You are an intent detection model. " + "Return attachment_summary as true if the user wants to talk about the document." + ) + + return settings + + +@pytest.fixture(name="sample_pdf_content") +def fixture_sample_pdf_content(): + """Create a dummy PDF content as BytesIO.""" + # This is a simple, valid one-page PDF content. + pdf_data = ( + b"%PDF-1.4\n1 0 obj\n<< /Type /Catalog /Pages 2 0 R >>\nendobj\n" + b"2 0 obj\n<< /Type /Pages /Kids [3 0 R] /Count 1 >>\nendobj\n" + b"3 0 obj\n<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] " + b"/Contents 4 0 R /Resources <<>> >>\nendobj\n" + b"4 0 obj\n<< /Length 35 >>\nstream\nBT /F1 24 Tf 100 700 Td (Hello PDF) " + b"Tj ET\nendstream\nendobj\n" + b"xref\n0 5\n0000000000 65535 f \n0000000010 00000 n \n0000000062 00000 n \n" + b"0000000118 00000 n \n0000000210 00000 n \n" + b"trailer\n<< /Size 5 /Root 1 0 R >>\nstartxref\n288\n%%EOF" + ) + return BytesIO(pdf_data) + + +@pytest.fixture(name="mock_albert_api") +def fixture_mock_albert_api(): + """Fixture to mock the Albert API endpoints.""" + # Mock collection creation + responses.post( + "https://albert.api.etalab.gouv.fr/v1/collections", + json={"id": "123", "name": "test-collection"}, + status=status.HTTP_200_OK, + ) + + # Mock PDF parsing + responses.post( + "https://albert.api.etalab.gouv.fr/v1/parse-beta", + json={ + "data": [ + { + "content": "This is the content of the PDF.", + "metadata": {"document_name": "sample.pdf"}, + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 20}, + }, + status=status.HTTP_200_OK, + ) + + # Mock document upload + responses.post( + "https://albert.api.etalab.gouv.fr/v1/documents", + json={"id": 456}, + status=status.HTTP_201_CREATED, + ) + + # Mock document search + responses.post( + "https://albert.api.etalab.gouv.fr/v1/search", + json={ + "data": [ + { + "method": "semantic", + "chunk": { + "id": 123, + "content": "This is the content of the PDF.", + "metadata": {"document_name": "sample.pdf"}, + }, + "score": 0.9, + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 20}, + }, + status=status.HTTP_200_OK, + ) + + +@pytest.fixture(name="mock_openai_stream") +@freeze_time("2025-07-25T10:36:35.297675Z") +def fixture_mock_openai_stream(): + """ + Fixture to mock the OpenAI stream response for document search queries. + """ + openai_stream = ( + "data: " + + json.dumps( + { + "id": "chatcmpl-1234567890", + "created": timezone.make_naive(timezone.now()).timestamp(), + "choices": [ + { + "delta": {"content": "From the document, I can see that "}, + "index": 0, + "finish_reason": None, + } + ], + "object": "chat.completion.chunk", + } + ) + + "\n\n" + "data: " + + json.dumps( + { + "id": "chatcmpl-1234567890", + "created": timezone.make_naive(timezone.now()).timestamp(), + "choices": [ + { + "delta": {"content": "it says 'Hello PDF'."}, + "index": 0, + "finish_reason": "stop", + } + ], + "object": "chat.completion.chunk", + "usage": { + "prompt_tokens": 150, + "completion_tokens": 25, + "total_tokens": 175, + }, + } + ) + + "\n\n" + "data: [DONE]\n\n" + ) + + async def mock_stream(): + for line in openai_stream.splitlines(keepends=True): + yield line.encode() + + route = respx.post("https://www.external-ai-service.com/chat/completions").mock( + return_value=httpx.Response(200, stream=mock_stream()) + ) + + return route + + +@pytest.fixture(name="mock_intent_detection_document") +def fixture_mock_intent_detection_document(): + """Fixture to mock the intent detection response for document summary.""" + intent_response = { + "id": "chatcmpl-intent-123", + "object": "chat.completion", + "created": int(timezone.make_naive(timezone.now()).timestamp()), + "model": "mini-model", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": '{"web_search": false, "attachment_summary": false}', + }, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 20, "completion_tokens": 5, "total_tokens": 25}, + } + + route = respx.post("https://www.mini-ai-service.com/chat/completions").mock( + return_value=httpx.Response(200, json=intent_response) + ) + + return route + + +@pytest.fixture(name="mock_intent_detection_document_summarize") +def fixture_mock_intent_detection_document_summarize(): + """Fixture to mock the intent detection response for document summary.""" + intent_response = { + "id": "chatcmpl-intent-123", + "object": "chat.completion", + "created": int(timezone.make_naive(timezone.now()).timestamp()), + "model": "mini-model", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": '{"web_search": false, "attachment_summary": true}', + }, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 20, "completion_tokens": 5, "total_tokens": 25}, + } + + route = respx.post("https://www.mini-ai-service.com/chat/completions").mock( + return_value=httpx.Response(200, json=intent_response) + ) + + return route + + +@responses.activate +@respx.mock +@freeze_time("2025-07-25T10:36:35.297675Z") +def test_post_conversation_with_document_upload_summarize( # pylint: disable=too-many-arguments,too-many-positional-arguments + api_client, + mock_albert_api, # pylint: disable=unused-argument + mock_openai_stream, # pylint: disable=unused-argument + mock_intent_detection_document_summarize, # pylint: disable=unused-argument + sample_pdf_content, +): + """ + Test POST to /api/v1/chats/{pk}/conversation/ with a PDF document + when user wants a summary of the document. + """ + chat_conversation = ChatConversationFactory() + api_client.force_authenticate(user=chat_conversation.owner) + + pdf_base64 = base64.b64encode(sample_pdf_content.read()).decode("utf-8") + message = UIMessage( + id="1", + role="user", + content="What does the document say?", + parts=[ + TextUIPart( + text="What does the document say?", + type="text", + ), + ], + experimental_attachments=[ + Attachment( + name="sample.pdf", + contentType="application/pdf", + url=f"data:application/pdf;base64,{pdf_base64}", + ) + ], + ) + + response = api_client.post( + f"/api/v1.0/chats/{chat_conversation.pk}/conversation/", + data={"messages": [message.model_dump(mode="json")]}, + format="json", + ) + + assert response.status_code == status.HTTP_200_OK + assert response.get("Content-Type") == "text/event-stream" + assert response.get("x-vercel-ai-data-stream") == "v1" + assert response.streaming + + # Wait for the streaming content to be fully received + response_content = b"".join(response.streaming_content).decode("utf-8") + assert response_content == '3:"attachment_summary_not_supported"\n' + + # Check that the conversation was updated + chat_conversation.refresh_from_db() + assert len(chat_conversation.messages) == 0 # might be improved in the future + assert len(chat_conversation.openai_messages) == 0 # might be improved in the future + + +@responses.activate +@respx.mock +@freeze_time("2025-07-25T10:36:35.297675Z") +def test_post_conversation_with_document_upload( # noqa:PLR0913 # pylint: disable=too-many-arguments,too-many-positional-arguments + api_client, + mock_albert_api, # pylint: disable=unused-argument + mock_openai_stream, # pylint: disable=unused-argument + mock_intent_detection_document, # pylint: disable=unused-argument + sample_pdf_content, + mock_uuid4, +): + """ + Test POST to /api/v1/chats/{pk}/conversation/ with a PDF document. + """ + chat_conversation = ChatConversationFactory() + api_client.force_authenticate(user=chat_conversation.owner) + + pdf_base64 = base64.b64encode(sample_pdf_content.read()).decode("utf-8") + message = UIMessage( + id="1", + role="user", + content="What does the document say?", + parts=[ + TextUIPart( + text="What does the document say?", + type="text", + ), + ], + experimental_attachments=[ + Attachment( + name="sample.pdf", + contentType="application/pdf", + url=f"data:application/pdf;base64,{pdf_base64}", + ) + ], + ) + + response = api_client.post( + f"/api/v1.0/chats/{chat_conversation.pk}/conversation/", + data={"messages": [message.model_dump(mode="json")]}, + format="json", + ) + + assert response.status_code == status.HTTP_200_OK + assert response.get("Content-Type") == "text/event-stream" + assert response.get("x-vercel-ai-data-stream") == "v1" + assert response.streaming + + # Wait for the streaming content to be fully received + response_content = b"".join(response.streaming_content).decode("utf-8") + assert response_content == ( + f'9:{{"toolCallId": "{mock_uuid4}", "toolName": ' + '"document_parsing", "args": {"documents": [{"identifier": "sample.pdf"}]}}\n' + f'a:{{"toolCallId": "{mock_uuid4}", "result": {{"state": "done"}}}}\n' + f'h:{{"source_type": "url", "id": "{mock_uuid4}", ' + '"url": "sample.pdf", "title": null, "providerMetadata": {}}\n' + '0:"From the document, I can see that "\n' + "0:\"it says 'Hello PDF'.\"\n" + 'd:{"finishReason": "stop", "usage": {"promptTokens": 150, ' + '"completionTokens": 25}}\n' + ) + + # Check that the conversation was updated + chat_conversation.refresh_from_db() + assert len(chat_conversation.messages) == 2 + assert chat_conversation.messages[0] == UIMessage( + id=str(mock_uuid4), + createdAt=timezone.now(), + content="What does the document say?", + reasoning=None, + experimental_attachments=None, + role="user", + annotations=None, + toolInvocations=None, + parts=[TextUIPart(type="text", text="What does the document say?")], + ) + assert chat_conversation.messages[1] == UIMessage( + id=str(mock_uuid4), + createdAt=timezone.now(), + content="From the document, I can see that it says 'Hello PDF'.", + reasoning=None, + experimental_attachments=None, + role="assistant", + annotations=None, + toolInvocations=None, + parts=[ + TextUIPart(type="text", text="From the document, I can see that it says 'Hello PDF'."), + SourceUIPart( + type="source", + source=LanguageModelV1Source( + source_type="url", + id=str(mock_uuid4), + url="sample.pdf", + title=None, + providerMetadata={}, + ), + ), + ], + ) + + assert len(chat_conversation.openai_messages) == 2 + assert chat_conversation.openai_messages[0] == { + "instructions": None, + "kind": "request", + "parts": [ + { + "content": "You are a helpful assistant. Escape formulas or any " + "math notation between `$$`, like `$$x^2 + y^2 = z^2$$` " + "or `$$C_l$$`. You can use Markdown to format your " + "answers. ", + "dynamic_ref": None, + "part_kind": "system-prompt", + "timestamp": "2025-07-25T10:36:35.297675Z", + }, + { + "content": [ + "Based on the following document contents:\n" + "\n" + " - From: sample.pdf:\n" + " content: This is the content of the PDF.\n" + "\n" + "\n" + "\n" + "\n" + "\n" + "Please answer the user's question: What does the " + "document say?" + ], + "part_kind": "user-prompt", + "timestamp": "2025-07-25T10:36:35.297675Z", + }, + ], + } + assert chat_conversation.openai_messages[1] == { + "kind": "response", + "model_name": "test-model", + "parts": [ + { + "content": "From the document, I can see that it says 'Hello PDF'.", + "part_kind": "text", + } + ], + "timestamp": "2025-07-25T10:36:35.297675Z", + "usage": { + "details": None, + "request_tokens": 150, + "requests": 1, + "response_tokens": 25, + "total_tokens": 175, + }, + "vendor_details": None, + "vendor_id": None, + } diff --git a/src/backend/chat/tests/views/chat/conversations/test_conversation_with_web_search_rag.py b/src/backend/chat/tests/views/chat/conversations/test_conversation_with_web_search_rag.py index 7fcaa26..e421c17 100644 --- a/src/backend/chat/tests/views/chat/conversations/test_conversation_with_web_search_rag.py +++ b/src/backend/chat/tests/views/chat/conversations/test_conversation_with_web_search_rag.py @@ -134,7 +134,7 @@ def fixture_mock_intent_detection_web_search(): "index": 0, "message": { "role": "assistant", - "content": '{"web_search": true}', + "content": '{"web_search": true, "attachment_summary": false}', }, "finish_reason": "stop", } diff --git a/src/backend/conftest.py b/src/backend/conftest.py new file mode 100644 index 0000000..1116038 --- /dev/null +++ b/src/backend/conftest.py @@ -0,0 +1,26 @@ +"""Global fixtures for the backend tests.""" + +import pytest +from urllib3.connectionpool import HTTPConnectionPool + + +@pytest.fixture(autouse=True) +def no_http_requests(monkeypatch): + """ + Prevents HTTP requests from being made during tests. + This is useful for tests that do not require actual HTTP requests + and helps to avoid network-related issues. + + Credits: https://blog.jerrycodes.com/no-http-requests/ + """ + + allowed_hosts = {"localhost"} + original_urlopen = HTTPConnectionPool.urlopen + + def urlopen_mock(self, method, url, *args, **kwargs): + if self.host in allowed_hosts: + return original_urlopen(self, method, url, *args, **kwargs) + + raise RuntimeError(f"The test was about to {method} {self.scheme}://{self.host}{url}") + + monkeypatch.setattr("urllib3.connectionpool.HTTPConnectionPool.urlopen", urlopen_mock) diff --git a/src/backend/conversations/settings.py b/src/backend/conversations/settings.py index 8458c5b..dc7dd65 100755 --- a/src/backend/conversations/settings.py +++ b/src/backend/conversations/settings.py @@ -485,7 +485,10 @@ class Base(Configuration): " - `web_search`: The user requests, explicitly or not, to have recent or precise" " information, or anything which would request to do some web research to add" " context before answering. Any semantic language like 'recent', 'latest information'" - " and similar should trigger a web search.\n", + " and similar should trigger a web search.\n" + " - `attachment_summary`: The user requests, explicitly or not, to have a summary of" + " a document, or any specific file. Any semantic language like 'summary', 'overview'," + " 'highlights', 'key points' and similar should trigger a document summary.\n", environ_name="AI_ROUTING_SYSTEM_PROMPT", environ_prefix=None, ) @@ -497,6 +500,18 @@ class Base(Configuration): environ_prefix=None, ) + # Documents + ALBERT_API_PARSE_TIMEOUT = values.PositiveIntegerValue( + default=120, # seconds + environ_name="ALBERT_API_PARSE_TIMEOUT", + environ_prefix=None, + ) + RAG_DOCUMENT_SEARCH_BACKEND = values.Value( + "chat.agent_rag.document_search.albert_api.AlbertRagDocumentSearch", + environ_name="RAG_DOCUMENT_SEARCH_BACKEND", + environ_prefix=None, + ) + # Web search RAG_WEB_SEARCH_BACKEND = values.Value( # "chat.agent_rag.web_search.albert_api.AlbertWebSearchManager", @@ -505,9 +520,12 @@ class Base(Configuration): environ_prefix=None, ) RAG_WEB_SEARCH_PROMPT_UPDATE = values.Value( - ''' -You are a subject-matter expert assistant. You are given web search result(s) or raw webpage text -that may contain navigation, menus, comments, category names, and other unrelated page details. + """ +You are a subject-matter expert assistant. +You are given : + - web search result(s) or raw webpage text that may contain navigation, menus, comments, + category names, and other unrelated page details. + - document(s) that may contain text, images, and metadata, formatted in markdown. **Your mission:** - Use ONLY the main informational content that directly and factually answers the user's explicit @@ -539,14 +557,12 @@ that may contain navigation, menus, comments, category names, and other unrelate navigation tools. - Invent summaries for structure or offer general commentary about the website. -WEB SEARCH RESULTS TO USE AS CONTEXT: -""" {search_results} -""" USER QUESTION: + {user_prompt} - ''', + """, environ_name="RAG_WEB_SEARCH_PROMPT_UPDATE", environ_prefix=None, ) @@ -795,6 +811,14 @@ class Test(Base): os.environ["OPENAI_AGENTS_DISABLE_TRACING"] = "true" + AI_BASE_URL = None + AI_API_KEY = None + AI_MODEL = None + + AI_ROUTING_MODEL_BASE_URL = None + AI_ROUTING_MODEL = None + AI_ROUTING_MODEL_API_KEY = None + def __init__(self): # pylint: disable=invalid-name self.INSTALLED_APPS += ["drf_spectacular_sidecar"] diff --git a/src/backend/pyproject.toml b/src/backend/pyproject.toml index 8f83362..d1528d0 100644 --- a/src/backend/pyproject.toml +++ b/src/backend/pyproject.toml @@ -50,6 +50,7 @@ dependencies = [ "jsonschema==4.24.0", "lxml==5.4.0", "markdown==3.8", + "markitdown[pptx,docx,xlsx,xls,pdf,outlook]==0.1.2", "mozilla-django-oidc==4.0.1", "nested-multipart-parser==1.5.0", "pydantic==2.11.7", @@ -80,6 +81,7 @@ dev = [ "pyfakefs==5.8.0", "pylint-django==2.6.1", "pylint==3.3.7", + "pylint-pydantic==0.3.5", "pytest-asyncio==1.1.0", "pytest-cov==6.2.1", "pytest-django==4.11.1",