From 8a11f9dab5b3412aab03544fccd688a2f6cab002 Mon Sep 17 00:00:00 2001 From: 0xallam Date: Sat, 25 Apr 2026 15:25:44 -0700 Subject: [PATCH] refactor(dedupe): route through MultiProvider + cache wrapper + retry policy MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ``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) --- strix/llm/dedupe.py | 106 +++++++++++++++++++++++----------- strix/run_config_factory.py | 8 ++- strix/tools/reporting/tool.py | 8 +-- 3 files changed, 80 insertions(+), 42 deletions(-) diff --git a/strix/llm/dedupe.py b/strix/llm/dedupe.py index 8d2be96..291d20a 100644 --- a/strix/llm/dedupe.py +++ b/strix/llm/dedupe.py @@ -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: diff --git a/strix/run_config_factory.py b/strix/run_config_factory.py index a6c9916..a2a7f50 100644 --- a/strix/run_config_factory.py +++ b/strix/run_config_factory.py @@ -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( diff --git a/strix/tools/reporting/tool.py b/strix/tools/reporting/tool.py index b2b8810..b61e674 100644 --- a/strix/tools/reporting/tool.py +++ b/strix/tools/reporting/tool.py @@ -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,