e4dcfc49aa
Tests / Import Check (Python 3.13) (push) Has been cancelled
Tests / Import Check (Python 3.14) (push) Has been cancelled
Tests / Python Tests (Python 3.11) (push) Has been cancelled
Tests / Python Tests (Python 3.12) (push) Has been cancelled
Tests / Python Tests (Python 3.14) (push) Has been cancelled
Tests / Test Summary (push) Has been cancelled
Tests / Lint and Format (push) Has been cancelled
Tests / Web Node Tests (push) Has been cancelled
Tests / Import Check (Python 3.11) (push) Has been cancelled
Tests / Import Check (Python 3.12) (push) Has been cancelled
Tests / Python Tests (Python 3.13) (push) Has been cancelled
546 lines
22 KiB
Python
546 lines
22 KiB
Python
"""Agentic context investigator.
|
|
|
|
The investigator is a read-only pre-pass that runs *before* the answer loop's
|
|
first LLM call. Given the user's request and the manifest of the turn's
|
|
attached sources, it freely interleaves the ``read_source`` tool to load the
|
|
sources it actually needs, follows leads across them, and produces one
|
|
objective, third-person investigation the answer loop consumes as grounding.
|
|
|
|
``read_source`` lives **only** here: the answer loop no longer mounts it, so
|
|
loading attached-source full text is wholly owned by this pre-pass. The
|
|
investigation it returns is the answer loop's window into the sources — it
|
|
never streams CONTENT events, so it can never be mistaken for the turn's
|
|
user-facing answer.
|
|
|
|
Two execution paths:
|
|
|
|
* **Native tool calling** (:meth:`_run_loop`) — the agentic investigation: a
|
|
bounded loop of LLM call → ``read_source`` dispatch → repeat, until the model
|
|
stops calling tools and writes its investigation.
|
|
* **Fallback** (:meth:`_single_pass`) — for providers without native tool
|
|
calling (or when the loop fails): the original dump-the-source-text-and-brief
|
|
single pass, modelled on :class:`~deeptutor.agents.notebook.analysis_agent`.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from contextlib import suppress
|
|
from dataclasses import dataclass, field
|
|
import logging
|
|
from typing import Any, Callable
|
|
|
|
from deeptutor.core.agentic import (
|
|
LLMClientConfig,
|
|
build_completion_kwargs,
|
|
build_openai_client,
|
|
can_use_native_tool_calling,
|
|
dispatch_tool_calls,
|
|
)
|
|
from deeptutor.core.context import UnifiedContext
|
|
from deeptutor.core.stream_bus import StreamBus
|
|
from deeptutor.core.trace import build_trace_metadata, merge_trace_metadata, new_call_id
|
|
from deeptutor.runtime.registry.tool_registry import get_tool_registry
|
|
from deeptutor.services.llm import clean_thinking_tags, get_llm_config, get_token_limit_kwargs
|
|
from deeptutor.services.llm import stream as llm_stream
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Trace identity. ``source="chat"`` keeps the pre-pass in the turn's existing
|
|
# activity lane; the dedicated stage gives it its own labelled group. The
|
|
# frontend keys its "Exploring your context…" status and "Context exploration"
|
|
# row header off ``call_kind="context_exploration"`` / ``stage`` — see
|
|
# ``web/components/chat/home/TracePanels.tsx``.
|
|
EXPLORE_STAGE = "context_exploration"
|
|
EXPLORE_SOURCE = "chat"
|
|
|
|
# Agentic-loop budget: at most this many LLM rounds; the model normally
|
|
# finishes earlier by writing its investigation without a tool call. The last
|
|
# round runs with tools disabled so it is always forced to finish.
|
|
MAX_LOOP_ROUNDS = 5
|
|
LOOP_MAX_TOKENS = 2000
|
|
|
|
# Single-pass fallback budgets. The same sources remain reachable via
|
|
# ``read_source`` in the loop path, so clipping here only bounds the fallback.
|
|
MAX_SOURCES = 12
|
|
CHARS_PER_SOURCE = 8000
|
|
TOTAL_CHARS = 48000
|
|
BRIEFING_MAX_TOKENS = 1400
|
|
|
|
# Maps a manifest source-id prefix to a human kind label (en, zh).
|
|
_KIND_BY_PREFIX: dict[str, tuple[str, str]] = {
|
|
"hs-": ("Conversation transcript", "对话记录"),
|
|
"nb-": ("Notebook record", "笔记本记录"),
|
|
"bk-": ("Book excerpt", "书籍节选"),
|
|
"qb-": ("Question-bank entry", "题库条目"),
|
|
"at-": ("Document", "文档"),
|
|
}
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class _CallResult:
|
|
text: str = ""
|
|
tool_calls: list[dict[str, Any]] = field(default_factory=list)
|
|
output_chars: int = 0
|
|
|
|
|
|
class ContextExplorer:
|
|
"""Investigate the turn's attached sources and return an objective briefing."""
|
|
|
|
def __init__(self, *, language: str, prompts: dict[str, Any]) -> None:
|
|
self.language = "zh" if str(language or "en").lower().startswith("zh") else "en"
|
|
self._prompts = prompts or {}
|
|
cfg = get_llm_config()
|
|
self.model = getattr(cfg, "model", None)
|
|
self.api_key = getattr(cfg, "api_key", None)
|
|
self.base_url = getattr(cfg, "base_url", None)
|
|
self.api_version = getattr(cfg, "api_version", None)
|
|
self.binding = getattr(cfg, "binding", None) or "openai"
|
|
self.extra_headers = getattr(cfg, "extra_headers", None) or {}
|
|
self.reasoning_effort = getattr(cfg, "reasoning_effort", None)
|
|
self.registry = get_tool_registry()
|
|
self._client_config = LLMClientConfig(
|
|
binding=self.binding,
|
|
model=self.model,
|
|
api_key=self.api_key,
|
|
base_url=self.base_url,
|
|
api_version=self.api_version,
|
|
extra_headers=self.extra_headers or None,
|
|
reasoning_effort=self.reasoning_effort,
|
|
)
|
|
|
|
async def investigate(
|
|
self,
|
|
*,
|
|
context: UnifiedContext,
|
|
stream: StreamBus,
|
|
usage: Any | None = None,
|
|
) -> str:
|
|
"""Run the pre-pass and return the investigation wrapped in its header.
|
|
|
|
Returns ``""`` when there is nothing to investigate or every path
|
|
fails, so the caller can simply skip injection.
|
|
"""
|
|
source_index = self._source_index(context)
|
|
if not source_index:
|
|
return ""
|
|
|
|
investigation = ""
|
|
if can_use_native_tool_calling(binding=self.binding, model=self.model):
|
|
try:
|
|
investigation = await self._run_loop(context, stream, source_index, usage)
|
|
except Exception:
|
|
logger.warning(
|
|
"context exploration loop failed; falling back to single pass",
|
|
exc_info=True,
|
|
)
|
|
investigation = ""
|
|
# Non-native providers, or a failed/empty loop, degrade to the robust
|
|
# dump-and-brief single pass so the answer loop is never left without
|
|
# grounding (it can no longer read sources itself).
|
|
if not investigation.strip():
|
|
investigation = await self._single_pass(context, stream, source_index, usage)
|
|
|
|
investigation = clean_thinking_tags(investigation, self.binding, self.model).strip()
|
|
if not investigation:
|
|
return ""
|
|
return f"{self._briefing_header()}\n\n{investigation}".strip()
|
|
|
|
# ---- agentic loop ----------------------------------------------------
|
|
|
|
async def _run_loop(
|
|
self,
|
|
context: UnifiedContext,
|
|
stream: StreamBus,
|
|
source_index: dict[str, str],
|
|
usage: Any | None,
|
|
) -> str:
|
|
system_prompt = self._t("loop.system")
|
|
user_template = self._t("loop.user_template")
|
|
if not system_prompt or not user_template:
|
|
logger.warning("explore_context loop prompts missing; using single pass")
|
|
return ""
|
|
|
|
messages: list[dict[str, Any]] = [
|
|
{"role": "system", "content": system_prompt},
|
|
{
|
|
"role": "user",
|
|
"content": user_template.format(
|
|
question=(context.user_message or "").strip() or "(empty)",
|
|
mode=str(context.active_capability or "chat"),
|
|
manifest=(context.source_manifest or "").strip() or "(none)",
|
|
),
|
|
},
|
|
]
|
|
tool_schemas = self._read_source_schemas(source_index)
|
|
augmenter = self._augmenter(source_index)
|
|
client = build_openai_client(self._client_config)
|
|
|
|
call_id = new_call_id("explore-context")
|
|
stage_meta = build_trace_metadata(
|
|
call_id=call_id,
|
|
phase=EXPLORE_STAGE,
|
|
label=self._status_exploring(),
|
|
call_kind="context_exploration",
|
|
trace_id=call_id,
|
|
trace_role="explore",
|
|
trace_group="stage",
|
|
)
|
|
chunk_meta = merge_trace_metadata(stage_meta, {"trace_kind": "llm_chunk"})
|
|
|
|
investigation = ""
|
|
total_in = 0
|
|
total_out = 0
|
|
async with stream.stage(EXPLORE_STAGE, source=EXPLORE_SOURCE, metadata=stage_meta):
|
|
await stream.progress(
|
|
self._status_exploring(),
|
|
source=EXPLORE_SOURCE,
|
|
stage=EXPLORE_STAGE,
|
|
metadata=merge_trace_metadata(
|
|
stage_meta, {"trace_kind": "call_status", "call_state": "running"}
|
|
),
|
|
)
|
|
for round_idx in range(MAX_LOOP_ROUNDS):
|
|
is_last = round_idx == MAX_LOOP_ROUNDS - 1
|
|
if is_last:
|
|
# Budget exhausted while still calling tools — force a
|
|
# tool-less finish from what has been gathered.
|
|
messages.append({"role": "user", "content": self._forced_finish_instruction()})
|
|
total_in += sum(_content_chars(m) for m in messages)
|
|
result = await self._call_llm(
|
|
client, messages, tool_schemas if not is_last else None, chunk_meta, stream
|
|
)
|
|
total_out += result.output_chars
|
|
if not result.tool_calls:
|
|
investigation = result.text
|
|
break
|
|
messages.append(_assistant_with_tool_calls(result.text, result.tool_calls))
|
|
dispatch = await dispatch_tool_calls(
|
|
tool_calls=result.tool_calls,
|
|
context=context,
|
|
stream=stream,
|
|
source=EXPLORE_SOURCE,
|
|
stage=EXPLORE_STAGE,
|
|
iteration_index=round_idx,
|
|
registry=self.registry,
|
|
kwarg_augmenter=augmenter,
|
|
tool_call_label=self._t("labels.tool_call", default="Tool call"),
|
|
trace_id_prefix="explore-context",
|
|
)
|
|
messages.extend(dispatch.tool_messages)
|
|
await stream.progress(
|
|
"",
|
|
source=EXPLORE_SOURCE,
|
|
stage=EXPLORE_STAGE,
|
|
metadata=merge_trace_metadata(
|
|
stage_meta, {"trace_kind": "call_status", "call_state": "complete"}
|
|
),
|
|
)
|
|
|
|
self._account_usage(usage, total_in, total_out, investigation)
|
|
return investigation
|
|
|
|
async def _call_llm(
|
|
self,
|
|
client: Any,
|
|
messages: list[dict[str, Any]],
|
|
tool_schemas: list[dict[str, Any]] | None,
|
|
chunk_meta: dict[str, Any],
|
|
stream: StreamBus,
|
|
) -> _CallResult:
|
|
"""One streamed LLM call. All output streams to the *thinking* channel
|
|
(never CONTENT — that is the answer loop's channel); the returned text
|
|
is the round's investigation when it carries no tool calls."""
|
|
kwargs: dict[str, Any] = {
|
|
"model": self.model,
|
|
"messages": messages,
|
|
"stream": True,
|
|
**build_completion_kwargs(
|
|
temperature=0.2,
|
|
model=self.model,
|
|
max_tokens=LOOP_MAX_TOKENS,
|
|
binding=self.binding,
|
|
reasoning_effort=self.reasoning_effort,
|
|
),
|
|
}
|
|
if tool_schemas:
|
|
kwargs["tools"] = tool_schemas
|
|
kwargs["tool_choice"] = "auto"
|
|
|
|
text_parts: list[str] = []
|
|
tool_acc: dict[int, dict[str, str]] = {}
|
|
output_chars = 0
|
|
response_stream = await client.chat.completions.create(**kwargs)
|
|
try:
|
|
async for chunk in response_stream:
|
|
choices = getattr(chunk, "choices", None) or []
|
|
if not choices:
|
|
continue
|
|
delta = getattr(choices[0], "delta", None)
|
|
if delta is None:
|
|
continue
|
|
reasoning = getattr(delta, "reasoning_content", None) or getattr(
|
|
delta, "reasoning", None
|
|
)
|
|
if reasoning:
|
|
output_chars += len(reasoning)
|
|
await stream.thinking(
|
|
reasoning, source=EXPLORE_SOURCE, stage=EXPLORE_STAGE, metadata=chunk_meta
|
|
)
|
|
content = getattr(delta, "content", None)
|
|
if content:
|
|
output_chars += len(content)
|
|
text_parts.append(content)
|
|
await stream.thinking(
|
|
content, source=EXPLORE_SOURCE, stage=EXPLORE_STAGE, metadata=chunk_meta
|
|
)
|
|
for tc in getattr(delta, "tool_calls", None) or []:
|
|
index = int(getattr(tc, "index", 0) or 0)
|
|
acc = tool_acc.setdefault(index, {"id": "", "name": "", "arguments": ""})
|
|
tcid = getattr(tc, "id", None)
|
|
if tcid:
|
|
acc["id"] += str(tcid)
|
|
fn = getattr(tc, "function", None)
|
|
if fn is None:
|
|
continue
|
|
name = getattr(fn, "name", None)
|
|
arguments = getattr(fn, "arguments", None)
|
|
if name:
|
|
acc["name"] += str(name)
|
|
output_chars += len(str(name))
|
|
if arguments:
|
|
acc["arguments"] += str(arguments)
|
|
output_chars += len(str(arguments))
|
|
finally:
|
|
close = getattr(response_stream, "close", None)
|
|
if callable(close):
|
|
with suppress(Exception):
|
|
await close()
|
|
|
|
tool_calls = [
|
|
{
|
|
"id": data.get("id") or f"call_{idx}",
|
|
"name": data.get("name", ""),
|
|
"arguments": data.get("arguments") or "{}",
|
|
}
|
|
for idx, data in sorted(tool_acc.items())
|
|
if data.get("name")
|
|
]
|
|
return _CallResult(
|
|
text="".join(text_parts), tool_calls=tool_calls, output_chars=output_chars
|
|
)
|
|
|
|
def _read_source_schemas(self, source_index: dict[str, str]) -> list[dict[str, Any]]:
|
|
schemas = self.registry.build_openai_schemas(["read_source"])
|
|
source_ids = sorted(source_index.keys())
|
|
for schema in schemas:
|
|
function = schema.get("function") if isinstance(schema, dict) else None
|
|
if not isinstance(function, dict):
|
|
continue
|
|
parameters = function.get("parameters")
|
|
if not isinstance(parameters, dict):
|
|
continue
|
|
properties = parameters.get("properties") or {}
|
|
if (
|
|
function.get("name") == "read_source"
|
|
and isinstance(properties.get("source_id"), dict)
|
|
and source_ids
|
|
):
|
|
properties["source_id"]["enum"] = source_ids
|
|
parameters["additionalProperties"] = False
|
|
return schemas
|
|
|
|
@staticmethod
|
|
def _augmenter(source_index: dict[str, str]) -> Callable[..., dict[str, Any]]:
|
|
def _augment(tool_name: str, args: dict[str, Any], _ctx: UnifiedContext) -> dict[str, Any]:
|
|
kwargs = dict(args)
|
|
if tool_name == "read_source":
|
|
kwargs["source_index"] = source_index
|
|
return kwargs
|
|
|
|
return _augment
|
|
|
|
# ---- single-pass fallback -------------------------------------------
|
|
|
|
async def _single_pass(
|
|
self,
|
|
context: UnifiedContext,
|
|
stream: StreamBus,
|
|
source_index: dict[str, str],
|
|
usage: Any | None,
|
|
) -> str:
|
|
sources_text = self._render_source_blocks(source_index)
|
|
if not sources_text:
|
|
return ""
|
|
system_prompt = self._t("system")
|
|
user_template = self._t("user_template")
|
|
if not system_prompt or not user_template:
|
|
logger.warning("explore_context single-pass prompts missing; skipping pre-pass")
|
|
return ""
|
|
user_prompt = user_template.format(
|
|
question=(context.user_message or "").strip() or "(empty)",
|
|
mode=str(context.active_capability or "chat"),
|
|
manifest=(context.source_manifest or "").strip() or "(none)",
|
|
sources=sources_text,
|
|
)
|
|
|
|
call_id = new_call_id("explore-context")
|
|
stage_meta = build_trace_metadata(
|
|
call_id=call_id,
|
|
phase=EXPLORE_STAGE,
|
|
label=self._status_exploring(),
|
|
call_kind="context_exploration",
|
|
trace_id=call_id,
|
|
trace_role="explore",
|
|
trace_group="stage",
|
|
)
|
|
chunk_meta = merge_trace_metadata(stage_meta, {"trace_kind": "llm_chunk"})
|
|
chunks: list[str] = []
|
|
async with stream.stage(EXPLORE_STAGE, source=EXPLORE_SOURCE, metadata=stage_meta):
|
|
await stream.progress(
|
|
self._status_exploring(),
|
|
source=EXPLORE_SOURCE,
|
|
stage=EXPLORE_STAGE,
|
|
metadata=merge_trace_metadata(
|
|
stage_meta, {"trace_kind": "call_status", "call_state": "running"}
|
|
),
|
|
)
|
|
try:
|
|
async for chunk in llm_stream(
|
|
prompt=user_prompt,
|
|
system_prompt=system_prompt,
|
|
model=self.model,
|
|
api_key=self.api_key,
|
|
base_url=self.base_url,
|
|
api_version=self.api_version,
|
|
binding=self.binding,
|
|
temperature=0.2,
|
|
**self._token_kwargs(BRIEFING_MAX_TOKENS),
|
|
):
|
|
if not chunk:
|
|
continue
|
|
chunks.append(chunk)
|
|
await stream.thinking(
|
|
chunk, source=EXPLORE_SOURCE, stage=EXPLORE_STAGE, metadata=chunk_meta
|
|
)
|
|
except Exception:
|
|
logger.warning("context exploration single pass failed", exc_info=True)
|
|
await stream.progress(
|
|
"",
|
|
source=EXPLORE_SOURCE,
|
|
stage=EXPLORE_STAGE,
|
|
metadata=merge_trace_metadata(
|
|
stage_meta, {"trace_kind": "call_status", "call_state": "complete"}
|
|
),
|
|
)
|
|
|
|
briefing = clean_thinking_tags("".join(chunks), self.binding, self.model).strip()
|
|
self._account_usage(usage, len(system_prompt) + len(user_prompt), len(briefing), briefing)
|
|
return briefing
|
|
|
|
def _render_source_blocks(self, source_index: dict[str, str]) -> str:
|
|
blocks: list[str] = []
|
|
total = 0
|
|
for sid, text in source_index.items():
|
|
body = str(text or "").strip()
|
|
if not body:
|
|
continue
|
|
if len(blocks) >= MAX_SOURCES or total >= TOTAL_CHARS:
|
|
break
|
|
remaining = TOTAL_CHARS - total
|
|
clipped = self._clip(body, min(CHARS_PER_SOURCE, remaining))
|
|
blocks.append(f"### [{sid}] ({self._kind_label(sid)})\n{clipped}")
|
|
total += len(clipped)
|
|
return "\n\n".join(blocks)
|
|
|
|
def _kind_label(self, sid: str) -> str:
|
|
for prefix, (en, zh) in _KIND_BY_PREFIX.items():
|
|
if sid.startswith(prefix):
|
|
return zh if self.language == "zh" else en
|
|
return "来源" if self.language == "zh" else "Source"
|
|
|
|
def _clip(self, text: str, limit: int) -> str:
|
|
text = (text or "").strip()
|
|
if limit <= 0:
|
|
return ""
|
|
if len(text) <= limit:
|
|
return text
|
|
note = "\n…(已截断)" if self.language == "zh" else "\n…(truncated)"
|
|
return text[:limit].rstrip() + note
|
|
|
|
# ---- shared helpers --------------------------------------------------
|
|
|
|
def _account_usage(
|
|
self, usage: Any | None, input_chars: int, output_chars: int, produced: str
|
|
) -> None:
|
|
if usage is None or not produced.strip():
|
|
return
|
|
try:
|
|
usage.add_estimated(input_chars=input_chars, output_chars=output_chars)
|
|
except Exception: # pragma: no cover - usage accounting is best-effort
|
|
logger.debug("explore_context usage accounting failed", exc_info=True)
|
|
|
|
@staticmethod
|
|
def _source_index(context: UnifiedContext) -> dict[str, str]:
|
|
idx = context.metadata.get("source_index")
|
|
if isinstance(idx, dict) and idx:
|
|
return {str(k): str(v) for k, v in idx.items()}
|
|
return {}
|
|
|
|
def _token_kwargs(self, max_tokens: int) -> dict[str, Any]:
|
|
if not self.model:
|
|
return {}
|
|
return get_token_limit_kwargs(self.model, max_tokens)
|
|
|
|
def _t(self, key: str, default: str = "") -> str:
|
|
value: Any = self._prompts
|
|
for part in key.split("."):
|
|
if not isinstance(value, dict) or part not in value:
|
|
return default
|
|
value = value[part]
|
|
return value.strip() if isinstance(value, str) else default
|
|
|
|
def _forced_finish_instruction(self) -> str:
|
|
return self._t(
|
|
"loop.forced_finish",
|
|
default=(
|
|
"Investigation budget reached. Stop calling tools and write your "
|
|
"objective, detailed investigation now from what you have gathered."
|
|
),
|
|
)
|
|
|
|
def _status_exploring(self) -> str:
|
|
return self._t(
|
|
"status.exploring",
|
|
default="上下文调查" if self.language == "zh" else "Context exploration",
|
|
)
|
|
|
|
def _briefing_header(self) -> str:
|
|
return self._t("briefing_header", default="[Context Investigation]")
|
|
|
|
|
|
def _content_chars(message: dict[str, Any]) -> int:
|
|
content = message.get("content")
|
|
if isinstance(content, str):
|
|
return len(content)
|
|
if isinstance(content, list):
|
|
return sum(len(str(part.get("text") or "")) for part in content if isinstance(part, dict))
|
|
return 0
|
|
|
|
|
|
def _assistant_with_tool_calls(content: str, tool_calls: list[dict[str, Any]]) -> dict[str, Any]:
|
|
return {
|
|
"role": "assistant",
|
|
"content": content or None,
|
|
"tool_calls": [
|
|
{
|
|
"id": tc["id"],
|
|
"type": "function",
|
|
"function": {"name": tc["name"], "arguments": tc.get("arguments") or "{}"},
|
|
}
|
|
for tc in tool_calls
|
|
],
|
|
}
|
|
|
|
|
|
__all__ = ["ContextExplorer", "EXPLORE_STAGE"]
|