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"]