diff --git a/src/backend/chat/clients/pydantic_ai.py b/src/backend/chat/clients/pydantic_ai.py index b946fb228..833b82597 100644 --- a/src/backend/chat/clients/pydantic_ai.py +++ b/src/backend/chat/clients/pydantic_ai.py @@ -156,6 +156,7 @@ from chat.tools.document_search_rag import add_document_rag_search_tool from chat.tools.document_summarize import document_summarize from chat.tools.self_documentation import build_self_documentation_payload +from chat.tools.url_fetch import add_url_fetch_tool, detect_urls, is_url_allowed from chat.vercel_ai_sdk.core import events_v4, events_v5 from chat.vercel_ai_sdk.encoder import CURRENT_EVENT_ENCODER_VERSION, EventEncoder @@ -310,6 +311,8 @@ def __init__( # pylint: disable=too-many-arguments,too-many-positional-argument ) self._web_search_tool_registered = False self._self_documentation_tool_registered = False + self._url_fetch_tool_registered = False + self._rag_tools_registered = False self.conversation_agent = ConversationAgent( model_hrid=self.model_hrid, @@ -521,12 +524,17 @@ def force_web_search_prompt() -> str: return True - async def _check_should_enable_rag(self, conversation_has_documents: bool) -> bool: - """Check if RAG should be enabled based on existing documents.""" + async def _check_should_enable_rag( + self, conversation_has_documents: bool, user_prompt: str = "" + ) -> bool: + """Check if RAG should be enabled based on existing documents or URLs in the prompt.""" if not self._is_document_upload_enabled: return False + if user_prompt and any(is_url_allowed(u) for u in detect_urls(user_prompt)): + return True + # Check for existing documents (any non-image attachment for this conversation) has_documents = await ( models.ChatConversationAttachment.objects.filter( @@ -776,6 +784,9 @@ async def _build_document_context_instruction(self) -> str: def _setup_rag_tools(self, document_context_instruction: str = "") -> None: """Register RAG-related tools and instructions on the conversation agent.""" + if self._rag_tools_registered: + return + self._rag_tools_registered = True add_document_rag_search_tool(self.conversation_agent) @self.conversation_agent.instructions @@ -804,6 +815,17 @@ async def summarize(ctx: RunContext, *args, **kwargs) -> ToolReturn: """Wrap the document_summarize tool to provide context and add the tool.""" return await document_summarize(ctx, *args, **kwargs) + def _setup_url_fetch_tool(self) -> None: + """Register url_fetch tool when RAG backend is configured.""" + if self._url_fetch_tool_registered: + return + if not self._is_document_upload_enabled: + return + if not getattr(settings, "RAG_DOCUMENT_SEARCH_BACKEND", None): + return + add_url_fetch_tool(self.conversation_agent) + self._url_fetch_tool_registered = True + def _setup_web_search_tool(self) -> None: """Register model-specific web search tool when configured.""" if self._web_search_tool_registered: @@ -1182,8 +1204,9 @@ async def _run_agent( # pylint: disable=too-many-locals self._setup_self_documentation_tool() self._setup_web_search_tool() self._setup_web_search(force_web_search) + self._setup_url_fetch_tool() - if await self._check_should_enable_rag(conversation_has_documents): + if await self._check_should_enable_rag(conversation_has_documents, user_prompt): document_context_instruction = await self._build_document_context_instruction() self._setup_rag_tools(document_context_instruction=document_context_instruction) diff --git a/src/backend/chat/tests/tools/test_url_fetch.py b/src/backend/chat/tests/tools/test_url_fetch.py new file mode 100644 index 000000000..c6099d907 --- /dev/null +++ b/src/backend/chat/tests/tools/test_url_fetch.py @@ -0,0 +1,614 @@ +"""Tests for the URL fetch tool.""" + +import asyncio +from unittest.mock import AsyncMock, MagicMock, Mock, patch + +import httpx +import pytest +from pydantic_ai import Agent, RunContext, RunUsage +from pydantic_ai.exceptions import ModelRetry + +from chat.tools.url_fetch import add_url_fetch_tool, detect_urls, is_url_allowed + + +@pytest.fixture(name="url_fetch_settings") +def url_fetch_settings_fixture(settings): + """Set up URL fetch settings for tests.""" + settings.URL_FETCH_BLOCKED_SCHEMES = ["http"] + settings.URL_FETCH_BLOCKED_HOSTS = ["localhost", "127.0.0.1", "0.0.0.0", "::1"] + settings.URL_FETCH_BLOCKED_TLDS = [".ru", ".cn", ".kp", ".ir"] + + +def test_settings_exist(url_fetch_settings, settings): + """URL fetch settings exist with expected defaults.""" + assert settings.URL_FETCH_BLOCKED_SCHEMES == ["http"] + assert "localhost" in settings.URL_FETCH_BLOCKED_HOSTS + assert ".ru" in settings.URL_FETCH_BLOCKED_TLDS + + +def test_detect_urls_simple(): + """detect_urls finds a plain HTTPS URL.""" + result = detect_urls("Check out https://example.com for details") + assert result == ["https://example.com"] + + +def test_detect_urls_multiple(): + """detect_urls finds multiple URLs.""" + result = detect_urls("See https://example.com and https://other.org/path?q=1") + assert result == ["https://example.com", "https://other.org/path?q=1"] + + +def test_detect_urls_http(): + """detect_urls also detects plain HTTP URLs.""" + result = detect_urls("Visit http://example.com now") + assert result == ["http://example.com"] + + +def test_detect_urls_empty(): + """detect_urls returns empty list when no URL present.""" + assert detect_urls("No URL here at all") == [] + + +def test_detect_urls_ignores_non_http(): + """detect_urls ignores non-HTTP(S) schemes.""" + assert detect_urls("ftp://example.com and file:///etc/passwd") == [] + + +def test_is_url_allowed_valid(url_fetch_settings): + """A standard HTTPS URL is allowed.""" + assert is_url_allowed("https://docs.numerique.gouv.fr/page") is True + + +def test_is_url_allowed_blocks_http_scheme(url_fetch_settings): + """HTTP URLs are blocked by scheme filter.""" + assert is_url_allowed("http://example.com/page") is False + + +def test_is_url_allowed_blocks_localhost(url_fetch_settings): + """localhost is blocked by host filter.""" + assert is_url_allowed("https://localhost/api") is False + + +def test_is_url_allowed_blocks_internal_ip(url_fetch_settings): + """127.0.0.1 is blocked by host filter.""" + assert is_url_allowed("https://127.0.0.1/secret") is False + + +def test_is_url_allowed_blocks_tld(url_fetch_settings): + """A blocked TLD is rejected.""" + assert is_url_allowed("https://example.ru/page") is False + + +def test_is_url_allowed_malformed(url_fetch_settings): + """Malformed URL returns False.""" + assert is_url_allowed("not-a-url") is False + + +def test_is_url_allowed_blocks_cloud_metadata(url_fetch_settings): + """Cloud metadata IP is blocked by IP range check.""" + assert is_url_allowed("https://169.254.169.254/latest/meta-data/") is False + + +def test_is_url_allowed_blocks_private_ip_range(url_fetch_settings): + """Private RFC-1918 IP is blocked by IP range check.""" + assert is_url_allowed("https://10.0.0.1/internal") is False + + +def test_is_url_allowed_blocks_private_ip_192(url_fetch_settings): + """Private 192.168.x.x IP is blocked.""" + assert is_url_allowed("https://192.168.1.100/admin") is False + + +@pytest.fixture(name="mock_rag_backend") +def mock_rag_backend_fixture(): + """Mock RAG backend class and instance.""" + backend_instance = MagicMock() + backend_instance.collection_id = "test-collection-123" + backend_instance.astore_document = AsyncMock() + backend_instance.asearch = AsyncMock() + backend_instance.parse_and_store_document = MagicMock(return_value="parsed content") + backend_instance.acreate_collection = AsyncMock(return_value="new-collection-id") + backend_class = MagicMock(return_value=backend_instance) + return backend_class, backend_instance + + +@pytest.fixture(name="mock_conversation") +def mock_conversation_fixture(): + """Mock conversation with collection_id.""" + conv = MagicMock() + conv.collection_id = "test-collection-123" + conv.pk = "conv-pk-1" + conv.asave = AsyncMock() + return conv + + +@pytest.fixture(name="mock_ctx") +def mock_ctx_fixture(mock_conversation): + """Mock RunContext for url_fetch tool tests.""" + ctx = Mock(spec=RunContext) + ctx.deps = Mock() + ctx.deps.conversation = mock_conversation + ctx.deps.user = Mock() + ctx.deps.user.sub = "user-sub-123" + ctx.deps.session = {} + ctx.usage = RunUsage(input_tokens=0, output_tokens=0) + ctx.retries = {} + ctx.max_retries = 2 + ctx.tool_name = "url_fetch" + user_msg = Mock() + user_msg.kind = "request" + user_prompt_part = Mock() + user_prompt_part.part_kind = "user-prompt" + user_prompt_part.content = "What does https://example.com say about Python?" + user_msg.parts = [user_prompt_part] + ctx.messages = [user_msg] + return ctx + + +@pytest.fixture(name="rag_search_results") +def rag_search_results_fixture(): + """Mock RAG search results.""" + result = MagicMock() + result.usage.prompt_tokens = 10 + result.usage.completion_tokens = 5 + chunk = MagicMock() + chunk.url = "https://example.com" + chunk.content = "Python is a great language." + result.data = [chunk] + return result + + +def _get_tool_fn(agent): + """Helper: extract the url_fetch tool function from a Pydantic AI agent.""" + return agent._function_toolset.tools["url_fetch"].function + + +@pytest.mark.asyncio +async def test_url_fetch_html_success(url_fetch_settings, mock_ctx, mock_rag_backend): + """url_fetch fetches HTML, stores via astore_document, returns confirmation.""" + backend_class, backend_instance = mock_rag_backend + + with ( + patch("chat.tools.url_fetch.import_string", return_value=backend_class), + patch("chat.tools.url_fetch.settings") as mock_settings, + patch("chat.tools.url_fetch.sync_to_async") as mock_sync_to_async, + patch("chat.tools.url_fetch.asyncio.to_thread", new_callable=AsyncMock), + patch( + "chat.tools.url_fetch.models.ChatConversationAttachment.objects.acreate", + new_callable=AsyncMock, + ), + ): + mock_settings.RAG_DOCUMENT_SEARCH_BACKEND = "some.backend" + mock_settings.URL_FETCH_BLOCKED_SCHEMES = ["http"] + mock_settings.URL_FETCH_BLOCKED_HOSTS = ["localhost", "127.0.0.1", "0.0.0.0", "::1"] + mock_settings.URL_FETCH_BLOCKED_TLDS = [".ru", ".cn", ".kp", ".ir"] + mock_sync_to_async.return_value = AsyncMock(return_value="Python is a great language.") + + get_response = Mock() + get_response.is_redirect = False + get_response.headers = {"content-type": "text/html; charset=utf-8"} + get_response.text = "Python is a great language." + get_response.raise_for_status = Mock() + + mock_client = AsyncMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + mock_client.get = AsyncMock(return_value=get_response) + + with patch("chat.tools.url_fetch.httpx.AsyncClient", return_value=mock_client): + agent = Agent("test") + add_url_fetch_tool(agent) + tool_fn = _get_tool_fn(agent) + assert tool_fn is not None + result = await tool_fn(mock_ctx, url="https://example.com") + + assert "indexed" in result.return_value + backend_instance.astore_document.assert_called_once() + backend_instance.asearch.assert_not_called() + + +@pytest.mark.asyncio +async def test_url_fetch_pdf_success(url_fetch_settings, mock_ctx, mock_rag_backend): + """url_fetch downloads PDF bytes, calls parse_and_store_document, returns confirmation.""" + backend_class, backend_instance = mock_rag_backend + + with ( + patch("chat.tools.url_fetch.import_string", return_value=backend_class), + patch("chat.tools.url_fetch.settings") as mock_settings, + patch("chat.tools.url_fetch.asyncio.to_thread", new_callable=AsyncMock) as mock_to_thread, + patch( + "chat.tools.url_fetch.models.ChatConversationAttachment.objects.acreate", + new_callable=AsyncMock, + ), + ): + mock_settings.RAG_DOCUMENT_SEARCH_BACKEND = "some.backend" + mock_settings.URL_FETCH_BLOCKED_SCHEMES = ["http"] + mock_settings.URL_FETCH_BLOCKED_HOSTS = ["localhost", "127.0.0.1", "0.0.0.0", "::1"] + mock_settings.URL_FETCH_BLOCKED_TLDS = [".ru", ".cn", ".kp", ".ir"] + # to_thread calls: save pdf, parse_and_store_document (returns parsed text), save md + mock_to_thread.side_effect = [None, "# Parsed PDF content", None] + + get_response = Mock() + get_response.is_redirect = False + get_response.headers = {"content-type": "application/pdf"} + get_response.content = b"%PDF-1.4 fake pdf content" + get_response.raise_for_status = Mock() + + mock_client = AsyncMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + mock_client.get = AsyncMock(return_value=get_response) + + with patch("chat.tools.url_fetch.httpx.AsyncClient", return_value=mock_client): + agent = Agent("test") + add_url_fetch_tool(agent) + tool_fn = _get_tool_fn(agent) + result = await tool_fn(mock_ctx, url="https://example.com/doc.pdf") + + assert "indexed" in result.return_value + assert mock_to_thread.call_count == 3 # save pdf, parse, save md + backend_instance.asearch.assert_not_called() + + +@pytest.mark.asyncio +async def test_url_fetch_blocked_url(url_fetch_settings, mock_ctx): + """Blocked URL returns error string (ModelCannotRetry caught by wrapper).""" + with patch("chat.tools.url_fetch.import_string"): + agent = Agent("test") + add_url_fetch_tool(agent) + tool_fn = _get_tool_fn(agent) + + result = await tool_fn(mock_ctx, url="http://localhost/evil") + assert "not allowed for security reasons" in result + + +@pytest.mark.asyncio +async def test_url_fetch_401_raises_cannot_retry(url_fetch_settings, mock_ctx): + """HTTP 401 returns auth error string (ModelCannotRetry caught by wrapper).""" + with patch("chat.tools.url_fetch.import_string"): + agent = Agent("test") + add_url_fetch_tool(agent) + tool_fn = _get_tool_fn(agent) + + mock_response = Mock() + mock_response.status_code = 401 + mock_client = AsyncMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + mock_client.get = AsyncMock( + side_effect=httpx.HTTPStatusError("401", request=Mock(), response=mock_response) + ) + with patch("chat.tools.url_fetch.httpx.AsyncClient", return_value=mock_client): + result = await tool_fn(mock_ctx, url="https://private.example.com/doc") + assert "requires authentication" in result + + +@pytest.mark.asyncio +async def test_url_fetch_403_raises_cannot_retry(url_fetch_settings, mock_ctx): + """HTTP 403 returns auth error string (ModelCannotRetry caught by wrapper).""" + with patch("chat.tools.url_fetch.import_string"): + agent = Agent("test") + add_url_fetch_tool(agent) + tool_fn = _get_tool_fn(agent) + + mock_response = Mock() + mock_response.status_code = 403 + mock_client = AsyncMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + mock_client.get = AsyncMock( + side_effect=httpx.HTTPStatusError("403", request=Mock(), response=mock_response) + ) + with patch("chat.tools.url_fetch.httpx.AsyncClient", return_value=mock_client): + result = await tool_fn(mock_ctx, url="https://private.example.com/doc") + assert "requires authentication" in result + + +@pytest.mark.asyncio +async def test_url_fetch_404_raises_cannot_retry(url_fetch_settings, mock_ctx): + """HTTP 404 returns not-found string (ModelCannotRetry caught by wrapper).""" + with patch("chat.tools.url_fetch.import_string"): + agent = Agent("test") + add_url_fetch_tool(agent) + tool_fn = _get_tool_fn(agent) + + mock_response = Mock() + mock_response.status_code = 404 + mock_client = AsyncMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + mock_client.get = AsyncMock( + side_effect=httpx.HTTPStatusError("404", request=Mock(), response=mock_response) + ) + with patch("chat.tools.url_fetch.httpx.AsyncClient", return_value=mock_client): + result = await tool_fn(mock_ctx, url="https://example.com/missing") + assert "was not found" in result + + +@pytest.mark.asyncio +async def test_url_fetch_5xx_raises_model_retry(url_fetch_settings, mock_ctx): + """HTTP 500 raises ModelRetry (retryable).""" + with patch("chat.tools.url_fetch.import_string"): + agent = Agent("test") + add_url_fetch_tool(agent) + tool_fn = _get_tool_fn(agent) + + mock_response = Mock() + mock_response.status_code = 500 + mock_client = AsyncMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + mock_client.get = AsyncMock( + side_effect=httpx.HTTPStatusError("500", request=Mock(), response=mock_response) + ) + with ( + patch("chat.tools.url_fetch.httpx.AsyncClient", return_value=mock_client), + pytest.raises(ModelRetry), + ): + await tool_fn(mock_ctx, url="https://example.com/broken") + + +@pytest.mark.asyncio +async def test_url_fetch_timeout_raises_model_retry(url_fetch_settings, mock_ctx): + """Network timeout raises ModelRetry (retryable).""" + with patch("chat.tools.url_fetch.import_string"): + agent = Agent("test") + add_url_fetch_tool(agent) + tool_fn = _get_tool_fn(agent) + + mock_client = AsyncMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + mock_client.get = AsyncMock(side_effect=httpx.TimeoutException("timeout")) + with ( + patch("chat.tools.url_fetch.httpx.AsyncClient", return_value=mock_client), + pytest.raises(ModelRetry, match="timed out"), + ): + await tool_fn(mock_ctx, url="https://slow.example.com/") + + +@pytest.mark.asyncio +async def test_url_fetch_unsupported_content_type(url_fetch_settings, mock_ctx): + """Unsupported content type returns error string (ModelCannotRetry caught by wrapper).""" + with patch("chat.tools.url_fetch.import_string"): + agent = Agent("test") + add_url_fetch_tool(agent) + tool_fn = _get_tool_fn(agent) + + for unsupported_ct in [ + "video/mp4", + "audio/mpeg", + "image/png", + "application/zip", + "application/octet-stream", + ]: + mock_response = Mock() + mock_response.is_redirect = False + mock_response.headers = {"content-type": unsupported_ct} + mock_response.raise_for_status = Mock() + mock_client = AsyncMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + mock_client.get = AsyncMock(return_value=mock_response) + with patch("chat.tools.url_fetch.httpx.AsyncClient", return_value=mock_client): + result = await tool_fn(mock_ctx, url="https://example.com/file") + assert "cannot analyse" in result + + +@pytest.mark.asyncio +async def test_url_fetch_empty_trafilatura_extraction( + url_fetch_settings, mock_ctx, mock_rag_backend +): + """Empty trafilatura extraction raises ModelCannotRetry.""" + backend_class, backend_instance = mock_rag_backend + + with ( + patch("chat.tools.url_fetch.import_string", return_value=backend_class), + patch("chat.tools.url_fetch.settings") as mock_settings, + patch("chat.tools.url_fetch.sync_to_async") as mock_sync_to_async, + ): + mock_settings.RAG_DOCUMENT_SEARCH_BACKEND = "some.backend" + mock_settings.URL_FETCH_BLOCKED_SCHEMES = ["http"] + mock_settings.URL_FETCH_BLOCKED_HOSTS = ["localhost", "127.0.0.1", "0.0.0.0", "::1"] + mock_settings.URL_FETCH_BLOCKED_TLDS = [".ru", ".cn", ".kp", ".ir"] + mock_sync_to_async.return_value = AsyncMock(return_value=None) + + get_response = Mock() + get_response.is_redirect = False + get_response.headers = {"content-type": "text/html"} + get_response.text = "" + get_response.raise_for_status = Mock() + + mock_client = AsyncMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + mock_client.get = AsyncMock(return_value=get_response) + + with patch("chat.tools.url_fetch.httpx.AsyncClient", return_value=mock_client): + agent = Agent("test") + add_url_fetch_tool(agent) + tool_fn = _get_tool_fn(agent) + + result = await tool_fn(mock_ctx, url="https://spa.example.com/") + assert "did not return readable content" in result + + +@pytest.mark.asyncio +async def test_url_fetch_creates_collection_when_missing(url_fetch_settings, mock_rag_backend): + """_ensure_collection creates collection when collection_id is empty.""" + backend_class, backend_instance = mock_rag_backend + backend_instance.collection_id = "" # no collection yet + + conv = MagicMock() + conv.collection_id = "" + conv.pk = "conv-pk-1" + conv.asave = AsyncMock() + + ctx = Mock(spec=RunContext) + ctx.deps = Mock() + ctx.deps.conversation = conv + ctx.deps.user = Mock() + ctx.deps.user.sub = "user-sub-123" + ctx.deps.session = {} + ctx.usage = RunUsage(input_tokens=0, output_tokens=0) + ctx.retries = {} + ctx.max_retries = 2 + ctx.tool_name = "url_fetch" + user_msg = Mock() + user_msg.kind = "request" + user_prompt_part = Mock() + user_prompt_part.part_kind = "user-prompt" + user_prompt_part.content = "Check https://example.com" + user_msg.parts = [user_prompt_part] + ctx.messages = [user_msg] + + with ( + patch("chat.tools.url_fetch.import_string", return_value=backend_class), + patch("chat.tools.url_fetch.settings") as mock_settings, + patch("chat.tools.url_fetch.sync_to_async") as mock_sync_to_async, + patch("chat.tools.url_fetch.asyncio.to_thread", new_callable=AsyncMock), + patch( + "chat.tools.url_fetch.models.ChatConversationAttachment.objects.acreate", + new_callable=AsyncMock, + ), + ): + mock_settings.RAG_DOCUMENT_SEARCH_BACKEND = "some.backend" + mock_settings.URL_FETCH_BLOCKED_SCHEMES = ["http"] + mock_settings.URL_FETCH_BLOCKED_HOSTS = ["localhost", "127.0.0.1", "0.0.0.0", "::1"] + mock_settings.URL_FETCH_BLOCKED_TLDS = [".ru", ".cn", ".kp", ".ir"] + mock_sync_to_async.return_value = AsyncMock(return_value="Python content.") + + get_response = Mock() + get_response.is_redirect = False + get_response.headers = {"content-type": "text/html"} + get_response.text = "Python content." + get_response.raise_for_status = Mock() + + mock_client = AsyncMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + mock_client.get = AsyncMock(return_value=get_response) + + with patch("chat.tools.url_fetch.httpx.AsyncClient", return_value=mock_client): + agent = Agent("test") + add_url_fetch_tool(agent) + tool_fn = _get_tool_fn(agent) + await tool_fn(ctx, url="https://example.com") + + backend_instance.acreate_collection.assert_called_once_with(name="conversation-conv-pk-1") + conv.asave.assert_called_once_with(update_fields=["collection_id", "updated_at"]) + + +@pytest.mark.asyncio +async def test_url_fetch_network_error_raises_model_retry(url_fetch_settings, mock_ctx): + """Generic network error raises ModelRetry (retryable).""" + with patch("chat.tools.url_fetch.import_string"): + agent = Agent("test") + add_url_fetch_tool(agent) + tool_fn = _get_tool_fn(agent) + + mock_client = AsyncMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + mock_client.get = AsyncMock(side_effect=httpx.ConnectError("connection refused")) + with ( + patch("chat.tools.url_fetch.httpx.AsyncClient", return_value=mock_client), + pytest.raises(ModelRetry, match="Network error"), + ): + await tool_fn(mock_ctx, url="https://down.example.com/") + + +@pytest.mark.asyncio +async def test_url_fetch_redirect_to_blocked_url(url_fetch_settings, mock_ctx): + """Redirect to a blocked URL returns error string.""" + with patch("chat.tools.url_fetch.import_string"): + agent = Agent("test") + add_url_fetch_tool(agent) + tool_fn = _get_tool_fn(agent) + + redirect_response = Mock() + redirect_response.is_redirect = True + redirect_response.headers = {"location": "http://169.254.169.254/meta-data/"} + redirect_response.raise_for_status = Mock() + + mock_client = AsyncMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + mock_client.get = AsyncMock(return_value=redirect_response) + + with patch("chat.tools.url_fetch.httpx.AsyncClient", return_value=mock_client): + result = await tool_fn(mock_ctx, url="https://example.com/redirect") + assert "not allowed" in result + + +def test_instructions_no_url_returns_empty(url_fetch_settings): + """Instructions closure returns empty string when no URL in latest message.""" + agent = Agent("test") + add_url_fetch_tool(agent) + + ctx = Mock() + user_msg = Mock() + user_msg.kind = "request" + part = Mock() + part.part_kind = "user-prompt" + part.content = "Tell me about Python" + user_msg.parts = [part] + ctx.messages = [user_msg] + + instructions_fn = None + for fn in agent._instructions: + if "url_fetch" in getattr(fn, "__name__", ""): + instructions_fn = fn + break + assert instructions_fn is not None + result = instructions_fn(ctx) + assert result == "" + + +def test_instructions_blocked_url_returns_empty(url_fetch_settings): + """Instructions closure returns empty string when URL is blocked.""" + agent = Agent("test") + add_url_fetch_tool(agent) + + ctx = Mock() + user_msg = Mock() + user_msg.kind = "request" + part = Mock() + part.part_kind = "user-prompt" + part.content = "Check http://localhost/admin" + user_msg.parts = [part] + ctx.messages = [user_msg] + + instructions_fn = None + for fn in agent._instructions: + if "url_fetch" in getattr(fn, "__name__", ""): + instructions_fn = fn + break + + result = instructions_fn(ctx) + assert result == "" + + +def test_instructions_injects_only_allowed_urls(url_fetch_settings): + """Instructions closure lists only the allowed URL when message has a blocked one too.""" + agent = Agent("test") + add_url_fetch_tool(agent) + + ctx = Mock() + user_msg = Mock() + user_msg.kind = "request" + part = Mock() + part.part_kind = "user-prompt" + part.content = "See https://example.com and http://localhost" + user_msg.parts = [part] + ctx.messages = [user_msg] + + instructions_fn = None + for fn in agent._instructions: + if "url_fetch" in getattr(fn, "__name__", ""): + instructions_fn = fn + break + + result = instructions_fn(ctx) + assert "https://example.com" in result + assert "localhost" not in result diff --git a/src/backend/chat/tools/descriptions.py b/src/backend/chat/tools/descriptions.py index 9c0b6a559..3f09e940c 100644 --- a/src/backend/chat/tools/descriptions.py +++ b/src/backend/chat/tools/descriptions.py @@ -86,3 +86,20 @@ "internet access, accepted files, or hosting, call the " "self_documentation tool before answering." ) + +URL_FETCH_TOOL_DESCRIPTION = """ +Fetch and read the content of a specific URL provided by the user. + +Use this tool ONLY when the user's message contains an explicit URL (starting with https://). +Do NOT use this tool for general searches — use web_search for that. +Do NOT call this tool for a URL that was mentioned in a previous turn unless the user explicitly asks again. + +The tool fetches the page and indexes it into the conversation's document store. +Once indexed, use the document search tool to retrieve relevant excerpts and answer the user's question. +""" + +URL_FETCH_SYSTEM_PROMPT = ( + "The user's message contains the following URLs you must fetch using url_fetch: `{urls}`. " + "Only fetch URLs from this list. " + "Do NOT call url_fetch for any URL not explicitly listed here." +) diff --git a/src/backend/chat/tools/url_fetch.py b/src/backend/chat/tools/url_fetch.py new file mode 100644 index 000000000..f43d7d385 --- /dev/null +++ b/src/backend/chat/tools/url_fetch.py @@ -0,0 +1,266 @@ +"""URL fetch tool for the conversation agent.""" + +import asyncio +import ipaddress +import logging +import re +import uuid +from io import BytesIO +from urllib.parse import urlparse + +from django.conf import settings +from django.core.files.storage import default_storage +from django.utils.module_loading import import_string + +import httpx +from asgiref.sync import sync_to_async +from pydantic_ai import Agent, RunContext +from pydantic_ai.exceptions import ModelRetry +from pydantic_ai.messages import ToolReturn +from trafilatura import extract + +from chat import models +from chat.tools.descriptions import URL_FETCH_SYSTEM_PROMPT, URL_FETCH_TOOL_DESCRIPTION +from chat.tools.exceptions import ModelCannotRetry +from chat.tools.utils import last_model_retry_soft_fail + +logger = logging.getLogger(__name__) + +# Regex: match http(s):// followed by non-whitespace characters +_URL_PATTERN = re.compile(r"https?://\S+") +# Strip trailing punctuation that is likely sentence punctuation, not part of the URL +_TRAILING_PUNCTUATION = re.compile(r"[.,;:!?\"'`\]})>]+$") + + +def detect_urls(text: str) -> list[str]: + """Extract all HTTP/HTTPS URLs from a string, stripping trailing punctuation.""" + return [_TRAILING_PUNCTUATION.sub("", u) for u in _URL_PATTERN.findall(text)] + + +def _is_private_address(hostname: str) -> bool: + """Return True if hostname is a private/link-local/loopback IP address.""" + try: + addr = ipaddress.ip_address(hostname) + return addr.is_private or addr.is_loopback or addr.is_link_local or addr.is_reserved + except ValueError: + # Not a bare IP address (it's a domain name) — let it through + # DNS resolution is not performed here to avoid blocking on DNS + return False + + +def is_url_allowed(url: str) -> bool: + """Return True if the URL passes all security filters.""" + try: + parsed = urlparse(url) + except ValueError: + return False + + if not parsed.scheme or not parsed.netloc: + return False + + hostname = parsed.hostname or "" + blocked = ( + parsed.scheme in settings.URL_FETCH_BLOCKED_SCHEMES + or hostname in settings.URL_FETCH_BLOCKED_HOSTS + or _is_private_address(hostname) + or any(hostname.endswith(tld) for tld in settings.URL_FETCH_BLOCKED_TLDS) + ) + return not blocked + + +def _get_last_user_text(ctx: RunContext) -> str: + """Extract the text of the latest user message from RunContext messages.""" + for msg in reversed(ctx.messages): + if msg.kind == "request": + for part in msg.parts: + if part.part_kind == "user-prompt": + content = part.content + if isinstance(content, str): + return content + parts = [] + for item in content: + if isinstance(item, str): + parts.append(item) + elif hasattr(item, "text"): + parts.append(item.text) + return " ".join(filter(None, parts)) + return "" + + +async def _ensure_collection(conversation, document_store_backend) -> object: + """Return a document store instance, creating the collection if needed.""" + document_store = document_store_backend(conversation.collection_id) + if not document_store.collection_id: + collection_id = await document_store.acreate_collection( + name=f"conversation-{conversation.pk}" + ) + conversation.collection_id = str(collection_id) + await conversation.asave(update_fields=["collection_id", "updated_at"]) + document_store = document_store_backend(conversation.collection_id) + return document_store + + +async def _fetch_url(url: str) -> httpx.Response: + """Fetch a URL with SSRF-safe redirect handling; raises ModelCannotRetry/ModelRetry on error.""" + try: + async with httpx.AsyncClient(timeout=30, follow_redirects=False) as client: + resp = await client.get(url) + if resp.is_redirect: + redirect_url = resp.headers.get("location", "") + if not is_url_allowed(redirect_url): + raise ModelCannotRetry( + f"The URL {url} redirects to a location that is not allowed " + "for security reasons." + ) + resp = await client.get(redirect_url, follow_redirects=False) + if resp.is_redirect: + raise ModelCannotRetry(f"Too many redirects for {url}.") + resp.raise_for_status() + return resp + except httpx.HTTPStatusError as exc: + status = exc.response.status_code + if status in (401, 403): + logger.warning("URL requires authentication: %s (HTTP %s)", url, status) + raise ModelCannotRetry( + f"The page at {url} requires authentication. I cannot access it." + ) from exc + if status == 404: + logger.warning("URL not found: %s (HTTP 404)", url) + raise ModelCannotRetry(f"The page at {url} was not found (404).") from exc + if status >= 500: + logger.warning("Server error fetching %s (HTTP %s)", url, status) + raise ModelRetry( + f"The server at {url} returned an error ({status}). Retrying..." + ) from exc + raise ModelCannotRetry(f"Failed to access {url} (HTTP {status}).") from exc + except httpx.TimeoutException as exc: + logger.warning("Timeout fetching %s", url) + raise ModelRetry(f"Request to {url} timed out. Retrying...") from exc + except httpx.HTTPError as exc: + logger.warning("HTTP error fetching %s: %s", url, type(exc).__name__) + raise ModelRetry( + f"Network error accessing {url}: {type(exc).__name__}. Retrying..." + ) from exc + + +async def _store_html_content( # noqa: PLR0913 + url: str, extracted: str, conversation, user, session: dict, document_store +) -> None: + """Save extracted text as a text/plain attachment and index it in the RAG store.""" + content_bytes = extracted.encode("utf-8") + key = f"{conversation.pk}/attachments/url-{uuid.uuid4()}.txt" + await asyncio.to_thread(default_storage.save, key, BytesIO(content_bytes)) + await models.ChatConversationAttachment.objects.acreate( + conversation=conversation, + uploaded_by=user, + key=key, + file_name=url[:255], + content_type="text/plain", + size=len(content_bytes), + upload_state=models.AttachmentStatus.READY, + ) + await document_store.astore_document(url, extracted, user_sub=user.sub, session=session) + + +async def _store_pdf_content( # noqa: PLR0913 + url: str, content: bytes, conversation, user, session: dict, document_store +) -> None: + """Save PDF as attachment, parse and index it, and create a markdown attachment.""" + pdf_key = f"{conversation.pk}/attachments/url-{uuid.uuid4()}.pdf" + await asyncio.to_thread(default_storage.save, pdf_key, BytesIO(content)) + await models.ChatConversationAttachment.objects.acreate( + conversation=conversation, + uploaded_by=user, + key=pdf_key, + file_name=url[:255], + content_type="application/pdf", + size=len(content), + upload_state=models.AttachmentStatus.READY, + ) + parsed_content = await asyncio.to_thread( + document_store.parse_and_store_document, + name=url, + content_type="application/pdf", + content=content, + user_sub=user.sub, + session=session, + ) + if parsed_content: + md_key = f"{conversation.pk}/attachments/url-{uuid.uuid4()}.md" + md_bytes = parsed_content.encode("utf-8") + await asyncio.to_thread(default_storage.save, md_key, BytesIO(md_bytes)) + await models.ChatConversationAttachment.objects.acreate( + conversation=conversation, + uploaded_by=user, + key=md_key, + file_name=f"{url}.md"[:255], + content_type="text/markdown", + size=len(md_bytes), + conversion_from=pdf_key, + upload_state=models.AttachmentStatus.READY, + ) + + +def add_url_fetch_tool(agent: Agent) -> None: + """Register the url_fetch tool and its instructions hint on the given agent.""" + + @agent.tool( + name="url_fetch", + retries=1, + description=URL_FETCH_TOOL_DESCRIPTION, + ) + @last_model_retry_soft_fail + async def url_fetch(ctx: RunContext, url: str) -> ToolReturn: + """ + Args: + ctx (RunContext): The run context containing the conversation. + url (str): The URL to fetch and index. + """ + logger.debug("Fetching URL: %s", url) + + if not is_url_allowed(url): + raise ModelCannotRetry(f"This URL is not allowed for security reasons ({url}).") + + resp = await _fetch_url(url) + mime = resp.headers.get("content-type", "").split(";")[0].strip() + + supported_text = {"text/html", "text/plain"} + supported_pdf = {"application/pdf"} + if mime not in supported_text | supported_pdf: + raise ModelCannotRetry( + f"The link points to a {mime} file. I cannot analyse this type of content yet." + ) + + document_store_backend = import_string(settings.RAG_DOCUMENT_SEARCH_BACKEND) + conversation = ctx.deps.conversation + user = ctx.deps.user + session = ctx.deps.session + document_store = await _ensure_collection(conversation, document_store_backend) + + if mime in supported_text: + extracted = await sync_to_async(extract)( + resp.text, include_comments=False, no_fallback=True + ) + if not extracted: + logger.warning("Empty extraction from %s (JS-rendered or auth-protected)", url) + raise ModelCannotRetry( + f"The page at {url} did not return readable content. " + "It may require authentication or be a JavaScript-rendered page." + ) + await _store_html_content(url, extracted, conversation, user, session, document_store) + else: + await _store_pdf_content(url, resp.content, conversation, user, session, document_store) + + return ToolReturn(return_value=f"Content from {url} has been indexed successfully.") + + @agent.instructions + def url_fetch_instructions(ctx: RunContext) -> str: + """Inject URL hint when the latest user message contains allowed URLs.""" + text = _get_last_user_text(ctx) + if not text: + return "" + allowed = [u for u in detect_urls(text) if is_url_allowed(u)] + if not allowed: + return "" + url_list = ", ".join(allowed) + return URL_FETCH_SYSTEM_PROMPT.format(urls=url_list) diff --git a/src/backend/conversations/settings.py b/src/backend/conversations/settings.py index 1b5c0d3e9..8cff51ba8 100755 --- a/src/backend/conversations/settings.py +++ b/src/backend/conversations/settings.py @@ -29,6 +29,7 @@ from chat.llm_configuration import cached_load_llm_configuration, load_llm_configuration from conversations.brave_settings import BraveSettings +from conversations.url_fetch_settings import UrlFetchSettings # Build paths inside the project like this: BASE_DIR / 'subdir'. BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) @@ -47,7 +48,7 @@ def get_release(): return "NA" # Default: not available -class Base(BraveSettings, Configuration): +class Base(BraveSettings, UrlFetchSettings, Configuration): """ This is the base configuration every configuration (aka environment) should inherit from. It is recommended to configure third-party applications by creating a configuration mixins in diff --git a/src/backend/conversations/url_fetch_settings.py b/src/backend/conversations/url_fetch_settings.py new file mode 100644 index 000000000..21f4a5f78 --- /dev/null +++ b/src/backend/conversations/url_fetch_settings.py @@ -0,0 +1,23 @@ +"""Django configuration mixin for URL fetch settings.""" + +from configurations import values + + +class UrlFetchSettings: + """Settings for the url_fetch tool.""" + + URL_FETCH_BLOCKED_SCHEMES = values.ListValue( + default=["http"], + environ_name="URL_FETCH_BLOCKED_SCHEMES", + environ_prefix=None, + ) + URL_FETCH_BLOCKED_HOSTS = values.ListValue( + default=["localhost", "127.0.0.1", "0.0.0.0", "::1"], # noqa: S104 + environ_name="URL_FETCH_BLOCKED_HOSTS", + environ_prefix=None, + ) + URL_FETCH_BLOCKED_TLDS = values.ListValue( + default=[".ru", ".cn", ".kp", ".ir"], + environ_name="URL_FETCH_BLOCKED_TLDS", + environ_prefix=None, + )