From a9068faefa4ec72f73c117d520b96407576fd77c Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 22 Jun 2026 13:13:51 +0100 Subject: [PATCH] feat: small mixes on mixes sources --- application/agents/agentic_agent.py | 22 +++++++++-- application/agents/classic_agent.py | 22 +++++++++-- application/agents/tools/internal_search.py | 6 +++ .../api/answer/services/stream_processor.py | 31 +++++++++++++++ tests/agents/test_classic_agent.py | 38 +++++++++++++++++++ tests/agents/test_internal_search_tool.py | 27 +++++++++++++ .../answer/test_stream_processor_exposure.py | 31 ++++++++++++++- 7 files changed, 170 insertions(+), 7 deletions(-) diff --git a/application/agents/agentic_agent.py b/application/agents/agentic_agent.py index fe08277a..31a7d1da 100644 --- a/application/agents/agentic_agent.py +++ b/application/agents/agentic_agent.py @@ -56,8 +56,24 @@ class AgenticAgent(BaseAgent): ) def _collect_internal_sources(self): - """Collect retrieved docs from the cached InternalSearchTool instance.""" + """Merge the cached InternalSearchTool's docs into ``retrieved_docs``, + deduped, preserving any pre-fetched docs so a mixed-exposure agent cites + both pre-fetched and tool-retrieved sources (not just the tool's).""" cache_key = f"internal_search:{INTERNAL_TOOL_ID}:{self.user or ''}" tool = self.tool_executor._loaded_tools.get(cache_key) - if tool and hasattr(tool, "retrieved_docs") and tool.retrieved_docs: - self.retrieved_docs = tool.retrieved_docs + if not (tool and getattr(tool, "retrieved_docs", None)): + return + + def _key(d): + if isinstance(d, dict): + return (d.get("source"), d.get("title"), d.get("text")) + return id(d) + + merged = list(self.retrieved_docs or []) + seen = {_key(d) for d in merged} + for doc in tool.retrieved_docs: + k = _key(doc) + if k not in seen: + seen.add(k) + merged.append(doc) + self.retrieved_docs = merged diff --git a/application/agents/classic_agent.py b/application/agents/classic_agent.py index 45621cd1..ea8a32b7 100644 --- a/application/agents/classic_agent.py +++ b/application/agents/classic_agent.py @@ -58,8 +58,24 @@ class ClassicAgent(BaseAgent): ) def _collect_internal_sources(self): - """Collect retrieved docs from the cached InternalSearchTool instance.""" + """Merge the cached InternalSearchTool's docs into ``retrieved_docs``, + deduped, preserving any pre-fetched docs so a mixed-exposure agent cites + both pre-fetched and tool-retrieved sources (not just the tool's).""" cache_key = f"internal_search:{INTERNAL_TOOL_ID}:{self.user or ''}" tool = self.tool_executor._loaded_tools.get(cache_key) - if tool and hasattr(tool, "retrieved_docs") and tool.retrieved_docs: - self.retrieved_docs = tool.retrieved_docs + if not (tool and getattr(tool, "retrieved_docs", None)): + return + + def _key(d): + if isinstance(d, dict): + return (d.get("source"), d.get("title"), d.get("text")) + return id(d) + + merged = list(self.retrieved_docs or []) + seen = {_key(d) for d in merged} + for doc in tool.retrieved_docs: + k = _key(doc) + if k not in seen: + seen.add(k) + merged.append(doc) + self.retrieved_docs = merged diff --git a/application/agents/tools/internal_search.py b/application/agents/tools/internal_search.py index d711b0de..1baf1eda 100644 --- a/application/agents/tools/internal_search.py +++ b/application/agents/tools/internal_search.py @@ -450,6 +450,10 @@ def add_internal_search_tool(tools_dict: Dict, retriever_config: Dict) -> None: has_dir = sources_have_directory_structure(source) internal_entry = build_internal_tool_entry(has_directory_structure=has_dir) + # The executor resolves a tool row by ``id``; the internal tool is synthetic + # (no DB row), so stamp its sentinel id or _get_or_load_tool drops it with + # ``tool_missing_row_id``. + internal_entry["id"] = INTERNAL_TOOL_ID internal_entry["config"] = build_internal_tool_config( **retriever_config, has_directory_structure=has_dir, @@ -471,6 +475,7 @@ def build_internal_tool_config( llm_name: str = None, api_key: str = None, decoded_token: Optional[Dict] = None, + request_id: Optional[str] = None, has_directory_structure: bool = False, ) -> Dict: """Build the config dict for InternalSearchTool.""" @@ -492,5 +497,6 @@ def build_internal_tool_config( "llm_name": llm_name or settings.LLM_PROVIDER, "api_key": api_key or settings.API_KEY, "decoded_token": decoded_token, + "request_id": request_id, "has_directory_structure": has_directory_structure, } diff --git a/application/api/answer/services/stream_processor.py b/application/api/answer/services/stream_processor.py index bd4b78c1..39c170f5 100644 --- a/application/api/answer/services/stream_processor.py +++ b/application/api/answer/services/stream_processor.py @@ -666,12 +666,43 @@ class StreamProcessor: active_docs = self.data["active_docs"] if active_docs and active_docs != "default": self.source = {"active_docs": active_docs} + self.all_sources = self._load_request_sources(active_docs) else: self.source = {} + self.all_sources = [] return self.source = {} self.all_sources = [] + def _load_request_sources(self, active_docs) -> list: + """Per-source list (with each source's retrieval config) for a non-agent + request, so per-source overrides (exposure, chunks, ...) are honored on + the default chat just like the agent path. Lenient read: a missing or + inaccessible source falls back to default config and never raises. + """ + owner = self.initial_user_id + ids = active_docs if isinstance(active_docs, list) else [active_docs] + sources = [] + for sid in ids: + if not sid or sid == "default": + continue + source_doc = None + if owner: + try: + with db_readonly() as conn: + source_doc = SourcesRepository(conn).get(str(sid), owner) + except Exception: + source_doc = None + sources.append( + { + "id": sid, + "retrieval": SourceConfig.parse( + (source_doc or {}).get("config") + ).retrieval, + } + ) + return sources + def _has_active_docs(self) -> bool: """Return True if a real document source is configured for retrieval.""" active_docs = self.source.get("active_docs") if self.source else None diff --git a/tests/agents/test_classic_agent.py b/tests/agents/test_classic_agent.py index 50a14230..a189b20e 100644 --- a/tests/agents/test_classic_agent.py +++ b/tests/agents/test_classic_agent.py @@ -290,6 +290,44 @@ class TestClassicAgentSearchExposure: agent._collect_internal_sources() assert [d["title"] for d in agent.retrieved_docs] == ["Doc"] + def test_collect_internal_sources_merges_with_prefetched( + self, agent_base_params, mock_llm_creator, mock_llm_handler_creator + ): + # Mixed exposure: a pre-fetched source's docs must survive alongside the + # tool-retrieved docs (not be overwritten), so both are cited. + retriever_config = {"source": {"active_docs": ["b"]}} + agent = ClassicAgent(retriever_config=retriever_config, **agent_base_params) + agent.retrieved_docs = [ + {"text": "Prefetched", "title": "Prefetched Doc", "source": "a"} + ] + mock_tool = Mock() + mock_tool.retrieved_docs = [ + {"text": "Found", "title": "Tool Doc", "source": "b"} + ] + cache_key = f"internal_search:{INTERNAL_TOOL_ID}:{agent.user or ''}" + agent.tool_executor._loaded_tools[cache_key] = mock_tool + + agent._collect_internal_sources() + assert [d["title"] for d in agent.retrieved_docs] == [ + "Prefetched Doc", + "Tool Doc", + ] + + def test_collect_internal_sources_dedupes( + self, agent_base_params, mock_llm_creator, mock_llm_handler_creator + ): + retriever_config = {"source": {"active_docs": ["b"]}} + agent = ClassicAgent(retriever_config=retriever_config, **agent_base_params) + dup = {"text": "X", "title": "Dup", "source": "b"} + agent.retrieved_docs = [dup] + mock_tool = Mock() + mock_tool.retrieved_docs = [dict(dup)] + cache_key = f"internal_search:{INTERNAL_TOOL_ID}:{agent.user or ''}" + agent.tool_executor._loaded_tools[cache_key] = mock_tool + + agent._collect_internal_sources() + assert [d["title"] for d in agent.retrieved_docs] == ["Dup"] + @pytest.mark.integration class TestClassicAgentIntegration: diff --git a/tests/agents/test_internal_search_tool.py b/tests/agents/test_internal_search_tool.py index 92a97541..e36ad5e1 100644 --- a/tests/agents/test_internal_search_tool.py +++ b/tests/agents/test_internal_search_tool.py @@ -12,6 +12,33 @@ from application.agents.tools.internal_search import ( ) +@pytest.mark.unit +class TestAddInternalSearchTool: + def test_stamps_id_so_executor_can_load(self): + # The executor resolves a tool by its row ``id``; the synthetic internal + # tool must carry the sentinel id or it is dropped (tool_missing_row_id). + config = { + "source": {"active_docs": ["s1"]}, + "retriever_name": "classic", + "chunks": 2, + } + tools: dict = {} + with patch( + "application.agents.tools.internal_search." + "sources_have_directory_structure", + return_value=False, + ): + add_internal_search_tool(tools, config) + assert INTERNAL_TOOL_ID in tools + assert tools[INTERNAL_TOOL_ID]["id"] == INTERNAL_TOOL_ID + assert tools[INTERNAL_TOOL_ID]["name"] == "internal_search" + + def test_no_active_docs_adds_nothing(self): + tools: dict = {} + add_internal_search_tool(tools, {"source": {}}) + assert tools == {} + + @pytest.mark.unit class TestInternalSearchToolSearch: diff --git a/tests/api/answer/test_stream_processor_exposure.py b/tests/api/answer/test_stream_processor_exposure.py index 7133d5cd..0fc35d8d 100644 --- a/tests/api/answer/test_stream_processor_exposure.py +++ b/tests/api/answer/test_stream_processor_exposure.py @@ -2,7 +2,7 @@ from __future__ import annotations -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest @@ -146,3 +146,32 @@ class TestBuildAgentExposure: assert pf_kwargs.get("exposure") == "prefetch" _, kwargs = sp.create_agent.call_args assert [e["id"] for e in kwargs["agentic_sources"]] == ["b"] + + +@pytest.mark.unit +class TestNonAgentSourceConfig: + """The non-agent chat path must also load per-source config so exposure + (and other per-source overrides) are honored, not just the agent path.""" + + def test_active_docs_loads_per_source_retrieval_config(self): + sp = StreamProcessor.__new__(StreamProcessor) + sp._agent_data = None + sp.data = {"active_docs": "s1"} + sp.initial_user_id = "user1" + sp.source = {} + sp.all_sources = [] + fake_source = { + "config": {"retrieval": {"exposure": "agentic_tool", "chunks": 7}} + } + with patch( + "application.api.answer.services.stream_processor.db_readonly" + ), patch( + "application.api.answer.services.stream_processor.SourcesRepository" + ) as repo: + repo.return_value.get.return_value = fake_source + sp._configure_source() + assert sp.source == {"active_docs": "s1"} + assert len(sp.all_sources) == 1 + assert sp.all_sources[0]["id"] == "s1" + assert sp.all_sources[0]["retrieval"].exposure == "agentic_tool" + assert sp.all_sources[0]["retrieval"].chunks == 7