add rag for long context retrived from urls

This commit is contained in:
camilleAND
2025-12-05 11:30:45 +01:00
parent db84c6906c
commit 7049782e70
3 changed files with 161 additions and 19 deletions
+13 -9
View File
@@ -507,11 +507,13 @@ 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
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
@@ -528,13 +530,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)
@@ -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(
+145 -10
View File
@@ -5,14 +5,22 @@ import random
import re
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
from chat import models
from chat.ai_sdk_types import TextUIPart
logger = logging.getLogger(__name__)
MAX_INLINE_CONTENT_CHARS = 8000
# Regex pattern to detect URLs
URL_PATTERN = re.compile(
r'http[s]?://(?:[a-zA-Z]|[0-9]|[$-_@.&+]|[!*\\(\\),]|(?:%[0-9a-fA-F][0-9a-fA-F]))+'
@@ -169,6 +177,63 @@ async def _get_with_retry(
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.
"""
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"])
# Parse and store in RAG backend (returns markdown)
# 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"
parsed_content = document_store.parse_and_store_document(
name=url,
content_type=content_type,
content=content_bytes,
)
# Create a markdown attachment so that the rest of the pipeline
# (document_search_rag, summarization, etc.) can see documents exist.
safe_name = slugify(url)[:100] or "document"
file_name = f"{safe_name}.md"
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=None,
)
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
async def fetch_url(ctx: RunContext, url: str) -> ToolReturn:
"""
Fetch content from a URL.
@@ -182,7 +247,10 @@ async def fetch_url(ctx: RunContext, url: str) -> ToolReturn:
ToolReturn: The fetched content from the URL.
"""
# Access the Django conversation object from the agent dependencies
conversation = getattr(getattr(ctx, "deps", None), "conversation", None)
deps = getattr(ctx, "deps", None)
conversation = getattr(deps, "conversation", None)
user = getattr(deps, "user", None)
urls = detect_url_in_conversation(conversation)
if url not in urls:
@@ -209,15 +277,42 @@ async def fetch_url(ctx: RunContext, url: str) -> ToolReturn:
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."
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.numerique.gouv.fr",
},
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[:20000], # Limit content
"source": "docs.numerique.gouv.fr"
"content": content[:MAX_INLINE_CONTENT_CHARS],
"source": "docs.numerique.gouv.fr",
}
)
except Exception as e:
@@ -229,15 +324,55 @@ async def fetch_url(ctx: RunContext, url: str) -> ToolReturn:
async with httpx.AsyncClient(timeout=30.0, follow_redirects=True) as client:
response = await _get_with_retry(client, url)
content = trafilatura.extract(response.text)
raw_text = response.text
extracted = trafilatura.extract(raw_text) or raw_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")
# 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": content[:20000], # Limit content to first 20000 chars
"content_type": response.headers.get("content-type", "unknown"),
}
"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)