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