Track SDK LLM usage
This commit is contained in:
@@ -30,6 +30,7 @@ if TYPE_CHECKING:
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
StreamEventSink = Callable[[str, Any], None]
|
||||
UsageSink = Callable[[str, str | None, Any], None]
|
||||
|
||||
|
||||
async def run_agent_loop(
|
||||
@@ -45,6 +46,7 @@ async def run_agent_loop(
|
||||
session: Session | None = None,
|
||||
start_parked: bool = False,
|
||||
event_sink: StreamEventSink | None = None,
|
||||
usage_sink: UsageSink | None = None,
|
||||
) -> RunResultBase | None:
|
||||
await coordinator.attach_runtime(
|
||||
agent_id,
|
||||
@@ -66,6 +68,7 @@ async def run_agent_loop(
|
||||
session=session,
|
||||
interactive=interactive,
|
||||
event_sink=event_sink,
|
||||
usage_sink=usage_sink,
|
||||
)
|
||||
else:
|
||||
result = await _run_noninteractive_until_lifecycle(
|
||||
@@ -78,6 +81,7 @@ async def run_agent_loop(
|
||||
max_turns=max_turns,
|
||||
session=session,
|
||||
event_sink=event_sink,
|
||||
usage_sink=usage_sink,
|
||||
)
|
||||
|
||||
if not interactive:
|
||||
@@ -101,6 +105,7 @@ async def run_agent_loop(
|
||||
session=session,
|
||||
interactive=interactive,
|
||||
event_sink=event_sink,
|
||||
usage_sink=usage_sink,
|
||||
)
|
||||
|
||||
|
||||
@@ -119,6 +124,7 @@ async def spawn_child_agent(
|
||||
skills: list[str],
|
||||
parent_history: list[Any],
|
||||
event_sink: StreamEventSink | None = None,
|
||||
usage_sink: UsageSink | None = None,
|
||||
) -> dict[str, Any]:
|
||||
parent_id = parent_ctx.get("agent_id")
|
||||
if not isinstance(parent_id, str):
|
||||
@@ -155,6 +161,7 @@ async def spawn_child_agent(
|
||||
parent_history=parent_history,
|
||||
),
|
||||
event_sink=event_sink,
|
||||
usage_sink=usage_sink,
|
||||
)
|
||||
|
||||
return {
|
||||
@@ -178,6 +185,7 @@ async def respawn_subagents(
|
||||
parent_ctx: dict[str, Any],
|
||||
root_id: str,
|
||||
event_sink: StreamEventSink | None = None,
|
||||
usage_sink: UsageSink | None = None,
|
||||
) -> None:
|
||||
"""Re-spawn subagent runners from a restored coordinator snapshot."""
|
||||
async with coordinator._lock:
|
||||
@@ -234,6 +242,7 @@ async def respawn_subagents(
|
||||
initial_input=[],
|
||||
start_parked=start_parked,
|
||||
event_sink=event_sink,
|
||||
usage_sink=usage_sink,
|
||||
)
|
||||
logger.info(
|
||||
"respawned %s (%s) parent=%s task_len=%d",
|
||||
@@ -259,6 +268,7 @@ async def _run_noninteractive_until_lifecycle(
|
||||
max_turns: int,
|
||||
session: Session | None,
|
||||
event_sink: StreamEventSink | None,
|
||||
usage_sink: UsageSink | None,
|
||||
) -> RunResultBase | None:
|
||||
"""Non-chat mode keeps running until finish_scan / agent_finish settles status."""
|
||||
result: RunResultBase | None = None
|
||||
@@ -278,6 +288,7 @@ async def _run_noninteractive_until_lifecycle(
|
||||
session=session,
|
||||
interactive=False,
|
||||
event_sink=event_sink,
|
||||
usage_sink=usage_sink,
|
||||
)
|
||||
|
||||
status = await _agent_status(coordinator, agent_id)
|
||||
@@ -322,6 +333,7 @@ async def _run_cycle(
|
||||
session: Session | None,
|
||||
interactive: bool,
|
||||
event_sink: StreamEventSink | None,
|
||||
usage_sink: UsageSink | None,
|
||||
) -> RunResultBase | None:
|
||||
try:
|
||||
await coordinator.mark_running(agent_id)
|
||||
@@ -344,6 +356,8 @@ async def _run_cycle(
|
||||
if stream.run_loop_exception is not None:
|
||||
raise stream.run_loop_exception
|
||||
finally:
|
||||
if usage_sink is not None:
|
||||
_emit_stream_usage(agent, agent_id, stream, usage_sink)
|
||||
await coordinator.detach_stream(agent_id, stream)
|
||||
except Exception as exc:
|
||||
if not interactive:
|
||||
@@ -463,6 +477,7 @@ async def _start_child_runner(
|
||||
initial_input: Any,
|
||||
start_parked: bool = False,
|
||||
event_sink: StreamEventSink | None = None,
|
||||
usage_sink: UsageSink | None = None,
|
||||
) -> None:
|
||||
session = open_agent_session(child_id, agents_db_path)
|
||||
sessions_to_close.append(session)
|
||||
@@ -486,7 +501,25 @@ async def _start_child_runner(
|
||||
session=session,
|
||||
start_parked=start_parked,
|
||||
event_sink=event_sink,
|
||||
usage_sink=usage_sink,
|
||||
),
|
||||
name=f"agent-{name}-{child_id}",
|
||||
)
|
||||
await coordinator.attach_runtime(child_id, task=task_handle)
|
||||
|
||||
|
||||
def _emit_stream_usage(
|
||||
agent: Any,
|
||||
agent_id: str,
|
||||
stream: RunResultBase,
|
||||
usage_sink: UsageSink,
|
||||
) -> None:
|
||||
context_wrapper = getattr(stream, "context_wrapper", None)
|
||||
usage = getattr(context_wrapper, "usage", None)
|
||||
if usage is None:
|
||||
return
|
||||
agent_name = getattr(agent, "name", None)
|
||||
try:
|
||||
usage_sink(agent_id, agent_name if isinstance(agent_name, str) else None, usage)
|
||||
except Exception:
|
||||
logger.exception("usage sink failed for %s", agent_id)
|
||||
|
||||
@@ -12,7 +12,7 @@ from strix.config.models import DEFAULT_MODEL_RETRY
|
||||
|
||||
|
||||
# Default max_turns budget passed to the SDK runner.
|
||||
DEFAULT_MAX_TURNS = 300
|
||||
DEFAULT_MAX_TURNS = 500
|
||||
|
||||
|
||||
def build_root_task(scan_config: dict[str, Any]) -> str:
|
||||
@@ -112,6 +112,7 @@ def make_model_settings(
|
||||
parallel_tool_calls=False,
|
||||
tool_choice="required",
|
||||
retry=DEFAULT_MODEL_RETRY,
|
||||
include_usage=True,
|
||||
)
|
||||
if reasoning_effort is not None:
|
||||
model_settings = model_settings.resolve(
|
||||
|
||||
@@ -36,6 +36,7 @@ from strix.core.inputs import (
|
||||
)
|
||||
from strix.core.paths import run_dir_for, runtime_state_dir
|
||||
from strix.core.sessions import open_agent_session
|
||||
from strix.report.state import get_global_report_state
|
||||
from strix.runtime import session_manager
|
||||
from strix.telemetry.logging import set_scan_id, setup_scan_logging
|
||||
|
||||
@@ -167,6 +168,17 @@ async def run_strix_scan(
|
||||
trace_include_sensitive_data=False,
|
||||
)
|
||||
|
||||
def usage_sink(agent_id: str, agent_name: str | None, usage: Any) -> None:
|
||||
report_state = get_global_report_state()
|
||||
if report_state is None:
|
||||
return
|
||||
report_state.record_sdk_usage(
|
||||
agent_id=agent_id,
|
||||
agent_name=agent_name,
|
||||
model=resolved_model,
|
||||
usage=usage,
|
||||
)
|
||||
|
||||
scope_context = build_scope_context(scan_config)
|
||||
|
||||
root_agent = build_strix_agent(
|
||||
@@ -205,6 +217,7 @@ async def run_strix_scan(
|
||||
max_turns=max_turns,
|
||||
interactive=interactive,
|
||||
event_sink=event_sink,
|
||||
usage_sink=usage_sink,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
@@ -236,6 +249,7 @@ async def run_strix_scan(
|
||||
parent_ctx=context,
|
||||
root_id=root_id,
|
||||
event_sink=event_sink,
|
||||
usage_sink=usage_sink,
|
||||
)
|
||||
|
||||
initial_input: Any = [] if is_resume else root_task
|
||||
@@ -276,6 +290,7 @@ async def run_strix_scan(
|
||||
session=root_session,
|
||||
start_parked=bool(interactive and is_resume and root_status != "running"),
|
||||
event_sink=event_sink,
|
||||
usage_sink=usage_sink,
|
||||
)
|
||||
except BaseException:
|
||||
logger.exception("Strix scan %s failed", scan_id)
|
||||
|
||||
@@ -47,6 +47,24 @@ def get_cvss_color(cvss_score: float) -> str:
|
||||
return "#6b7280"
|
||||
|
||||
|
||||
def format_token_count(count: float | None) -> str:
|
||||
value = int(count or 0)
|
||||
if value >= 1_000_000:
|
||||
return f"{value / 1_000_000:.1f}M"
|
||||
if value >= 1_000:
|
||||
return f"{value / 1_000:.1f}K"
|
||||
return str(value)
|
||||
|
||||
|
||||
def format_cost(cost: float | None) -> str:
|
||||
value = float(cost or 0.0)
|
||||
if value < 0.0001:
|
||||
return f"${value:.6f}"
|
||||
if value < 1:
|
||||
return f"${value:.4f}"
|
||||
return f"${value:.2f}"
|
||||
|
||||
|
||||
def format_vulnerability_report(report: dict[str, Any]) -> Text: # noqa: PLR0915
|
||||
"""Format a vulnerability report for CLI display with all rich fields."""
|
||||
field_style = "bold #4ade80"
|
||||
@@ -235,6 +253,71 @@ def _build_vulnerability_stats(stats_text: Text, report_state: Any) -> None:
|
||||
stats_text.append("\n")
|
||||
|
||||
|
||||
def _llm_usage(report_state: Any) -> dict[str, Any]:
|
||||
if hasattr(report_state, "get_total_llm_usage"):
|
||||
usage = report_state.get_total_llm_usage()
|
||||
return usage if isinstance(usage, dict) else {}
|
||||
usage = getattr(report_state, "run_record", {}).get("llm_usage")
|
||||
return usage if isinstance(usage, dict) else {}
|
||||
|
||||
|
||||
def _int_stat(usage: dict[str, Any], key: str) -> int:
|
||||
try:
|
||||
return max(0, int(usage.get(key) or 0))
|
||||
except (TypeError, ValueError):
|
||||
return 0
|
||||
|
||||
|
||||
def _float_stat(usage: dict[str, Any], key: str) -> float:
|
||||
try:
|
||||
value = float(usage.get(key) or 0.0)
|
||||
except (TypeError, ValueError):
|
||||
return 0.0
|
||||
return value if value > 0 else 0.0
|
||||
|
||||
|
||||
def _detail_value(usage: dict[str, Any], detail_key: str, value_key: str) -> int:
|
||||
details = usage.get(detail_key)
|
||||
if isinstance(details, list):
|
||||
details = details[0] if details and isinstance(details[0], dict) else {}
|
||||
if not isinstance(details, dict):
|
||||
return 0
|
||||
return _int_stat(details, value_key)
|
||||
|
||||
|
||||
def _build_llm_usage_stats(
|
||||
stats_text: Text,
|
||||
report_state: Any,
|
||||
*,
|
||||
include_requests: bool = False,
|
||||
) -> None:
|
||||
usage = _llm_usage(report_state)
|
||||
if not usage or _int_stat(usage, "total_tokens") <= 0:
|
||||
return
|
||||
|
||||
input_tokens = _int_stat(usage, "input_tokens")
|
||||
output_tokens = _int_stat(usage, "output_tokens")
|
||||
cached_tokens = _detail_value(usage, "input_tokens_details", "cached_tokens")
|
||||
reasoning_tokens = _detail_value(usage, "output_tokens_details", "reasoning_tokens")
|
||||
cost = _float_stat(usage, "cost")
|
||||
|
||||
stats_text.append("LLM Usage ", style="bold cyan")
|
||||
if include_requests:
|
||||
stats_text.append(f"{_int_stat(usage, 'requests')} req", style="white")
|
||||
stats_text.append(" | ", style="dim white")
|
||||
stats_text.append(f"In: {format_token_count(input_tokens)}", style="white")
|
||||
if cached_tokens > 0:
|
||||
stats_text.append(f" ({format_token_count(cached_tokens)} cached)", style="dim white")
|
||||
stats_text.append(" | ", style="dim white")
|
||||
stats_text.append(f"Out: {format_token_count(output_tokens)}", style="white")
|
||||
if reasoning_tokens > 0:
|
||||
stats_text.append(f" ({format_token_count(reasoning_tokens)} reasoning)", style="dim white")
|
||||
if cost > 0:
|
||||
stats_text.append(" | ", style="dim white")
|
||||
stats_text.append(f"Cost: {format_cost(cost)}", style="white")
|
||||
stats_text.append("\n")
|
||||
|
||||
|
||||
def build_final_stats_text(report_state: Any) -> Text:
|
||||
"""Build final stats from Strix-owned scan artifacts."""
|
||||
stats_text = Text()
|
||||
@@ -242,6 +325,7 @@ def build_final_stats_text(report_state: Any) -> Text:
|
||||
return stats_text
|
||||
|
||||
_build_vulnerability_stats(stats_text, report_state)
|
||||
_build_llm_usage_stats(stats_text, report_state, include_requests=True)
|
||||
|
||||
return stats_text
|
||||
|
||||
@@ -284,6 +368,8 @@ def build_live_stats_text(report_state: Any) -> Text:
|
||||
|
||||
stats_text.append("\n")
|
||||
|
||||
_build_llm_usage_stats(stats_text, report_state)
|
||||
|
||||
return stats_text
|
||||
|
||||
|
||||
@@ -295,6 +381,16 @@ def build_tui_stats_text(report_state: Any) -> Text:
|
||||
model = load_settings().llm.model or "unknown"
|
||||
stats_text.append(str(model), style="white")
|
||||
|
||||
usage = _llm_usage(report_state)
|
||||
if usage and _int_stat(usage, "total_tokens") > 0:
|
||||
stats_text.append("\n")
|
||||
stats_text.append("Tokens: ", style="bold white")
|
||||
stats_text.append(format_token_count(_int_stat(usage, "total_tokens")), style="white")
|
||||
cost = _float_stat(usage, "cost")
|
||||
if cost > 0:
|
||||
stats_text.append(" Cost: ", style="bold white")
|
||||
stats_text.append(format_cost(cost), style="white")
|
||||
|
||||
caido_url = getattr(report_state, "caido_url", None)
|
||||
if caido_url:
|
||||
stats_text.append("\n")
|
||||
|
||||
+12
-2
@@ -17,6 +17,7 @@ from strix.config.models import (
|
||||
configure_sdk_model_defaults,
|
||||
normalize_model_name,
|
||||
)
|
||||
from strix.report.state import get_global_report_state
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -187,11 +188,12 @@ async def check_duplicate(
|
||||
)
|
||||
|
||||
configure_sdk_model_defaults(settings)
|
||||
model = MultiProvider().get_model(normalize_model_name(model_name))
|
||||
resolved_model = normalize_model_name(model_name)
|
||||
model = MultiProvider().get_model(resolved_model)
|
||||
response = await model.get_response(
|
||||
system_instructions=DEDUPE_SYSTEM_PROMPT,
|
||||
input=user_msg,
|
||||
model_settings=ModelSettings(retry=DEFAULT_MODEL_RETRY),
|
||||
model_settings=ModelSettings(retry=DEFAULT_MODEL_RETRY, include_usage=True),
|
||||
tools=[],
|
||||
output_schema=None,
|
||||
handoffs=[],
|
||||
@@ -200,6 +202,14 @@ async def check_duplicate(
|
||||
conversation_id=None,
|
||||
prompt=None,
|
||||
)
|
||||
report_state = get_global_report_state()
|
||||
if report_state is not None:
|
||||
report_state.record_sdk_usage(
|
||||
agent_id="dedupe",
|
||||
agent_name="dedupe",
|
||||
model=resolved_model,
|
||||
usage=response.usage,
|
||||
)
|
||||
content = _extract_text(response)
|
||||
if not content:
|
||||
return {
|
||||
|
||||
@@ -6,7 +6,10 @@ from pathlib import Path
|
||||
from typing import Any, Optional
|
||||
from uuid import uuid4
|
||||
|
||||
from agents.usage import Usage
|
||||
|
||||
from strix.core.paths import run_dir_for
|
||||
from strix.report.usage import LLMUsageLedger
|
||||
from strix.report.writer import (
|
||||
read_run_record,
|
||||
write_executive_report,
|
||||
@@ -51,6 +54,7 @@ class ReportState:
|
||||
|
||||
self.scan_results: dict[str, Any] | None = None
|
||||
self.scan_config: dict[str, Any] | None = None
|
||||
self._llm_usage = LLMUsageLedger()
|
||||
self.run_record: dict[str, Any] = {
|
||||
"run_id": self.run_id,
|
||||
"run_name": self.run_name,
|
||||
@@ -58,6 +62,7 @@ class ReportState:
|
||||
"end_time": None,
|
||||
"status": "running",
|
||||
"targets_info": [],
|
||||
"llm_usage": self._build_llm_usage_record(),
|
||||
}
|
||||
self._run_dir: Path | None = None
|
||||
self._saved_vuln_ids: set[str] = set()
|
||||
@@ -104,6 +109,7 @@ class ReportState:
|
||||
if isinstance(scan_results, dict):
|
||||
self.scan_results = scan_results
|
||||
self.final_scan_result = self._format_final_scan_result(scan_results)
|
||||
self._hydrate_llm_usage(data.get("llm_usage"))
|
||||
logger.info("report state hydrated run.json from %s", run_dir)
|
||||
|
||||
json_path = run_dir / "vulnerabilities.json"
|
||||
@@ -206,6 +212,26 @@ class ReportState:
|
||||
def get_existing_vulnerabilities(self) -> list[dict[str, Any]]:
|
||||
return list(self.vulnerability_reports)
|
||||
|
||||
def record_sdk_usage(
|
||||
self,
|
||||
*,
|
||||
agent_id: str,
|
||||
usage: Usage | None,
|
||||
agent_name: str | None = None,
|
||||
model: str | None = None,
|
||||
) -> None:
|
||||
"""Record SDK-native token usage for one completed model run/cycle."""
|
||||
if self._llm_usage.record(
|
||||
agent_id=agent_id,
|
||||
agent_name=agent_name,
|
||||
model=model,
|
||||
usage=usage,
|
||||
):
|
||||
self.save_run_data()
|
||||
|
||||
def get_total_llm_usage(self) -> dict[str, Any]:
|
||||
return dict(self.run_record.get("llm_usage") or self._build_llm_usage_record())
|
||||
|
||||
def update_scan_final_fields(
|
||||
self,
|
||||
executive_summary: str,
|
||||
@@ -264,6 +290,7 @@ class ReportState:
|
||||
self.run_record["end_time"] = self.end_time
|
||||
self.run_record["status"] = status
|
||||
|
||||
self._sync_llm_usage_record()
|
||||
self._save_artifacts()
|
||||
|
||||
def cleanup(self, status: str = "stopped") -> None:
|
||||
@@ -304,3 +331,13 @@ class ReportState:
|
||||
logger.info("Essential scan data saved to: %s", run_dir)
|
||||
except (OSError, RuntimeError):
|
||||
logger.exception("Failed to save scan data")
|
||||
|
||||
def _sync_llm_usage_record(self) -> None:
|
||||
self.run_record["llm_usage"] = self._build_llm_usage_record()
|
||||
|
||||
def _build_llm_usage_record(self) -> dict[str, Any]:
|
||||
return self._llm_usage.to_record()
|
||||
|
||||
def _hydrate_llm_usage(self, raw_usage: Any) -> None:
|
||||
self._llm_usage.hydrate(raw_usage)
|
||||
self._sync_llm_usage_record()
|
||||
|
||||
@@ -0,0 +1,234 @@
|
||||
"""SDK-native LLM usage aggregation for scan reports."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
from agents.usage import Usage, deserialize_usage, serialize_usage
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class LLMUsageLedger:
|
||||
"""Aggregate SDK ``Usage`` objects and attach best-effort cost estimates."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._total_usage = Usage()
|
||||
self._agent_usage: dict[str, Usage] = {}
|
||||
self._agent_metadata: dict[str, dict[str, str]] = {}
|
||||
self._total_cost = 0.0
|
||||
self._agent_cost: dict[str, float] = {}
|
||||
|
||||
def record(
|
||||
self,
|
||||
*,
|
||||
agent_id: str,
|
||||
usage: Usage | None,
|
||||
agent_name: str | None = None,
|
||||
model: str | None = None,
|
||||
) -> bool:
|
||||
if usage is None or not _usage_has_activity(usage):
|
||||
return False
|
||||
|
||||
normalized_agent_id = str(agent_id or "unknown")
|
||||
estimated_cost = _estimate_litellm_cost(usage, model)
|
||||
|
||||
self._total_usage.add(usage)
|
||||
self._agent_usage.setdefault(normalized_agent_id, Usage()).add(usage)
|
||||
|
||||
metadata = self._agent_metadata.setdefault(normalized_agent_id, {})
|
||||
if agent_name:
|
||||
metadata["agent_name"] = agent_name
|
||||
if model:
|
||||
metadata["model"] = model
|
||||
|
||||
if estimated_cost is not None:
|
||||
self._total_cost += estimated_cost
|
||||
self._agent_cost[normalized_agent_id] = (
|
||||
self._agent_cost.get(normalized_agent_id, 0.0) + estimated_cost
|
||||
)
|
||||
|
||||
return True
|
||||
|
||||
def to_record(self) -> dict[str, Any]:
|
||||
record = serialize_usage(self._total_usage)
|
||||
record["cost"] = _round_cost(self._total_cost)
|
||||
record["cost_source"] = "litellm_estimate"
|
||||
record["agents"] = []
|
||||
|
||||
for agent_id in sorted(self._agent_usage):
|
||||
usage = self._agent_usage[agent_id]
|
||||
metadata = self._agent_metadata.get(agent_id, {})
|
||||
agent_record = serialize_usage(usage)
|
||||
agent_record.update(
|
||||
{
|
||||
"agent_id": agent_id,
|
||||
"agent_name": metadata.get("agent_name") or agent_id,
|
||||
"model": metadata.get("model"),
|
||||
"cost": _round_cost(self._agent_cost.get(agent_id, 0.0)),
|
||||
"cost_source": "litellm_estimate",
|
||||
}
|
||||
)
|
||||
record["agents"].append(agent_record)
|
||||
|
||||
return record
|
||||
|
||||
def hydrate(self, raw_usage: Any) -> None:
|
||||
self._total_usage = Usage()
|
||||
self._agent_usage.clear()
|
||||
self._agent_metadata.clear()
|
||||
self._total_cost = 0.0
|
||||
self._agent_cost.clear()
|
||||
|
||||
if not isinstance(raw_usage, dict):
|
||||
return
|
||||
|
||||
try:
|
||||
self._total_usage = deserialize_usage(raw_usage)
|
||||
except Exception:
|
||||
logger.exception("Failed to hydrate aggregate llm_usage from run.json")
|
||||
self._total_usage = Usage()
|
||||
|
||||
self._total_cost = _float_or_zero(raw_usage.get("cost"))
|
||||
agents = raw_usage.get("agents") or []
|
||||
if not isinstance(agents, list):
|
||||
return
|
||||
|
||||
for raw_agent in agents:
|
||||
if not isinstance(raw_agent, dict):
|
||||
continue
|
||||
agent_id = str(raw_agent.get("agent_id") or "").strip()
|
||||
if not agent_id:
|
||||
continue
|
||||
try:
|
||||
self._agent_usage[agent_id] = deserialize_usage(raw_agent)
|
||||
except Exception:
|
||||
logger.exception("Failed to hydrate llm_usage for agent %s", agent_id)
|
||||
self._agent_usage[agent_id] = Usage()
|
||||
|
||||
metadata: dict[str, str] = {}
|
||||
agent_name = raw_agent.get("agent_name")
|
||||
model = raw_agent.get("model")
|
||||
if isinstance(agent_name, str) and agent_name:
|
||||
metadata["agent_name"] = agent_name
|
||||
if isinstance(model, str) and model:
|
||||
metadata["model"] = model
|
||||
self._agent_metadata[agent_id] = metadata
|
||||
self._agent_cost[agent_id] = _float_or_zero(raw_agent.get("cost"))
|
||||
|
||||
|
||||
def _usage_has_activity(usage: Usage) -> bool:
|
||||
return bool(
|
||||
usage.requests
|
||||
or usage.input_tokens
|
||||
or usage.output_tokens
|
||||
or usage.total_tokens
|
||||
or usage.request_usage_entries
|
||||
)
|
||||
|
||||
|
||||
def _estimate_litellm_cost(usage: Usage, model: str | None) -> float | None:
|
||||
litellm_model = _litellm_model_name(model)
|
||||
if not litellm_model:
|
||||
return None
|
||||
|
||||
entries = list(usage.request_usage_entries)
|
||||
if not entries:
|
||||
return _estimate_litellm_entry_cost(usage, litellm_model)
|
||||
|
||||
total = 0.0
|
||||
estimated_any = False
|
||||
for entry in entries:
|
||||
cost = _estimate_litellm_entry_cost(entry, litellm_model)
|
||||
if cost is None:
|
||||
continue
|
||||
total += cost
|
||||
estimated_any = True
|
||||
|
||||
return total if estimated_any else None
|
||||
|
||||
|
||||
def _estimate_litellm_entry_cost(entry: Any, model: str) -> float | None:
|
||||
prompt_tokens = _int_or_zero(getattr(entry, "input_tokens", 0))
|
||||
completion_tokens = _int_or_zero(getattr(entry, "output_tokens", 0))
|
||||
total_tokens = _int_or_zero(getattr(entry, "total_tokens", 0))
|
||||
if total_tokens <= 0:
|
||||
total_tokens = prompt_tokens + completion_tokens
|
||||
if total_tokens <= 0:
|
||||
return None
|
||||
|
||||
usage_payload: dict[str, Any] = {
|
||||
"prompt_tokens": prompt_tokens,
|
||||
"completion_tokens": completion_tokens,
|
||||
"total_tokens": total_tokens,
|
||||
}
|
||||
prompt_details = _details_to_dict(getattr(entry, "input_tokens_details", None))
|
||||
completion_details = _details_to_dict(getattr(entry, "output_tokens_details", None))
|
||||
if prompt_details:
|
||||
usage_payload["prompt_tokens_details"] = prompt_details
|
||||
if completion_details:
|
||||
usage_payload["completion_tokens_details"] = completion_details
|
||||
|
||||
try:
|
||||
from litellm import completion_cost
|
||||
|
||||
cost = completion_cost(
|
||||
completion_response={
|
||||
"model": model.split("/", 1)[-1],
|
||||
"usage": usage_payload,
|
||||
},
|
||||
model=model,
|
||||
)
|
||||
except Exception: # noqa: BLE001 - LiteLLM raises plain Exception for unknown model prices.
|
||||
logger.debug("LiteLLM cost estimate unavailable for model %s", model, exc_info=True)
|
||||
return None
|
||||
|
||||
return cost if isinstance(cost, int | float) and cost >= 0 else None
|
||||
|
||||
|
||||
def _litellm_model_name(model: str | None) -> str | None:
|
||||
if not model:
|
||||
return None
|
||||
normalized = model.strip()
|
||||
for prefix in ("litellm/", "any-llm/", "openai/"):
|
||||
if normalized.startswith(prefix):
|
||||
normalized = normalized.removeprefix(prefix)
|
||||
break
|
||||
return normalized or None
|
||||
|
||||
|
||||
def _details_to_dict(details: Any) -> dict[str, Any]:
|
||||
if details is None:
|
||||
return {}
|
||||
if isinstance(details, list):
|
||||
for item in details:
|
||||
result = _details_to_dict(item)
|
||||
if result:
|
||||
return result
|
||||
return {}
|
||||
if hasattr(details, "model_dump"):
|
||||
return _details_to_dict(details.model_dump())
|
||||
if not isinstance(details, dict):
|
||||
return {}
|
||||
return {str(k): v for k, v in details.items() if v is not None}
|
||||
|
||||
|
||||
def _int_or_zero(value: Any) -> int:
|
||||
try:
|
||||
return max(0, int(value or 0))
|
||||
except (TypeError, ValueError):
|
||||
return 0
|
||||
|
||||
|
||||
def _float_or_zero(value: Any) -> float:
|
||||
try:
|
||||
result = float(value or 0.0)
|
||||
except (TypeError, ValueError):
|
||||
return 0.0
|
||||
return result if result >= 0 else 0.0
|
||||
|
||||
|
||||
def _round_cost(cost: float) -> float:
|
||||
return round(max(0.0, cost), 10)
|
||||
@@ -131,6 +131,20 @@ def end(report_state: "ReportState", exit_reason: str = "completed") -> None:
|
||||
except (ValueError, TypeError, AttributeError):
|
||||
pass
|
||||
|
||||
llm_props: dict[str, int | float] = {}
|
||||
try:
|
||||
usage = report_state.get_total_llm_usage()
|
||||
if isinstance(usage, dict):
|
||||
llm_props = {
|
||||
"llm_requests": int(usage.get("requests") or 0),
|
||||
"llm_input_tokens": int(usage.get("input_tokens") or 0),
|
||||
"llm_output_tokens": int(usage.get("output_tokens") or 0),
|
||||
"llm_tokens": int(usage.get("total_tokens") or 0),
|
||||
"llm_cost": float(usage.get("cost") or 0.0),
|
||||
}
|
||||
except (TypeError, ValueError, AttributeError):
|
||||
pass
|
||||
|
||||
_send(
|
||||
"scan_ended",
|
||||
{
|
||||
@@ -139,6 +153,7 @@ def end(report_state: "ReportState", exit_reason: str = "completed") -> None:
|
||||
"duration_seconds": round(duration),
|
||||
"vulnerabilities_total": len(report_state.vulnerability_reports),
|
||||
**{f"vulnerabilities_{k}": v for k, v in vulnerabilities_counts.items()},
|
||||
**llm_props,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user