✨(chat) add document RAG on document uploaded by user
This allow use to upload documents to ask question on them. The summarize feature is willingly disabled for now, as it is not the same process.
This commit is contained in:
Quentin BEY
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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,
|
||||
),
|
||||
)
|
||||
@@ -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)
|
||||
@@ -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):
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
),
|
||||
),
|
||||
]
|
||||
@@ -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",
|
||||
)
|
||||
|
||||
@@ -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
|
||||
@@ -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")
|
||||
@@ -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,
|
||||
},
|
||||
|
||||
+457
@@ -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,
|
||||
}
|
||||
+1
-1
@@ -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",
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
@@ -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"]
|
||||
|
||||
@@ -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",
|
||||
|
||||
Reference in New Issue
Block a user