use pydanticAI messages + refacto + settings timeout + correct source for downstream summarize
This commit is contained in:
@@ -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:
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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}},
|
||||
)
|
||||
|
||||
|
||||
@@ -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 l’interroger, tu dois utiliser l’outil `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 l’interroger, tu dois utiliser l’outil `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
|
||||
|
||||
@@ -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.
|
||||
|
||||
Reference in New Issue
Block a user