Track LLM token usage in reports

src/skillspector/llm_analyzer_base.py

@@ -35,7 +35,12 @@
 from langchain_core.messages import BaseMessage
 from pydantic import BaseModel, Field, field_validator
 
-from skillspector.llm_utils import get_chat_model
+from skillspector.llm_utils import (
+    LLMTokenUsage,
+    empty_token_usage,
+    extract_token_usage,
+    get_chat_model,
+)
 from skillspector.logging_config import get_logger
 from skillspector.model_info import get_max_input_tokens
 from skillspector.models import Finding
@@ -114,6 +119,12 @@ class LLMAnalysisResult(BaseModel):
     findings: list[LLMFinding] = Field(default_factory=list)
 
 
+def _add_token_usage(total: LLMTokenUsage, usage: LLMTokenUsage) -> None:
+    total["input_tokens"] += usage["input_tokens"]
+    total["output_tokens"] += usage["output_tokens"]
+    total["total_tokens"] += usage["total_tokens"]
+
+
 def estimate_tokens(text: str) -> int:
     """Approximate token count from character length."""
     return len(text) // CHARS_PER_TOKEN
@@ -275,8 +286,36 @@ def __init__(self, base_prompt: str, model: str):
         self._input_budget = get_max_input_tokens(model)
         self._llm = get_chat_model(model=model)
         self._structured_llm = (
-            self._llm.with_structured_output(self.response_schema) if self.response_schema else None
+            self._llm.with_structured_output(self.response_schema, include_raw=True)
+            if self.response_schema
+            else None
         )
+        self._llm_usage = empty_token_usage()
+
+    @property
+    def llm_usage(self) -> LLMTokenUsage:
+        """Cumulative token usage from the most recent batch run."""
+        return dict(self._llm_usage)  # type: ignore[return-value]
+
+    def _reset_llm_usage(self) -> None:
+        self._llm_usage = empty_token_usage()
+
+    def _record_usage_from_raw(self, raw: object) -> None:
+        _add_token_usage(self._llm_usage, extract_token_usage(raw))
+
+    def _unwrap_structured_response(self, response: object) -> object:
+        if not isinstance(response, dict) or not {"raw", "parsed", "parsing_error"} <= set(
+            response
+        ):
+            return response
+        raw = response.get("raw")
+        self._record_usage_from_raw(raw)
+        parsing_error = response.get("parsing_error")
+        if parsing_error is not None:
+            if isinstance(parsing_error, BaseException):
+                raise parsing_error
+            raise ValueError(str(parsing_error))
+        return response.get("parsed")
 
     # -- Batching -----------------------------------------------------------
 
@@ -376,6 +415,7 @@ def run_batches(
         :meth:`parse_response` returns :class:`Finding` objects; subclasses may
         return dicts or other types.
         """
+        self._reset_llm_usage()
         results: list[tuple[Batch, list]] = []
         for batch in batches:
             prompt = self.build_prompt(batch, **kwargs)
@@ -386,9 +426,11 @@ def run_batches(
                 len(batch.findings),
             )
             if self._structured_llm:
-                response = self._structured_llm.invoke(prompt)
+                response = self._unwrap_structured_response(self._structured_llm.invoke(prompt))
             else:
-                response = _message_text(self._llm.invoke(prompt))
+                raw_response = self._llm.invoke(prompt)
+                self._record_usage_from_raw(raw_response)
+                response = _message_text(raw_response)
             logger.debug("LLM response for %s", batch.file_label)
             parsed = self.parse_response(response, batch)
             results.append((batch, parsed))
@@ -417,6 +459,7 @@ async def arun_batches(
 
         The return type mirrors :meth:`run_batches`.
         """
+        self._reset_llm_usage()
         sem = asyncio.Semaphore(max_concurrency)
 
         async def _process(batch: Batch) -> tuple[Batch, list]:
@@ -429,9 +472,13 @@ async def _process(batch: Batch) -> tuple[Batch, list]:
                     len(batch.findings),
                 )
                 if self._structured_llm:
-                    response = await self._structured_llm.ainvoke(prompt)
+                    response = self._unwrap_structured_response(
+                        await self._structured_llm.ainvoke(prompt)
+                    )
                 else:
-                    response = _message_text(await self._llm.ainvoke(prompt))
+                    raw_response = await self._llm.ainvoke(prompt)
+                    self._record_usage_from_raw(raw_response)
+                    response = _message_text(raw_response)
                 logger.debug("LLM response for %s", batch.file_label)
                 return (batch, self.parse_response(response, batch))
 

src/skillspector/llm_utils.py

@@ -36,9 +36,10 @@
 
 import asyncio
 import json
-from typing import NoReturn
+from typing import NoReturn, TypedDict
 
 from langchain_core.language_models.chat_models import BaseChatModel
+from pydantic import BaseModel
 
 from skillspector.model_info import get_max_input_tokens, get_max_output_tokens
 from skillspector.providers import (
@@ -53,6 +54,32 @@
 from skillspector.providers.openai import OpenAIProvider
 
 
+class LLMTokenUsage(TypedDict):
+    """Provider-normalized token usage for LLM calls."""
+
+    input_tokens: int
+    output_tokens: int
+    total_tokens: int
+
+
+def empty_token_usage() -> LLMTokenUsage:
+    return {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0}
+
+
+def extract_token_usage(raw: object) -> LLMTokenUsage:
+    usage = getattr(raw, "usage_metadata", None) or {}
+    if not isinstance(usage, dict):
+        return empty_token_usage()
+    input_tokens = int(usage.get("input_tokens") or usage.get("prompt_tokens") or 0)
+    output_tokens = int(usage.get("output_tokens") or usage.get("completion_tokens") or 0)
+    total_tokens = int(usage.get("total_tokens") or input_tokens + output_tokens)
+    return {
+        "input_tokens": input_tokens,
+        "output_tokens": output_tokens,
+        "total_tokens": total_tokens,
+    }
+
+
 def _resolve_llm_credentials() -> tuple[str, str | None]:
     """Return ``(api_key, base_url)`` resolved from the environment.
 
@@ -161,11 +188,19 @@ class _StructuredAgentCLIModel:
     ``complete()``, then parses and validates the response into *schema*.
     """
 
-    def __init__(self, provider: object, model: str, max_output_tokens: int, schema: type) -> None:
+    def __init__(
+        self,
+        provider: object,
+        model: str,
+        max_output_tokens: int,
+        schema: type[BaseModel],
+        include_raw: bool = False,
+    ) -> None:
         self._provider = provider
         self._model = model
         self._max_output_tokens = max_output_tokens
         self._schema = schema
+        self._include_raw = include_raw
 
     def _augment(self, prompt: str) -> str:
         schema_json = json.dumps(self._schema.model_json_schema(), indent=2)
@@ -182,7 +217,17 @@ def invoke(self, prompt: str) -> object:
             model=self._model,
             max_output_tokens=self._max_output_tokens,
         )
-        return self._schema.model_validate(_extract_json_object(raw))
+        try:
+            parsed = self._schema.model_validate(_extract_json_object(raw))
+            parsing_error = None
+        except Exception as exc:
+            parsed = None
+            parsing_error = exc
+            if not self._include_raw:
+                raise
+        if self._include_raw:
+            return {"raw": _AgentCLIMessage(raw), "parsed": parsed, "parsing_error": parsing_error}
+        return parsed
 
     async def ainvoke(self, prompt: str) -> object:
         return await asyncio.to_thread(self.invoke, prompt)
@@ -227,9 +272,11 @@ def invoke(self, prompt: str) -> _AgentCLIMessage:
     async def ainvoke(self, prompt: str) -> _AgentCLIMessage:
         return await asyncio.to_thread(self.invoke, prompt)
 
-    def with_structured_output(self, schema: type) -> _StructuredAgentCLIModel:
+    def with_structured_output(
+        self, schema: type[BaseModel], *, include_raw: bool = False
+    ) -> _StructuredAgentCLIModel:
         return _StructuredAgentCLIModel(
-            self._provider, self._model, self._max_output_tokens, schema
+            self._provider, self._model, self._max_output_tokens, schema, include_raw=include_raw
         )
 
 
@@ -272,7 +319,17 @@ def chat_completion(prompt: str, *, model: str | None = None) -> str:
     which normalise content blocks to a single string) and falls back to
     ``.content`` for the CLI adapter's ``_AgentCLIMessage``.
     """
+    content, _usage = chat_completion_with_usage(prompt, model=model)
+    return content
+
+
+def chat_completion_with_usage(
+    prompt: str, *, model: str | None = None
+) -> tuple[str, LLMTokenUsage]:
+    """Request a chat completion and return its content and token usage."""
     response = get_chat_model(model=model).invoke(prompt)
     if hasattr(response, "text"):
-        return response.text  # type: ignore[union-attr]
-    return response.content or ""  # type: ignore[union-attr]
+        content = response.text
+    else:
+        content = response.content or ""
+    return content, extract_token_usage(response)

src/skillspector/nodes/analyzers/mcp_tool_poisoning.py

@@ -23,7 +23,7 @@
 import re
 import unicodedata
 
-from skillspector.llm_utils import chat_completion
+from skillspector.llm_utils import chat_completion_with_usage, empty_token_usage
 from skillspector.models import Finding
 from skillspector.state import (
     AnalyzerNodeResponse,
@@ -691,6 +691,7 @@ def _check_tp4(state: SkillspectorState) -> tuple[list[Finding], LLMCallRecord |
     as a degraded LLM stage. See :func:`skillspector.state.llm_call_record`.
     """
     attempted = False
+    usage = empty_token_usage()
     try:
         manifest: dict = state.get("manifest") or {}
         description = manifest.get("description")
@@ -762,7 +763,7 @@ def _check_tp4(state: SkillspectorState) -> tuple[list[Finding], LLMCallRecord |
 }}"""
 
         attempted = True
-        response = chat_completion(prompt, model=model)
+        response, usage = chat_completion_with_usage(prompt, model=model)
 
         # Parse JSON — handle optional ```json code blocks
         json_text = response.strip()
@@ -776,7 +777,7 @@ def _check_tp4(state: SkillspectorState) -> tuple[list[Finding], LLMCallRecord |
                 json_text = json_text.rstrip()[:-3].rstrip()
 
         result = json.loads(json_text)
-        ok_record = llm_call_record(ANALYZER_ID, ok=True)
+        ok_record = llm_call_record(ANALYZER_ID, ok=True, **usage)
 
         if not result.get("is_mismatch"):
             return [], ok_record
@@ -818,7 +819,7 @@ def _check_tp4(state: SkillspectorState) -> tuple[list[Finding], LLMCallRecord |
         # Only record a failure if the LLM call was actually attempted; a failure
         # before the call (e.g. building the prompt) is not an LLM-stage failure.
         if attempted:
-            return [], llm_call_record(ANALYZER_ID, ok=False, error=str(exc))
+            return [], llm_call_record(ANALYZER_ID, ok=False, error=str(exc), **usage)
         return [], None
 
 

src/skillspector/nodes/analyzers/semantic_developer_intent.py

@@ -32,6 +32,11 @@
 ANALYZER_ID = "semantic_developer_intent"
 logger = get_logger(__name__)
 
+
+class _NoUsage:
+    llm_usage = {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0}
+
+
 ANALYZER_PROMPT = """\
 You are a developer-intent auditor for AI agent skills.  Your job is to
 detect mismatches between what a skill *claims* to do (its manifest and
@@ -179,12 +184,22 @@ def node(state: SkillspectorState) -> AnalyzerNodeResponse:
         results = asyncio.run(analyzer.arun_batches(batches))
         findings = analyzer.collect_findings(results)
         logger.info("%s: %d findings", ANALYZER_ID, len(findings))
-        return {"findings": findings, "llm_call_log": [llm_call_record(ANALYZER_ID, ok=True)]}
+        return {
+            "findings": findings,
+            "llm_call_log": [llm_call_record(ANALYZER_ID, ok=True, **analyzer.llm_usage)],
+        }
     except ValueError:
         raise
     except Exception as exc:
         logger.warning("%s failed: %s", ANALYZER_ID, exc)
         return {
             "findings": [],
-            "llm_call_log": [llm_call_record(ANALYZER_ID, ok=False, error=str(exc))],
+            "llm_call_log": [
+                llm_call_record(
+                    ANALYZER_ID,
+                    ok=False,
+                    error=str(exc),
+                    **locals().get("analyzer", _NoUsage()).llm_usage,
+                )
+            ],
         }

src/skillspector/nodes/analyzers/semantic_quality_policy.py

@@ -32,6 +32,11 @@
 ANALYZER_ID = "semantic_quality_policy"
 logger = get_logger(__name__)
 
+
+class _NoUsage:
+    llm_usage = {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0}
+
+
 ANALYZER_PROMPT = """\
 You are a quality and safety auditor for AI agent skills.  Your job is to
 review a single skill file and report findings that fall into the categories
@@ -148,12 +153,22 @@ def node(state: SkillspectorState) -> AnalyzerNodeResponse:
         results = asyncio.run(analyzer.arun_batches(batches))
         findings = analyzer.collect_findings(results)
         logger.info("%s: %d findings", ANALYZER_ID, len(findings))
-        return {"findings": findings, "llm_call_log": [llm_call_record(ANALYZER_ID, ok=True)]}
+        return {
+            "findings": findings,
+            "llm_call_log": [llm_call_record(ANALYZER_ID, ok=True, **analyzer.llm_usage)],
+        }
     except ValueError:
         raise
     except Exception as exc:
         logger.warning("%s failed: %s", ANALYZER_ID, exc)
         return {
             "findings": [],
-            "llm_call_log": [llm_call_record(ANALYZER_ID, ok=False, error=str(exc))],
+            "llm_call_log": [
+                llm_call_record(
+                    ANALYZER_ID,
+                    ok=False,
+                    error=str(exc),
+                    **locals().get("analyzer", _NoUsage()).llm_usage,
+                )
+            ],
         }

src/skillspector/nodes/analyzers/semantic_security_discovery.py

@@ -27,6 +27,11 @@
 ANALYZER_ID = "semantic_security_discovery"
 logger = get_logger(__name__)
 
+
+class _NoUsage:
+    llm_usage = {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0}
+
+
 ANALYZER_PROMPT = """\
 You are a security analyzer for AI agent skill files. Your task is to identify \
 **intent and attack-phrasing risks** — issues that evade regex/static detection because \
@@ -90,14 +95,22 @@ def node(state: SkillspectorState) -> AnalyzerNodeResponse:
         results = analyzer.run_batches(batches)
         findings = analyzer.collect_findings(results)
         logger.info("%s: %d findings", ANALYZER_ID, len(findings))
-        return {"findings": findings, "llm_call_log": [llm_call_record(ANALYZER_ID, ok=True)]}
+        return {
+            "findings": findings,
+            "llm_call_log": [llm_call_record(ANALYZER_ID, ok=True, **analyzer.llm_usage)],
+        }
     except ValidationError as exc:
         # Malformed LLM response — degrade gracefully rather than crashing the graph
         logger.warning("%s: LLM returned malformed response: %s", ANALYZER_ID, exc)
         return {
             "findings": [],
             "llm_call_log": [
-                llm_call_record(ANALYZER_ID, ok=False, error=f"malformed LLM response: {exc}")
+                llm_call_record(
+                    ANALYZER_ID,
+                    ok=False,
+                    error=f"malformed LLM response: {exc}",
+                    **locals().get("analyzer", _NoUsage()).llm_usage,
+                )
             ],
         }
     except ValueError:
@@ -106,5 +119,12 @@ def node(state: SkillspectorState) -> AnalyzerNodeResponse:
         logger.warning("%s failed: %s", ANALYZER_ID, exc)
         return {
             "findings": [],
-            "llm_call_log": [llm_call_record(ANALYZER_ID, ok=False, error=str(exc))],
+            "llm_call_log": [
+                llm_call_record(
+                    ANALYZER_ID,
+                    ok=False,
+                    error=str(exc),
+                    **locals().get("analyzer", _NoUsage()).llm_usage,
+                )
+            ],
         }