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:
+72
-34
@@ -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:
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user