use pydanticAI messages + refacto + settings timeout + correct source for downstream summarize

This commit is contained in:
camilleAND
2025-12-11 14:53:19 +01:00
parent 1ffbd2370e
commit 8448340b5c
6 changed files with 120 additions and 131 deletions
+8 -16
View File
@@ -76,7 +76,7 @@ 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 URL_PATTERN, detect_url_in_conversation, fetch_url
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
@@ -391,26 +391,18 @@ 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)
# Check if URL is present in current message or conversation, and add fetch_url tool dynamically
# Check current message first (most recent)
has_url_in_current_message = any(
URL_PATTERN.search(part.text) if isinstance(part, TextUIPart) else False
for part in messages[-1].parts or []
) or (URL_PATTERN.search(messages[-1].content) if messages[-1].content else False)
# Also check conversation history
has_url_in_conversation = detect_url_in_conversation(self.conversation)
# Check if fetch_url tool already exists (might be in configuration)
fetch_url_exists = "fetch_url" in self.conversation_agent._function_toolset.tools # pylint: disable=protected-access
if (has_url_in_current_message or has_url_in_conversation) and not fetch_url_exists:
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}
@@ -507,7 +499,7 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
.aexists()
)
should_enable_rag = has_not_pdf_docs or has_url_in_current_message or has_url_in_conversation
should_enable_rag = has_not_pdf_docs or has_url_in_conversation
document_urls = []
if not conversation_has_documents and not should_enable_rag:
+11 -3
View File
@@ -57,7 +57,7 @@ def test_detect_url_in_conversation_with_ui_messages(conversation):
"parts": [{"type": "text", "text": "Check this: https://example.com/page"}],
}
]
urls = detect_url_in_conversation(conversation)
urls = detect_url_in_conversation(conversation.ui_messages)
assert "https://example.com/page" in urls
@@ -74,7 +74,7 @@ def test_detect_url_in_conversation_multiple_urls(conversation):
],
}
]
urls = detect_url_in_conversation(conversation)
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
@@ -85,7 +85,7 @@ def test_detect_url_in_conversation_no_urls(conversation):
conversation.ui_messages = [
{"role": "user", "parts": [{"type": "text", "text": "No URL here"}]}
]
urls = detect_url_in_conversation(conversation)
urls = detect_url_in_conversation(conversation.ui_messages)
assert urls == []
@@ -102,6 +102,7 @@ async def test_fetch_url_not_detected_in_conversation(mocked_context):
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")
@@ -117,6 +118,7 @@ async def test_fetch_url_docs_numerique_gouv_fr_success(mocked_context):
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"
@@ -143,6 +145,7 @@ async def test_fetch_url_docs_numerique_gouv_fr_large_content(mocked_context):
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
@@ -199,6 +202,7 @@ async def test_fetch_url_wikipedia_html(mocked_context):
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 = """
@@ -239,6 +243,7 @@ async def test_fetch_url_arxiv_pdf(mocked_context):
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"
@@ -300,6 +305,7 @@ async def test_fetch_url_http_error(mocked_context):
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))
@@ -318,6 +324,7 @@ async def test_fetch_url_timeout(mocked_context):
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"))
@@ -336,6 +343,7 @@ async def test_fetch_url_docs_numerique_gouv_fr_empty_content(mocked_context):
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(
-7
View File
@@ -18,13 +18,6 @@ def get_pydantic_tools_by_name(name: str) -> Tool:
"""Get a tool by its name."""
tool_dict = {
"get_current_weather": Tool(get_current_weather, takes_ctx=False),
# Note: fetch_url is added dynamically in pydantic_ai.py when URL is detected
# It's kept here for reference but won't be used via prepare
"fetch_url": Tool(
fetch_url,
takes_ctx=True,
max_retries=2,
),
"web_search_brave": Tool(
web_search_brave,
takes_ctx=True,
+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}},
)
+85 -103
View File
@@ -3,20 +3,21 @@
import logging
import random
import re
from io import BytesIO
import httpx
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
import trafilatura
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__)
@@ -100,19 +101,16 @@ def _extract_text_from_message(message) -> str:
return ' '.join(text_parts)
def detect_url_in_conversation(conversation) -> list[str]:
def detect_url_in_conversation(messages=None) -> list[str]:
"""
Detect URLs present in the conversation messages.
Args:
conversation: The ChatConversation instance.
messages: Iterable of UIMessage/dict messages (latest payload).
Returns:
list[str]: List of unique URLs found in the conversation.
"""
if not conversation:
return []
found_urls = set()
def extract_urls_from_messages(messages):
@@ -126,14 +124,8 @@ def detect_url_in_conversation(conversation) -> list[str]:
matches = URL_PATTERN.findall(text_content)
found_urls.update(matches)
# Get URLs from ui_messages and messages if present
if hasattr(conversation, 'ui_messages') and conversation.ui_messages:
extract_urls_from_messages(conversation.ui_messages)
if hasattr(conversation, 'messages') and conversation.messages:
extract_urls_from_messages(conversation.messages)
if found_urls:
logger.info("URL detected in messages: %s", found_urls)
if messages:
extract_urls_from_messages(messages)
return list(found_urls)
@@ -196,16 +188,7 @@ async def _store_in_rag_and_attachments(
Returns the markdown content stored, mainly to allow generating a short preview.
"""
document_store_backend = import_string(settings.RAG_DOCUMENT_SEARCH_BACKEND)
document_store = document_store_backend(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-{conversation.pk}",
)
conversation.collection_id = str(collection_id)
await conversation.asave(update_fields=["collection_id", "updated_at"])
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
@@ -223,15 +206,12 @@ async def _store_in_rag_and_attachments(
# However, AlbertRagBackend.store_document uses the same name for both filename and metadata.
# We try to pass the original URL to store_document, hoping the storage endpoint is more
# robust than the parser endpoint regarding filenames.
parsed_content = document_store.parse_document(
parsed_content = await store_document_in_rag(
conversation=conversation,
name=safe_rag_name,
content_type=content_type,
content=BytesIO(content_bytes),
)
document_store.store_document(
name=url,
content=parsed_content
content=content_bytes,
store_name=url,
)
# Create a markdown attachment so that the rest of the pipeline
@@ -239,17 +219,16 @@ async def _store_in_rag_and_attachments(
file_name = f"{safe_rag_name}.md"
key = f"{conversation.pk}/attachments/{file_name}"
md_attachment = await models.ChatConversationAttachment.objects.acreate(
await create_markdown_attachment(
conversation=conversation,
uploaded_by=user,
key=key,
user=user,
file_name=file_name,
content_type="text/markdown",
conversion_from=None,
parsed_content=parsed_content,
key=key,
# Keep track of the original URL so downstream tools (e.g. summarize)
# can surface a clickable source instead of the slugified filename.
conversion_from=url,
)
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"])
return parsed_content
@@ -271,86 +250,89 @@ async def fetch_url(ctx: RunContext, url: str) -> ToolReturn:
conversation = getattr(deps, "conversation", None)
user = getattr(deps, "user", None)
urls = detect_url_in_conversation(conversation)
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 url not in 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",
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
if DOCS_HOST in url and "/docs/" in url:
# Use regex to extract the document ID
m = re.search(r'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=30.0, 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}},
)
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,
"content": content[:MAX_INLINE_CONTENT_CHARS],
"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}},
)
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=30.0, follow_redirects=True) as client:
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)
extracted = trafilatura.extract(response.text) or response.text
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:
extracted = trafilatura.extract(response.text) or response.text
# For large or binary/PDF content, store in RAG instead of returning everything inline.
if (
conversation
+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.