Compare commits

...

14 Commits

Author SHA1 Message Date
camilleAND 3e1934ce9c add error handling for document storage 2025-12-11 16:14:15 +01:00
camilleAND 4c255845f9 async trafilatura 2025-12-11 15:59:15 +01:00
camilleAND 6f51b62392 delete non needed comment 2025-12-11 15:48:32 +01:00
camilleAND 92d4938e8e add message for streaming tool on front 2025-12-11 15:44:25 +01:00
camilleAND 6513420658 refacto document upload 2025-12-11 15:43:59 +01:00
camilleAND 8448340b5c use pydanticAI messages + refacto + settings timeout + correct source for downstream summarize 2025-12-11 14:53:19 +01:00
camilleAND 1ffbd2370e improve regexp for url 2025-12-05 16:33:39 +01:00
camilleAND 8a97448ff2 remove literal repetition 2025-12-05 16:26:19 +01:00
camilleAND f50453b629 add ignore for my dev default version 2025-12-05 15:55:21 +01:00
camilleAND 5fce562459 add test for fetch url 2025-12-05 13:10:16 +01:00
camilleAND dc0200099d add pdf gestion for fetch_url + rag, add url deduplicates in sources 2025-12-05 12:57:31 +01:00
camilleAND 7049782e70 add rag for long context retrived from urls 2025-12-05 11:30:45 +01:00
camilleAND db84c6906c 📝(new tool): add security, enable fetch only for detected urls 2025-12-05 11:30:45 +01:00
camilleAND 54fd920c0c 📝(new tool): add fetch_url, specific fetching for docs urls 2025-12-05 11:30:45 +01:00
12 changed files with 1084 additions and 53 deletions
+3
View File
@@ -79,3 +79,6 @@ db.sqlite3
# Docker compose override
compose.override.yml
# LLM configuration
src/backend/conversations/configuration/llm/default_dev.json
@@ -18,7 +18,7 @@ class DocumentConverter:
*,
name: str,
content_type: str,
content: bytes,
content: bytes | BytesIO,
) -> str:
"""
Convert a document to Markdown format.
@@ -28,9 +28,18 @@ class DocumentConverter:
Args:
name (str): The name of the document.
content_type (str): The MIME type of the document (e.g., "application/pdf").
content (bytes): The content of the document as bytes.
content (bytes | BytesIO): The content of the document as bytes or BytesIO.
"""
return self._convert(BytesIO(content), file_extension=os.path.splitext(name)[1])
# Handle both bytes and BytesIO
if isinstance(content, BytesIO):
# Read the BytesIO to bytes, then create a new BytesIO for the converter
content_bytes = content.read()
content_io = BytesIO(content_bytes)
else:
# content is already bytes
content_io = BytesIO(content)
return self._convert(content_io, file_extension=os.path.splitext(name)[1])
def _convert(self, document: BytesIO, file_extension: str) -> str:
"""
+56 -44
View File
@@ -13,7 +13,6 @@ import logging
import time
import uuid
from contextlib import AsyncExitStack, ExitStack
from io import BytesIO
from typing import Dict, List, Optional, Tuple
from django.conf import settings
@@ -22,7 +21,6 @@ from django.core.cache import cache
from django.core.exceptions import ImproperlyConfigured
from django.core.files.storage import default_storage
from django.db.models import Q
from django.utils.module_loading import import_string
from asgiref.sync import sync_to_async
from langfuse import get_client
@@ -63,6 +61,7 @@ from chat.agents.local_media_url_processors import (
from chat.ai_sdk_types import (
LanguageModelV1Source,
SourceUIPart,
TextUIPart,
UIMessage,
)
from chat.clients.async_to_sync import convert_async_generator_to_sync
@@ -71,10 +70,12 @@ from chat.clients.pydantic_ui_message_converter import (
model_message_to_ui_message,
ui_message_to_user_content,
)
from chat.document_storage import store_document_with_attachment
from chat.mcp_servers import get_mcp_servers
from chat.tools.document_generic_search_rag import add_document_rag_search_tool_from_setting
from chat.tools.document_search_rag import add_document_rag_search_tool
from chat.tools.document_summarize import document_summarize
from chat.tools.fetch_url import detect_url_in_conversation, fetch_url
from chat.vercel_ai_sdk.core import events_v4, events_v5
from chat.vercel_ai_sdk.encoder import EventEncoder
@@ -249,19 +250,10 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
):
raise ValueError("Document URL does not belong to the conversation.")
document_store_backend = import_string(settings.RAG_DOCUMENT_SEARCH_BACKEND)
document_store = document_store_backend(self.conversation.collection_id)
if not document_store.collection_id:
# Create a new collection for the conversation
collection_id = document_store.create_collection(
name=f"conversation-{self.conversation.pk}",
)
self.conversation.collection_id = str(collection_id)
await self.conversation.asave(update_fields=["collection_id", "updated_at"])
for document in documents:
key = None
document_data = None
if isinstance(document, DocumentUrl):
if document.url.startswith("/media-key/"):
# Local file, retrieve from object storage
@@ -272,33 +264,31 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
# Retrieve the document data
with default_storage.open(key, "rb") as file:
document_data = file.read()
parsed_content = document_store.parse_and_store_document(
name=document.identifier,
content_type=document.media_type,
content=document_data,
)
else:
# Remote URL
raise ValueError("External document URL are not accepted yet.")
else:
parsed_content = document_store.parse_and_store_document(
name=document.identifier,
content_type=document.media_type,
content=document.data,
)
document_data = document.data
# Convert BytesIO to bytes if needed
if hasattr(document_data, 'read'):
document_data = document_data.read()
if not document.media_type.startswith("text/"):
md_attachment = await models.ChatConversationAttachment.objects.acreate(
conversation=self.conversation,
uploaded_by=self.user,
key=key or f"{self.conversation.pk}/attachments/{document.identifier}.md",
file_name=f"{document.identifier}.md",
content_type="text/markdown",
conversion_from=key, # might be None
)
default_storage.save(md_attachment.key, BytesIO(parsed_content.encode("utf8")))
md_attachment.upload_state = models.AttachmentStatus.READY
await md_attachment.asave(update_fields=["upload_state", "updated_at"])
# Use the shared document storage utility
create_attachment = not document.media_type.startswith("text/")
attachment_key = None
if create_attachment:
attachment_key = key or f"{self.conversation.pk}/attachments/{document.identifier}.md"
await store_document_with_attachment(
conversation=self.conversation,
user=self.user,
name=document.identifier,
content_type=document.media_type,
content=document_data,
create_attachment=create_attachment,
conversion_from=key, # might be None
attachment_key=attachment_key,
)
def prepare_prompt( # noqa: PLR0912 # pylint: disable=too-many-branches
self, message: UIMessage
@@ -389,6 +379,19 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
langfuse.update_current_trace(
input=user_prompt if self._store_analytics else "REDACTED"
)
# Check conversation history or provided messages
urls_in_conversation = detect_url_in_conversation(messages)
has_url_in_conversation = bool(urls_in_conversation)
if has_url_in_conversation:
# Add fetch_url tool dynamically if URL is detected and tool doesn't exist yet
@self.conversation_agent.tool(name="fetch_url", retries=2)
@functools.wraps(fetch_url)
async def fetch_url_tool(ctx: RunContext, url: str) -> ToolReturn:
"""Wrap the fetch_url tool to provide context and add the tool."""
ctx.deps.messages = messages
return await fetch_url(ctx, url)
usage = {"promptTokens": 0, "completionTokens": 0}
@@ -484,11 +487,13 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
.aexists()
)
should_enable_rag = has_not_pdf_docs or has_url_in_conversation
document_urls = []
if not conversation_has_documents and not has_not_pdf_docs:
if not conversation_has_documents and not should_enable_rag:
# No documents to process
pass
elif has_not_pdf_docs:
elif should_enable_rag:
add_document_rag_search_tool(self.conversation_agent)
@self.conversation_agent.instructions
@@ -505,13 +510,15 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
)
# Inform the model (system-level) that documents are attached and available
@self.conversation_agent.system_prompt
def attached_documents_note() -> str:
return (
"[Internal context] User documents are attached to this conversation. "
"Do not request re-upload of documents; consider them already available "
"via the internal store."
)
# Only if we actually have documents (not just URL), to avoid hallucination
if has_not_pdf_docs:
@self.conversation_agent.system_prompt
def attached_documents_note() -> str:
return (
"[Internal context] User documents are attached to this conversation. "
"Do not request re-upload of documents; consider them already available "
"via the internal store."
)
@self.conversation_agent.tool(name="summarize", retries=2)
@functools.wraps(document_summarize)
@@ -549,6 +556,7 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
_final_output_from_tool = None
_ui_sources = []
_added_source_urls = set()
# Help Mistral to prevent `Unexpected role 'user' after role 'tool'` error.
if history and history[-1].kind == "request":
@@ -661,6 +669,10 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
sources := event.result.metadata.get("sources")
):
for source_url in sources:
# Skip if we've already added this URL to avoid duplicates
if source_url in _added_source_urls:
continue
_added_source_urls.add(source_url)
url_source = LanguageModelV1Source(
sourceType="url",
id=str(uuid.uuid4()),
+207
View File
@@ -0,0 +1,207 @@
"""
Utilities for storing documents in RAG backend and creating attachments.
This module provides shared functionality for processing documents and storing them
in the RAG backend, as well as creating markdown attachments for non-text documents.
"""
import logging
from io import BytesIO
from typing import Optional
from django.conf import settings
from django.core.files.base import ContentFile
from django.core.files.storage import default_storage
from django.utils.module_loading import import_string
from django.utils.text import slugify
from chat import models
logger = logging.getLogger(__name__)
async def ensure_collection_exists(conversation) -> None:
"""
Ensure that a document collection exists for the conversation.
Creates a new collection if one doesn't exist and updates the conversation.
Args:
conversation: The ChatConversation instance.
"""
document_store_backend = import_string(settings.RAG_DOCUMENT_SEARCH_BACKEND)
document_store = document_store_backend(conversation.collection_id)
if not document_store.collection_id:
collection_id = document_store.create_collection(
name=f"conversation-{conversation.pk}",
)
conversation.collection_id = str(collection_id)
await conversation.asave(update_fields=["collection_id", "updated_at"])
async def store_document_in_rag(
conversation,
name: str,
content_type: str,
content: bytes | BytesIO,
store_name: Optional[str] = None,
) -> str:
"""
Parse and store a document in the RAG backend.
Args:
conversation: The ChatConversation instance.
name: The name/identifier to use for parsing (should be filesystem-safe).
content_type: The MIME type of the document.
content: The document content as bytes or BytesIO.
store_name: Optional name to use for storing (for metadata/citations).
If None, uses `name`.
Returns:
str: The parsed markdown content.
"""
await ensure_collection_exists(conversation)
document_store_backend = import_string(settings.RAG_DOCUMENT_SEARCH_BACKEND)
document_store = document_store_backend(conversation.collection_id)
# Normalize content to bytes first, then create a fresh BytesIO
# The backend expects BytesIO in its signature, but internally passes it directly
# to convert_raw which expects bytes. We ensure content is always bytes first,
# then create a fresh BytesIO for the backend (which it needs for PDFs).
if isinstance(content, BytesIO):
# Read the BytesIO content to bytes
content_bytes = content.read()
# Reset position if possible (though we create a new BytesIO anyway)
content.seek(0) if hasattr(content, 'seek') else None
elif isinstance(content, bytes):
content_bytes = content
else:
raise TypeError(f"content must be bytes or BytesIO, got {type(content)}")
# Create a fresh BytesIO from the bytes for the backend
# The backend needs BytesIO for PDFs (file-like object), but will pass it
# directly to convert_raw for non-PDFs (which expects bytes).
# This is a limitation of the backend that we can't fix.
content_io = BytesIO(content_bytes)
# Parse the document
parsed_content = document_store.parse_document(
name=name,
content_type=content_type,
content=content_io,
)
# Store the document (use store_name if provided, otherwise use name)
document_store.store_document(
name=store_name or name,
content=parsed_content,
)
return parsed_content
async def create_markdown_attachment(
conversation,
user,
file_name: str,
parsed_content: str,
key: Optional[str] = None,
conversion_from: Optional[str] = None,
) -> models.ChatConversationAttachment:
"""
Create a markdown attachment for a parsed document.
Args:
conversation: The ChatConversation instance.
user: The user who uploaded/created the document.
file_name: The name of the markdown file to create.
parsed_content: The markdown content to store.
key: Optional storage key. If None, generates from conversation.pk and file_name.
conversion_from: Optional key of the original file if this is a conversion.
Returns:
ChatConversationAttachment: The created attachment instance.
"""
if key is None:
key = f"{conversation.pk}/attachments/{file_name}"
md_attachment = await models.ChatConversationAttachment.objects.acreate(
conversation=conversation,
uploaded_by=user,
key=key,
file_name=file_name,
content_type="text/markdown",
conversion_from=conversion_from,
)
try:
default_storage.save(key, ContentFile(parsed_content.encode("utf8")))
md_attachment.upload_state = models.AttachmentStatus.READY
await md_attachment.asave(update_fields=["upload_state", "updated_at"])
except Exception as exc:
logger.error("Failed to save markdown attachment to storage: %s", exc)
await md_attachment.adelete()
raise
return md_attachment
async def store_document_with_attachment(
conversation,
user,
name: str,
content_type: str,
content: bytes | BytesIO,
store_name: Optional[str] = None,
create_attachment: bool = True,
conversion_from: Optional[str] = None,
attachment_key: Optional[str] = None,
) -> tuple[str, Optional[models.ChatConversationAttachment]]:
"""
Store a document in RAG and optionally create a markdown attachment.
This is a convenience function that combines store_document_in_rag and
create_markdown_attachment. It handles the common workflow of storing
non-text documents.
Args:
conversation: The ChatConversation instance.
user: The user who uploaded/created the document.
name: The name/identifier to use for parsing (should be filesystem-safe).
content_type: The MIME type of the document.
content: The document content as bytes or BytesIO.
store_name: Optional name to use for storing (for metadata/citations).
If None, uses `name`.
create_attachment: Whether to create a markdown attachment.
Defaults to True for non-text content types.
conversion_from: Optional key of the original file if this is a conversion.
attachment_key: Optional custom key for the attachment. If None, generates
from conversation.pk and file_name.
Returns:
tuple[str, Optional[ChatConversationAttachment]]: The parsed content and
the created attachment (if created).
"""
parsed_content = await store_document_in_rag(
conversation=conversation,
name=name,
content_type=content_type,
content=content,
store_name=store_name,
)
attachment = None
if create_attachment and not content_type.startswith("text/"):
file_name = f"{name}.md"
attachment = await create_markdown_attachment(
conversation=conversation,
user=user,
file_name=file_name,
parsed_content=parsed_content,
key=attachment_key,
conversion_from=conversion_from,
)
return parsed_content, attachment
@@ -0,0 +1,361 @@
"""Tests for the fetch_url tool."""
from unittest.mock import AsyncMock, Mock, patch
import httpx
import pytest
import respx
from pydantic_ai import RunContext, RunUsage
from chat.factories import ChatConversationFactory
from chat.tools.fetch_url import detect_url_in_conversation, fetch_url
from core.factories import UserFactory
pytestmark = pytest.mark.django_db()
@pytest.fixture(autouse=True)
def fetch_url_settings(settings):
"""Define settings for fetch_url tests."""
settings.RAG_DOCUMENT_SEARCH_BACKEND = (
"chat.agent_rag.document_rag_backends.albert_rag_backend.AlbertRagBackend"
)
settings.ALBERT_API_URL = "https://albert.api.etalab.gouv.fr"
settings.ALBERT_API_KEY = "test-albert-key"
settings.ALBERT_API_TIMEOUT = 30
settings.ALBERT_API_PARSE_TIMEOUT = 60
@pytest.fixture(name="mocked_context")
def fixture_mocked_context(conversation, user):
"""Fixture for a mocked RunContext with conversation and user."""
mock_ctx = Mock(spec=RunContext)
mock_ctx.usage = RunUsage(input_tokens=0, output_tokens=0)
mock_ctx.deps = Mock()
mock_ctx.deps.conversation = conversation
mock_ctx.deps.user = user
return mock_ctx
@pytest.fixture(name="conversation")
def fixture_conversation():
"""Create a test conversation."""
return ChatConversationFactory()
@pytest.fixture(name="user")
def fixture_user():
"""Create a test user."""
return UserFactory()
def test_detect_url_in_conversation_with_ui_messages(conversation):
"""Test URL detection in ui_messages."""
conversation.ui_messages = [
{
"role": "user",
"parts": [{"type": "text", "text": "Check this: https://example.com/page"}],
}
]
urls = detect_url_in_conversation(conversation.ui_messages)
assert "https://example.com/page" in urls
def test_detect_url_in_conversation_multiple_urls(conversation):
"""Test URL detection with multiple URLs."""
conversation.ui_messages = [
{
"role": "user",
"parts": [
{
"type": "text",
"text": "See https://example.com/1 and https://example.com/2",
}
],
}
]
urls = detect_url_in_conversation(conversation.ui_messages)
assert len(urls) == 2
assert "https://example.com/1" in urls
assert "https://example.com/2" in urls
def test_detect_url_in_conversation_no_urls(conversation):
"""Test URL detection when no URLs are present."""
conversation.ui_messages = [
{"role": "user", "parts": [{"type": "text", "text": "No URL here"}]}
]
urls = detect_url_in_conversation(conversation.ui_messages)
assert urls == []
def test_detect_url_in_conversation_empty_conversation():
"""Test URL detection with None conversation."""
urls = detect_url_in_conversation(None)
assert urls == []
@pytest.mark.asyncio
@respx.mock
async def test_fetch_url_not_detected_in_conversation(mocked_context):
"""Test fetch_url when URL is not detected in conversation."""
mocked_context.deps.conversation.ui_messages = [
{"role": "user", "parts": [{"type": "text", "text": "Hello"}]}
]
mocked_context.deps.messages = mocked_context.deps.conversation.ui_messages
result = await fetch_url(mocked_context, "https://example.com")
assert result.return_value["error"] == "URL not detected in conversation"
assert "not detected" in result.content
@pytest.mark.asyncio
@respx.mock
async def test_fetch_url_docs_numerique_gouv_fr_success(mocked_context):
"""Test fetch_url with docs.numerique.gouv.fr URL."""
url = "https://docs.numerique.gouv.fr/docs/1ef86abf-f7e0-46ce-b6c7-8be8b8af4c3d/"
mocked_context.deps.conversation.ui_messages = [
{"role": "user", "parts": [{"type": "text", "text": f"Check {url}"}]}
]
mocked_context.deps.messages = mocked_context.deps.conversation.ui_messages
# Mock the Docs API response
docs_api_url = "https://docs.numerique.gouv.fr/api/v1.0/documents/1ef86abf-f7e0-46ce-b6c7-8be8b8af4c3d/content/?content_format=markdown"
respx.get(docs_api_url).mock(
return_value=httpx.Response(
status_code=200,
json={"content": "# Test Document\n\nThis is test content."},
)
)
result = await fetch_url(mocked_context, url)
assert result.return_value["url"] == url
assert result.return_value["source"] == "docs.numerique.gouv.fr"
assert "content" in result.return_value
assert "# Test Document" in result.return_value["content"]
@pytest.mark.asyncio
@respx.mock
async def test_fetch_url_docs_numerique_gouv_fr_large_content(mocked_context):
"""Test fetch_url with docs.numerique.gouv.fr when content is large."""
url = "https://docs.numerique.gouv.fr/docs/1ef86abf-f7e0-46ce-b6c7-8be8b8af4c3d/"
mocked_context.deps.conversation.ui_messages = [
{"role": "user", "parts": [{"type": "text", "text": f"Check {url}"}]}
]
mocked_context.deps.messages = mocked_context.deps.conversation.ui_messages
# Create large content (> 8000 chars)
large_content = "# Large Document\n\n" + "x" * 10000
docs_api_url = "https://docs.numerique.gouv.fr/api/v1.0/documents/1ef86abf-f7e0-46ce-b6c7-8be8b8af4c3d/content/?content_format=markdown"
respx.get(docs_api_url).mock(
return_value=httpx.Response(
status_code=200,
json={"content": large_content},
)
)
# Mock Albert API document storage
respx.post("https://albert.api.etalab.gouv.fr/v1/documents").mock(
return_value=httpx.Response(
status_code=200,
json={"id": 456},
)
)
# Mock RAG backend
with patch("chat.tools.fetch_url.import_string") as mock_import, patch(
"chat.tools.fetch_url.models.ChatConversationAttachment.objects.acreate", new_callable=AsyncMock
) as mock_attachment_create, patch("chat.tools.fetch_url.default_storage.save") as mock_storage:
mock_backend = Mock()
mock_backend.collection_id = None # Will trigger collection creation
mock_backend.create_collection = Mock(return_value="123")
mock_backend.parse_document = Mock(return_value=large_content)
mock_backend.store_document = Mock()
mock_import.return_value = Mock(return_value=mock_backend)
# Mock conversation.asave for collection_id update
mocked_context.deps.conversation.asave = AsyncMock()
# Mock attachment creation
mock_attachment = Mock()
mock_attachment.upload_state = None
mock_attachment.asave = AsyncMock()
mock_attachment_create.return_value = mock_attachment
result = await fetch_url(mocked_context, url)
assert result.return_value["url"] == url
assert result.return_value["stored_in_rag"] is True
assert "content_preview" in result.return_value
assert result.metadata["sources"] == {url}
@pytest.mark.asyncio
@respx.mock
async def test_fetch_url_wikipedia_html(mocked_context):
"""Test fetch_url with Wikipedia HTML page."""
url = "https://fr.wikipedia.org/wiki/%C3%89lectron"
mocked_context.deps.conversation.ui_messages = [
{"role": "user", "parts": [{"type": "text", "text": f"Read {url}"}]}
]
mocked_context.deps.messages = mocked_context.deps.conversation.ui_messages
# Mock Wikipedia HTML response
html_content = """
<html>
<head><title>Électron - Wikipédia</title></head>
<body>
<h1>Électron</h1>
<p>L'électron est une particule élémentaire.</p>
</body>
</html>
"""
respx.get(url).mock(
return_value=httpx.Response(
status_code=200,
content=html_content.encode("utf-8"),
headers={"content-type": "text/html; charset=utf-8"},
)
)
# Mock trafilatura extraction
with patch("chat.tools.fetch_url.trafilatura.extract") as mock_extract:
mock_extract.return_value = "Électron\n\nL'électron est une particule élémentaire."
result = await fetch_url(mocked_context, url)
assert result.return_value["url"] == url
assert result.return_value["status_code"] == 200
assert "content" in result.return_value
assert "Électron" in result.return_value["content"]
assert result.metadata["sources"] == {url}
@pytest.mark.asyncio
@respx.mock
async def test_fetch_url_arxiv_pdf(mocked_context):
"""Test fetch_url with arXiv PDF."""
url = "https://arxiv.org/pdf/1706.08595"
mocked_context.deps.conversation.ui_messages = [
{"role": "user", "parts": [{"type": "text", "text": f"Read {url}"}]}
]
mocked_context.deps.messages = mocked_context.deps.conversation.ui_messages
# Mock PDF response
pdf_content = b"%PDF-1.4\n1 0 obj\n<< /Type /Catalog >>\nendobj\n"
respx.get(url).mock(
return_value=httpx.Response(
status_code=200,
content=pdf_content,
headers={"content-type": "application/pdf"},
)
)
# Mock Albert API parse endpoint
parsed_content = "# PDF Content\n\nExtracted text from PDF."
respx.post("https://albert.api.etalab.gouv.fr/v1/parse-beta").mock(
return_value=httpx.Response(
status_code=200,
json={"data": [{"content": parsed_content}]},
)
)
# Mock Albert API document storage
respx.post("https://albert.api.etalab.gouv.fr/v1/documents").mock(
return_value=httpx.Response(
status_code=200,
json={"id": 456},
)
)
# Mock RAG backend for PDF storage
with patch("chat.tools.fetch_url.import_string") as mock_import, patch(
"chat.tools.fetch_url.models.ChatConversationAttachment.objects.acreate", new_callable=AsyncMock
) as mock_attachment_create, patch("chat.tools.fetch_url.default_storage.save") as mock_storage:
mock_backend = Mock()
mock_backend.collection_id = "123"
mock_backend.parse_document = Mock(return_value=parsed_content)
mock_backend.store_document = Mock()
mock_import.return_value = Mock(return_value=mock_backend)
# Mock attachment creation
mock_attachment = Mock()
mock_attachment.upload_state = None
mock_attachment.asave = AsyncMock()
mock_attachment_create.return_value = mock_attachment
result = await fetch_url(mocked_context, url)
assert result.return_value["url"] == url
assert result.return_value["stored_in_rag"] is True
assert "content_preview" in result.return_value
assert result.return_value["content_type"] == "application/pdf"
assert result.metadata["sources"] == {url}
@pytest.mark.asyncio
@respx.mock
async def test_fetch_url_http_error(mocked_context):
"""Test fetch_url with HTTP error."""
url = "https://example.com/error"
mocked_context.deps.conversation.ui_messages = [
{"role": "user", "parts": [{"type": "text", "text": f"Check {url}"}]}
]
mocked_context.deps.messages = mocked_context.deps.conversation.ui_messages
respx.get(url).mock(return_value=httpx.Response(status_code=404))
result = await fetch_url(mocked_context, url)
assert result.return_value["url"] == url
assert "error" in result.return_value
assert "404" in result.return_value["error"]
@pytest.mark.asyncio
@respx.mock
async def test_fetch_url_timeout(mocked_context):
"""Test fetch_url with timeout."""
url = "https://example.com/slow"
mocked_context.deps.conversation.ui_messages = [
{"role": "user", "parts": [{"type": "text", "text": f"Check {url}"}]}
]
mocked_context.deps.messages = mocked_context.deps.conversation.ui_messages
respx.get(url).mock(side_effect=httpx.TimeoutException("Request timed out"))
result = await fetch_url(mocked_context, url)
assert result.return_value["url"] == url
assert "error" in result.return_value
assert "Timeout" in result.return_value["error"]
@pytest.mark.asyncio
@respx.mock
async def test_fetch_url_docs_numerique_gouv_fr_empty_content(mocked_context):
"""Test fetch_url with docs.numerique.gouv.fr when content is empty."""
url = "https://docs.numerique.gouv.fr/docs/1ef86abf-f7e0-46ce-b6c7-8be8b8af4c3d/"
mocked_context.deps.conversation.ui_messages = [
{"role": "user", "parts": [{"type": "text", "text": f"Check {url}"}]}
]
mocked_context.deps.messages = mocked_context.deps.conversation.ui_messages
docs_api_url = "https://docs.numerique.gouv.fr/api/v1.0/documents/1ef86abf-f7e0-46ce-b6c7-8be8b8af4c3d/content/?content_format=markdown"
respx.get(docs_api_url).mock(
return_value=httpx.Response(
status_code=200,
json={"content": ""},
)
)
result = await fetch_url(mocked_context, url)
assert result.return_value["url"] == url
assert result.return_value["error"] == "Content empty or private"
assert "n'est pas public" in result.content
+1
View File
@@ -3,6 +3,7 @@
from pydantic_ai import Tool, ToolDefinition
from .fake_current_weather import get_current_weather
from .fetch_url import fetch_url
from .web_seach_albert_rag import web_search_albert_rag
from .web_search_brave import web_search_brave, web_search_brave_with_document_backend
from .web_search_tavily import web_search_tavily
@@ -26,6 +26,9 @@ def add_document_rag_search_tool(agent: Agent) -> None:
document_store = document_store_backend(ctx.deps.conversation.collection_id)
if not ctx.deps.conversation.collection_id:
return ToolReturn(return_value=[], content="No documents to search in yet.")
rag_results = document_store.search(query)
ctx.usage += RunUsage(
+9 -2
View File
@@ -23,7 +23,14 @@ logger = logging.getLogger(__name__)
def read_document_content(doc):
"""Read document content asynchronously."""
with default_storage.open(doc.key) as f:
return doc.file_name, f.read().decode("utf-8")
# Prefer original URL when the attachment comes from a fetch_url ingestion,
# fallback to the stored filename otherwise.
source_name = (
doc.conversion_from
if doc.conversion_from and doc.conversion_from.startswith(("http://", "https://"))
else doc.file_name
)
return source_name, f.read().decode("utf-8")
async def summarize_chunk(idx, chunk, total_chunks, summarization_agent, ctx):
@@ -173,7 +180,7 @@ async def document_summarize( # pylint: disable=too-many-locals
logger.debug("[summarize] MERGE response<= %s", final_summary)
return ToolReturn(
return_value=final_summary,
return_value=final_summary + "\n Copy paste this summary to the user.",
metadata={"sources": {doc[0] for doc in documents}},
)
+398
View File
@@ -0,0 +1,398 @@
"""Tool to fetch content from a URL detected in the conversation."""
import logging
import random
import re
import httpx
from asgiref.sync import sync_to_async
from django.conf import settings
from django.utils.text import slugify
from pydantic_ai import RunContext
from pydantic_ai.messages import ToolReturn
import trafilatura
from chat import models
from chat.ai_sdk_types import TextUIPart
from chat.document_storage import (
create_markdown_attachment,
ensure_collection_exists,
store_document_in_rag,
)
logger = logging.getLogger(__name__)
MAX_INLINE_CONTENT_CHARS = 8000
# Host for Docs.numerique.gouv.fr
DOCS_HOST = "docs.numerique.gouv.fr"
# Regex pattern to detect URLs
# Note: This is a permissive pattern for detection in free text, not strict validation
URL_PATTERN = re.compile(
r'https?://(?:[a-zA-Z0-9\-._~:/?#\[\]@!$&\'()*+,;=%]|(?:%[0-9a-fA-F]{2}))+'
)
def _get_headers() -> dict:
"""
Return a random set of HTTP headers for each request.
For now this only randomizes the User-Agent, but we can easily extend this
list with more header variants (Accept-Language, Referer, etc.) if needed.
"""
headers_pool = [
{
"User-Agent": (
"Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
"AppleWebKit/537.36 (KHTML, like Gecko) "
"Chrome/120.0.0.0 Safari/537.36"
)
},
{
"User-Agent": (
"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) "
"AppleWebKit/605.1.15 (KHTML, like Gecko) "
"Version/17.1 Safari/605.1.15"
)
},
{
"User-Agent": (
"Mozilla/5.0 (X11; Ubuntu; Linux x86_64; rv:121.0) "
"Gecko/20100101 Firefox/121.0"
)
},
]
return random.choice(headers_pool)
def _extract_text_from_message(message) -> str:
"""
Extract all text content from a message.
Args:
message: A message object (UIMessage or dict).
Returns:
str: All text content concatenated.
"""
text_parts = []
# Handle UIMessage objects (Pydantic models)
if hasattr(message, 'parts'):
for part in message.parts or []:
if isinstance(part, TextUIPart) and part.text:
text_parts.append(part.text)
# Also check the deprecated content field
if hasattr(message, 'content') and message.content:
text_parts.append(message.content)
# Handle dict-based messages (JSON deserialized)
elif isinstance(message, dict):
# Check parts
parts = message.get('parts', [])
for part in parts:
if isinstance(part, dict) and part.get('type') == 'text':
text = part.get('text', '')
if text:
text_parts.append(text)
# Check deprecated content field
content = message.get('content', '')
if content:
text_parts.append(content)
return ' '.join(text_parts)
def detect_url_in_conversation(messages=None) -> list[str]:
"""
Detect URLs present in the conversation messages.
Args:
messages: Iterable of UIMessage/dict messages (latest payload).
Returns:
list[str]: List of unique URLs found in the conversation.
"""
found_urls = set()
def extract_urls_from_messages(messages):
if not messages:
return
for message in messages:
if not message:
continue
text_content = _extract_text_from_message(message)
if text_content and URL_PATTERN.search(text_content):
matches = URL_PATTERN.findall(text_content)
found_urls.update(matches)
if messages:
extract_urls_from_messages(messages)
return list(found_urls)
async def _get_with_retry(
client: httpx.AsyncClient,
url: str,
max_attempts: int = 3,
) -> httpx.Response:
"""
Perform a GET request with randomized headers and a simple retry strategy.
We retry once on header-related HTTP status codes (e.g. 403, 429), each time
using a new random header set. Other errors are propagated immediately.
"""
last_exception: httpx.HTTPStatusError | None = None
for attempt in range(max_attempts):
headers = _get_headers()
try:
response = await client.get(url, headers=headers)
response.raise_for_status()
return response
except httpx.HTTPStatusError as exc:
last_exception = exc
status_code = exc.response.status_code
# Only retry on codes that are likely related to headers / rate limits.
should_retry = status_code in (403, 429)
is_last_attempt = attempt >= max_attempts - 1
logger.debug(
"HTTP error %s for URL %s on attempt %s with headers %s (retry=%s)",
status_code,
url,
attempt + 1,
headers,
should_retry and not is_last_attempt,
)
if (not should_retry) or is_last_attempt:
raise
# Should not be reached, but keeps type-checkers happy.
if last_exception is not None:
raise last_exception
raise RuntimeError("Unexpected state in _get_with_retry")
async def _store_in_rag_and_attachments(
conversation,
user,
url: str,
content_bytes: bytes,
content_type: str,
) -> str:
"""
Store the fetched document into the RAG backend and create a markdown attachment.
Returns the markdown content stored, mainly to allow generating a short preview.
"""
await ensure_collection_exists(conversation)
# Force content_type to "application/pdf" if it seems to be a PDF but the header was weird
# This ensures AlbertRagBackend uses the PDF parser
if "pdf" in content_type or url.lower().endswith(".pdf"):
content_type = "application/pdf"
# Use a safe filename (slugified) for the RAG backend to avoid API errors with URLs.
safe_rag_name = slugify(url)[:100] or "document"
# We must split parsing and storing to handle the filename vs metadata issue:
# 1. Parsing needs a safe filename (no slashes) to avoid 500 errors from the API.
# 2. Storing needs the original URL in metadata so citations are correct.
parsed_content = await store_document_in_rag(
conversation=conversation,
name=safe_rag_name,
content_type=content_type,
content=content_bytes,
store_name=url,
)
# Create a markdown attachment so that the rest of the pipeline
# (document_search_rag, summarization, etc.) can see documents exist.
file_name = f"{safe_rag_name}.md"
key = f"{conversation.pk}/attachments/{file_name}"
await create_markdown_attachment(
conversation=conversation,
user=user,
file_name=file_name,
parsed_content=parsed_content,
key=key,
# Keep track of the original URL for downstream tools (e.g. summarize)
conversion_from=url,
)
return parsed_content
async def fetch_url(ctx: RunContext, url: str) -> ToolReturn:
"""
Fetch content from a URL.
When an URL is detected and you need to fetch content from it, you should use this tool.
Args:
ctx (RunContext): The run context containing the conversation.
url (str): The URL to fetch content from.
Returns:
ToolReturn: The fetched content from the URL.
"""
# Access the Django conversation object from the agent dependencies
deps = getattr(ctx, "deps", None)
conversation = getattr(deps, "conversation", None)
user = getattr(deps, "user", None)
messages_for_detection = getattr(deps, "messages", None)
urls = detect_url_in_conversation(messages_for_detection)
logger.info("URLs authorized (extracted from messages): %s", urls)
# If messages are provided, enforce URL presence; otherwise skip the check.
if messages_for_detection is not None and url not in urls:
return ToolReturn(
return_value={"url": url, "error": "URL not detected in conversation", "content" : f"URL {url} not detected in conversation"},
)
try:
# Special handling for docs.numerique.gouv.fr
m = re.search(r"https?://(?:www\.)?docs\.numerique\.gouv\.fr/docs/([^/?#]+)", url)
if m:
docs_id = m.group(1)
url_transformed = f"https://{DOCS_HOST}/api/v1.0/documents/{docs_id}/content/?content_format=markdown"
try:
async with httpx.AsyncClient(timeout=settings.FETCH_URL_TIMEOUT, follow_redirects=True) as client:
response = await _get_with_retry(client, url_transformed)
data = response.json()
content = data.get('content', '')
if not content:
return ToolReturn(
return_value={"url": url, "error": "Content empty or private", "content": "Ce document Docs n'est pas public ou est vide."},
)
# If the Docs content is very large, route it through RAG instead of
# returning everything inline.
if conversation and user and len(content) > MAX_INLINE_CONTENT_CHARS:
parsed = await _store_in_rag_and_attachments(
conversation=conversation,
user=user,
url=url,
content_bytes=content.encode("utf-8"),
content_type="text/markdown",
)
preview = parsed[:MAX_INLINE_CONTENT_CHARS]
return ToolReturn(
return_value={
"url": url,
"original_url": url,
"stored_in_rag": True,
"content_preview": preview,
"source": DOCS_HOST,
"content":(
"Le contenu de ce document est volumineux et a été indexé dans "
"la base de documents de la conversation. "
"Pour linterroger, tu dois utiliser loutil `document_search_rag` "
"avec une requête précise décrivant ce que tu cherches dans ce document."
)
},
metadata={"sources": {url}},
)
return ToolReturn(
return_value={
"url": url,
"original_url": url,
"content": content[:MAX_INLINE_CONTENT_CHARS],
"source": DOCS_HOST,
}
)
except Exception as e:
logger.warning("Error fetching Docs content %s: %s", url, e)
return ToolReturn(
return_value={"url": url, "error": str(e), "content": "Ce document Docs n'est pas public ou une erreur est survenue."},
)
async with httpx.AsyncClient(timeout=settings.FETCH_URL_TIMEOUT, follow_redirects=True) as client:
response = await _get_with_retry(client, url)
content_type_header = response.headers.get("content-type", "unknown")
content_type = content_type_header.split(";", 1)[0].strip().lower()
is_binary_like = not content_type.startswith("text/")
is_pdf = "pdf" in content_type or url.lower().endswith(".pdf")
# Avoid trafilatura on PDFs
if is_pdf:
extracted = ""
else:
# Run trafilatura.extract in a thread to avoid blocking the event loop
extracted = await sync_to_async(trafilatura.extract)(response.text) or response.text
# For large or binary/PDF content, store in RAG instead of returning everything inline.
if (
conversation
and user
and (is_binary_like or is_pdf or len(extracted) > MAX_INLINE_CONTENT_CHARS)
):
parsed = await _store_in_rag_and_attachments(
conversation=conversation,
user=user,
url=url,
content_bytes=response.content,
content_type=content_type,
)
preview = parsed[:MAX_INLINE_CONTENT_CHARS]
return ToolReturn(
return_value={
"url": url,
"status_code": response.status_code,
"stored_in_rag": True,
"content_preview": preview,
"content_type": content_type_header,
"content":(
"Le contenu de cette ressource est volumineux ou de type fichier "
"(par exemple PDF). Il a été indexé dans la base de documents de la "
"conversation. Pour lexploiter, tu dois utiliser loutil "
"`document_search_rag` avec une requête précise décrivant les "
"informations que tu souhaites retrouver."
)
},
metadata={"sources": {url}},
)
# Small textual content: keep the existing behaviour with inline content.
return ToolReturn(
return_value={
"url": url,
"status_code": response.status_code,
"content": extracted[:MAX_INLINE_CONTENT_CHARS],
"content_type": content_type_header,
},
metadata={"sources": {url}},
)
except httpx.HTTPStatusError as e:
logger.warning("HTTP error fetching %s: %s", url, e)
return ToolReturn(
return_value={
"url": url,
"error": f"HTTP {e.response.status_code}: {str(e)}",
}
)
except httpx.TimeoutException as e:
logger.warning("Timeout fetching %s: %s", url, e)
return ToolReturn(
return_value={
"url": url,
"error": f"Timeout: {str(e)}",
}
)
except Exception as e:
logger.exception("Error fetching %s: %s", url, e)
return ToolReturn(
return_value={
"url": url,
"error": f"Error: {str(e)}",
}
)
+7
View File
@@ -841,6 +841,13 @@ USER QUESTION:
environ_prefix=None,
)
# Fetch URL
FETCH_URL_TIMEOUT = values.PositiveIntegerValue(
default=5, # seconds
environ_name="FETCH_URL_TIMEOUT",
environ_prefix=None,
)
# Logging
# We want to make it easy to log to console but by default we log production
# to Sentry and don't want to log to console.
@@ -791,20 +791,39 @@ export const Chat = ({
<Loader />
<Text $variation="600" $size="md">
{(() => {
const toolInvocation = message.parts?.find(
// Find the tool invocation that is currently running (not completed)
const toolInvocations = message.parts?.filter(
(part) =>
part.type === 'tool-invocation' &&
part.toolInvocation.toolName !==
'document_parsing',
);
) || [];
// Find the last tool invocation that is not yet completed
const activeToolInvocation = [...toolInvocations]
.reverse()
.find(
(part) =>
part.type === 'tool-invocation' &&
part.toolInvocation.state !== 'result',
);
if (
toolInvocation?.type ===
activeToolInvocation?.type ===
'tool-invocation' &&
toolInvocation.toolInvocation.toolName ===
activeToolInvocation.toolInvocation.toolName ===
'summarize'
) {
return t('Summarizing...');
}
if (
activeToolInvocation?.type ===
'tool-invocation' &&
activeToolInvocation.toolInvocation.toolName ===
'fetch_url'
) {
return t('Fetching URL...');
}
return t('Search...');
})()}
</Text>
@@ -104,6 +104,7 @@
"Start conversation": "Entamer la conversation",
"Stop": "Stop",
"Summarizing...": "Résumé en cours...",
"Fetching URL...": "Visite du lien en cours...",
"The Assistant is a sovereign conversational AI designed for public servants. It helps you save time on daily tasks like rephrasing, summarising, translating, or searching information. Your data never leaves France and is stored on secure, state-compliant infrastructures. It is never used for commercial purposes.": "L'Assistant est une IA souveraine conçue pour les fonctionnaires. Il vous permet de gagner du temps sur des tâches quotidiennes telles que la reformulation, le résumé, la traduction ou la recherche d'informations. Vos données ne quittent jamais la France et sont stockées sur des infrastructures sûres et conformes à l'état et ne sont jamais utilisées à des fins commerciales.",
"The Assistant is in Beta": "L'Assistant est en Bêta",
"The conversation has been deleted.": "La conversation a été supprimée.",
@@ -236,6 +237,7 @@
"Start conversation": "Begin een gesprek",
"Stop": "Stop",
"Summarizing...": "Samenvatten...",
"Fetching URL...": "URL ophalen...",
"The Assistant is a sovereign conversational AI designed for public servants. It helps you save time on daily tasks like rephrasing, summarising, translating, or searching information. Your data never leaves France and is stored on secure, state-compliant infrastructures. It is never used for commercial purposes.": "De Assistent is een soevereine conversationele AI, ontworpen voor ambtenaren. Het helpt je tijd te besparen bij dagelijkse taken zoals het herformuleren, samenvatten, vertalen of zoeken van informatie. Je gegevens verlaten het land nooit en worden opgeslagen op beveiligde, door de overheid goedgekeurde infrastructuren. Ze worden nooit gebruikt voor commerciële doeleinden.",
"The Assistant is in Beta": "De Assistent is in bèta",
"The conversation has been deleted.": "Het gesprek is verwijderd.",
@@ -368,6 +370,7 @@
"Start conversation": "Начать беседу",
"Stop": "Остановить",
"Summarizing...": "Обобщение...",
"Fetching URL...": "Получение URL...",
"The Assistant is a sovereign conversational AI designed for public servants. It helps you save time on daily tasks like rephrasing, summarising, translating, or searching information. Your data never leaves France and is stored on secure, state-compliant infrastructures. It is never used for commercial purposes.": "Помощник - собеседник на основе ИИ для государственных служащих. Он поможет вам сэкономить время на ежедневных задачах, таких как перефразирование, обобщение, перевод или поиск информации. Ваши данные никогда не покидают Францию и хранятся в охраняемой государственной инфраструктуре, которая никогда не используется в коммерческих целях.",
"The Assistant is in Beta": "Помощник находится на этапе Бета-версии",
"The conversation has been deleted.": "Беседа была удалена.",
@@ -500,6 +503,7 @@
"Start conversation": "Почати розмову",
"Stop": "Зупинити",
"Summarizing...": "Узагальнення...",
"Fetching URL...": "Отримання URL...",
"The Assistant is a sovereign conversational AI designed for public servants. It helps you save time on daily tasks like rephrasing, summarising, translating, or searching information. Your data never leaves France and is stored on secure, state-compliant infrastructures. It is never used for commercial purposes.": "Помічник - це розмовний ШІ, призначений для державних службовців. Він допоможе вам зберегти час в таких щоденних завданнях, як рефразування, узагальнення, переклад або пошукова інформація. Ваші дані ніколи не покидають Францію та зберігаються на захищеній державній інфраструктурі. Вони ніколи не використовуються для комерційних цілей.",
"The Assistant is in Beta": "Помічник у бета-версії",
"The conversation has been deleted.": "Розмова була видалена.",