refactor(dedupe): route through MultiProvider + cache wrapper + retry policy

``check_duplicate`` was calling ``litellm.completion(...)`` directly
via ``resolve_llm_config()``, bypassing every layer the main agent
loop runs through:

- :class:`MultiProvider` (so ``anthropic/...`` aliases never went
  through :class:`AnthropicCachingLitellmModel` and missed the
  ``cache_control`` patching on the system prompt — 4x cost on
  repeated dedupe calls within the same scan).
- :data:`DEFAULT_RETRY` (no retry on 429s / network blips — the
  caller's broad except-and-fallback was hiding this).

Switch to the SDK's :meth:`Model.get_response` directly: same model
selection, same retry policy, same cache wrapper. Extract assistant
text from ``ModelResponse.output`` via the canonical
``ResponseOutputMessage`` walk.

``check_duplicate`` is now async — drops the ``asyncio.to_thread``
indirection in ``_do_create``. Validation logic is fast-sync; running
it on the event loop is fine.

Drive-by: rename ``_DEFAULT_RETRY`` → ``DEFAULT_RETRY`` in
``run_config_factory`` so the dedupe path can reuse the same constant
without reaching into a private name.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
0xallam
2026-04-25 15:25:44 -07:00
parent b3f7cfd040
commit 8a11f9dab5
3 changed files with 80 additions and 42 deletions
+72 -34
View File
@@ -1,10 +1,28 @@
"""LLM-based vulnerability-report deduplication.
Routes through the same :class:`MultiProvider` (so ``anthropic/...``
models pick up :class:`AnthropicCachingLitellmModel`'s cache_control
patching) and :data:`DEFAULT_RETRY` policy as the main agent loop —
no parallel litellm code path.
"""
from __future__ import annotations
import json
import logging
from typing import Any
from typing import TYPE_CHECKING, Any
import litellm
from agents.model_settings import ModelSettings
from agents.models.interface import ModelTracing
from openai.types.responses import ResponseOutputMessage
from strix.config.config import resolve_llm_config
from strix.config.config import Config
from strix.llm.multi_provider_setup import build_multi_provider
from strix.run_config_factory import DEFAULT_RETRY
if TYPE_CHECKING:
from agents.items import ModelResponse
logger = logging.getLogger(__name__)
@@ -126,7 +144,25 @@ def _parse_dedupe_response(content: str) -> dict[str, Any]:
}
def check_duplicate(
def _extract_text(response: ModelResponse) -> str:
"""Concatenate ``output_text`` fragments across every message item.
The SDK returns OpenAI Responses-API-shaped output; for a plain
chat-completion the assistant message has a list of content parts,
each of which carries a ``.text`` attribute we can pull verbatim.
"""
parts: list[str] = []
for item in response.output:
if not isinstance(item, ResponseOutputMessage):
continue
for chunk in item.content:
text = getattr(chunk, "text", None)
if text:
parts.append(text)
return "".join(parts)
async def check_duplicate(
candidate: dict[str, Any], existing_reports: list[dict[str, Any]]
) -> dict[str, Any]:
if not existing_reports:
@@ -138,39 +174,39 @@ def check_duplicate(
}
try:
model_name = Config.get("strix_llm")
if not model_name:
return {
"is_duplicate": False,
"duplicate_id": "",
"confidence": 0.0,
"reason": "STRIX_LLM not configured; skipping dedupe check",
}
candidate_cleaned = _prepare_report_for_comparison(candidate)
existing_cleaned = [_prepare_report_for_comparison(r) for r in existing_reports]
comparison_data = {"candidate": candidate_cleaned, "existing_reports": existing_cleaned}
model_name, api_key, api_base = resolve_llm_config()
litellm_model: str | None = model_name
user_msg = (
f"Compare this candidate vulnerability against existing reports:\n\n"
f"{json.dumps(comparison_data, indent=2)}\n\n"
f"Respond with ONLY the JSON object described in the system prompt."
)
messages = [
{"role": "system", "content": DEDUPE_SYSTEM_PROMPT},
{
"role": "user",
"content": (
f"Compare this candidate vulnerability against existing reports:\n\n"
f"{json.dumps(comparison_data, indent=2)}\n\n"
f"Respond with ONLY the JSON object described in the system prompt."
),
},
]
completion_kwargs: dict[str, Any] = {
"model": litellm_model,
"messages": messages,
"timeout": 120,
}
if api_key:
completion_kwargs["api_key"] = api_key
if api_base:
completion_kwargs["api_base"] = api_base
response = litellm.completion(**completion_kwargs)
content = response.choices[0].message.content
model = build_multi_provider().get_model(model_name)
response = await model.get_response(
system_instructions=DEDUPE_SYSTEM_PROMPT,
input=user_msg,
model_settings=ModelSettings(retry=DEFAULT_RETRY),
tools=[],
output_schema=None,
handoffs=[],
tracing=ModelTracing.DISABLED,
previous_response_id=None,
conversation_id=None,
prompt=None,
)
content = _extract_text(response)
if not content:
return {
"is_duplicate": False,
@@ -182,8 +218,10 @@ def check_duplicate(
result = _parse_dedupe_response(content)
logger.info(
f"Deduplication check: is_duplicate={result['is_duplicate']}, "
f"confidence={result['confidence']}, reason={result['reason'][:100]}"
"Deduplication check: is_duplicate=%s, confidence=%.2f, reason=%s",
result["is_duplicate"],
result["confidence"],
result["reason"][:100],
)
except Exception as e:
+5 -3
View File
@@ -33,8 +33,10 @@ STRIX_DEFAULT_MAX_TURNS = 300
# Retry: 5 attempts with ``min(90, 2*2^n)`` backoff. 4xx auth/validation
# errors are excluded from the retryable status list — they can't be
# fixed by retrying and should fail fast.
_DEFAULT_RETRY = ModelRetrySettings(
# fixed by retrying and should fail fast. Public so the dedupe path
# (and any other one-shot LLM call outside ``Runner.run``) reuses the
# same policy.
DEFAULT_RETRY = ModelRetrySettings(
max_retries=5,
backoff=ModelRetryBackoffSettings(
initial_delay=2.0,
@@ -82,7 +84,7 @@ def make_run_config(
base_settings = ModelSettings(
parallel_tool_calls=False,
tool_choice="required",
retry=_DEFAULT_RETRY,
retry=DEFAULT_RETRY,
)
if reasoning_effort is not None:
base_settings = base_settings.resolve(
+3 -5
View File
@@ -2,7 +2,6 @@
from __future__ import annotations
import asyncio
import json
import logging
import re
@@ -152,7 +151,7 @@ _REQUIRED_FIELDS = {
}
def _do_create( # noqa: PLR0912
async def _do_create( # noqa: PLR0912
*,
title: str,
description: str,
@@ -238,7 +237,7 @@ def _do_create( # noqa: PLR0912
"endpoint": endpoint,
"method": method,
}
dedupe = check_duplicate(candidate, existing)
dedupe = await check_duplicate(candidate, existing)
if dedupe.get("is_duplicate"):
duplicate_id = dedupe.get("duplicate_id", "")
duplicate_title = next(
@@ -397,8 +396,7 @@ async def create_vulnerability_report(
``fix_before`` (verbatim source), ``fix_after`` (suggested
replacement).
"""
result = await asyncio.to_thread(
_do_create,
result = await _do_create(
title=title,
description=description,
impact=impact,