mirror of
https://github.com/norandom/OpenUltraSAST.git
synced 2026-10-16 03:49:26 +00:00
feat(learning-harness): chat endpoint independent of embeddings, with the DeepSeek adapter
Task 1.2. Detectors, the classifier's model tier, the proposer and the judge resolve one chat endpoint: an explicit override first, then the scripted client flag, then an endpoint named in the config, then DeepSeek, then OpenRouter. Prove-order retrieval keeps its own embedding endpoint untouched, because the chat provider we default to offers none. The adapter handles the provider's documented quirks: thinking is on by default there and silently ignores temperature, so it is disabled explicitly; the reasoning field is kept on the reply for the caller to replay across tool turns; JSON mode has no schema and occasionally returns empty content, so one retry is built in, never a loop; token log-probabilities are requested only when the caller wants them, and their mean is what lets the classifier abstain. Cost comes from the provider's own usage fields, cache hits priced separately, and a model with no recorded price reports no cost rather than a made-up one. The endpoint record deliberately holds no key, so it is safe to journal. Co-Authored-By: Claude Opus 5 (1M context) <[email protected]> Claude-Session: https://claude.ai/code/session_01T7XD7x5KonMdmGRkp63EQG
This commit is contained in:
Marius Ciepluch
co-authored by
Claude Opus 5
Claude Opus 5
parent
289648338c
commit
944df86d12
@@ -33,6 +33,9 @@ class ModelConfig:
|
||||
hunter: str | None = None
|
||||
verifier: str | None = None
|
||||
patcher: str | None = None
|
||||
judge: str | None = None # second, independent judgement before anything is published as proven
|
||||
chat_base_url: str | None = None # chat endpoint; independent of the embedding endpoint
|
||||
chat_api_key_env: str | None = None # environment variable holding that endpoint's key
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
@@ -135,6 +138,19 @@ class VariantsConfig:
|
||||
max_mechanisms: int = 500
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class LearningConfig:
|
||||
# learning-harness: the classified detector set and its rounds.
|
||||
families_path: str | None = None # taxonomy file; the shipped ruleset when unset
|
||||
configs_dir: str = ".openultrasast/learning/configs"
|
||||
k_runs: int = 5 # runs per pair; a single run misses most real changes
|
||||
round_cost_cap_usd: float = 10.0
|
||||
minibatch: int = 8 # train pairs in a round's first stage
|
||||
detector_model: str | None = None
|
||||
judge_model: str | None = None
|
||||
cutoff_date: str | None = None # rows fixed after this date are reported as a separate slice
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ObligationsConfig:
|
||||
# authorization-obligations: the obligation checker in MAP (standard/deep only).
|
||||
@@ -160,6 +176,7 @@ class ResolvedConfig:
|
||||
regress: RegressConfig = RegressConfig()
|
||||
variants: VariantsConfig = VariantsConfig()
|
||||
obligations: ObligationsConfig = ObligationsConfig()
|
||||
learning: LearningConfig = LearningConfig()
|
||||
runs_dir: str = ".openultrasast/runs"
|
||||
|
||||
|
||||
@@ -186,6 +203,7 @@ def load_config(config_path: Path | None = None) -> ResolvedConfig:
|
||||
regress=_load_regress(data.get("regress", {})),
|
||||
variants=_load_variants(data.get("variants", {})),
|
||||
obligations=_load_obligations(data.get("obligations", {})),
|
||||
learning=_load_learning(data.get("learning", {})),
|
||||
runs_dir=os.environ.get("OPENULTRASAST_RUNS_DIR", ".openultrasast/runs"),
|
||||
)
|
||||
|
||||
@@ -210,6 +228,9 @@ def _load_models(value: object) -> ModelConfig:
|
||||
hunter=_string(data.get("hunter")),
|
||||
verifier=_string(data.get("verifier")),
|
||||
patcher=_string(data.get("patcher")),
|
||||
judge=_string(data.get("judge")),
|
||||
chat_base_url=_string(data.get("chat_base_url")),
|
||||
chat_api_key_env=_string(data.get("chat_api_key_env")),
|
||||
)
|
||||
|
||||
|
||||
@@ -357,6 +378,20 @@ def _load_variants(value: object) -> VariantsConfig:
|
||||
)
|
||||
|
||||
|
||||
def _load_learning(value: object) -> LearningConfig:
|
||||
data = _section(value)
|
||||
return LearningConfig(
|
||||
families_path=_string(data.get("families_path")),
|
||||
configs_dir=_string(data.get("configs_dir")) or LearningConfig.configs_dir,
|
||||
k_runs=_int_value(data.get("k_runs"), LearningConfig.k_runs),
|
||||
round_cost_cap_usd=_float_value(data.get("round_cost_cap_usd"), LearningConfig.round_cost_cap_usd),
|
||||
minibatch=_int_value(data.get("minibatch"), LearningConfig.minibatch),
|
||||
detector_model=_string(data.get("detector_model")),
|
||||
judge_model=_string(data.get("judge_model")),
|
||||
cutoff_date=_string(data.get("cutoff_date")),
|
||||
)
|
||||
|
||||
|
||||
def _load_obligations(value: object) -> ObligationsConfig:
|
||||
data = _section(value)
|
||||
return ObligationsConfig(
|
||||
|
||||
@@ -0,0 +1,215 @@
|
||||
"""The chat endpoint, resolved independently of embeddings (learning-harness, Req 6.6).
|
||||
|
||||
Detectors, the classifier's model tier, the proposer and the judge all talk to one chat endpoint;
|
||||
prove-order retrieval keeps its own embedding endpoint, because the chat provider we default to has
|
||||
none. Resolution order is explicit override, the scripted client flag, an endpoint named in the
|
||||
config, DeepSeek, then OpenRouter. ``DeepSeekChatClient`` adapts the provider's quirks: thinking is
|
||||
on by default there and silently ignores ``temperature``, so it is disabled explicitly; the reasoning
|
||||
field is kept on the reply for the caller to replay across tool turns; JSON mode has no schema and
|
||||
occasionally returns empty content, so one retry is built in. Prices are data, and an unpriced model
|
||||
reports no cost rather than a made-up one.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import Literal
|
||||
|
||||
from ..config import ResolvedConfig
|
||||
from ..provider.openrouter import OpenRouterChatClient, OpenRouterError
|
||||
from ..tool_hunter import CLIENT_ENV, ChatClient, ChatResponse, chat_response_from_message, scripted_hunter_client
|
||||
|
||||
DEEPSEEK_BASE_URL = "https://api.deepseek.com" # no version suffix: the platform rejects /v1
|
||||
DEEPSEEK_KEY_ENV = "DEEPSEEK_API_KEY"
|
||||
DEEPSEEK_BASE_ENV = "DEEPSEEK_BASE_URL"
|
||||
OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1"
|
||||
DEFAULT_DETECTOR_MODEL = "deepseek-v4-flash"
|
||||
DEFAULT_JUDGE_MODEL = "deepseek-v4-pro"
|
||||
Provider = Literal["deepseek", "openrouter", "custom", "scripted"]
|
||||
_SCRIPTED_FLAGS = frozenset({"scripted", "script", "dump", "dump-only", "unsafe", "unsafe-snippet"})
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Prices:
|
||||
"""USD per million tokens. DeepSeek platform pricing, peak rate, recorded 2026-09-06."""
|
||||
|
||||
cache_hit_per_m: float
|
||||
input_per_m: float
|
||||
output_per_m: float
|
||||
|
||||
|
||||
_PRICES: dict[str, Prices] = {
|
||||
"deepseek-v4-flash": Prices(cache_hit_per_m=0.014, input_per_m=0.44, output_per_m=1.32),
|
||||
"deepseek-v4-flash-vision-exp": Prices(cache_hit_per_m=0.014, input_per_m=0.44, output_per_m=1.32),
|
||||
"deepseek-v4-pro": Prices(cache_hit_per_m=0.044, input_per_m=1.32, output_per_m=3.96),
|
||||
}
|
||||
|
||||
|
||||
def price_of(model: str) -> Prices | None:
|
||||
return _PRICES.get(model.strip())
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ChatEndpoint:
|
||||
"""What the run should record about where the model calls went. Never holds the key."""
|
||||
|
||||
provider: Provider
|
||||
base_url: str
|
||||
thinking: bool # the provider's models reason by default, so the adapter disables it explicitly
|
||||
prices: Prices | None = None
|
||||
|
||||
def cost(self, usage: Mapping[str, object] | None) -> float | None:
|
||||
"""Dollars for one call from the provider's own usage fields, or None when the model has no recorded price."""
|
||||
if self.prices is None or not isinstance(usage, Mapping):
|
||||
return None
|
||||
cache_hit = _tokens(usage.get("prompt_cache_hit_tokens"))
|
||||
miss = usage.get("prompt_cache_miss_tokens")
|
||||
input_tokens = _tokens(miss) if miss is not None else max(_tokens(usage.get("prompt_tokens")) - cache_hit, 0)
|
||||
output_tokens = _tokens(usage.get("completion_tokens"))
|
||||
return (
|
||||
cache_hit * self.prices.cache_hit_per_m + input_tokens * self.prices.input_per_m + output_tokens * self.prices.output_per_m
|
||||
) / 1_000_000
|
||||
|
||||
|
||||
class DeepSeekChatClient:
|
||||
"""Adapt an OpenAI-shaped chat client to what detectors and the classifier need."""
|
||||
|
||||
def __init__(self, client: OpenRouterChatClient, *, disable_thinking: bool = True, endpoint: ChatEndpoint | None = None) -> None:
|
||||
self._client = client
|
||||
self._disable_thinking = disable_thinking
|
||||
self.endpoint = endpoint
|
||||
self.usage: list[dict[str, object]] = [] # one entry per call, for the round's cost meter
|
||||
|
||||
def complete(
|
||||
self,
|
||||
*,
|
||||
model: str,
|
||||
messages: list[dict[str, object]],
|
||||
tools: list[dict[str, object]],
|
||||
timeout_seconds: int = 60,
|
||||
json_object: bool = False,
|
||||
logprobs: bool = False,
|
||||
) -> ChatResponse:
|
||||
response = self._call(
|
||||
model=model, messages=messages, tools=tools, timeout_seconds=timeout_seconds, json_object=json_object, logprobs=logprobs
|
||||
)
|
||||
if json_object and not (response.content or "").strip():
|
||||
response = self._call(
|
||||
model=model, messages=messages, tools=tools, timeout_seconds=timeout_seconds, json_object=json_object, logprobs=logprobs
|
||||
) # documented behaviour: JSON mode occasionally returns empty content. One retry, never a loop.
|
||||
return response
|
||||
|
||||
def cost_usd(self) -> float:
|
||||
"""Recorded cost of the calls made so far; calls whose model has no price contribute nothing."""
|
||||
if self.endpoint is None:
|
||||
return 0.0
|
||||
return sum(self.endpoint.cost(entry) or 0.0 for entry in self.usage)
|
||||
|
||||
def _call(
|
||||
self,
|
||||
*,
|
||||
model: str,
|
||||
messages: Sequence[Mapping[str, object]],
|
||||
tools: Sequence[Mapping[str, object]],
|
||||
timeout_seconds: int,
|
||||
json_object: bool,
|
||||
logprobs: bool,
|
||||
) -> ChatResponse:
|
||||
extra: dict[str, object] = {}
|
||||
if self._disable_thinking:
|
||||
extra["thinking"] = {"type": "disabled"} # otherwise temperature is silently ignored
|
||||
if json_object:
|
||||
extra["response_format"] = {"type": "json_object"}
|
||||
if logprobs:
|
||||
extra["logprobs"] = True
|
||||
extra["top_logprobs"] = 1
|
||||
raw = getattr(self._client, "complete_chat_raw", None)
|
||||
if callable(raw):
|
||||
payload = raw(
|
||||
model=model, messages=list(messages), tools=list(tools) or None, timeout_seconds=timeout_seconds, extra_body=extra
|
||||
)
|
||||
usage = payload.get("usage") if isinstance(payload, dict) else None
|
||||
if isinstance(usage, dict):
|
||||
self.usage.append(dict(usage))
|
||||
message = _message_of(payload)
|
||||
else: # a stub or a client without the raw call
|
||||
message = self._client.complete_chat(
|
||||
model=model, messages=list(messages), tools=list(tools) or None, timeout_seconds=timeout_seconds, extra_body=extra
|
||||
)
|
||||
return chat_response_from_message(message)
|
||||
|
||||
|
||||
def resolve_chat_endpoint(config: ResolvedConfig, *, override: ChatClient | None = None) -> tuple[ChatClient, ChatEndpoint] | None:
|
||||
"""The chat client for detectors, classifier, proposer and judge, with what to record about it."""
|
||||
if override is not None:
|
||||
return override, ChatEndpoint(provider="scripted", base_url="", thinking=False)
|
||||
flag = os.environ.get(CLIENT_ENV, "").strip().lower()
|
||||
if flag in _SCRIPTED_FLAGS:
|
||||
return scripted_hunter_client(flag), ChatEndpoint(provider="scripted", base_url="", thinking=False)
|
||||
model = resolve_models(config)[0]
|
||||
configured = config.models.chat_base_url
|
||||
key_env = config.models.chat_api_key_env
|
||||
if configured or key_env:
|
||||
key = os.environ.get(key_env or DEEPSEEK_KEY_ENV, "")
|
||||
base_url = configured or DEEPSEEK_BASE_URL
|
||||
return _deepseek_client(key, base_url, model, provider=_provider_of(base_url))
|
||||
deepseek_key = os.environ.get(DEEPSEEK_KEY_ENV)
|
||||
if deepseek_key:
|
||||
return _deepseek_client(deepseek_key, os.environ.get(DEEPSEEK_BASE_ENV, DEEPSEEK_BASE_URL), model, provider="deepseek")
|
||||
if os.environ.get("OPENROUTER_API_KEY"):
|
||||
try:
|
||||
client = OpenRouterChatClient.from_env()
|
||||
except OpenRouterError:
|
||||
return None
|
||||
endpoint = ChatEndpoint(provider="openrouter", base_url=client.base_url, thinking=False, prices=price_of(model))
|
||||
return DeepSeekChatClient(client, disable_thinking=False, endpoint=endpoint), endpoint
|
||||
return None
|
||||
|
||||
|
||||
def resolve_models(config: ResolvedConfig) -> tuple[str, str]:
|
||||
"""(detector model, judge model). `[learning]` wins, then `[models]`, then the shipped defaults."""
|
||||
detector = config.learning.detector_model or config.models.hunter or DEFAULT_DETECTOR_MODEL
|
||||
judge = config.learning.judge_model or config.models.judge or DEFAULT_JUDGE_MODEL
|
||||
return detector, judge
|
||||
|
||||
|
||||
def _deepseek_client(key: str, base_url: str, model: str, *, provider: Provider) -> tuple[ChatClient, ChatEndpoint]:
|
||||
endpoint = ChatEndpoint(provider=provider, base_url=base_url.rstrip("/"), thinking=True, prices=price_of(model))
|
||||
client = OpenRouterChatClient(api_key=key, base_url=endpoint.base_url)
|
||||
return DeepSeekChatClient(client, endpoint=endpoint), endpoint
|
||||
|
||||
|
||||
def _provider_of(base_url: str) -> Provider:
|
||||
host = base_url.lower()
|
||||
if "deepseek" in host:
|
||||
return "deepseek"
|
||||
return "openrouter" if "openrouter" in host else "custom"
|
||||
|
||||
|
||||
def _message_of(payload: object) -> Mapping[str, object]:
|
||||
if isinstance(payload, Mapping):
|
||||
choices = payload.get("choices")
|
||||
if isinstance(choices, Sequence) and choices and isinstance(choices[0], Mapping):
|
||||
message = choices[0].get("message")
|
||||
if isinstance(message, Mapping):
|
||||
return message
|
||||
return {}
|
||||
|
||||
|
||||
def _tokens(value: object) -> int:
|
||||
return int(value) if isinstance(value, int | float) else 0
|
||||
|
||||
|
||||
__all__ = [
|
||||
"DEEPSEEK_BASE_URL",
|
||||
"DEFAULT_DETECTOR_MODEL",
|
||||
"DEFAULT_JUDGE_MODEL",
|
||||
"ChatEndpoint",
|
||||
"DeepSeekChatClient",
|
||||
"Prices",
|
||||
"price_of",
|
||||
"resolve_chat_endpoint",
|
||||
"resolve_models",
|
||||
]
|
||||
@@ -68,10 +68,29 @@ class OpenRouterChatClient:
|
||||
messages: Sequence[Mapping[str, object]],
|
||||
tools: Sequence[Mapping[str, object]] | None = None,
|
||||
timeout_seconds: int = 60,
|
||||
extra_body: Mapping[str, object] | None = None,
|
||||
) -> dict[str, object]:
|
||||
"""The assistant message of the reply, provider fields included (a reasoning field is kept for the caller to replay)."""
|
||||
payload = self.complete_chat_raw(
|
||||
model=model, messages=messages, tools=tools, timeout_seconds=timeout_seconds, extra_body=extra_body
|
||||
)
|
||||
return _extract_message(payload)
|
||||
|
||||
def complete_chat_raw(
|
||||
self,
|
||||
*,
|
||||
model: str,
|
||||
messages: Sequence[Mapping[str, object]],
|
||||
tools: Sequence[Mapping[str, object]] | None = None,
|
||||
timeout_seconds: int = 60,
|
||||
extra_body: Mapping[str, object] | None = None,
|
||||
) -> dict[str, object]:
|
||||
"""The whole response payload, so a caller can read `usage` for cost accounting."""
|
||||
body: dict[str, object] = {"model": model, "messages": list(messages), "temperature": 0}
|
||||
if tools:
|
||||
body["tools"] = list(tools)
|
||||
for key, value in (extra_body or {}).items():
|
||||
body[key] = value
|
||||
payload = json.dumps(body).encode()
|
||||
request = urllib.request.Request(
|
||||
f"{self.base_url.rstrip('/')}/chat/completions",
|
||||
@@ -88,10 +107,9 @@ class OpenRouterChatClient:
|
||||
return cast(dict[str, object], json.loads(response.read().decode()))
|
||||
|
||||
try:
|
||||
response_payload = call_with_retry(_do, attempts=self.max_attempts, base_delay=self.retry_base_delay)
|
||||
return call_with_retry(_do, attempts=self.max_attempts, base_delay=self.retry_base_delay)
|
||||
except (urllib.error.URLError, TimeoutError, json.JSONDecodeError) as exc:
|
||||
raise OpenRouterError(f"OpenRouter chat request failed: {exc}") from exc
|
||||
return _extract_message(response_payload)
|
||||
|
||||
def complete_json(self, *, model: str, messages: list[dict[str, str]], timeout_seconds: int = 60) -> object:
|
||||
message = self.complete_chat(model=model, messages=messages, timeout_seconds=timeout_seconds)
|
||||
|
||||
@@ -101,6 +101,8 @@ class ToolCall:
|
||||
class ChatResponse:
|
||||
content: str | None = None
|
||||
tool_calls: tuple[ToolCall, ...] = ()
|
||||
reasoning: str | None = None # provider chain-of-thought; replayed on tool turns, never shown or scored
|
||||
mean_logprob: float | None = None # averaged token log-probability when the caller asked for it
|
||||
|
||||
|
||||
class ChatClient(Protocol):
|
||||
@@ -198,7 +200,24 @@ def chat_response_from_message(message: Mapping[str, object]) -> ChatResponse:
|
||||
parsed = _tool_call_from_openrouter(item, index)
|
||||
if parsed is not None:
|
||||
calls.append(parsed)
|
||||
return ChatResponse(content=text, tool_calls=tuple(calls))
|
||||
reasoning = message.get("reasoning_content") or message.get("reasoning")
|
||||
return ChatResponse(
|
||||
content=text,
|
||||
tool_calls=tuple(calls),
|
||||
reasoning=reasoning if isinstance(reasoning, str) and reasoning else None,
|
||||
mean_logprob=_mean_logprob(message.get("logprobs")),
|
||||
)
|
||||
|
||||
|
||||
def _mean_logprob(payload: object) -> float | None:
|
||||
"""The mean token log-probability of a reply, used by the classifier to abstain when it is unsure."""
|
||||
if not isinstance(payload, Mapping):
|
||||
return None
|
||||
tokens = payload.get("content")
|
||||
if not isinstance(tokens, Sequence) or isinstance(tokens, str | bytes):
|
||||
return None
|
||||
values = [float(item["logprob"]) for item in tokens if isinstance(item, Mapping) and isinstance(item.get("logprob"), int | float)]
|
||||
return sum(values) / len(values) if values else None
|
||||
|
||||
|
||||
def _tool_call_from_openrouter(item: object, index: int) -> ToolCall | None:
|
||||
|
||||
@@ -0,0 +1,162 @@
|
||||
"""learning-harness task 1.2: the chat endpoint resolves independently of embeddings (offline)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import urllib.request
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from openultrasast.config import ResolvedConfig
|
||||
from openultrasast.provider import openrouter as provider
|
||||
|
||||
|
||||
class _Recorder:
|
||||
"""Stands in for OpenRouterChatClient.complete_chat and records what the adapter asked for."""
|
||||
|
||||
def __init__(self, replies: list[dict[str, object]]) -> None:
|
||||
self.replies = replies
|
||||
self.calls: list[dict[str, Any]] = []
|
||||
|
||||
def complete_chat(self, **kwargs: Any) -> dict[str, object]:
|
||||
self.calls.append(kwargs)
|
||||
return self.replies[min(len(self.calls) - 1, len(self.replies) - 1)]
|
||||
|
||||
|
||||
def _endpoint_env(monkeypatch: pytest.MonkeyPatch, **values: str | None) -> None:
|
||||
for name in ("DEEPSEEK_API_KEY", "DEEPSEEK_BASE_URL", "OPENROUTER_API_KEY", "OPENROUTER_BASE_URL", "OPENULTRASAST_HUNTER_CLIENT"):
|
||||
monkeypatch.delenv(name, raising=False)
|
||||
for name, value in values.items():
|
||||
if value is not None:
|
||||
monkeypatch.setenv(name, value)
|
||||
|
||||
|
||||
def test_deepseek_wins_over_openrouter_and_never_touches_the_embedding_endpoint(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
from openultrasast.learning.endpoint import DEEPSEEK_BASE_URL, resolve_chat_endpoint
|
||||
|
||||
_endpoint_env(monkeypatch, DEEPSEEK_API_KEY="ds-key", OPENROUTER_API_KEY="or-key", OPENROUTER_BASE_URL="https://openrouter.ai/api/v1")
|
||||
resolved = resolve_chat_endpoint(ResolvedConfig())
|
||||
assert resolved is not None
|
||||
_client, endpoint = resolved
|
||||
assert endpoint.provider == "deepseek" and endpoint.base_url == DEEPSEEK_BASE_URL
|
||||
assert not endpoint.base_url.endswith("/v1") # DeepSeek rejects the version suffix
|
||||
assert "api_key" not in {field for field in endpoint.__dict__} # the key stays on the client, never in a journaled record
|
||||
embedding = provider.OpenRouterEmbeddingClient.from_env()
|
||||
assert embedding.base_url == "https://openrouter.ai/api/v1" and embedding.api_key == "or-key"
|
||||
|
||||
|
||||
def test_resolution_order_override_then_scripted_then_openrouter_then_none(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
from openultrasast.learning.endpoint import resolve_chat_endpoint
|
||||
from openultrasast.tool_hunter import ChatResponse, ScriptedChatClient
|
||||
|
||||
injected = ScriptedChatClient([ChatResponse(content="{}")])
|
||||
_endpoint_env(monkeypatch, DEEPSEEK_API_KEY="ds-key")
|
||||
resolved = resolve_chat_endpoint(ResolvedConfig(), override=injected)
|
||||
assert resolved is not None and resolved[0] is injected and resolved[1].provider == "scripted"
|
||||
|
||||
_endpoint_env(monkeypatch, DEEPSEEK_API_KEY="ds-key", OPENULTRASAST_HUNTER_CLIENT="scripted")
|
||||
resolved = resolve_chat_endpoint(ResolvedConfig())
|
||||
assert resolved is not None and resolved[1].provider == "scripted"
|
||||
|
||||
_endpoint_env(monkeypatch, OPENROUTER_API_KEY="or-key")
|
||||
resolved = resolve_chat_endpoint(ResolvedConfig())
|
||||
assert resolved is not None and resolved[1].provider == "openrouter" and resolved[1].base_url.endswith("/api/v1")
|
||||
|
||||
_endpoint_env(monkeypatch)
|
||||
assert resolve_chat_endpoint(ResolvedConfig()) is None
|
||||
|
||||
|
||||
def test_the_config_can_point_the_chat_endpoint_somewhere_else(monkeypatch: pytest.MonkeyPatch, tmp_path: Any) -> None:
|
||||
from openultrasast.config import load_config
|
||||
from openultrasast.learning.endpoint import resolve_chat_endpoint, resolve_models
|
||||
|
||||
_endpoint_env(monkeypatch, MY_KEY="local-key")
|
||||
monkeypatch.setenv("MY_KEY", "local-key")
|
||||
(tmp_path / "ousast.toml").write_text(
|
||||
'[models]\njudge = "deepseek-v4-pro"\nchat_base_url = "http://localhost:11434/v1"\nchat_api_key_env = "MY_KEY"\n\n'
|
||||
"[learning]\nk_runs = 3\nround_cost_cap_usd = 2.5\nminibatch = 4\n"
|
||||
)
|
||||
config = load_config(tmp_path / "ousast.toml")
|
||||
assert config.learning.k_runs == 3 and config.learning.round_cost_cap_usd == 2.5 and config.learning.minibatch == 4
|
||||
resolved = resolve_chat_endpoint(config)
|
||||
assert resolved is not None and resolved[1].base_url == "http://localhost:11434/v1"
|
||||
detector, judge = resolve_models(config)
|
||||
assert detector == "deepseek-v4-flash" and judge == "deepseek-v4-pro"
|
||||
|
||||
|
||||
def test_the_adapter_disables_thinking_asks_for_json_and_keeps_the_reasoning_field() -> None:
|
||||
from openultrasast.learning.endpoint import DeepSeekChatClient
|
||||
|
||||
recorder = _Recorder([{"content": '{"intent": "public"}', "reasoning_content": "step one"}])
|
||||
client = DeepSeekChatClient(recorder) # type: ignore[arg-type]
|
||||
response = client.complete(model="deepseek-v4-flash", messages=[{"role": "user", "content": "json please"}], tools=[], json_object=True)
|
||||
assert response.content == '{"intent": "public"}' and response.reasoning == "step one"
|
||||
body = recorder.calls[0]["extra_body"]
|
||||
assert body["thinking"] == {"type": "disabled"} # otherwise temperature is silently ignored
|
||||
assert body["response_format"] == {"type": "json_object"} and "logprobs" not in body
|
||||
|
||||
|
||||
def test_the_adapter_retries_once_on_empty_json_content_then_gives_up() -> None:
|
||||
from openultrasast.learning.endpoint import DeepSeekChatClient
|
||||
|
||||
recorder = _Recorder([{"content": ""}, {"content": '{"ok": true}'}])
|
||||
client = DeepSeekChatClient(recorder) # type: ignore[arg-type]
|
||||
assert client.complete(model="m", messages=[], tools=[], json_object=True).content == '{"ok": true}'
|
||||
assert len(recorder.calls) == 2
|
||||
empty = _Recorder([{"content": ""}, {"content": ""}])
|
||||
client = DeepSeekChatClient(empty) # type: ignore[arg-type]
|
||||
assert client.complete(model="m", messages=[], tools=[], json_object=True).content == ""
|
||||
assert len(empty.calls) == 2 # one retry, never a loop
|
||||
|
||||
|
||||
def test_log_probabilities_are_requested_only_when_asked_for() -> None:
|
||||
from openultrasast.learning.endpoint import DeepSeekChatClient
|
||||
|
||||
recorder = _Recorder([{"content": "{}", "logprobs": {"content": [{"logprob": -0.1}, {"logprob": -0.3}]}}])
|
||||
client = DeepSeekChatClient(recorder) # type: ignore[arg-type]
|
||||
response = client.complete(model="m", messages=[], tools=[], json_object=True, logprobs=True)
|
||||
assert recorder.calls[0]["extra_body"]["logprobs"] is True
|
||||
assert response.mean_logprob is not None and -0.3 < response.mean_logprob < -0.1
|
||||
|
||||
|
||||
def test_cost_comes_from_the_providers_own_usage_fields() -> None:
|
||||
from openultrasast.learning.endpoint import ChatEndpoint, price_of
|
||||
|
||||
flash = ChatEndpoint(provider="deepseek", base_url="https://api.deepseek.com", thinking=False, prices=price_of("deepseek-v4-flash"))
|
||||
cost = flash.cost({"prompt_cache_hit_tokens": 1_000_000, "prompt_cache_miss_tokens": 1_000_000, "completion_tokens": 1_000_000})
|
||||
assert cost is not None and abs(cost - (0.014 + 0.44 + 1.32)) < 1e-9
|
||||
openrouter_shape = flash.cost({"prompt_tokens": 1_000_000, "completion_tokens": 0})
|
||||
assert openrouter_shape is not None and abs(openrouter_shape - 0.44) < 1e-9 # no cache split reported: all miss
|
||||
unpriced = ChatEndpoint(provider="openrouter", base_url="x", thinking=False, prices=price_of("some/unknown-model"))
|
||||
assert unpriced.cost({"prompt_tokens": 10}) is None # never claim a cost we cannot source
|
||||
|
||||
|
||||
def test_the_underlying_client_forwards_extra_body_and_keeps_usage(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
sent: dict[str, object] = {}
|
||||
|
||||
class _Response:
|
||||
def __enter__(self) -> _Response:
|
||||
return self
|
||||
|
||||
def __exit__(self, *_: object) -> None:
|
||||
return None
|
||||
|
||||
def read(self) -> bytes:
|
||||
return json.dumps(
|
||||
{"choices": [{"message": {"content": "hi", "reasoning_content": "because"}}], "usage": {"prompt_tokens": 7}}
|
||||
).encode()
|
||||
|
||||
def _urlopen(request: Any, timeout: int = 0) -> _Response: # noqa: ARG001
|
||||
sent.update(json.loads(request.data.decode()))
|
||||
return _Response()
|
||||
|
||||
monkeypatch.setattr(urllib.request, "urlopen", _urlopen)
|
||||
client = provider.OpenRouterChatClient(api_key="k", base_url="https://api.deepseek.com")
|
||||
payload = client.complete_chat_raw(
|
||||
model="deepseek-v4-flash", messages=[{"role": "user", "content": "x"}], extra_body={"thinking": {"type": "disabled"}}
|
||||
)
|
||||
assert sent["thinking"] == {"type": "disabled"} and sent["temperature"] == 0
|
||||
assert payload["usage"] == {"prompt_tokens": 7}
|
||||
message = provider._extract_message(payload)
|
||||
assert message["reasoning_content"] == "because" # retained for the caller to replay
|
||||
Reference in New Issue
Block a user