soc_ai/agent/targeted_investigator.py
Rama main · réplica de lectura del repositorio original.
"""Phase D — targeted investigator.
Replaces the multi-tool LLM-driven investigator agent for the synth-first
pipeline. The synth round-1 emits a TargetedGap naming ONE tool + exact
args. This module dispatches that tool deterministically — no LLM call,
no tool selection, just argument unpacking + dispatch.
Why deterministic? The 30B was burning 6K-15K reasoning-trace tokens per
investigator turn. Phase D removes the LLM from the dispatch loop
entirely; the synth (120B) decides what to call, and we just call it.
The bytes saved are the long-tail timeout fix.
"""
from __future__ import annotations
import inspect
import json
import logging
from collections.abc import Callable
from typing import Any, NoReturn, cast
from soc_ai.agent.toolset import (
PHASE_D_TOOLS,
_clamp_tool_result,
_egress_tool_idents,
strip_synth_markers,
)
from soc_ai.agent.triage import TargetedGap
_LOGGER = logging.getLogger(__name__)
# Per-tool argument ceilings, mirroring the clamps register_read_tools enforces
# inline in the interactive path (toolset.py). Phase-D dispatch calls the tool
# functions DIRECTLY, bypassing those wrappers, so we must re-apply the SAME
# caps here — otherwise an oversized max_results/k (a synth mistake, or a
# prompt-injected steer from attacker-influenceable alert data) pulls thousands
# of full docs back and blows the round-2 synth prompt past the model's context.
# Keep in sync with the min(...) clamps in soc_ai/agent/toolset.py.
_PHASE_D_ARG_CEILINGS: dict[str, dict[str, int]] = {
"t_query_events_oql": {"max_results": 25},
"t_query_zeek_logs": {"max_results": 25},
"t_query_cases": {"max_results": 10},
"t_query_detections": {"max_results": 10},
"t_get_playbooks": {"max_results": 10},
"t_lookup_runbook": {"k": 5},
# The pcap window drives one ssh + tcpdump pass per ring file it overlaps on
# the production sensor, so an oversized window is sensor load, not just
# prompt bytes.
"t_get_pcap": {"window_minutes": 60},
}
def _clamp_arg_ceilings(tool_name: str, tool_args: dict[str, Any]) -> dict[str, Any]:
"""Return ``tool_args`` with any capped arg clamped to its per-tool ceiling.
Returns the SAME dict object untouched when nothing needs clamping (so the
common path stays byte-identical); otherwise a shallow copy with the
offending ints lowered. Never mutates the caller's ``gap.tool_args``.
"""
ceilings = _PHASE_D_ARG_CEILINGS.get(tool_name)
if not ceilings:
return tool_args
clamped: dict[str, Any] | None = None
for arg, cap in ceilings.items():
val = tool_args.get(arg)
# bool is an int subclass — exclude it so a stray True/False isn't "clamped".
if isinstance(val, int) and not isinstance(val, bool) and val > cap:
if clamped is None:
clamped = dict(tool_args)
clamped[arg] = cap
return clamped if clamped is not None else tool_args
def build_targeted_investigator_prompt(gap: TargetedGap) -> str:
"""Stub prompt — kept around for audit logs / Phase D alternatives.
The default Phase D dispatch path is deterministic (no LLM call); this
prompt is only used if a future variant wants to A/B against an LLM-
driven Phase D where the 30B confirms the synth's tool selection.
"""
return (
"You are a targeted-tool executor. The synthesizer needs ONE specific\n"
"piece of evidence to finalize a triage verdict. Run this ONE tool call\n"
"and return the raw result. Do not call any other tool.\n\n"
f"Question: {gap.question}\n"
f"Tool: {gap.tool_name}\n"
f"Args: {json.dumps(gap.tool_args)}\n"
f"Why: {gap.why_this_matters}\n"
)
async def run_targeted_investigation(
gap: TargetedGap,
*,
ctx: Any,
) -> dict[str, Any] | str:
"""Dispatch the single tool the synth requested. Return its raw result.
On any error (unknown tool, dispatch failure, tool exception), return
a string message describing the failure — synth round 2 sees the
failure and can still emit a verdict.
"""
try:
return await _dispatch_named_tool(gap.tool_name, gap.tool_args, ctx)
except Exception as e:
_LOGGER.exception("targeted investigation dispatch failed")
# This string lands verbatim in the round-2 synth prompt; an ES error
# can name the failing index, so it passes the same model boundary the
# success path does (_clamp_tool_result strips it there).
return strip_synth_markers(f"targeted dispatch error: {type(e).__name__}: {e}")
def _raise_arg_type_error(tool_name: str, args: dict[str, Any], cause: TypeError) -> NoReturn:
"""Raise a structured TypeError so run_targeted_investigation's error
channel produces an actionable message that names the tool, the offending
args, and the exception — letting synth round 2 re-call with corrected
types instead of silently missing the evidence.
"""
# The message lands in the round-2 synth prompt, so a huge argument dump
# is bounded. The bracket marker tells the model the repr is cut.
args_repr = repr(args)
if len(args_repr) > 400:
args_repr = args_repr[:400] + " [truncated]"
raise TypeError(
f"tool {tool_name} rejected arguments {args_repr}: {cause}"
" — re-call with corrected argument types"
) from cause
def _dispatch_table() -> dict[str, Callable[..., Any]]:
"""The Phase-D tool_name → callable table. Imported lazily to dodge cycles.
Its key set is drift-tested against :data:`PHASE_D_TOOLS`
(tests/test_toolset.py) — add new Phase-D tools in BOTH places.
"""
from soc_ai.tools.crawl_page import crawl_page # noqa: PLC0415
from soc_ai.tools.enrichment import enrich_domain, enrich_hash, enrich_ip # noqa: PLC0415
from soc_ai.tools.get_event_raw import get_event_raw # noqa: PLC0415
from soc_ai.tools.get_pcap import get_pcap_facts # noqa: PLC0415
from soc_ai.tools.get_playbooks import get_playbooks # noqa: PLC0415
from soc_ai.tools.get_rule_content import get_rule_content # noqa: PLC0415
from soc_ai.tools.lookup_runbook import lookup_runbook # noqa: PLC0415
from soc_ai.tools.query_cases import query_cases # noqa: PLC0415
from soc_ai.tools.query_detections import query_detections # noqa: PLC0415
from soc_ai.tools.query_events import query_events_oql # noqa: PLC0415
from soc_ai.tools.query_zeek import query_zeek_logs # noqa: PLC0415
from soc_ai.tools.web_search import web_search # noqa: PLC0415
return {
"t_enrich_ip": enrich_ip,
"t_enrich_domain": enrich_domain,
"t_enrich_hash": enrich_hash,
"t_query_zeek_logs": query_zeek_logs,
"t_query_events_oql": query_events_oql,
"t_query_cases": query_cases,
"t_query_detections": query_detections,
"t_get_rule_content": get_rule_content,
"t_get_event_raw": get_event_raw,
"t_get_playbooks": get_playbooks,
"t_lookup_runbook": lookup_runbook,
"t_get_pcap": get_pcap_facts,
"t_web_search": web_search,
"t_crawl_page": crawl_page,
}
async def _dispatch_named_tool(
tool_name: str,
tool_args: dict[str, Any],
ctx: Any,
) -> dict[str, Any]:
"""Validate tool_name against PHASE_D_TOOLS, then invoke it with kwargs."""
if tool_name not in PHASE_D_TOOLS:
raise ValueError(f"unknown tool {tool_name!r}")
# The table's key set == PHASE_D_TOOLS (drift-tested), so this lookup
# cannot miss for a validated name.
fn = _dispatch_table()[tool_name]
# Re-apply the interactive path's per-tool arg ceilings (toolset.py). The
# tool functions here are called directly, so their wrapper clamps (e.g.
# max_results=min(..., 25)) don't run — do it before dispatch.
tool_args = _clamp_arg_ceilings(tool_name, tool_args)
base_kwargs: dict[str, Any] = {"settings": ctx.settings}
if tool_name in {"t_enrich_ip", "t_enrich_domain", "t_enrich_hash"}:
base_kwargs["misp"] = getattr(ctx, "misp", None)
base_kwargs["blocklist"] = getattr(ctx, "blocklist", None)
if tool_name == "t_enrich_ip":
base_kwargs["maxmind"] = getattr(ctx, "maxmind", None)
base_kwargs["cloud"] = getattr(ctx, "cloud", None)
elif tool_name in {
"t_query_zeek_logs",
"t_query_events_oql",
"t_query_cases",
"t_query_detections",
"t_get_rule_content",
"t_get_event_raw",
"t_get_playbooks",
}:
base_kwargs["elastic"] = ctx.elastic
base_kwargs["auth"] = getattr(ctx, "auth", None)
# Center the windowed query tools (query_events_oql / query_zeek_logs) on
# the alert's @timestamp, exactly as the interactive wrappers do
# (toolset.py). Without it these fall to their now-relative branch
# (_build_time_filter with time_anchor=None -> [now-Nm, now]) and return
# zero hits for any batch/older alert. Family members that don't accept
# `time_anchor` (query_cases / query_detections / get_rule_content /
# get_event_raw / get_playbooks) drop it via the signature filter below —
# the same mechanism that already drops the injected `auth`.
base_kwargs["time_anchor"] = getattr(ctx, "default_time_anchor", None)
# Thread the synth-doc visibility scope exactly as the interactive
# wrappers do (toolset.py passes include_synth=ctx.include_synth on
# every query tool). Without it query_events_oql / query_zeek_logs fall
# to their prod default (include_synth=False), which injects `must_not
# exists synth.scenario_id` — so an eval-mode Phase-D query
# STRUCTURALLY cannot return the scenario's own planted documents: it
# returns zero hits and round-2 concludes there is nothing there.
# (t_get_event_raw fetches by _id with no filter, which is why traces
# still showed planted docs arriving by that route.) The default False
# keeps the prod kill-switch for any context without an explicit
# opt-in; family members without the kwarg drop it via the signature
# filter below, like `auth` and `time_anchor`.
base_kwargs["include_synth"] = getattr(ctx, "include_synth", False)
elif tool_name == "t_lookup_runbook":
# Operator-runbook search hits the local store, not ES — inject the
# session factory instead of elastic/auth. ``settings`` is unused by
# this tool, so drop it too (its signature has no **kwargs).
base_kwargs = {"db_sessionmaker": getattr(ctx, "db_sessionmaker", None)}
elif tool_name == "t_get_pcap":
# get_pcap_facts needs only settings (already in base_kwargs) +
# alert_ts from the context time anchor.
base_kwargs["alert_ts"] = getattr(ctx, "default_time_anchor", None)
elif tool_name in {"t_web_search", "t_crawl_page"}:
# The online egress guards recognise internal identifiers by the
# EFFECTIVE sets (env config unioned with the DB-discovered hosts and
# suffixes), exactly as the interactive wrappers pass them. Left to
# the raw settings tuples, a discovered internal name in the synth's
# args would go out to the search engine or be crawled server-side.
base_kwargs["suffixes"], base_kwargs["extra_hosts"] = await _egress_tool_idents(ctx)
# Drop any injected base kwarg the target signature does not accept. The
# injection tables above are keyed by tool NAME, not by inspecting each
# signature, so a dependency that only some tools in a family take (e.g.
# ``auth`` — the SO web-API session — which the ES-query family does NOT
# accept, they are elastic-only) would otherwise raise an unconditional
# ``TypeError: unexpected keyword argument`` on every real dispatch. Honor
# ``**kwargs``: a function that declares VAR_KEYWORD accepts anything.
_sig = inspect.signature(fn)
if not any(p.kind == inspect.Parameter.VAR_KEYWORD for p in _sig.parameters.values()):
base_kwargs = {k: v for k, v in base_kwargs.items() if k in _sig.parameters}
# Lenient kwarg dispatch. Some reasoning models can hallucinate
# tool-arg names (e.g. ``'社区ID'``,
# ``'filter'``, ``'dataset'`` — all observed). Rather than
# fail-stop the whole Phase D dispatch on the first unknown kwarg,
# drop the unknown keys and retry once. Honors VAR_KEYWORD if the
# target function accepts ``**kwargs``.
try:
raw = await fn(**tool_args, **base_kwargs)
except TypeError as exc:
sig = inspect.signature(fn)
has_var_keyword = any(
p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values()
)
accepted = {n for n in sig.parameters if n not in base_kwargs}
if has_var_keyword:
# Function accepts **kwargs — the TypeError must be something
# else (wrong type for an accepted arg, missing positional).
_raise_arg_type_error(tool_name, tool_args, exc)
dropped = sorted(set(tool_args) - accepted)
if not dropped:
# No unknown kwargs — the error is a wrong-typed value for a valid
# kwarg name. Produce an actionable structured message so synth
# round 2 can re-call with corrected types instead of silently
# missing the evidence.
_raise_arg_type_error(tool_name, tool_args, exc)
_LOGGER.warning(
"targeted dispatch on %s: dropping hallucinated kwargs %s and retrying (was: %s)",
tool_name,
dropped,
exc,
)
cleaned_args = {k: v for k, v in tool_args.items() if k in accepted}
try:
raw = await fn(**cleaned_args, **base_kwargs)
except TypeError as type_err:
# Wrong-typed value for a valid kwarg survives the retry.
# Same structured error path.
_raise_arg_type_error(tool_name, cleaned_args, type_err)
if hasattr(raw, "model_dump"):
# `raw` is a pydantic model here; model_dump(mode="json") yields a dict.
result = cast("dict[str, Any]", raw.model_dump(mode="json"))
elif isinstance(raw, dict):
result = raw
else:
result = {"result": raw}
# Clamp to the per-tool budget, exactly as the interactive wrappers do,
# before this result is embedded verbatim in the round-2 synth prompt.
return _clamp_tool_result(result)
__all__ = [
"build_targeted_investigator_prompt",
"run_targeted_investigation",
]