mirror of
https://github.com/usestrix/strix.git
synced 2026-08-16 09:26:39 +02:00
feat(dedupe): add dedicated deduplication model (#823)
Co-authored-by: oyasumi <oyasumi@kantilabs.xyz> Co-authored-by: Ahmed Allam <ahmed39652003@gmail.com>
This commit is contained in:
co-authored by
oyasumi
Ahmed Allam
parent
d4f4697533
commit
08126eb518
@@ -35,6 +35,31 @@ Configure Strix using environment variables or a config file.
|
|||||||
Timeout in seconds for memory compression operations (context summarization).
|
Timeout in seconds for memory compression operations (context summarization).
|
||||||
</ParamField>
|
</ParamField>
|
||||||
|
|
||||||
|
### Dedicated deduplication model
|
||||||
|
|
||||||
|
Finding deduplication is a cheap, structured classification task. By default it
|
||||||
|
runs on the main model, but you can route it to a smaller/cheaper model without
|
||||||
|
affecting the agents that do the actual testing.
|
||||||
|
|
||||||
|
<ParamField path="STRIX_DEDUPE_MODEL" type="string">
|
||||||
|
Model used to judge whether a candidate finding duplicates an existing report.
|
||||||
|
Falls back to `STRIX_LLM` when unset.
|
||||||
|
</ParamField>
|
||||||
|
|
||||||
|
<ParamField path="DEDUPE_LLM_API_KEY" type="string">
|
||||||
|
Optional provider key for the deduplication model.
|
||||||
|
</ParamField>
|
||||||
|
|
||||||
|
<ParamField path="DEDUPE_LLM_API_BASE" type="string">
|
||||||
|
Optional custom API base URL for the deduplication model. Use when the dedupe
|
||||||
|
model runs on a different endpoint than the main model.
|
||||||
|
</ParamField>
|
||||||
|
|
||||||
|
<ParamField path="STRIX_DEDUPE_REASONING_EFFORT" type="string">
|
||||||
|
Reasoning effort for the deduplication model. Defaults to the model's own
|
||||||
|
baseline when unset.
|
||||||
|
</ParamField>
|
||||||
|
|
||||||
## Optional Features
|
## Optional Features
|
||||||
|
|
||||||
<ParamField path="PERPLEXITY_API_KEY" type="string">
|
<ParamField path="PERPLEXITY_API_KEY" type="string">
|
||||||
|
|||||||
@@ -259,6 +259,9 @@ ignore = [
|
|||||||
# a runtime ``Callable`` annotation on ``vulnerability_found_callback``.
|
# a runtime ``Callable`` annotation on ``vulnerability_found_callback``.
|
||||||
"strix/report/state.py" = ["TC003", "PLR0912", "PLR0915", "E501", "PERF401", "PLC0415"]
|
"strix/report/state.py" = ["TC003", "PLR0912", "PLR0915", "E501", "PERF401", "PLC0415"]
|
||||||
"strix/report/usage.py" = ["PLC0415"]
|
"strix/report/usage.py" = ["PLC0415"]
|
||||||
|
# Lazy import of strix.config.models avoids a circular dependency between the
|
||||||
|
# report pipeline and the config layer.
|
||||||
|
"strix/report/dedupe.py" = ["PLC0415"]
|
||||||
"strix/telemetry/logging.py" = ["PLC0415"]
|
"strix/telemetry/logging.py" = ["PLC0415"]
|
||||||
"strix/config/models.py" = ["PLC0415"]
|
"strix/config/models.py" = ["PLC0415"]
|
||||||
# Heavy inference deps (httpx, openai) imported lazily so auth-status checks
|
# Heavy inference deps (httpx, openai) imported lazily so auth-status checks
|
||||||
|
|||||||
@@ -17,6 +17,7 @@ from strix.config.loader import (
|
|||||||
persist_current,
|
persist_current,
|
||||||
)
|
)
|
||||||
from strix.config.settings import (
|
from strix.config.settings import (
|
||||||
|
DedupeSettings,
|
||||||
IntegrationSettings,
|
IntegrationSettings,
|
||||||
LlmSettings,
|
LlmSettings,
|
||||||
RuntimeSettings,
|
RuntimeSettings,
|
||||||
@@ -26,6 +27,7 @@ from strix.config.settings import (
|
|||||||
|
|
||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
|
"DedupeSettings",
|
||||||
"IntegrationSettings",
|
"IntegrationSettings",
|
||||||
"LlmSettings",
|
"LlmSettings",
|
||||||
"RuntimeSettings",
|
"RuntimeSettings",
|
||||||
|
|||||||
@@ -43,6 +43,18 @@ class LlmSettings(BaseSettings):
|
|||||||
timeout: int = Field(default=300, alias="LLM_TIMEOUT")
|
timeout: int = Field(default=300, alias="LLM_TIMEOUT")
|
||||||
|
|
||||||
|
|
||||||
|
class DedupeSettings(BaseSettings):
|
||||||
|
model_config = _BASE_CONFIG
|
||||||
|
|
||||||
|
model: str | None = Field(default=None, alias="STRIX_DEDUPE_MODEL")
|
||||||
|
reasoning_effort: ReasoningEffort | None = Field(
|
||||||
|
default=None,
|
||||||
|
alias="STRIX_DEDUPE_REASONING_EFFORT",
|
||||||
|
)
|
||||||
|
api_key: str | None = Field(default=None, alias="DEDUPE_LLM_API_KEY")
|
||||||
|
api_base: str | None = Field(default=None, alias="DEDUPE_LLM_API_BASE")
|
||||||
|
|
||||||
|
|
||||||
class RuntimeSettings(BaseSettings):
|
class RuntimeSettings(BaseSettings):
|
||||||
model_config = _BASE_CONFIG
|
model_config = _BASE_CONFIG
|
||||||
|
|
||||||
@@ -85,6 +97,7 @@ class Settings(BaseSettings):
|
|||||||
model_config = _BASE_CONFIG
|
model_config = _BASE_CONFIG
|
||||||
|
|
||||||
llm: LlmSettings = Field(default_factory=LlmSettings)
|
llm: LlmSettings = Field(default_factory=LlmSettings)
|
||||||
|
dedupe: DedupeSettings = Field(default_factory=DedupeSettings)
|
||||||
runtime: RuntimeSettings = Field(default_factory=RuntimeSettings)
|
runtime: RuntimeSettings = Field(default_factory=RuntimeSettings)
|
||||||
telemetry: TelemetrySettings = Field(default_factory=TelemetrySettings)
|
telemetry: TelemetrySettings = Field(default_factory=TelemetrySettings)
|
||||||
integrations: IntegrationSettings = Field(default_factory=IntegrationSettings)
|
integrations: IntegrationSettings = Field(default_factory=IntegrationSettings)
|
||||||
|
|||||||
@@ -394,6 +394,33 @@ async def warm_up_llm(show_model_warning: bool = True) -> None:
|
|||||||
)
|
)
|
||||||
logger.info("LLM warm-up succeeded for model %s", (llm.model or "").strip())
|
logger.info("LLM warm-up succeeded for model %s", (llm.model or "").strip())
|
||||||
|
|
||||||
|
if settings.dedupe.model:
|
||||||
|
from strix.report.dedupe import _dedupe_extra_args
|
||||||
|
|
||||||
|
dedupe_model = settings.dedupe.model.strip()
|
||||||
|
raw_model = dedupe_model
|
||||||
|
deduper = StrixProvider().get_model(dedupe_model)
|
||||||
|
# Match the runtime path: send the dedupe key/endpoint per call so a
|
||||||
|
# separate-provider dedupe model authenticates during warm-up too.
|
||||||
|
deduper_extra = _dedupe_extra_args(settings.dedupe)
|
||||||
|
deduper_settings = ModelSettings(extra_args=deduper_extra or None)
|
||||||
|
await asyncio.wait_for(
|
||||||
|
deduper.get_response(
|
||||||
|
system_instructions="You are a helpful assistant.",
|
||||||
|
input="Reply with just 'OK'.",
|
||||||
|
model_settings=deduper_settings,
|
||||||
|
tools=[],
|
||||||
|
output_schema=None,
|
||||||
|
handoffs=[],
|
||||||
|
tracing=ModelTracing.DISABLED,
|
||||||
|
previous_response_id=None,
|
||||||
|
conversation_id=None,
|
||||||
|
prompt=None,
|
||||||
|
),
|
||||||
|
timeout=llm.timeout,
|
||||||
|
)
|
||||||
|
logger.info("LLM warm-up succeeded for dedupe model %s", dedupe_model)
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.exception("LLM warm-up failed")
|
logger.exception("LLM warm-up failed")
|
||||||
error_text = Text()
|
error_text = Text()
|
||||||
|
|||||||
+42
-8
@@ -13,20 +13,55 @@ from openai.types.responses import ResponseOutputMessage
|
|||||||
|
|
||||||
from strix.config import load_settings
|
from strix.config import load_settings
|
||||||
from strix.config.models import (
|
from strix.config.models import (
|
||||||
DEFAULT_MODEL_RETRY,
|
|
||||||
StrixProvider,
|
StrixProvider,
|
||||||
configure_sdk_model_defaults,
|
configure_sdk_model_defaults,
|
||||||
request_timeout_extra_args,
|
|
||||||
)
|
)
|
||||||
|
from strix.core.inputs import make_model_settings
|
||||||
from strix.report.state import get_global_report_state
|
from strix.report.state import get_global_report_state
|
||||||
|
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from agents.items import ModelResponse
|
from agents.items import ModelResponse
|
||||||
|
|
||||||
|
from strix.config.settings import DedupeSettings
|
||||||
|
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
def _dedupe_extra_args(dedupe: DedupeSettings) -> dict[str, str]:
|
||||||
|
"""Per-call credential + endpoint for the dedupe model.
|
||||||
|
|
||||||
|
Provider env vars and the global base URL are process-wide, so a
|
||||||
|
shared-provider dedupe key or a distinct dedupe endpoint can't be installed
|
||||||
|
globally without clobbering (or being clobbered by) the main model's
|
||||||
|
config. Passing them per call keeps the two apart. Only applies when a
|
||||||
|
dedicated dedupe model is configured.
|
||||||
|
"""
|
||||||
|
if not dedupe.model:
|
||||||
|
return {}
|
||||||
|
extra: dict[str, str] = {}
|
||||||
|
if dedupe.api_key and dedupe.api_key.strip():
|
||||||
|
extra["api_key"] = dedupe.api_key.strip()
|
||||||
|
if dedupe.api_base and dedupe.api_base.strip():
|
||||||
|
extra["api_base"] = dedupe.api_base.strip()
|
||||||
|
return extra
|
||||||
|
|
||||||
|
|
||||||
|
def _dedupe_model_settings(
|
||||||
|
dedupe: DedupeSettings, model_name: str, request_timeout: float | None
|
||||||
|
) -> ModelSettings:
|
||||||
|
settings = make_model_settings(
|
||||||
|
dedupe.reasoning_effort,
|
||||||
|
model_name=model_name,
|
||||||
|
force_required_tool_choice=False,
|
||||||
|
request_timeout=request_timeout,
|
||||||
|
)
|
||||||
|
extra = _dedupe_extra_args(dedupe)
|
||||||
|
if extra:
|
||||||
|
settings = settings.resolve(ModelSettings(extra_args=extra))
|
||||||
|
return settings
|
||||||
|
|
||||||
DEDUPE_SYSTEM_PROMPT = """You are an expert vulnerability report deduplication judge.
|
DEDUPE_SYSTEM_PROMPT = """You are an expert vulnerability report deduplication judge.
|
||||||
Your task is to determine if a candidate vulnerability report describes the SAME vulnerability
|
Your task is to determine if a candidate vulnerability report describes the SAME vulnerability
|
||||||
as any existing report.
|
as any existing report.
|
||||||
@@ -286,13 +321,14 @@ async def check_duplicate(
|
|||||||
|
|
||||||
try:
|
try:
|
||||||
settings = load_settings()
|
settings = load_settings()
|
||||||
model_name = settings.llm.model
|
dedupe = settings.dedupe
|
||||||
|
model_name = (dedupe.model or "").strip() or settings.llm.model
|
||||||
if not model_name:
|
if not model_name:
|
||||||
return {
|
return {
|
||||||
"is_duplicate": False,
|
"is_duplicate": False,
|
||||||
"duplicate_id": "",
|
"duplicate_id": "",
|
||||||
"confidence": 0.0,
|
"confidence": 0.0,
|
||||||
"reason": "STRIX_LLM not configured; skipping dedupe check",
|
"reason": "No LLM model configured; skipping dedupe check",
|
||||||
}
|
}
|
||||||
|
|
||||||
candidate_cleaned = _prepare_report_for_comparison(candidate)
|
candidate_cleaned = _prepare_report_for_comparison(candidate)
|
||||||
@@ -311,10 +347,8 @@ async def check_duplicate(
|
|||||||
response = await model.get_response(
|
response = await model.get_response(
|
||||||
system_instructions=DEDUPE_SYSTEM_PROMPT,
|
system_instructions=DEDUPE_SYSTEM_PROMPT,
|
||||||
input=user_msg,
|
input=user_msg,
|
||||||
model_settings=ModelSettings(
|
model_settings=_dedupe_model_settings(
|
||||||
retry=DEFAULT_MODEL_RETRY,
|
dedupe, resolved_model, settings.llm.timeout
|
||||||
include_usage=True,
|
|
||||||
extra_args=request_timeout_extra_args(settings.llm.timeout),
|
|
||||||
),
|
),
|
||||||
tools=[],
|
tools=[],
|
||||||
output_schema=None,
|
output_schema=None,
|
||||||
|
|||||||
@@ -0,0 +1,94 @@
|
|||||||
|
"""Tests for the dedicated deduplication model configuration."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import json
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
|
from strix.config import loader
|
||||||
|
from strix.config.settings import DedupeSettings
|
||||||
|
from strix.report.dedupe import _dedupe_model_settings
|
||||||
|
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
|
||||||
|
def test_dedupe_key_sent_per_call_not_via_global_env() -> None:
|
||||||
|
dedupe = DedupeSettings(STRIX_DEDUPE_MODEL="deepseek/cheap", DEDUPE_LLM_API_KEY="dedupe-key")
|
||||||
|
settings = _dedupe_model_settings(dedupe, "deepseek/cheap", 300)
|
||||||
|
# The key rides on the request, so a shared-provider main key can't clobber
|
||||||
|
# it (and vice versa) through the global provider env var.
|
||||||
|
assert (settings.extra_args or {})["api_key"] == "dedupe-key"
|
||||||
|
|
||||||
|
|
||||||
|
def test_dedupe_settings_omit_api_key_when_unset() -> None:
|
||||||
|
dedupe = DedupeSettings(STRIX_DEDUPE_MODEL="deepseek/cheap")
|
||||||
|
settings = _dedupe_model_settings(dedupe, "deepseek/cheap", 300)
|
||||||
|
assert "api_key" not in (settings.extra_args or {})
|
||||||
|
assert "api_base" not in (settings.extra_args or {})
|
||||||
|
|
||||||
|
|
||||||
|
def test_dedupe_endpoint_sent_per_call() -> None:
|
||||||
|
dedupe = DedupeSettings(
|
||||||
|
STRIX_DEDUPE_MODEL="openai/cheap",
|
||||||
|
DEDUPE_LLM_API_KEY="dedupe-key",
|
||||||
|
DEDUPE_LLM_API_BASE="https://dedupe.example/v1",
|
||||||
|
)
|
||||||
|
settings = _dedupe_model_settings(dedupe, "openai/cheap", 300)
|
||||||
|
# A distinct dedupe endpoint rides on the request instead of the
|
||||||
|
# process-wide base URL, so it can't clobber the main model's endpoint.
|
||||||
|
assert (settings.extra_args or {})["api_base"] == "https://dedupe.example/v1"
|
||||||
|
assert (settings.extra_args or {})["api_key"] == "dedupe-key"
|
||||||
|
|
||||||
|
|
||||||
|
def test_dedupe_defaults_are_empty() -> None:
|
||||||
|
settings = DedupeSettings()
|
||||||
|
assert settings.model is None
|
||||||
|
assert settings.reasoning_effort is None
|
||||||
|
assert settings.api_key is None
|
||||||
|
|
||||||
|
|
||||||
|
def test_dedupe_model_read_from_env(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||||
|
monkeypatch.setenv("STRIX_DEDUPE_MODEL", "deepseek/deepseek-v4-flash")
|
||||||
|
monkeypatch.setenv("STRIX_DEDUPE_REASONING_EFFORT", "low")
|
||||||
|
|
||||||
|
settings = DedupeSettings()
|
||||||
|
|
||||||
|
assert settings.model == "deepseek/deepseek-v4-flash"
|
||||||
|
assert settings.reasoning_effort == "low"
|
||||||
|
|
||||||
|
|
||||||
|
def test_config_file_loads_dedupe_model(
|
||||||
|
tmp_path: Path,
|
||||||
|
monkeypatch: pytest.MonkeyPatch,
|
||||||
|
) -> None:
|
||||||
|
for key in ("STRIX_LLM", "STRIX_DEDUPE_MODEL", "STRIX_DEDUPE_REASONING_EFFORT"):
|
||||||
|
monkeypatch.delenv(key, raising=False)
|
||||||
|
path = tmp_path / "config.json"
|
||||||
|
path.write_text(
|
||||||
|
json.dumps(
|
||||||
|
{
|
||||||
|
"env": {
|
||||||
|
"STRIX_LLM": "openai/root",
|
||||||
|
"STRIX_DEDUPE_MODEL": "deepseek/cheap",
|
||||||
|
"STRIX_DEDUPE_REASONING_EFFORT": "minimal",
|
||||||
|
}
|
||||||
|
}
|
||||||
|
),
|
||||||
|
encoding="utf-8",
|
||||||
|
)
|
||||||
|
loader._cached = None
|
||||||
|
loader._override = path
|
||||||
|
try:
|
||||||
|
settings = loader.load_settings()
|
||||||
|
finally:
|
||||||
|
loader._cached = None
|
||||||
|
loader._override = None
|
||||||
|
|
||||||
|
assert settings.dedupe.model == "deepseek/cheap"
|
||||||
|
assert settings.dedupe.reasoning_effort == "minimal"
|
||||||
|
# Main model stays independent of the dedupe override.
|
||||||
|
assert settings.llm.model == "openai/root"
|
||||||
Reference in New Issue
Block a user