From babc067aa4d11d4f56068bab54c30ab227804289 Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 22 Jun 2026 10:40:33 +0100 Subject: [PATCH] feat: remove old chunk management --- application/agents/classic_agent.py | 36 +++++- application/agents/tools/internal_search.py | 1 + .../api/answer/services/stream_processor.py | 41 ++++++- application/retriever/classic_rag.py | 5 +- application/retriever/dispatcher.py | 3 + application/retriever/stages/prescreen.py | 8 +- frontend/src/agents/NewAgent.tsx | 26 ----- .../src/agents/workflow/WorkflowBuilder.tsx | 5 +- frontend/src/locale/en.json | 6 +- frontend/src/settings/General.tsx | 24 ---- .../settings/components/RetrievalOptions.tsx | 20 ++-- tests/agents/test_classic_agent.py | 110 ++++++++++++++++++ .../answer/test_stream_processor_exposure.py | 48 ++++++-- tests/retriever/test_prescreen.py | 24 ++++ tests/test_retriever.py | 7 ++ 15 files changed, 276 insertions(+), 88 deletions(-) diff --git a/application/agents/classic_agent.py b/application/agents/classic_agent.py index d6dced9e..45621cd1 100644 --- a/application/agents/classic_agent.py +++ b/application/agents/classic_agent.py @@ -1,14 +1,34 @@ import logging -from typing import Dict, Generator +from typing import Dict, Generator, Optional from application.agents.base import BaseAgent +from application.agents.tools.internal_search import ( + INTERNAL_TOOL_ID, + add_internal_search_tool, +) from application.logging import LogContext logger = logging.getLogger(__name__) class ClassicAgent(BaseAgent): - """A simplified agent with clear execution flow""" + """A simplified agent with clear execution flow. + + Pre-fetches ``prefetch`` sources into the prompt and, when a + ``retriever_config`` is supplied, also exposes ``agentic_tool`` sources + via the internal_search tool. With no ``retriever_config`` (every source + at the default ``prefetch`` exposure) no search tool is added and behavior + is identical to plain pre-fetch. + """ + + def __init__( + self, + retriever_config: Optional[Dict] = None, + *args, + **kwargs, + ): + super().__init__(*args, **kwargs) + self.retriever_config = retriever_config or {} def _gen_inner( self, query: str, log_context: LogContext @@ -16,6 +36,8 @@ class ClassicAgent(BaseAgent): """Core generator function for ClassicAgent execution flow""" tools_dict = self.tool_executor.get_tools() + if self.retriever_config: + add_internal_search_tool(tools_dict, self.retriever_config) self._prepare_tools(tools_dict) messages = self._build_messages(self.prompt, query) @@ -25,9 +47,19 @@ class ClassicAgent(BaseAgent): llm_response, tools_dict, messages, log_context ) + if self.retriever_config: + self._collect_internal_sources() + yield {"sources": self.retrieved_docs} yield {"tool_calls": self._get_truncated_tool_calls()} log_context.stacks.append( {"component": "agent", "data": {"tool_calls": self.tool_calls.copy()}} ) + + def _collect_internal_sources(self): + """Collect retrieved docs from the cached InternalSearchTool instance.""" + 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 diff --git a/application/agents/tools/internal_search.py b/application/agents/tools/internal_search.py index a622cce4..d711b0de 100644 --- a/application/agents/tools/internal_search.py +++ b/application/agents/tools/internal_search.py @@ -45,6 +45,7 @@ class InternalSearchTool(Tool): llm_name=self.config.get("llm_name", settings.LLM_PROVIDER), api_key=self.config.get("api_key", settings.API_KEY), decoded_token=self.config.get("decoded_token"), + request_id=self.config.get("request_id"), ) def _legacy_classic(): diff --git a/application/api/answer/services/stream_processor.py b/application/api/answer/services/stream_processor.py index ad7ec3cf..bd4b78c1 100644 --- a/application/api/answer/services/stream_processor.py +++ b/application/api/answer/services/stream_processor.py @@ -187,6 +187,26 @@ class StreamProcessor: tools_data = self.pre_fetch_tools() return self.create_agent(tools_data=tools_data) + # Classic agents (D11): partition sources by exposure. Pre-fetch the + # ``prefetch`` subset into the prompt and expose the ``agentic_tool`` + # subset via the internal_search tool. ``agentic_sources`` is empty when + # no source opts into ``agentic_tool`` (the default) or when no + # per-source detail is known (single-source / no-config requests). In + # that case fall back to the unscoped pre-fetch and add no search tool — + # behavior is byte-identical to today's classic. + _, agentic_sources = self._exposure_partition() + if agentic_sources: + docs_together, docs_list = self.pre_fetch_docs( + question, exposure="prefetch" + ) + tools_data = self.pre_fetch_tools() + return self.create_agent( + docs_together=docs_together, + docs=docs_list, + tools_data=tools_data, + agentic_sources=agentic_sources, + ) + docs_together, docs_list = self.pre_fetch_docs(question) tools_data = self.pre_fetch_tools() return self.create_agent( @@ -900,6 +920,7 @@ class StreamProcessor: user_api_key=self.agent_config["user_api_key"], agent_id=self.agent_id, decoded_token=self.decoded_token, + request_id=self.data.get("request_id"), ) def _legacy_classic(): @@ -1330,7 +1351,10 @@ class StreamProcessor: "tool_executor": tool_executor, } - if agent_key in ("agentic", "research") and retriever_config: + # Restore the search-tool config on resume. Classic agents carry one + # only when they had ``agentic_tool`` sources; a default classic agent + # serializes an empty config (falsy), so its behavior is unchanged. + if retriever_config and agent_key in ("classic", "agentic", "research"): agent_kwargs["retriever_config"] = retriever_config agent = AgentCreator.create_agent(agent_key, **agent_kwargs) @@ -1498,11 +1522,15 @@ class StreamProcessor: } # Type-specific kwargs - if agent_type in ("agentic", "research"): - # D11: when an ``agentic_tool`` subset is supplied, scope the search - # tool to it; otherwise the tool exposes every source (today's - # behavior). ``tool_sources`` drives both the source dict and the - # per-source dispatch list the InternalSearchTool uses. + # D11: agentic/research always carry a retriever_config; classic carries + # one only when an ``agentic_tool`` subset is supplied. A default classic + # agent (``agentic_sources is None``) gets NO retriever_config, so + # ClassicAgent adds no internal_search tool and stays today's behavior. + if agent_type in ("agentic", "research") or agentic_sources: + # When an ``agentic_tool`` subset is supplied, scope the search tool + # to it; otherwise (agentic/research only) the tool exposes every + # source (today's behavior). ``tool_sources`` drives both the source + # dict and the per-source dispatch list the InternalSearchTool uses. tool_sources = ( agentic_sources if agentic_sources is not None @@ -1537,6 +1565,7 @@ class StreamProcessor: "llm_name": provider or settings.LLM_PROVIDER, "api_key": system_api_key, "decoded_token": self.decoded_token, + "request_id": self.data.get("request_id"), } elif agent_type == "workflow": diff --git a/application/retriever/classic_rag.py b/application/retriever/classic_rag.py index f155347b..a95e7750 100644 --- a/application/retriever/classic_rag.py +++ b/application/retriever/classic_rag.py @@ -24,6 +24,7 @@ class ClassicRAG(BaseRetriever): decoded_token=None, model_user_id=None, defer_rephrase=False, + request_id=None, ): self.original_question = source.get("question", "") self.chat_history = chat_history if chat_history is not None else [] @@ -62,8 +63,10 @@ class ClassicRAG(BaseRetriever): model_user_id=self.model_user_id, ) # Query-rephrase LLM is a side channel — tag it so its rows - # land as ``source='rag_condense'`` in cost-attribution. + # land as ``source='rag_condense'`` in cost-attribution, and stamp + # the originating request so the rows correlate to it. self.llm._token_usage_source = "rag_condense" + self.llm._request_id = request_id if "active_docs" in source and source["active_docs"] is not None: if isinstance(source["active_docs"], list): diff --git a/application/retriever/dispatcher.py b/application/retriever/dispatcher.py index 9391f4ee..0efdd742 100644 --- a/application/retriever/dispatcher.py +++ b/application/retriever/dispatcher.py @@ -60,6 +60,7 @@ class Dispatcher(BaseRetriever): api_key=settings.API_KEY, decoded_token=None, model_user_id=None, + request_id=None, sources: Optional[List[Dict[str, Any]]] = None, stages: Optional[List[Stage]] = None, ): @@ -89,6 +90,7 @@ class Dispatcher(BaseRetriever): api_key=api_key, decoded_token=decoded_token, model_user_id=model_user_id, + request_id=request_id, ) self.source = source or {} self.original_question = self.source.get("question", "") @@ -243,6 +245,7 @@ class Dispatcher(BaseRetriever): decoded_token=self._ctor_kwargs.get("decoded_token"), agent_id=self._ctor_kwargs.get("agent_id"), model_user_id=self._ctor_kwargs.get("model_user_id"), + request_id=self._ctor_kwargs.get("request_id"), ) return list(self.stages) + prescreen diff --git a/application/retriever/stages/prescreen.py b/application/retriever/stages/prescreen.py index 7d89995c..4050784a 100644 --- a/application/retriever/stages/prescreen.py +++ b/application/retriever/stages/prescreen.py @@ -55,6 +55,7 @@ class PreScreenStage: decoded_token: Optional[Dict[str, Any]] = None, agent_id: Optional[str] = None, model_user_id: Optional[str] = None, + request_id: Optional[str] = None, ): """Build the stage. @@ -76,6 +77,7 @@ class PreScreenStage: self.decoded_token = decoded_token self.agent_id = agent_id self.model_user_id = model_user_id + self.request_id = request_id def _resolve_model(self) -> Optional[str]: """Use the configured model, else fall back to the request model.""" @@ -92,8 +94,10 @@ class PreScreenStage: agent_id=self.agent_id, model_user_id=self.model_user_id, ) - # Tag rows so the screening calls land as a distinct cost source. + # Tag rows so the screening calls land as a distinct cost source, and + # stamp the originating request so the rows correlate to it. llm._token_usage_source = "rag_prescreen" + llm._request_id = self.request_id return llm @staticmethod @@ -183,6 +187,7 @@ def build_prescreen_stages( decoded_token: Optional[Dict[str, Any]] = None, agent_id: Optional[str] = None, model_user_id: Optional[str] = None, + request_id: Optional[str] = None, ) -> List[Stage]: """Build prescreen stages from a group's per-source retrieval configs. @@ -216,6 +221,7 @@ def build_prescreen_stages( decoded_token=decoded_token, agent_id=agent_id, model_user_id=model_user_id, + request_id=request_id, ) ) return stages diff --git a/frontend/src/agents/NewAgent.tsx b/frontend/src/agents/NewAgent.tsx index 5f2b6244..af60283c 100644 --- a/frontend/src/agents/NewAgent.tsx +++ b/frontend/src/agents/NewAgent.tsx @@ -159,10 +159,8 @@ export default function NewAgent({ mode }: { mode: 'new' | 'edit' | 'draft' }) { trackChanges: false, }, }; - const chunks = ['0', '2', '4', '6', '8', '10']; const agentTypes = [ { label: t('agents.form.agentTypes.classic'), value: 'classic' }, - { label: 'Agentic', value: 'agentic' }, { label: 'Research', value: 'research' }, ]; @@ -984,30 +982,6 @@ export default function NewAgent({ mode }: { mode: 'new' | 'edit' | 'draft' }) { } /> -
- -
diff --git a/frontend/src/agents/workflow/WorkflowBuilder.tsx b/frontend/src/agents/workflow/WorkflowBuilder.tsx index caff5b00..e58566ce 100644 --- a/frontend/src/agents/workflow/WorkflowBuilder.tsx +++ b/frontend/src/agents/workflow/WorkflowBuilder.tsx @@ -81,7 +81,7 @@ import type { Model } from '../../models/types'; const PRIMARY_ACTION_SPINNER_DELAY_MS = 180; interface AgentNodeConfig { - agent_type: 'classic' | 'agentic' | 'research'; + agent_type: 'classic' | 'research'; llm_name?: string; model_id?: string; system_prompt: string; @@ -1840,9 +1840,6 @@ function WorkflowBuilderInner() { Classic - - Agentic - Research diff --git a/frontend/src/locale/en.json b/frontend/src/locale/en.json index fb315b62..3bd8400a 100644 --- a/frontend/src/locale/en.json +++ b/frontend/src/locale/en.json @@ -137,7 +137,7 @@ "retrieval": { "title": "Retrieval", "note": "These changes take effect immediately, no re-ingest needed.", - "tag": "live · no re-ingest", + "tag": "no re-ingest", "chunks": "Chunks to retrieve (top-k)", "scoreThreshold": "Score threshold", "scoreThresholdPlaceholder": "None", @@ -150,9 +150,9 @@ "exposure": "Search exposure", "exposures": { "prefetch": "Pre-fetch into context", - "agentic_tool": "Agentic search tool" + "agentic_tool": "On-demand search tool" }, - "exposureHint": "Only affects agentic/research agents. Classic agents always pre-fetch." + "exposureHint": "Pre-fetch this source into the prompt, or let the agent search it on demand as a tool." }, "prescreen": { "enable": "Enable LLM prescreen", diff --git a/frontend/src/settings/General.tsx b/frontend/src/settings/General.tsx index a77c8ee2..be2bd3bd 100644 --- a/frontend/src/settings/General.tsx +++ b/frontend/src/settings/General.tsx @@ -12,10 +12,8 @@ import { } from '../components/ui/select'; import { useDarkTheme } from '../hooks'; import { - selectChunks, selectPrompt, selectPrompts, - setChunks, setModalStateDeleteConv, setPrompt, setPrompts, @@ -41,9 +39,7 @@ export default function General() { { label: '繁體中文（臺灣）', value: 'zhTW' }, { label: 'Русский', value: 'ru' }, ]; - const chunks = ['0', '2', '4', '6', '8', '10']; const prompts = useSelector(selectPrompts); - const selectedChunks = useSelector(selectChunks); const [isDarkTheme, toggleTheme] = useDarkTheme(); const [selectedTheme, setSelectedTheme] = React.useState( isDarkTheme ? 'Dark' : 'Light', @@ -73,26 +69,6 @@ export default function General() { setPrompts={(newPrompts) => dispatch(setPrompts(newPrompts))} />
-
- - -
diff --git a/tests/agents/test_classic_agent.py b/tests/agents/test_classic_agent.py index ac5ed833..50a14230 100644 --- a/tests/agents/test_classic_agent.py +++ b/tests/agents/test_classic_agent.py @@ -2,6 +2,7 @@ from unittest.mock import Mock import pytest from application.agents.classic_agent import ClassicAgent +from application.agents.tools.internal_search import INTERNAL_TOOL_ID @pytest.fixture @@ -16,6 +17,15 @@ def _no_tools(monkeypatch): ) +@pytest.fixture +def _no_dir_structure(monkeypatch): + """Stub the DB-backed directory-structure lookup used by the search tool.""" + monkeypatch.setattr( + "application.agents.tools.internal_search.sources_have_directory_structure", + lambda source: False, + ) + + @pytest.mark.unit class TestClassicAgent: @@ -181,6 +191,106 @@ class TestClassicAgent: assert "tool_calls" in agent_logs[0]["data"] +@pytest.mark.unit +class TestClassicAgentSearchExposure: + """ClassicAgent honors per-source exposure via an internal_search tool.""" + + def test_default_registers_no_internal_search( + self, + agent_base_params, + mock_llm, + mock_llm_handler, + mock_llm_creator, + mock_llm_handler_creator, + _no_tools, + log_context, + ): + # No retriever_config (every source at default prefetch) → no search + # tool is added and pre-fetched sources are preserved (unchanged). + mock_llm.gen_stream = Mock(return_value=iter(["Answer"])) + + def mock_handler(*args, **kwargs): + yield "Processed" + + mock_llm_handler.process_message_flow = Mock(side_effect=mock_handler) + + agent = ClassicAgent(**agent_base_params) + assert agent.retriever_config == {} + + captured = {} + original_prepare = agent._prepare_tools + + def _spy(tools_dict): + captured["tools_dict"] = tools_dict + return original_prepare(tools_dict) + + agent._prepare_tools = _spy + prefetched = [{"text": "pre", "title": "Pre", "source": "s"}] + agent.retrieved_docs = list(prefetched) + + list(agent._gen_inner("q", log_context)) + + assert INTERNAL_TOOL_ID not in captured["tools_dict"] + # Pre-fetched sources untouched (no _collect_internal_sources overwrite). + assert agent.retrieved_docs == prefetched + + def test_with_agentic_tool_source_registers_internal_search( + self, + agent_base_params, + mock_llm, + mock_llm_handler, + mock_llm_creator, + mock_llm_handler_creator, + _no_tools, + _no_dir_structure, + log_context, + ): + # A retriever_config (agentic_tool subset) → the internal_search tool is + # registered so the LLM can search that subset on demand. + mock_llm.gen_stream = Mock(return_value=iter(["Answer"])) + + def mock_handler(*args, **kwargs): + yield "Processed" + + mock_llm_handler.process_message_flow = Mock(side_effect=mock_handler) + + retriever_config = { + "source": {"active_docs": ["b"]}, + "retriever_name": "classic", + "chunks": 2, + "sources": [{"id": "b", "retrieval": None}], + } + agent = ClassicAgent(retriever_config=retriever_config, **agent_base_params) + + captured = {} + original_prepare = agent._prepare_tools + + def _spy(tools_dict): + captured["tools_dict"] = tools_dict + return original_prepare(tools_dict) + + agent._prepare_tools = _spy + + list(agent._gen_inner("q", log_context)) + + assert INTERNAL_TOOL_ID in captured["tools_dict"] + assert captured["tools_dict"][INTERNAL_TOOL_ID]["name"] == "internal_search" + + def test_collect_internal_sources_surfaces_tool_docs( + 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) + + mock_tool = Mock() + mock_tool.retrieved_docs = [{"text": "Found", "title": "Doc", "source": "t"}] + 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] == ["Doc"] + + @pytest.mark.integration class TestClassicAgentIntegration: diff --git a/tests/api/answer/test_stream_processor_exposure.py b/tests/api/answer/test_stream_processor_exposure.py index 8979b92e..7133d5cd 100644 --- a/tests/api/answer/test_stream_processor_exposure.py +++ b/tests/api/answer/test_stream_processor_exposure.py @@ -98,19 +98,51 @@ class TestBuildAgentExposure: _, kwargs = sp.create_agent.call_args assert [e["id"] for e in kwargs["agentic_sources"]] == ["b"] - def test_classic_agent_ignores_exposure(self): - # Classic agents pre-fetch all and never partition by exposure. + def _classic_processor(self): sp = _processor() sp.agent_config = {"agent_type": "classic"} - sp.all_sources = [ - {"id": "a", "retrieval": RetrievalConfig(exposure="agentic_tool")}, - ] sp.initialize = MagicMock() - sp.pre_fetch_docs = MagicMock(return_value=("t", ["d"])) + sp.pre_fetch_docs = MagicMock(return_value=("docs_together", ["doc"])) sp.pre_fetch_tools = MagicMock(return_value=None) sp.create_agent = MagicMock(return_value="AGENT") - sp.build_agent("q") - # Classic pre-fetch is called with no exposure scoping (all sources). + return sp + + def test_classic_default_prefetches_all_no_partition(self): + # Default classic (all prefetch / no config): byte-identical to today — + # unscoped pre-fetch and no agentic_sources, so no search tool is added. + sp = self._classic_processor() + sp.all_sources = [ + {"id": "a", "retrieval": RetrievalConfig()}, + {"id": "b", "retrieval": None}, + ] + result = sp.build_agent("q") + assert result == "AGENT" sp.pre_fetch_docs.assert_called_once_with("q") _, kwargs = sp.create_agent.call_args assert "agentic_sources" not in kwargs + + def test_classic_no_per_source_detail_is_today(self): + # Single-source / no-config requests carry no per-source detail; classic + # must fall back to the unscoped pre-fetch (no exposure scoping). + sp = self._classic_processor() + sp.all_sources = [] + sp.source = {"active_docs": ["a"]} + sp.build_agent("q") + sp.pre_fetch_docs.assert_called_once_with("q") + _, kwargs = sp.create_agent.call_args + assert "agentic_sources" not in kwargs + + def test_classic_mixed_prefetches_and_scopes_tool(self): + # Classic with an agentic_tool source now pre-fetches only the prefetch + # subset and exposes the agentic_tool subset via the search tool. + sp = self._classic_processor() + sp.all_sources = [ + {"id": "a", "retrieval": RetrievalConfig(exposure="prefetch")}, + {"id": "b", "retrieval": RetrievalConfig(exposure="agentic_tool")}, + ] + sp.build_agent("q") + sp.pre_fetch_docs.assert_called_once() + _, pf_kwargs = sp.pre_fetch_docs.call_args + assert pf_kwargs.get("exposure") == "prefetch" + _, kwargs = sp.create_agent.call_args + assert [e["id"] for e in kwargs["agentic_sources"]] == ["b"] diff --git a/tests/retriever/test_prescreen.py b/tests/retriever/test_prescreen.py index 5c4e0960..bc0e81bf 100644 --- a/tests/retriever/test_prescreen.py +++ b/tests/retriever/test_prescreen.py @@ -148,6 +148,30 @@ class TestModelResolution: stage([{"text": "x"}], {"query": "q"}) assert captured["model_id"] == "cheap" + def test_request_id_and_source_stamped_on_llm(self): + config = PreScreenConfig(candidate_k=1, batch_size=1, max_keep=1, model=None) + created = {} + + def fake_create_llm(*args, **kwargs): + llm = Mock(gen=Mock(return_value='{"keep": [0]}'), model_id="m") + created["llm"] = llm + return llm + + with patch( + "application.retriever.stages.prescreen.LLMCreator.create_llm", + side_effect=fake_create_llm, + ): + stage = PreScreenStage( + config, + llm_name="openai", + api_key="k", + model_id="m", + request_id="req-123", + ) + stage([{"text": "x"}], {"query": "q"}) + assert created["llm"]._request_id == "req-123" + assert created["llm"]._token_usage_source == "rag_prescreen" + @pytest.mark.unit class TestStageBuilders: diff --git a/tests/test_retriever.py b/tests/test_retriever.py index 7d5baed5..b9334415 100644 --- a/tests/test_retriever.py +++ b/tests/test_retriever.py @@ -126,6 +126,13 @@ class TestClassicRAGInit: assert rag.chunks == 2 assert rag.vectorstores == [] + def test_request_id_and_source_stamped_on_rephrase_llm( + self, _patch_llm_creator + ): + _make_rag(request_id="req-123") + assert _patch_llm_creator._request_id == "req-123" + assert _patch_llm_creator._token_usage_source == "rag_condense" + def test_active_docs_as_list(self, _patch_llm_creator): rag = _make_rag(source={"question": "q", "active_docs": ["a", "b"]}) assert rag.vectorstores == ["a", "b"]