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 json
import logging 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__) 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]] candidate: dict[str, Any], existing_reports: list[dict[str, Any]]
) -> dict[str, Any]: ) -> dict[str, Any]:
if not existing_reports: if not existing_reports:
@@ -138,39 +174,39 @@ def check_duplicate(
} }
try: 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) candidate_cleaned = _prepare_report_for_comparison(candidate)
existing_cleaned = [_prepare_report_for_comparison(r) for r in existing_reports] existing_cleaned = [_prepare_report_for_comparison(r) for r in existing_reports]
comparison_data = {"candidate": candidate_cleaned, "existing_reports": existing_cleaned} comparison_data = {"candidate": candidate_cleaned, "existing_reports": existing_cleaned}
model_name, api_key, api_base = resolve_llm_config() user_msg = (
litellm_model: str | None = model_name 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 = [ model = build_multi_provider().get_model(model_name)
{"role": "system", "content": DEDUPE_SYSTEM_PROMPT}, response = await model.get_response(
{ system_instructions=DEDUPE_SYSTEM_PROMPT,
"role": "user", input=user_msg,
"content": ( model_settings=ModelSettings(retry=DEFAULT_RETRY),
f"Compare this candidate vulnerability against existing reports:\n\n" tools=[],
f"{json.dumps(comparison_data, indent=2)}\n\n" output_schema=None,
f"Respond with ONLY the JSON object described in the system prompt." handoffs=[],
), tracing=ModelTracing.DISABLED,
}, previous_response_id=None,
] conversation_id=None,
prompt=None,
completion_kwargs: dict[str, Any] = { )
"model": litellm_model, content = _extract_text(response)
"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
if not content: if not content:
return { return {
"is_duplicate": False, "is_duplicate": False,
@@ -182,8 +218,10 @@ def check_duplicate(
result = _parse_dedupe_response(content) result = _parse_dedupe_response(content)
logger.info( logger.info(
f"Deduplication check: is_duplicate={result['is_duplicate']}, " "Deduplication check: is_duplicate=%s, confidence=%.2f, reason=%s",
f"confidence={result['confidence']}, reason={result['reason'][:100]}" result["is_duplicate"],
result["confidence"],
result["reason"][:100],
) )
except Exception as e: 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 # 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 # errors are excluded from the retryable status list — they can't be
# fixed by retrying and should fail fast. # fixed by retrying and should fail fast. Public so the dedupe path
_DEFAULT_RETRY = ModelRetrySettings( # (and any other one-shot LLM call outside ``Runner.run``) reuses the
# same policy.
DEFAULT_RETRY = ModelRetrySettings(
max_retries=5, max_retries=5,
backoff=ModelRetryBackoffSettings( backoff=ModelRetryBackoffSettings(
initial_delay=2.0, initial_delay=2.0,
@@ -82,7 +84,7 @@ def make_run_config(
base_settings = ModelSettings( base_settings = ModelSettings(
parallel_tool_calls=False, parallel_tool_calls=False,
tool_choice="required", tool_choice="required",
retry=_DEFAULT_RETRY, retry=DEFAULT_RETRY,
) )
if reasoning_effort is not None: if reasoning_effort is not None:
base_settings = base_settings.resolve( base_settings = base_settings.resolve(
+3 -5
View File
@@ -2,7 +2,6 @@
from __future__ import annotations from __future__ import annotations
import asyncio
import json import json
import logging import logging
import re import re
@@ -152,7 +151,7 @@ _REQUIRED_FIELDS = {
} }
def _do_create( # noqa: PLR0912 async def _do_create( # noqa: PLR0912
*, *,
title: str, title: str,
description: str, description: str,
@@ -238,7 +237,7 @@ def _do_create( # noqa: PLR0912
"endpoint": endpoint, "endpoint": endpoint,
"method": method, "method": method,
} }
dedupe = check_duplicate(candidate, existing) dedupe = await check_duplicate(candidate, existing)
if dedupe.get("is_duplicate"): if dedupe.get("is_duplicate"):
duplicate_id = dedupe.get("duplicate_id", "") duplicate_id = dedupe.get("duplicate_id", "")
duplicate_title = next( duplicate_title = next(
@@ -397,8 +396,7 @@ async def create_vulnerability_report(
``fix_before`` (verbatim source), ``fix_after`` (suggested ``fix_before`` (verbatim source), ``fix_after`` (suggested
replacement). replacement).
""" """
result = await asyncio.to_thread( result = await _do_create(
_do_create,
title=title, title=title,
description=description, description=description,
impact=impact, impact=impact,