From ec5615dcc1010deb06d2eba2a866cd15f5d899cc Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Sat, 24 Oct 2026 12:38:45 +0800 Subject: [PATCH] feat: enhance retrieval query handling with agentic mode routing - Updated the `use_agentic` parameter behavior to allow for a query router to choose between classic and agent retrieval modes. - Modified documentation and environment examples to reflect the new routing logic. - Introduced a router decision structure to capture the mode chosen by the query router. - Adjusted retrieval execution logic to incorporate the new routing mechanism, ensuring seamless integration of agentic and classic retrieval methods. --- AGENTS.md | 6 +- apps/api/.env.example | 5 +- apps/api/app/api/v1/routes/retrieval.py | 13 +- .../test_retrieval_query_router_contract.py | 299 ++++++++++++++++++ apps/worker/.env.example | 5 +- .../retrieval/agent_explore/__init__.py | 7 +- .../execution/response_projection.py | 2 + .../services/retrieval/execution/routes.py | 40 ++- .../services/retrieval/routing/__init__.py | 1 + .../services/retrieval/routing/classifier.py | 80 +++++ .../services/retrieval/routing/config.py | 24 ++ .../services/retrieval/routing/prompt.py | 10 + .../services/retrieval/routing/types.py | 28 ++ .../tests/test_query_router_classifier.py | 177 +++++++++++ 14 files changed, 683 insertions(+), 14 deletions(-) create mode 100644 apps/api/tests/contract/test_retrieval_query_router_contract.py create mode 100644 packages/shared-python/shared/services/retrieval/routing/__init__.py create mode 100644 packages/shared-python/shared/services/retrieval/routing/classifier.py create mode 100644 packages/shared-python/shared/services/retrieval/routing/config.py create mode 100644 packages/shared-python/shared/services/retrieval/routing/prompt.py create mode 100644 packages/shared-python/shared/services/retrieval/routing/types.py create mode 100644 packages/shared-python/shared/tests/test_query_router_classifier.py diff --git a/AGENTS.md b/AGENTS.md index 587b3a673..1011436d2 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -102,8 +102,8 @@ flowchart TB subgraph RETRIEVE["⑤ Retrieval (shared)"] Query["POST /v1|/v2 retrieval/query"] --> Pipeline["run_retrieval_query"] - Pipeline --> Classic["classic_topk / small_corpus (use_agentic=False)"] - Pipeline --> Explore["agent_explore (use_agentic unset/true)"] + Pipeline --> Classic["classic_topk / small_corpus"] + Pipeline --> Explore["agent_explore"] Explore --> Cursor["AGENT_EXPLORE_HARNESS=cursor_sdk"] Explore --> Own["AGENT_EXPLORE_HARNESS=openai"] Classic --> Channels["map_unit_discovery: path+content BM25 -> RRF"] @@ -542,7 +542,7 @@ Core retrieval internals are grouped by ownership: ### Retrieval Switch and Agent Modes -Per-request `use_agentic` only turns agent retrieval on or off: `False` → classic top-K (map-unit BM25); `None`/`True` → agent_explore. When agent retrieval is on, `AGENT_EXPLORE_HARNESS` selects Cursor (`cursor_sdk`) or Knowhere's own harness (`openai`). +Per-request `use_agentic`: `False` → classic top-K (map-unit BM25); `True` → agent_explore; unset (`None`) → a query router chooses classic or agent. When agent retrieval is on, `AGENT_EXPLORE_HARNESS` selects Cursor (`cursor_sdk`) or Knowhere's own harness (`openai`). #### Classic Mode (map-unit BM25) diff --git a/apps/api/.env.example b/apps/api/.env.example index c1721c4ba..06114302e 100644 --- a/apps/api/.env.example +++ b/apps/api/.env.example @@ -103,10 +103,13 @@ ARK_API_KEY= # answer_text is always empty. # use_agentic on the request turns agent retrieval on or off: # false → no agent; classic map-unit BM25 -# omit/true → agent on; AGENT_EXPLORE_HARNESS selects the mode +# true → agent on; AGENT_EXPLORE_HARNESS selects the mode +# omit → query router chooses classic or agent # cursor_sdk → Cursor (needs CURSOR_API_KEY) # openai → Knowhere's own harness (DS_KEY / ALI_API_KEYS) # AGENT_EXPLORE_HARNESS=cursor_sdk +# RETRIEVAL_QUERY_ROUTER_ENABLED=true +# RETRIEVAL_QUERY_ROUTER_MODEL= # CURSOR_API_KEY= # only when harness is cursor_sdk # AGENT_EXPLORE_CURSOR_MODEL=composer-2.5 # Classic BM25 may use Postgres FTS prefilter: diff --git a/apps/api/app/api/v1/routes/retrieval.py b/apps/api/app/api/v1/routes/retrieval.py index 319a143c4..569872734 100644 --- a/apps/api/app/api/v1/routes/retrieval.py +++ b/apps/api/app/api/v1/routes/retrieval.py @@ -73,8 +73,9 @@ class RetrievalQueryRequest(BaseModel): None, description=( "Turns agent retrieval on or off. False uses classic map-unit " - "BM25. Unset or true turns agent retrieval on; " - "AGENT_EXPLORE_HARNESS then selects Cursor (cursor_sdk) or " + "BM25. True turns agent retrieval on. Unset (None) lets the " + "query router choose classic or agent. When agent retrieval is " + "on, AGENT_EXPLORE_HARNESS then selects Cursor (cursor_sdk) or " "Knowhere's own harness (openai)." ), ) @@ -156,6 +157,14 @@ class RetrievalQueryResponse(BaseModel): "observation/decision/result schema." ), ) + router_decision: dict[str, Any] | None = Field( + default=None, + description=( + "Present when use_agentic was unset and the query router chose " + "agent or classic. Fields: mode, reason, source, model_name, " + "tokens_used." + ), + ) async def execute_retrieval_query( diff --git a/apps/api/tests/contract/test_retrieval_query_router_contract.py b/apps/api/tests/contract/test_retrieval_query_router_contract.py new file mode 100644 index 000000000..1eb3cc0ed --- /dev/null +++ b/apps/api/tests/contract/test_retrieval_query_router_contract.py @@ -0,0 +1,299 @@ +from __future__ import annotations + +from collections.abc import Callable +from contextlib import AbstractAsyncContextManager +from typing import Any, cast +from uuid import uuid4 + +import pytest +from httpx import AsyncClient +from pytest import MonkeyPatch + +from shared.services.retrieval.routing.types import QueryMode, RouterDecision +from tests.contract.test_retrieval_classic_map_unit_contract import _publish_document +from tests.contract.test_retrieval_contract import ( + _episode_keeping_refs, + _patch_harness, + _seed_retrieval_document, +) + + +def _decision(*, mode: QueryMode, source: str = "llm") -> RouterDecision: + return RouterDecision( + mode=mode, + reason="contract", + source=source, + model_name="test-router-model", + tokens_used=2, + ) + + +async def _publish_above_top_k(*, namespace: str) -> dict[str, str]: + identifier = uuid4().hex[:8] + chunk_id = f"hit-{identifier}" + target = await _publish_document( + namespace=namespace, + source_file_name="target.pdf", + chunks=[ + { + "chunk_id": chunk_id, + "type": "text", + "content": "unique alpha ranking marker", + "path": "target.pdf/Root/Hit/body", + "order": 1, + "metadata": {}, + } + ], + ) + await _publish_document( + namespace=namespace, + source_file_name="filler.pdf", + chunks=[ + { + "chunk_id": f"filler-{identifier}", + "type": "text", + "content": "unrelated filler content", + "path": "filler.pdf/Root/Noise/body", + "order": 1, + "metadata": {}, + } + ], + ) + return { + **target, + "chunk_id": chunk_id, + "section_path": "Root / Hit / body", + } + + +def _patch_classifier( + monkeypatch: MonkeyPatch, + decision: RouterDecision, + calls: list[str], +) -> None: + from shared.services.retrieval.execution import routes as retrieval_routes + + async def fake_classify(query: str) -> RouterDecision: + calls.append(query) + return decision + + monkeypatch.setattr(retrieval_routes, "classify_query_mode", fake_classify) + + +def _patch_classifier_boom(monkeypatch: MonkeyPatch, calls: list[str], message: str) -> None: + from shared.services.retrieval.execution import routes as retrieval_routes + + async def boom(_query: str) -> RouterDecision: + calls.append("called") + raise AssertionError(message) + + monkeypatch.setattr(retrieval_routes, "classify_query_mode", boom) + + +@pytest.mark.asyncio +async def test_unset_use_agentic_classic_decision_uses_classic_route( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + monkeypatch: MonkeyPatch, +) -> None: + namespace = f"router-classic-{uuid4().hex[:8]}" + calls: list[str] = [] + + async with developer_api_client_factory() as api_client: + await _publish_above_top_k(namespace=namespace) + _patch_classifier(monkeypatch, _decision(mode=QueryMode.CLASSIC), calls) + response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": namespace, + "query": "unique alpha ranking", + "top_k": 1, + }, + ) + + assert response.status_code == 200 + body = cast(dict[str, Any], response.json()) + assert calls == ["unique alpha ranking"] + assert body["router_used"] == "classic_topk" + assert body["router_decision"] == { + "mode": "classic", + "reason": "contract", + "source": "llm", + "model_name": "test-router-model", + "tokens_used": 2, + } + + +@pytest.mark.asyncio +async def test_unset_use_agentic_agent_decision_uses_agent_explore( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + monkeypatch: MonkeyPatch, +) -> None: + namespace = f"router-agent-{uuid4().hex[:8]}" + calls: list[str] = [] + + async with developer_api_client_factory() as api_client: + target = await _publish_above_top_k(namespace=namespace) + _patch_classifier(monkeypatch, _decision(mode=QueryMode.AGENT), calls) + _patch_harness( + monkeypatch, + _episode_keeping_refs( + documents=[target], + notes="unique alpha ranking marker", + ), + ) + response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": namespace, + "query": "unique alpha ranking", + "top_k": 1, + }, + ) + + assert response.status_code == 200 + body = cast(dict[str, Any], response.json()) + assert calls == ["unique alpha ranking"] + assert body["router_used"] == "agent_explore" + assert body["router_decision"]["mode"] == "agent" + assert body["router_decision"]["source"] == "llm" + + +@pytest.mark.asyncio +async def test_explicit_use_agentic_does_not_call_classifier( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + monkeypatch: MonkeyPatch, +) -> None: + namespace = f"router-explicit-{uuid4().hex[:8]}" + calls: list[str] = [] + + async with developer_api_client_factory() as api_client: + target = await _publish_above_top_k(namespace=namespace) + _patch_classifier_boom( + monkeypatch, + calls, + "classifier must not run when use_agentic is set", + ) + classic = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": namespace, + "query": "unique alpha ranking", + "top_k": 1, + "use_agentic": False, + }, + ) + _patch_harness( + monkeypatch, + _episode_keeping_refs(documents=[target], notes="explicit agent"), + ) + agent = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": namespace, + "query": "unique alpha ranking", + "top_k": 1, + "use_agentic": True, + }, + ) + + assert classic.status_code == 200 + assert agent.status_code == 200 + classic_body = cast(dict[str, Any], classic.json()) + agent_body = cast(dict[str, Any], agent.json()) + assert calls == [] + assert classic_body["router_used"] == "classic_topk" + assert classic_body.get("router_decision") is None + assert agent_body["router_used"] == "agent_explore" + assert agent_body.get("router_decision") is None + + +@pytest.mark.asyncio +async def test_disabled_router_keeps_unset_use_agentic_on_agent( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + monkeypatch: MonkeyPatch, +) -> None: + namespace = f"router-disabled-{uuid4().hex[:8]}" + calls: list[str] = [] + + async with developer_api_client_factory() as api_client: + target = await _publish_above_top_k(namespace=namespace) + monkeypatch.setattr( + "shared.services.retrieval.routing.config.QUERY_ROUTER_ENABLED", + False, + ) + _patch_classifier_boom( + monkeypatch, + calls, + "classifier must not run when the router is disabled", + ) + _patch_harness( + monkeypatch, + _episode_keeping_refs(documents=[target], notes="disabled router"), + ) + response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": namespace, + "query": "unique alpha ranking", + "top_k": 1, + }, + ) + + assert response.status_code == 200 + body = cast(dict[str, Any], response.json()) + assert calls == [] + assert body["router_used"] == "agent_explore" + assert body["router_decision"] == { + "mode": "agent", + "reason": "", + "source": "disabled", + "model_name": "", + "tokens_used": 0, + } + + +@pytest.mark.asyncio +async def test_small_corpus_skips_query_router( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + monkeypatch: MonkeyPatch, +) -> None: + namespace = f"router-small-{uuid4().hex[:8]}" + calls: list[str] = [] + + async with developer_api_client_factory() as api_client: + await _seed_retrieval_document( + user_id="local-dev-user", + namespace=namespace, + source_file_name="only.pdf", + section_path="only/section", + content="only chunk", + ) + _patch_classifier_boom( + monkeypatch, + calls, + "classifier must not run on the small-corpus path", + ) + response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": namespace, + "query": "only chunk", + "top_k": 10, + }, + ) + + assert response.status_code == 200 + body = cast(dict[str, Any], response.json()) + assert calls == [] + assert body["router_used"] == "small_corpus_all" + assert body.get("router_decision") is None diff --git a/apps/worker/.env.example b/apps/worker/.env.example index dfce1d370..fe58ed661 100644 --- a/apps/worker/.env.example +++ b/apps/worker/.env.example @@ -97,10 +97,13 @@ ARK_API_KEY= # answer_text is always empty. # use_agentic on the request turns agent retrieval on or off: # false → no agent; classic map-unit BM25 -# omit/true → agent on; AGENT_EXPLORE_HARNESS selects the mode +# true → agent on; AGENT_EXPLORE_HARNESS selects the mode +# omit → query router chooses classic or agent # cursor_sdk → Cursor (needs CURSOR_API_KEY) # openai → Knowhere's own harness (DS_KEY / ALI_API_KEYS) # AGENT_EXPLORE_HARNESS=cursor_sdk +# RETRIEVAL_QUERY_ROUTER_ENABLED=true +# RETRIEVAL_QUERY_ROUTER_MODEL= # CURSOR_API_KEY= # only when harness is cursor_sdk # AGENT_EXPLORE_CURSOR_MODEL=composer-2.5 # Classic BM25 may use Postgres FTS prefilter: diff --git a/packages/shared-python/shared/services/retrieval/agent_explore/__init__.py b/packages/shared-python/shared/services/retrieval/agent_explore/__init__.py index d313f765d..206ed702c 100644 --- a/packages/shared-python/shared/services/retrieval/agent_explore/__init__.py +++ b/packages/shared-python/shared/services/retrieval/agent_explore/__init__.py @@ -1,8 +1,9 @@ """In-process tool-calling agentic retrieval route. -Default agentic path when ``use_agentic`` is unset or true. Runs the -``agent_tools`` corpus registry through an LLM tool-calling loop. Map-nav is -archived under ``deprecated/mapnav/`` and is not a live route. +Agentic path when ``use_agentic`` is true, or when it is unset and the query +router selects agent. Runs the ``agent_tools`` corpus registry through an LLM +tool-calling loop. Map-nav is archived under ``deprecated/mapnav/`` and is +not a live route. Which provider runs that loop (Cursor SDK or OpenAI-compatible) is the ``AGENT_EXPLORE_HARNESS`` switch resolved via ``resolve_harness()``. diff --git a/packages/shared-python/shared/services/retrieval/execution/response_projection.py b/packages/shared-python/shared/services/retrieval/execution/response_projection.py index 8878d82cf..b3019f6a3 100644 --- a/packages/shared-python/shared/services/retrieval/execution/response_projection.py +++ b/packages/shared-python/shared/services/retrieval/execution/response_projection.py @@ -38,6 +38,8 @@ async def project_public_retrieval_response(response: dict[str, Any]) -> dict[st public_response['failure_reason'] = response['failure_reason'] if response.get('decision_trace') is not None: public_response['decision_trace'] = response['decision_trace'] + if response.get('router_decision') is not None: + public_response['router_decision'] = response['router_decision'] projected_rows = await enrich_rows_with_retrieval_asset_url( response.get('results', []), diff --git a/packages/shared-python/shared/services/retrieval/execution/routes.py b/packages/shared-python/shared/services/retrieval/execution/routes.py index 6cba8f9f8..e0cd1a5ae 100644 --- a/packages/shared-python/shared/services/retrieval/execution/routes.py +++ b/packages/shared-python/shared/services/retrieval/execution/routes.py @@ -19,6 +19,9 @@ from shared.services.retrieval.hydration.result_assembly import ( assemble_retrieval_results, ) +from shared.services.retrieval.routing import config as query_router_config +from shared.services.retrieval.routing.classifier import classify_query_mode +from shared.services.retrieval.routing.types import QueryMode, RouterDecision from shared.services.retrieval.search.map_unit_discovery import map_unit_discovery from shared.services.retrieval.search.ranking import rank_retrieval_candidates from shared.services.retrieval.search.scoped_corpus import ( @@ -77,10 +80,39 @@ async def run_retrieval_route( if small_corpus_outcome is not None: return small_corpus_outcome - # Explicit False → classic map-unit BM25 top-K. None/True → agent_explore. - if context.use_agentic is False: - return await _run_classic_topk_route(context) - return await _run_agent_explore_route(context) + use_agentic = context.use_agentic + decision: RouterDecision | None = None + if use_agentic is None: + use_agentic, decision = await _resolve_auto_route(context) + + if use_agentic is False: + outcome = await _run_classic_topk_route(context) + else: + outcome = await _run_agent_explore_route(context) + + if decision is not None: + outcome.response["router_decision"] = decision.to_dict() + return outcome + + +async def _resolve_auto_route( + context: RetrievalRouteContext, +) -> tuple[bool, RouterDecision]: + if not query_router_config.QUERY_ROUTER_ENABLED: + decision = RouterDecision( + mode=QueryMode.AGENT, + reason="", + source="disabled", + model_name="", + tokens_used=0, + ) + else: + decision = await classify_query_mode(context.query) + logger.info( + f"retrieval query router mode={decision.mode.value} " + f"source={decision.source} reason={decision.reason}" + ) + return decision.mode == QueryMode.AGENT, decision async def _try_run_small_corpus_route( diff --git a/packages/shared-python/shared/services/retrieval/routing/__init__.py b/packages/shared-python/shared/services/retrieval/routing/__init__.py new file mode 100644 index 000000000..9049724e6 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/routing/__init__.py @@ -0,0 +1 @@ +"""Query-mode router: classify an unset ``use_agentic`` query as agent or classic.""" diff --git a/packages/shared-python/shared/services/retrieval/routing/classifier.py b/packages/shared-python/shared/services/retrieval/routing/classifier.py new file mode 100644 index 000000000..4519f47dd --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/routing/classifier.py @@ -0,0 +1,80 @@ +from __future__ import annotations + +import asyncio +import json +from typing import Any + +from loguru import logger + +from shared.services.retrieval.routing import config as routing_config +from shared.services.retrieval.routing.prompt import QUERY_ROUTER_SYSTEM_PROMPT +from shared.services.retrieval.routing.types import QueryMode, RouterDecision + + +def _resolve_client_and_model() -> tuple[Any, str]: + from shared.services.ai.llm_overrides import resolve_text + from shared.services.ai.openai_compatible_client_sync import get_openai_client + + requested = routing_config.QUERY_ROUTER_MODEL + effective_model, api_key, api_url = resolve_text(requested) + model = effective_model or requested + client = get_openai_client(model=model, api_key=api_key, api_url=api_url) + return client, model + + +def _fallback_decision(*, model_name: str, tokens_used: int = 0) -> RouterDecision: + return RouterDecision( + mode=QueryMode.AGENT, + reason="", + source="fallback", + model_name=model_name, + tokens_used=tokens_used, + ) + + +def _parse_decision( + content: str, *, model_name: str, tokens_used: int +) -> RouterDecision: + try: + parsed = json.loads(content) + except ValueError: + return _fallback_decision(model_name=model_name, tokens_used=tokens_used) + if not isinstance(parsed, dict): + return _fallback_decision(model_name=model_name, tokens_used=tokens_used) + mode_raw = str(parsed.get("mode") or "").strip().lower() + if mode_raw not in {QueryMode.AGENT.value, QueryMode.CLASSIC.value}: + return _fallback_decision(model_name=model_name, tokens_used=tokens_used) + return RouterDecision( + mode=QueryMode(mode_raw), + reason=str(parsed.get("reason") or ""), + source="llm", + model_name=model_name, + tokens_used=tokens_used, + ) + + +async def classify_query_mode(query: str) -> RouterDecision: + model_name = "" + try: + client, model_name = _resolve_client_and_model() + content, usage = await asyncio.to_thread( + client.chat_completion_with_usage, + [ + {"role": "system", "content": QUERY_ROUTER_SYSTEM_PROMPT}, + {"role": "user", "content": query}, + ], + model=model_name, + temperature=0.0, + max_tokens=routing_config.QUERY_ROUTER_MAX_TOKENS, + timeout=routing_config.QUERY_ROUTER_TIMEOUT_SECONDS, + response_format={"type": "json_object"}, + ) + tokens_used = int((usage or {}).get("total_tokens", 0) or 0) + return _parse_decision( + str(content or ""), + model_name=model_name, + tokens_used=tokens_used, + ) + except Exception as exc: + logger.warning("retrieval query router fallback to agent: {}", exc) + return _fallback_decision(model_name=model_name) diff --git a/packages/shared-python/shared/services/retrieval/routing/config.py b/packages/shared-python/shared/services/retrieval/routing/config.py new file mode 100644 index 000000000..203dae277 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/routing/config.py @@ -0,0 +1,24 @@ +from __future__ import annotations + +import os + +from shared.services.retrieval.agent_explore.config import AGENT_EXPLORE_MODEL + +QUERY_ROUTER_ENABLED_ENV = "RETRIEVAL_QUERY_ROUTER_ENABLED" +QUERY_ROUTER_MODEL_ENV = "RETRIEVAL_QUERY_ROUTER_MODEL" + +QUERY_ROUTER_TIMEOUT_SECONDS = 10 +QUERY_ROUTER_MAX_TOKENS = 256 + + +def _env_flag(name: str, default: bool) -> bool: + raw = os.environ.get(name) + if raw is None or not str(raw).strip(): + return default + return str(raw).strip().lower() in {"1", "true", "yes", "on"} + + +QUERY_ROUTER_ENABLED = _env_flag(QUERY_ROUTER_ENABLED_ENV, True) +QUERY_ROUTER_MODEL = ( + os.environ.get(QUERY_ROUTER_MODEL_ENV, "").strip() or AGENT_EXPLORE_MODEL +) diff --git a/packages/shared-python/shared/services/retrieval/routing/prompt.py b/packages/shared-python/shared/services/retrieval/routing/prompt.py new file mode 100644 index 000000000..e536ccc1f --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/routing/prompt.py @@ -0,0 +1,10 @@ +QUERY_ROUTER_SYSTEM_PROMPT = """\ +You classify one retrieval query into a search mode. Reply with json only. + +Modes: +- agent: the query needs multi-hop reasoning, cross-document synthesis, comparison, aggregation, causal or how/why analysis, or an open-ended investigation. +- classic: the query is looking up a single fact, a specific value, a named detail, or a keyword-style locate request such as what X is or where to find Y. + +Return a json object: +{"mode":"agent"|"classic","reason":"at most 30 characters"} +""" diff --git a/packages/shared-python/shared/services/retrieval/routing/types.py b/packages/shared-python/shared/services/retrieval/routing/types.py new file mode 100644 index 000000000..c0b17e9d9 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/routing/types.py @@ -0,0 +1,28 @@ +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum +from typing import Any + + +class QueryMode(str, Enum): + AGENT = "agent" + CLASSIC = "classic" + + +@dataclass(frozen=True) +class RouterDecision: + mode: QueryMode + reason: str + source: str + model_name: str + tokens_used: int + + def to_dict(self) -> dict[str, Any]: + return { + "mode": self.mode.value, + "reason": self.reason, + "source": self.source, + "model_name": self.model_name, + "tokens_used": self.tokens_used, + } diff --git a/packages/shared-python/shared/tests/test_query_router_classifier.py b/packages/shared-python/shared/tests/test_query_router_classifier.py new file mode 100644 index 000000000..da575cbcb --- /dev/null +++ b/packages/shared-python/shared/tests/test_query_router_classifier.py @@ -0,0 +1,177 @@ +"""Unit tests for the retrieval query-mode classifier and public projection.""" + +from __future__ import annotations + +import os +from typing import Any + +os.environ.setdefault("DATABASE_URL", "postgresql+asyncpg://test:test@localhost/test") +os.environ.setdefault("TMP_PATH", "/tmp/knowhere-test") +os.environ.setdefault("S3_BUCKET_NAME", "test-uploads") +os.environ.setdefault("S3_ACCESS_KEY_ID", "test") +os.environ.setdefault("S3_SECRET_ACCESS_KEY", "test") +os.environ.setdefault("S3_TEMP_PATH", "/tmp") + +import pytest + +from shared.core.config import settings +from shared.services.retrieval.execution.response_projection import ( + project_public_retrieval_response, +) +from shared.services.retrieval.routing.classifier import classify_query_mode +from shared.services.retrieval.routing.types import QueryMode + + +class _FakeClient: + def __init__(self, content: str, usage: dict[str, int] | None = None) -> None: + self.content = content + self.usage = usage or {"total_tokens": 4} + self.calls: list[dict[str, Any]] = [] + + def chat_completion_with_usage(self, messages: object, **kwargs: object) -> tuple[str, dict[str, int]]: + self.calls.append({"messages": messages, **kwargs}) + return self.content, self.usage + + +def _patch_client(monkeypatch: pytest.MonkeyPatch, client: _FakeClient) -> None: + monkeypatch.setattr( + "shared.services.retrieval.routing.classifier._resolve_client_and_model", + lambda: (client, "test-router-model"), + ) + + +@pytest.mark.asyncio +async def test_classifier_parses_classic_json(monkeypatch: pytest.MonkeyPatch) -> None: + client = _FakeClient('{"mode":"classic","reason":"single fact"}') + _patch_client(monkeypatch, client) + + decision = await classify_query_mode("What is the design flood of dam A?") + + assert decision.mode is QueryMode.CLASSIC + assert decision.source == "llm" + assert decision.reason == "single fact" + assert decision.model_name == "test-router-model" + assert decision.tokens_used == 4 + assert client.calls[0]["temperature"] == 0.0 + assert client.calls[0]["response_format"] == {"type": "json_object"} + assert client.calls[0]["model"] == "test-router-model" + + +@pytest.mark.asyncio +async def test_classifier_parses_agent_json_and_keeps_long_reason( + monkeypatch: pytest.MonkeyPatch, +) -> None: + reason = "cross-document comparison of every listed project" + client = _FakeClient(f'{{"mode":"agent","reason":"{reason}"}}') + _patch_client(monkeypatch, client) + + decision = await classify_query_mode( + "Compare the flood-control conclusions across these reports" + ) + + assert decision.mode is QueryMode.AGENT + assert decision.source == "llm" + assert decision.reason == reason + + +@pytest.mark.asyncio +async def test_classifier_invalid_mode_falls_back_to_agent( + monkeypatch: pytest.MonkeyPatch, +) -> None: + client = _FakeClient('{"mode":"hybrid","reason":"neither"}') + _patch_client(monkeypatch, client) + + decision = await classify_query_mode("anything") + + assert decision.mode is QueryMode.AGENT + assert decision.source == "fallback" + assert decision.reason == "" + assert decision.tokens_used == 4 + + +@pytest.mark.asyncio +async def test_classifier_invalid_json_falls_back_to_agent( + monkeypatch: pytest.MonkeyPatch, +) -> None: + client = _FakeClient("not-json") + _patch_client(monkeypatch, client) + + decision = await classify_query_mode("anything") + + assert decision.mode is QueryMode.AGENT + assert decision.source == "fallback" + assert decision.tokens_used == 4 + + +@pytest.mark.asyncio +async def test_classifier_exception_falls_back_to_agent( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class _BoomClient: + def chat_completion_with_usage(self, *_args: object, **_kwargs: object) -> tuple[str, dict[str, int]]: + raise TimeoutError("router timed out") + + monkeypatch.setattr( + "shared.services.retrieval.routing.classifier._resolve_client_and_model", + lambda: (_BoomClient(), "test-router-model"), + ) + + decision = await classify_query_mode("anything") + + assert decision.mode is QueryMode.AGENT + assert decision.source == "fallback" + assert decision.model_name == "test-router-model" + assert decision.tokens_used == 0 + + +@pytest.mark.asyncio +async def test_llm_mock_non_json_falls_back_to_agent( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(settings, "LLM_MOCK_ENABLED", True) + + decision = await classify_query_mode("Find the value of X") + + assert decision.mode is QueryMode.AGENT + assert decision.source == "fallback" + + +@pytest.mark.asyncio +async def test_projection_passes_router_decision() -> None: + public = await project_public_retrieval_response( + { + "namespace": "ns", + "query": "q", + "router_used": "classic_topk", + "router_decision": { + "mode": "classic", + "reason": "single fact", + "source": "llm", + "model_name": "test-router-model", + "tokens_used": 4, + }, + "results": [], + } + ) + + assert public["router_decision"] == { + "mode": "classic", + "reason": "single fact", + "source": "llm", + "model_name": "test-router-model", + "tokens_used": 4, + } + + +@pytest.mark.asyncio +async def test_projection_omits_router_decision_when_absent() -> None: + public = await project_public_retrieval_response( + { + "namespace": "ns", + "query": "q", + "router_used": "small_corpus_all", + "results": [], + } + ) + + assert "router_decision" not in public