mirror of
https://github.com/usestrix/strix.git
synced 2026-11-09 05:19:39 -10:00
Compare commits
7
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
8faf47de96 | ||
|
|
fb60e104f6 | ||
|
|
8267d7b311 | ||
|
|
8685c53b55 | ||
|
|
b4113e662f | ||
|
|
93dae3db07 | ||
|
|
1e1339e494 |
No files matched your search
@@ -49,6 +49,14 @@ Configure Strix using environment variables or a config file.
|
||||
Timeout in seconds for memory compression operations (context summarization).
|
||||
</ParamField>
|
||||
|
||||
<ParamField path="STRIX_OPENROUTER_STICKY_SESSIONS" default="false" type="boolean">
|
||||
Send a per-agent `session_id` on OpenRouter requests, so each agent's calls stay on one upstream provider and its prompt cache carries over between turns. When unset, OpenRouter routes every request freely.
|
||||
</ParamField>
|
||||
|
||||
<ParamField path="STRIX_CACHE_BLOCK_TOKENS" default="128" type="integer">
|
||||
Token block size that providers cache prompts in. A turn counts as a cache miss in the run report only when the cached tokens fall at least this many tokens short of the previous prompt and the prompt did not shrink. Lower it for providers with smaller blocks (DeepSeek and GLM use 64; vLLM defaults to 16).
|
||||
</ParamField>
|
||||
|
||||
### Dedicated deduplication model
|
||||
|
||||
Finding deduplication is a cheap, structured classification task. By default it
|
||||
|
||||
+32
-4
@@ -8,6 +8,7 @@ import inspect
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import AsyncGenerator
|
||||
from typing import TYPE_CHECKING, Any, cast
|
||||
|
||||
@@ -746,6 +747,11 @@ def _configure_litellm_compatibility() -> None:
|
||||
_install_openrouter_stream_cost_capture()
|
||||
|
||||
|
||||
# Agent ids are 8 hex characters and can repeat across runs; the session id
|
||||
# OpenRouter pins a provider to must not, so each agent gets its own UUID.
|
||||
_OPENROUTER_SESSION_IDS: dict[str, str] = {}
|
||||
|
||||
|
||||
def _install_openrouter_stream_cost_capture() -> None:
|
||||
"""Preserve OpenRouter's per-stream cost, which LiteLLM drops when streaming.
|
||||
|
||||
@@ -764,14 +770,16 @@ def _install_openrouter_stream_cost_capture() -> None:
|
||||
OpenrouterConfig,
|
||||
)
|
||||
|
||||
from strix.report.state import streamed_openrouter_costs
|
||||
from strix.report.state import record_openrouter_provider, streamed_openrouter_costs
|
||||
|
||||
class _StrixOpenRouterStreamingHandler(OpenRouterChatCompletionStreamingHandler):
|
||||
def chunk_parser(self, chunk: dict[str, Any]) -> Any:
|
||||
stream = super().chunk_parser(chunk)
|
||||
streamed_openrouter_costs.remember(
|
||||
chunk.get("id") or getattr(stream, "id", None), chunk.get("usage")
|
||||
)
|
||||
usage = chunk.get("usage")
|
||||
response_id = chunk.get("id") or getattr(stream, "id", None)
|
||||
streamed_openrouter_costs.remember(response_id, usage)
|
||||
if usage:
|
||||
record_openrouter_provider(chunk.get("provider"), usage)
|
||||
return stream
|
||||
|
||||
class _StrixOpenrouterConfig(OpenrouterConfig):
|
||||
@@ -784,6 +792,26 @@ def _install_openrouter_stream_cost_capture() -> None:
|
||||
json_mode=json_mode,
|
||||
)
|
||||
|
||||
def transform_response(self, *args: Any, **kwargs: Any) -> Any:
|
||||
# Non-streamed replies (LLM_DISABLE_STREAMING) skip the chunk parser.
|
||||
response = super().transform_response(*args, **kwargs)
|
||||
raw_response = kwargs.get("raw_response", args[1] if len(args) > 1 else None)
|
||||
with contextlib.suppress(Exception):
|
||||
body = raw_response.json() # type: ignore[union-attr]
|
||||
if body.get("usage"):
|
||||
record_openrouter_provider(body.get("provider"), body["usage"])
|
||||
return response
|
||||
|
||||
def transform_request(self, *args: Any, **kwargs: Any) -> dict[str, Any]:
|
||||
# Pin each agent's calls to one upstream provider so its prompt cache
|
||||
# survives between turns.
|
||||
body = super().transform_request(*args, **kwargs)
|
||||
agent_id = request_log.current_call_context().agent_id
|
||||
if agent_id and load_settings().llm.openrouter_sticky_sessions:
|
||||
session_id = _OPENROUTER_SESSION_IDS.setdefault(agent_id, str(uuid.uuid4()))
|
||||
body.setdefault("session_id", session_id)
|
||||
return body
|
||||
|
||||
# LiteLLM's provider-config factory reads litellm.OpenrouterConfig at call
|
||||
# time, so overriding the attribute is enough for the subclass to take
|
||||
# effect. (type: ignore — mypy rejects reassigning a class attribute.)
|
||||
|
||||
@@ -58,6 +58,14 @@ class LlmSettings(BaseSettings):
|
||||
default=True,
|
||||
alias="STRIX_PROMPT_CACHE",
|
||||
)
|
||||
# Providers cache prompts in fixed-size token blocks, so a fully cached prompt
|
||||
# can read back up to a block short. 128 covers the largest common size
|
||||
# (OpenAI; DeepSeek and GLM use 64, vLLM defaults to 16).
|
||||
cache_block_tokens: int = Field(default=128, ge=1, alias="STRIX_CACHE_BLOCK_TOKENS")
|
||||
openrouter_sticky_sessions: bool = Field(
|
||||
default=False,
|
||||
alias="STRIX_OPENROUTER_STICKY_SESSIONS",
|
||||
)
|
||||
disable_streaming: bool = Field(
|
||||
default=False,
|
||||
alias="LLM_DISABLE_STREAMING",
|
||||
|
||||
@@ -657,6 +657,35 @@ class ReportState:
|
||||
def record_observed_llm_cost(self, cost: float) -> None:
|
||||
self._llm_usage.record_observed_cost(cost)
|
||||
|
||||
def record_llm_provider(
|
||||
self,
|
||||
provider: str,
|
||||
*,
|
||||
agent_id: str | None,
|
||||
input_tokens: int,
|
||||
cached_tokens: int,
|
||||
cost: float,
|
||||
) -> None:
|
||||
self._llm_usage.record_provider(
|
||||
provider,
|
||||
agent_id=agent_id,
|
||||
input_tokens=input_tokens,
|
||||
cached_tokens=cached_tokens,
|
||||
cost=cost,
|
||||
cache_block_tokens=load_settings().llm.cache_block_tokens,
|
||||
)
|
||||
|
||||
def get_process_llm_providers(self) -> dict[str, dict[str, float]]:
|
||||
"""Per-provider usage since this process started, like get_process_llm_usage."""
|
||||
baseline = self._telemetry_llm_usage_baseline.get("providers") or {}
|
||||
providers: dict[str, dict[str, float]] = {}
|
||||
for name, tally in (self._llm_usage.to_record().get("providers") or {}).items():
|
||||
before = baseline.get(name) or {}
|
||||
delta = {key: max(0, value - _number(before.get(key))) for key, value in tally.items()}
|
||||
if delta["requests"]:
|
||||
providers[name] = delta
|
||||
return providers
|
||||
|
||||
def get_total_llm_usage(self) -> dict[str, Any]:
|
||||
return dict(self.run_record.get("llm_usage") or self._build_llm_usage_record())
|
||||
|
||||
@@ -990,6 +1019,30 @@ class StreamedOpenRouterCosts:
|
||||
streamed_openrouter_costs = StreamedOpenRouterCosts()
|
||||
|
||||
|
||||
def record_openrouter_provider(provider: Any, usage: Any) -> None:
|
||||
"""Tally which upstream provider served a stream, from its final usage chunk.
|
||||
|
||||
OpenRouter spreads one model across many providers whose prices, quantization
|
||||
and prompt caching differ, so this is what shows where a scan's tokens went.
|
||||
"""
|
||||
# Deferred: request_log pulls in the agents SDK, which strix.report must not import.
|
||||
from strix.llm.request_log import current_call_context
|
||||
|
||||
report_state = get_global_report_state()
|
||||
if report_state is None or not isinstance(usage, dict):
|
||||
return
|
||||
details = usage.get("prompt_tokens_details")
|
||||
report_state.record_llm_provider(
|
||||
provider if isinstance(provider, str) and provider else "unknown",
|
||||
agent_id=current_call_context().agent_id,
|
||||
input_tokens=int(_number(usage.get("prompt_tokens"))),
|
||||
cached_tokens=int(_number(details.get("cached_tokens")))
|
||||
if isinstance(details, dict)
|
||||
else 0,
|
||||
cost=openrouter_stream_cost(usage) or 0.0,
|
||||
)
|
||||
|
||||
|
||||
def litellm_cost_callback(
|
||||
kwargs: Any,
|
||||
completion_response: Any,
|
||||
|
||||
@@ -6,6 +6,7 @@ import logging
|
||||
from typing import Any
|
||||
|
||||
from agents.usage import Usage, deserialize_usage, serialize_usage
|
||||
from pydantic import BaseModel, TypeAdapter, ValidationError
|
||||
|
||||
from strix.report.pricing import resolve_litellm_model
|
||||
|
||||
@@ -13,6 +14,22 @@ from strix.report.pricing import resolve_litellm_model
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ProviderUsage(BaseModel):
|
||||
"""Running totals for one upstream provider OpenRouter routed calls to."""
|
||||
|
||||
requests: int = 0
|
||||
input_tokens: int = 0
|
||||
cached_tokens: int = 0
|
||||
cost: float = 0.0
|
||||
# Calls that didn't find the agent's whole previous prompt cached, and the
|
||||
# previous-prompt tokens they had to pay for again.
|
||||
cache_misses: int = 0
|
||||
missed_tokens: int = 0
|
||||
|
||||
|
||||
_PROVIDER_USAGE = TypeAdapter(dict[str, ProviderUsage])
|
||||
|
||||
|
||||
class LLMUsageLedger:
|
||||
"""Aggregate SDK ``Usage`` objects and attach best-effort cost estimates."""
|
||||
|
||||
@@ -23,6 +40,10 @@ class LLMUsageLedger:
|
||||
self._observed_cost = 0.0
|
||||
self._estimated_cost = 0.0
|
||||
self._has_observed_cost = False
|
||||
# Keyed by upstream provider name, e.g. "Z.AI" or "DeepInfra".
|
||||
self._providers: dict[str, ProviderUsage] = {}
|
||||
# Each agent's last prompt size, which its next call should find cached.
|
||||
self._last_input_tokens: dict[str, int] = {}
|
||||
# When True, tokens are still tracked but cost stays $0 — the run is on a
|
||||
# model subscription, so there is no metered per-token charge to report.
|
||||
self.zero_cost = False
|
||||
@@ -62,6 +83,33 @@ class LLMUsageLedger:
|
||||
self._observed_cost += float(cost)
|
||||
self._has_observed_cost = True
|
||||
|
||||
def record_provider(
|
||||
self,
|
||||
provider: str,
|
||||
*,
|
||||
agent_id: str | None,
|
||||
input_tokens: int,
|
||||
cached_tokens: int,
|
||||
cost: float,
|
||||
cache_block_tokens: int,
|
||||
) -> None:
|
||||
tally = self._providers.setdefault(provider, ProviderUsage())
|
||||
tally.requests += 1
|
||||
tally.input_tokens += input_tokens
|
||||
tally.cached_tokens += cached_tokens
|
||||
if agent_id:
|
||||
previous = self._last_input_tokens.get(agent_id, 0)
|
||||
# The previous prompt is a prefix of this one, so all of it but a
|
||||
# partial last block should read back cached. A shrinking prompt means
|
||||
# compaction rewrote it, so a miss is expected.
|
||||
missed = previous - cached_tokens
|
||||
if input_tokens >= previous and missed >= cache_block_tokens:
|
||||
tally.cache_misses += 1
|
||||
tally.missed_tokens += missed
|
||||
self._last_input_tokens[agent_id] = input_tokens
|
||||
if not self.zero_cost:
|
||||
tally.cost = _round_cost(tally.cost + cost)
|
||||
|
||||
@property
|
||||
def total_cost(self) -> float:
|
||||
if self.zero_cost:
|
||||
@@ -71,6 +119,7 @@ class LLMUsageLedger:
|
||||
def to_record(self) -> dict[str, Any]:
|
||||
record = serialize_usage(self._total_usage)
|
||||
record["cost"] = self.total_cost
|
||||
record["providers"] = {name: tally.model_dump() for name, tally in self._providers.items()}
|
||||
record["agents"] = []
|
||||
|
||||
agent_tokens = {aid: _resolve_total_tokens(u) for aid, u in self._agent_usage.items()}
|
||||
@@ -102,10 +151,16 @@ class LLMUsageLedger:
|
||||
self._observed_cost = 0.0
|
||||
self._estimated_cost = 0.0
|
||||
self._has_observed_cost = False
|
||||
self._providers = {}
|
||||
|
||||
if not isinstance(raw_usage, dict):
|
||||
return
|
||||
|
||||
try:
|
||||
self._providers = _PROVIDER_USAGE.validate_python(raw_usage.get("providers") or {})
|
||||
except ValidationError:
|
||||
logger.exception("Failed to hydrate llm_usage providers from run.json")
|
||||
|
||||
try:
|
||||
self._total_usage = deserialize_usage(raw_usage)
|
||||
except Exception:
|
||||
|
||||
@@ -118,6 +118,7 @@ def end(report_state: "ReportState", exit_reason: str = "completed") -> None:
|
||||
}
|
||||
except (TypeError, ValueError, AttributeError):
|
||||
pass
|
||||
providers = report_state.get_process_llm_providers()
|
||||
|
||||
report_state.posthog_scan_ended_sent = _send(
|
||||
"scan_ended",
|
||||
@@ -129,6 +130,7 @@ def end(report_state: "ReportState", exit_reason: str = "completed") -> None:
|
||||
"vulnerabilities_total": len(report_state.vulnerability_reports),
|
||||
**{f"vulnerabilities_{k}": v for k, v in vulnerabilities_counts.items()},
|
||||
**llm_props,
|
||||
**({"llm_providers": providers} if providers else {}),
|
||||
"skills": get_loaded_skill_names(),
|
||||
},
|
||||
)
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import urllib.parse
|
||||
from typing import TYPE_CHECKING, Any
|
||||
@@ -120,6 +121,7 @@ def end(report_state: ReportState, exit_reason: str = "completed") -> None:
|
||||
}
|
||||
except (TypeError, ValueError, AttributeError):
|
||||
pass
|
||||
providers = report_state.get_process_llm_providers()
|
||||
|
||||
report_state.scarf_scan_ended_sent = _send(
|
||||
"scan_ended",
|
||||
@@ -132,6 +134,8 @@ def end(report_state: ReportState, exit_reason: str = "completed") -> None:
|
||||
"vulnerabilities_total": len(report_state.vulnerability_reports),
|
||||
**{f"vulnerabilities_{k}": v for k, v in vulnerabilities_counts.items()},
|
||||
**llm_props,
|
||||
# Query params are flat, so the per-provider tally travels as JSON.
|
||||
**({"llm_providers": json.dumps(providers)} if providers else {}),
|
||||
"skills": ",".join(get_loaded_skill_names()),
|
||||
},
|
||||
)
|
||||
|
||||
+133
-1
@@ -2,9 +2,12 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, patch
|
||||
from typing import TYPE_CHECKING, Any
|
||||
from unittest.mock import MagicMock, call, patch
|
||||
|
||||
import httpx
|
||||
import litellm
|
||||
import pytest
|
||||
from litellm.types.utils import LlmProviders
|
||||
@@ -14,6 +17,7 @@ from strix.config.models import (
|
||||
_configure_litellm_compatibility,
|
||||
_install_openrouter_stream_cost_capture,
|
||||
)
|
||||
from strix.llm import request_log
|
||||
from strix.report.state import (
|
||||
ReportState,
|
||||
litellm_cost_callback,
|
||||
@@ -21,6 +25,11 @@ from strix.report.state import (
|
||||
set_global_report_state,
|
||||
streamed_openrouter_costs,
|
||||
)
|
||||
from strix.report.usage import LLMUsageLedger
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
@@ -257,3 +266,126 @@ def test_openrouter_stream_handler_records_cost() -> None:
|
||||
assert streamed_openrouter_costs.take(SimpleNamespace(id="gen-stream")) == pytest.approx(
|
||||
0.0035055
|
||||
)
|
||||
|
||||
|
||||
def test_openrouter_tallies_provider() -> None:
|
||||
_install_openrouter_stream_cost_capture()
|
||||
config = ProviderConfigManager.get_provider_chat_config(
|
||||
model="z-ai/glm-5.3", provider=LlmProviders.OPENROUTER
|
||||
)
|
||||
assert config is not None
|
||||
handler = config.get_model_response_iterator(streaming_response=iter([]), sync_stream=True)
|
||||
report_state = MagicMock()
|
||||
usage = {
|
||||
"prompt_tokens": 1000,
|
||||
"completion_tokens": 10,
|
||||
"cost": 0.002,
|
||||
"prompt_tokens_details": {"cached_tokens": 900},
|
||||
}
|
||||
with patch("strix.report.state.get_global_report_state", return_value=report_state):
|
||||
handler.chunk_parser(
|
||||
{
|
||||
"id": "gen-a",
|
||||
"created": 1,
|
||||
"model": "z-ai/glm-5.3",
|
||||
"provider": "Together",
|
||||
"choices": [{"index": 0, "delta": {"content": None}}],
|
||||
"usage": usage,
|
||||
}
|
||||
)
|
||||
# Non-streamed replies (LLM_DISABLE_STREAMING) carry the same fields.
|
||||
reply = {
|
||||
"choices": [{"message": {"role": "assistant"}}],
|
||||
"provider": "Together",
|
||||
"usage": usage,
|
||||
}
|
||||
config.transform_response(
|
||||
"z-ai/glm-5.3",
|
||||
httpx.Response(200, json=reply),
|
||||
litellm.ModelResponse(),
|
||||
MagicMock(),
|
||||
{},
|
||||
[],
|
||||
{},
|
||||
{},
|
||||
None,
|
||||
)
|
||||
|
||||
tally = call("Together", agent_id=None, input_tokens=1000, cached_tokens=900, cost=0.002)
|
||||
assert report_state.record_llm_provider.call_args_list == [tally, tally]
|
||||
|
||||
|
||||
def test_provider_tally_survives_run_record_round_trip() -> None:
|
||||
ledger = LLMUsageLedger()
|
||||
for input_tokens, cached_tokens, cost in [(1000, 900, 0.002), (500, 0, 0.001)]:
|
||||
ledger.record_provider(
|
||||
"Together",
|
||||
agent_id=None,
|
||||
input_tokens=input_tokens,
|
||||
cached_tokens=cached_tokens,
|
||||
cost=cost,
|
||||
cache_block_tokens=128,
|
||||
)
|
||||
|
||||
restored = LLMUsageLedger()
|
||||
restored.hydrate(ledger.to_record())
|
||||
|
||||
assert restored.to_record()["providers"] == {
|
||||
"Together": {
|
||||
"requests": 2,
|
||||
"input_tokens": 1500,
|
||||
"cached_tokens": 900,
|
||||
"cost": 0.003,
|
||||
"cache_misses": 0,
|
||||
"missed_tokens": 0,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def test_provider_tally_counts_cache_misses_per_agent() -> None:
|
||||
ledger = LLMUsageLedger()
|
||||
calls = [
|
||||
("Z.AI", "a1", 1000, 0), # first call: nothing to miss
|
||||
("Z.AI", "a1", 1200, 960), # 40 short of the previous 1000: within a block
|
||||
("DeepInfra", "a1", 1500, 200), # 1000 of the previous 1200 lost
|
||||
("Z.AI", "a2", 800, 0), # another agent's first call
|
||||
("Z.AI", "a1", 600, 0), # prompt shrank: compaction, not a miss
|
||||
]
|
||||
for provider, agent_id, input_tokens, cached_tokens in calls:
|
||||
ledger.record_provider(
|
||||
provider,
|
||||
agent_id=agent_id,
|
||||
input_tokens=input_tokens,
|
||||
cached_tokens=cached_tokens,
|
||||
cost=0.0,
|
||||
cache_block_tokens=128,
|
||||
)
|
||||
|
||||
providers = ledger.to_record()["providers"]
|
||||
assert providers["DeepInfra"]["cache_misses"] == 1
|
||||
assert providers["DeepInfra"]["missed_tokens"] == 1000
|
||||
assert providers["Z.AI"]["cache_misses"] == 0
|
||||
|
||||
|
||||
def test_openrouter_request_carries_agent_session_id() -> None:
|
||||
_install_openrouter_stream_cost_capture()
|
||||
config = ProviderConfigManager.get_provider_chat_config(
|
||||
model="moonshotai/kimi-k3", provider=LlmProviders.OPENROUTER
|
||||
)
|
||||
assert config is not None
|
||||
messages: list[AllMessageValues] = [{"role": "user", "content": "hi"}]
|
||||
|
||||
def body() -> dict[str, Any]:
|
||||
return config.transform_request("moonshotai/kimi-k3", messages, {}, {}, {})
|
||||
|
||||
assert "session_id" not in body()
|
||||
token = request_log.bind_call_context("a1b2c3d4", "root")
|
||||
try:
|
||||
assert "session_id" not in body()
|
||||
with patch("strix.config.models.load_settings") as settings:
|
||||
settings.return_value.llm.openrouter_sticky_sessions = True
|
||||
session_id = body()["session_id"]
|
||||
assert str(uuid.UUID(session_id)) == session_id
|
||||
assert body()["session_id"] == session_id
|
||||
finally:
|
||||
request_log.reset_call_context(token)
|
||||
Reference in new issue
Block a user