prompture 1.0.53.dev2__tar.gz → 1.0.54.dev1__tar.gz

prompture-1.0.54.dev1/.claude/skills/update-pricing/SKILL.md

@@ -0,0 +1,164 @@
+---
+name: update-pricing
+description: Update LLM model pricing in Prompture. Pricing resolves through a pluggable source registry — by default local KB JSON files (primary, curated) then models.dev (fallback). Use when model prices change, new models launch, models.dev data is stale, or you need to add a custom pricing source.
+metadata:
+  author: prompture
+  version: "4.0"
+---
+
+# Update Model Pricing
+
+## How Pricing Resolves
+
+Pricing flows through a **pluggable source registry** at `prompture/infra/pricing.py`. `get_model_rates(provider, model_id)` walks the registered sources in priority order (lowest first) and returns the first dict that includes both `input` and `output` per-1M rates.
+
+Default registration (controlled by `Settings.pricing_source`):
+
+| Setting value | Sources registered (priority order) |
+|---------------|-------------------------------------|
+| `local_first` (default) | `LocalKBPricingSource` (priority 0) → `ModelsDevPricingSource` (priority 100) |
+| `local_only` | `LocalKBPricingSource` only |
+| `models_dev_only` | `ModelsDevPricingSource` only |
+
+**Use `local_first` (default)** when you want curated rates with automatic fallback. **`local_only`** for offline/locked-down envs or when you don't trust models.dev. **`models_dev_only`** as a temporary opt-out while local files are unmaintained.
+
+`CostMixin._calculate_cost()` consumes the resolved dict — `input`, `output`, optional `cache_read`, `cache_write`, `reasoning` rates per 1M tokens. Cached prompt tokens are billed at `cache_read` (falls back to `input` rate when not published).
+
+## Local KB layout (the primary source)
+
+Per-provider JSON files in `prompture/infra/rates/` carry both **capabilities** and **pricing** for each model. Adding a `cost` block makes that model use local pricing; omitting it falls through to models.dev.
+
+```json
+{
+  "gpt-5.5": {
+    "supports_temperature": false,
+    "supports_tool_use": true,
+    "supports_structured_output": true,
+    "supports_vision": true,
+    "is_reasoning": true,
+    "context_window": 1050000,
+    "max_output_tokens": 128000,
+    "modalities_input": ["text", "image"],
+    "modalities_output": ["text"],
+    "api_type": "openai",
+    "tokens_param": "max_completion_tokens",
+    "cost": {
+      "input": 5.00,
+      "output": 30.00,
+      "cache_read": 0.50
+    }
+  }
+}
+```
+
+**`cost` block keys** (USD per 1M tokens):
+- `input` — required for cost to compute
+- `output` — required
+- `cache_read` — optional. Discount rate for cached prompt tokens (Anthropic / OpenAI / Kimi all support this).
+- `cache_write` — optional. Anthropic only — premium rate for tokens written to prompt cache.
+- `reasoning` — optional. For models that bill reasoning tokens separately.
+
+## Vendor pricing pages (verify before editing)
+
+**Always cross-check rates against the vendor's official page**, not just models.dev — feeds get stale.
+
+| Provider | Pricing page |
+|----------|--------------|
+| OpenAI | https://developers.openai.com/api/docs/pricing |
+| Anthropic (Claude) | https://platform.claude.com/docs/en/about-claude/pricing |
+| Moonshot (Kimi) | https://platform.kimi.ai/docs (search "Model Pricing") |
+| xAI (Grok) | https://docs.x.ai/docs/models |
+| Google (Gemini) | https://ai.google.dev/gemini-api/docs/pricing |
+| Groq | https://groq.com/pricing |
+| OpenRouter | https://openrouter.ai/models |
+| Moonshot K2 | https://platform.moonshot.ai/docs/pricing/chat |
+| Z.ai (GLM) | https://docs.z.ai/guides/llm/pricing |
+
+When a vendor publishes "cache hit" / "cache read" pricing, encode it as `cache_read`. When they publish a "cache write" / "cache creation" rate (Anthropic only), encode it as `cache_write`.
+
+## When to update what
+
+| Scenario | Action |
+|----------|--------|
+| Vendor changed prices for a model already in the KB | Update the `cost` block in `rates/{provider}.json` |
+| New model from existing provider | Add entry with capabilities + `cost` block |
+| New provider | Create `rates/{provider}.json`, add to `PROVIDER_MAP` in `model_rates.py`, wire descriptor in `provider_descriptors.py` |
+| Model not in KB but in models.dev | Nothing — fallback already covers it |
+| You want to override what models.dev says | Add the `cost` block to local KB; it wins |
+| You want to lock pricing to local-only | Set `pricing_source=local_only` in env / `.env` |
+
+## Adding a custom pricing source
+
+For internal price feeds, vendor APIs the registry doesn't cover, or per-environment overrides:
+
+```python
+from prompture.infra.pricing import register_pricing_source
+
+class MyInternalFeed:
+    name = "internal_feed"
+    priority = 50  # between local_kb (0) and models_dev (100)
+
+    def get_rates(self, provider: str, model_id: str) -> dict[str, float] | None:
+        # Return {"input": ..., "output": ..., "cache_read": ...} or None
+        return _query_my_internal_service(provider, model_id)
+
+register_pricing_source(MyInternalFeed())
+```
+
+Sources are ordered by `(priority, name)`. The first to return a dict containing both `input` and `output` wins. Exceptions inside `get_rates` are swallowed and the next source is tried.
+
+## Refreshing the models.dev cache
+
+```python
+from prompture.infra.model_rates import refresh_rates_cache
+refresh_rates_cache(force=True)
+```
+
+Or delete the cache file:
+```bash
+rm ~/.prompture/cache/models_dev.json
+```
+
+## Provider name → models.dev key mapping
+
+Lives in `prompture/infra/model_rates.py` (`PROVIDER_MAP`):
+
+```python
+{"openai": "openai", "claude": "anthropic", "google": "google", "groq": "groq",
+ "grok": "xai", "azure": "azure", "openrouter": "openrouter",
+ "moonshot": "moonshotai", "zai": "zai"}
+```
+
+## Free / local drivers (always $0)
+
+`ollama_driver.py`, `lmstudio_driver.py`, `local_http_driver.py`, `airllm_driver.py`, `hugging_driver.py`.
+
+## Steps for adding / updating a model price
+
+1. **Open** the vendor's pricing page (table above) and copy the current rates.
+2. **Edit** `prompture/infra/rates/{provider}.json` — find the model entry (or add a new one with capabilities) and add/update the `cost` block.
+3. **Verify** with `python -c "from prompture.infra.model_rates import get_model_rates; print(get_model_rates('openai', 'gpt-5.5'))"`. The dict should include exactly the keys you wrote.
+4. **Run tests**: `pytest tests/test_pricing.py tests/test_cached_tokens.py -q`
+5. **For long-running processes** that imported the LocalKBPricingSource before the edit: the in-memory KB is loaded once at startup. Either restart, or call `LocalKBPricingSource().reload()` on the registered instance.
+
+## Verification commands
+
+```bash
+# Check resolved rates for a model
+python -c "from prompture.infra.model_rates import get_model_rates; print(get_model_rates('openai', 'gpt-5.5'))"
+
+# Inspect which sources are registered
+python -c "from prompture.infra.pricing import get_pricing_sources; print([(s.name, s.priority) for s in get_pricing_sources()])"
+
+# Check that capabilities still load
+python -c "from prompture.infra.model_rates import get_model_capabilities; print(get_model_capabilities('openai', 'gpt-5.5'))"
+
+# Run pricing + cost tests
+pytest tests/test_pricing.py tests/test_cached_tokens.py tests/test_model_rates.py -q
+```
+
+## Side effects of editing the KB
+
+- `prompture/infra/discovery.py` reads KB model IDs via `get_kb_models_for_provider()` for static detection (no API key needed).
+- `CostMixin._get_model_config()` still reads `tokens_param` / `supports_temperature` from the KB capabilities.
+- A model with a local `cost` block now appears in cost calculations regardless of models.dev cache state — useful when models.dev lags vendor announcements.

{prompture-1.0.53.dev2 → prompture-1.0.54.dev1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: prompture
-Version: 1.0.53.dev2
+Version: 1.0.54.dev1
 Summary: Ask LLMs to return structured JSON and run cross-model tests. API-first.
 Author-email: Juan Denis <juan@vene.co>
 License-Expression: MIT

{prompture-1.0.53.dev2 → prompture-1.0.54.dev1}/prompture/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '1.0.53.dev2'
-__version_tuple__ = version_tuple = (1, 0, 53, 'dev2')
+__version__ = version = '1.0.54.dev1'
+__version_tuple__ = version_tuple = (1, 0, 54, 'dev1')
 
 __commit_id__ = commit_id = None

{prompture-1.0.53.dev2 → prompture-1.0.54.dev1}/prompture/drivers/async_base.py

@@ -232,6 +232,8 @@ class AsyncDriver(ABC):
                 prompt_tokens=meta.get("prompt_tokens", 0),
                 completion_tokens=meta.get("completion_tokens", 0),
                 total_tokens=meta.get("total_tokens", 0),
+                cached_prompt_tokens=meta.get("cached_prompt_tokens", 0),
+                cache_creation_tokens=meta.get("cache_creation_tokens", 0),
                 cost=meta.get("cost", 0.0),
                 elapsed_ms=elapsed_ms,
                 status=status,

{prompture-1.0.53.dev2 → prompture-1.0.54.dev1}/prompture/drivers/async_claude_driver.py

@@ -20,6 +20,7 @@ from .claude_driver import (
     _build_anthropic_meta,
     _build_anthropic_stream_done,
     _convert_tools_to_anthropic,
+    _extract_anthropic_cache_tokens,
     _extract_anthropic_system_and_messages,
     _extract_anthropic_text_and_tool_calls,
 )
@@ -104,8 +105,14 @@ class AsyncClaudeDriver(CostMixin, AsyncDriver):
         if not text and reasoning_content:
             text = reasoning_content
 
+        cache_read, cache_create = _extract_anthropic_cache_tokens(resp.usage)
         total_cost = self._calculate_cost(
-            "claude", model, resp.usage.input_tokens, resp.usage.output_tokens
+            "claude",
+            model,
+            resp.usage.input_tokens + cache_read + cache_create,
+            resp.usage.output_tokens,
+            cached_tokens=cache_read,
+            cache_creation_tokens=cache_create,
         )
         meta = _build_anthropic_meta(resp, model, total_cost)
 
@@ -157,8 +164,14 @@ class AsyncClaudeDriver(CostMixin, AsyncDriver):
 
         resp = await client.messages.create(**kwargs)
 
+        cache_read, cache_create = _extract_anthropic_cache_tokens(resp.usage)
         total_cost = self._calculate_cost(
-            "claude", model, resp.usage.input_tokens, resp.usage.output_tokens
+            "claude",
+            model,
+            resp.usage.input_tokens + cache_read + cache_create,
+            resp.usage.output_tokens,
+            cached_tokens=cache_read,
+            cache_creation_tokens=cache_create,
         )
         meta = _build_anthropic_meta(resp, model, total_cost)
 
@@ -205,8 +218,10 @@ class AsyncClaudeDriver(CostMixin, AsyncDriver):
 
         full_text = ""
         full_reasoning = ""
-        prompt_tokens = 0
+        base_input = 0
         completion_tokens = 0
+        cache_read = 0
+        cache_create = 0
 
         async with client.messages.stream(**kwargs) as stream:
             async for event in stream:
@@ -228,9 +243,19 @@ class AsyncClaudeDriver(CostMixin, AsyncDriver):
                     elif event.type == "message_start" and hasattr(event, "message"):
                         usage = getattr(event.message, "usage", None)
                         if usage:
-                            prompt_tokens = getattr(usage, "input_tokens", 0)
+                            base_input = getattr(usage, "input_tokens", 0)
+                            cache_read, cache_create = _extract_anthropic_cache_tokens(usage)
 
-        total_cost = self._calculate_cost("claude", model, prompt_tokens, completion_tokens)
+        prompt_tokens = base_input + cache_read + cache_create
+        total_cost = self._calculate_cost(
+            "claude",
+            model,
+            prompt_tokens,
+            completion_tokens,
+            cached_tokens=cache_read,
+            cache_creation_tokens=cache_create,
+        )
         yield _build_anthropic_stream_done(
-            model, full_text, full_reasoning, prompt_tokens, completion_tokens, total_cost
+            model, full_text, full_reasoning, prompt_tokens, completion_tokens, total_cost,
+            cached_prompt_tokens=cache_read, cache_creation_tokens=cache_create,
         )

{prompture-1.0.53.dev2 → prompture-1.0.54.dev1}/prompture/drivers/async_google_driver.py

@@ -64,11 +64,18 @@ class AsyncGoogleDriver(CostMixin, AsyncDriver):
     def _extract_usage_metadata(self, response: Any, messages: list[dict[str, Any]]) -> dict[str, Any]:
         """Extract token counts from response, falling back to character estimation."""
         usage = getattr(response, "usage_metadata", None)
+        cached_prompt_tokens = 0
         if usage:
             prompt_tokens = getattr(usage, "prompt_token_count", 0) or 0
             completion_tokens = getattr(usage, "candidates_token_count", 0) or 0
             total_tokens = getattr(usage, "total_token_count", 0) or (prompt_tokens + completion_tokens)
-            cost = self._calculate_cost("google", self.model, prompt_tokens, completion_tokens)
+            raw_cached = getattr(usage, "cached_content_token_count", 0) or 0
+            if isinstance(raw_cached, (int, float)):
+                cached_prompt_tokens = int(raw_cached)
+            cost = self._calculate_cost(
+                "google", self.model, prompt_tokens, completion_tokens,
+                cached_tokens=cached_prompt_tokens,
+            )
         else:
             # Fallback: estimate from character counts
             total_prompt_chars = 0
@@ -92,6 +99,7 @@ class AsyncGoogleDriver(CostMixin, AsyncDriver):
             "prompt_tokens": prompt_tokens,
             "completion_tokens": completion_tokens,
             "total_tokens": total_tokens,
+            "cached_prompt_tokens": cached_prompt_tokens,
             "cost": round(cost, 6),
         }
 

{prompture-1.0.53.dev2 → prompture-1.0.54.dev1}/prompture/drivers/async_grok_driver.py

@@ -81,17 +81,24 @@ class AsyncGrokDriver(CostMixin, AsyncDriver):
             except Exception as e:
                 raise RuntimeError(f"Grok API request failed: {e!s}") from e
 
+        from .moonshot_driver import _extract_moonshot_cached_tokens
+
         usage = resp.get("usage", {})
         prompt_tokens = usage.get("prompt_tokens", 0)
         completion_tokens = usage.get("completion_tokens", 0)
         total_tokens = usage.get("total_tokens", 0)
+        cached_prompt_tokens = _extract_moonshot_cached_tokens(usage)
 
-        total_cost = self._calculate_cost("grok", model, prompt_tokens, completion_tokens)
+        total_cost = self._calculate_cost(
+            "grok", model, prompt_tokens, completion_tokens,
+            cached_tokens=cached_prompt_tokens,
+        )
 
         meta = {
             "prompt_tokens": prompt_tokens,
             "completion_tokens": completion_tokens,
             "total_tokens": total_tokens,
+            "cached_prompt_tokens": cached_prompt_tokens,
             "cost": round(total_cost, 6),
             "raw_response": resp,
             "model_name": model,
@@ -159,16 +166,23 @@ class AsyncGrokDriver(CostMixin, AsyncDriver):
             except Exception as e:
                 raise RuntimeError(f"Grok API request failed: {e!s}") from e
 
+        from .moonshot_driver import _extract_moonshot_cached_tokens
+
         usage = resp.get("usage", {})
         prompt_tokens = usage.get("prompt_tokens", 0)
         completion_tokens = usage.get("completion_tokens", 0)
         total_tokens = usage.get("total_tokens", 0)
-        total_cost = self._calculate_cost("grok", model, prompt_tokens, completion_tokens)
+        cached_prompt_tokens = _extract_moonshot_cached_tokens(usage)
+        total_cost = self._calculate_cost(
+            "grok", model, prompt_tokens, completion_tokens,
+            cached_tokens=cached_prompt_tokens,
+        )
 
         meta = {
             "prompt_tokens": prompt_tokens,
             "completion_tokens": completion_tokens,
             "total_tokens": total_tokens,
+            "cached_prompt_tokens": cached_prompt_tokens,
             "cost": round(total_cost, 6),
             "raw_response": resp,
             "model_name": model,

{prompture-1.0.53.dev2 → prompture-1.0.54.dev1}/prompture/drivers/async_groq_driver.py

@@ -75,17 +75,24 @@ class AsyncGroqDriver(CostMixin, AsyncDriver):
 
         resp = await self.client.chat.completions.create(**kwargs)
 
+        from .openai_driver import _extract_openai_cached_tokens
+
         usage = getattr(resp, "usage", None)
         prompt_tokens = getattr(usage, "prompt_tokens", 0)
         completion_tokens = getattr(usage, "completion_tokens", 0)
         total_tokens = getattr(usage, "total_tokens", 0)
+        cached_prompt_tokens = _extract_openai_cached_tokens(usage)
 
-        total_cost = self._calculate_cost("groq", model, prompt_tokens, completion_tokens)
+        total_cost = self._calculate_cost(
+            "groq", model, prompt_tokens, completion_tokens,
+            cached_tokens=cached_prompt_tokens,
+        )
 
         meta = {
             "prompt_tokens": prompt_tokens,
             "completion_tokens": completion_tokens,
             "total_tokens": total_tokens,
+            "cached_prompt_tokens": cached_prompt_tokens,
             "cost": round(total_cost, 6),
             "raw_response": resp.model_dump(),
             "model_name": model,
@@ -137,16 +144,23 @@ class AsyncGroqDriver(CostMixin, AsyncDriver):
 
         resp = await self.client.chat.completions.create(**kwargs)
 
+        from .openai_driver import _extract_openai_cached_tokens
+
         usage = getattr(resp, "usage", None)
         prompt_tokens = getattr(usage, "prompt_tokens", 0)
         completion_tokens = getattr(usage, "completion_tokens", 0)
         total_tokens = getattr(usage, "total_tokens", 0)
-        total_cost = self._calculate_cost("groq", model, prompt_tokens, completion_tokens)
+        cached_prompt_tokens = _extract_openai_cached_tokens(usage)
+        total_cost = self._calculate_cost(
+            "groq", model, prompt_tokens, completion_tokens,
+            cached_tokens=cached_prompt_tokens,
+        )
 
         meta = {
             "prompt_tokens": prompt_tokens,
             "completion_tokens": completion_tokens,
             "total_tokens": total_tokens,
+            "cached_prompt_tokens": cached_prompt_tokens,
             "cost": round(total_cost, 6),
             "raw_response": resp.model_dump(),
             "model_name": model,

{prompture-1.0.53.dev2 → prompture-1.0.54.dev1}/prompture/drivers/async_moonshot_driver.py

@@ -20,7 +20,7 @@ import httpx
 from ..infra.cost_mixin import CostMixin, prepare_strict_schema
 from .async_base import AsyncDriver
 from .base import _parse_tool_arguments
-from .moonshot_driver import MoonshotDriver
+from .moonshot_driver import MoonshotDriver, _extract_moonshot_cached_tokens
 
 logger = logging.getLogger("prompture.drivers.moonshot")
 
@@ -142,6 +142,7 @@ class AsyncMoonshotDriver(CostMixin, AsyncDriver):
         prompt_tokens = usage.get("prompt_tokens", 0)
         completion_tokens = usage.get("completion_tokens", 0)
         total_tokens = usage.get("total_tokens", 0)
+        cached_prompt_tokens = _extract_moonshot_cached_tokens(usage)
 
         message = resp["choices"][0]["message"]
         text = message.get("content") or ""
@@ -189,6 +190,7 @@ class AsyncMoonshotDriver(CostMixin, AsyncDriver):
                     prompt_tokens += fb_usage.get("prompt_tokens", 0)
                     completion_tokens = fb_usage.get("completion_tokens", 0)
                     total_tokens = prompt_tokens + completion_tokens
+                    cached_prompt_tokens += _extract_moonshot_cached_tokens(fb_usage)
                     resp = fb_resp
                     fb_message = fb_resp["choices"][0]["message"]
                     text = fb_message.get("content") or ""
@@ -196,12 +198,16 @@ class AsyncMoonshotDriver(CostMixin, AsyncDriver):
                     if not text and reasoning_content:
                         text = reasoning_content
 
-        total_cost = self._calculate_cost("moonshot", model, prompt_tokens, completion_tokens)
+        total_cost = self._calculate_cost(
+            "moonshot", model, prompt_tokens, completion_tokens,
+            cached_tokens=cached_prompt_tokens,
+        )
 
         meta = {
             "prompt_tokens": prompt_tokens,
             "completion_tokens": completion_tokens,
             "total_tokens": total_tokens,
+            "cached_prompt_tokens": cached_prompt_tokens,
             "cost": round(total_cost, 6),
             "raw_response": resp,
             "model_name": model,
@@ -280,12 +286,17 @@ class AsyncMoonshotDriver(CostMixin, AsyncDriver):
         prompt_tokens = usage.get("prompt_tokens", 0)
         completion_tokens = usage.get("completion_tokens", 0)
         total_tokens = usage.get("total_tokens", 0)
-        total_cost = self._calculate_cost("moonshot", model, prompt_tokens, completion_tokens)
+        cached_prompt_tokens = _extract_moonshot_cached_tokens(usage)
+        total_cost = self._calculate_cost(
+            "moonshot", model, prompt_tokens, completion_tokens,
+            cached_tokens=cached_prompt_tokens,
+        )
 
         meta = {
             "prompt_tokens": prompt_tokens,
             "completion_tokens": completion_tokens,
             "total_tokens": total_tokens,
+            "cached_prompt_tokens": cached_prompt_tokens,
             "cost": round(total_cost, 6),
             "raw_response": resp,
             "model_name": model,
@@ -365,6 +376,7 @@ class AsyncMoonshotDriver(CostMixin, AsyncDriver):
         full_reasoning = ""
         prompt_tokens = 0
         completion_tokens = 0
+        cached_prompt_tokens = 0
 
         async with (
             httpx.AsyncClient() as client,
@@ -392,6 +404,7 @@ class AsyncMoonshotDriver(CostMixin, AsyncDriver):
                 if usage:
                     prompt_tokens = usage.get("prompt_tokens", 0)
                     completion_tokens = usage.get("completion_tokens", 0)
+                    cached_prompt_tokens = _extract_moonshot_cached_tokens(usage)
 
                 choices = chunk.get("choices", [])
                 if choices:
@@ -407,7 +420,10 @@ class AsyncMoonshotDriver(CostMixin, AsyncDriver):
                         yield {"type": "delta", "text": content}
 
         total_tokens = prompt_tokens + completion_tokens
-        total_cost = self._calculate_cost("moonshot", model, prompt_tokens, completion_tokens)
+        total_cost = self._calculate_cost(
+            "moonshot", model, prompt_tokens, completion_tokens,
+            cached_tokens=cached_prompt_tokens,
+        )
 
         done_chunk: dict[str, Any] = {
             "type": "done",
@@ -416,6 +432,7 @@ class AsyncMoonshotDriver(CostMixin, AsyncDriver):
                 "prompt_tokens": prompt_tokens,
                 "completion_tokens": completion_tokens,
                 "total_tokens": total_tokens,
+                "cached_prompt_tokens": cached_prompt_tokens,
                 "cost": round(total_cost, 6),
                 "raw_response": {},
                 "model_name": model,

{prompture-1.0.53.dev2 → prompture-1.0.54.dev1}/prompture/drivers/async_openai_driver.py

@@ -19,6 +19,7 @@ from .openai_driver import (
     _build_openai_base_kwargs,
     _build_openai_json_mode_response_format,
     _build_openai_stream_done,
+    _extract_openai_cached_tokens,
     _extract_openai_meta,
     _extract_openai_tool_calls,
 )
@@ -90,9 +91,13 @@ class AsyncOpenAIDriver(CostMixin, AsyncDriver):
 
         resp = await self.client.chat.completions.create(**kwargs)
 
-        prompt_tokens = getattr(getattr(resp, "usage", None), "prompt_tokens", 0)
-        completion_tokens = getattr(getattr(resp, "usage", None), "completion_tokens", 0)
-        total_cost = self._calculate_cost("openai", model, prompt_tokens, completion_tokens)
+        usage = getattr(resp, "usage", None)
+        prompt_tokens = getattr(usage, "prompt_tokens", 0)
+        completion_tokens = getattr(usage, "completion_tokens", 0)
+        cached_prompt_tokens = _extract_openai_cached_tokens(usage)
+        total_cost = self._calculate_cost(
+            "openai", model, prompt_tokens, completion_tokens, cached_tokens=cached_prompt_tokens
+        )
         meta = _extract_openai_meta(resp, model, total_cost)
 
         text = resp.choices[0].message.content
@@ -126,9 +131,13 @@ class AsyncOpenAIDriver(CostMixin, AsyncDriver):
 
         resp = await self.client.chat.completions.create(**kwargs)
 
-        prompt_tokens = getattr(getattr(resp, "usage", None), "prompt_tokens", 0)
-        completion_tokens = getattr(getattr(resp, "usage", None), "completion_tokens", 0)
-        total_cost = self._calculate_cost("openai", model, prompt_tokens, completion_tokens)
+        usage = getattr(resp, "usage", None)
+        prompt_tokens = getattr(usage, "prompt_tokens", 0)
+        completion_tokens = getattr(usage, "completion_tokens", 0)
+        cached_prompt_tokens = _extract_openai_cached_tokens(usage)
+        total_cost = self._calculate_cost(
+            "openai", model, prompt_tokens, completion_tokens, cached_tokens=cached_prompt_tokens
+        )
         meta = _extract_openai_meta(resp, model, total_cost)
 
         choice = resp.choices[0]
@@ -177,12 +186,14 @@ class AsyncOpenAIDriver(CostMixin, AsyncDriver):
         full_text = ""
         prompt_tokens = 0
         completion_tokens = 0
+        cached_prompt_tokens = 0
 
         async for chunk in stream:
             # Usage comes in the final chunk
             if getattr(chunk, "usage", None):
                 prompt_tokens = chunk.usage.prompt_tokens or 0
                 completion_tokens = chunk.usage.completion_tokens or 0
+                cached_prompt_tokens = _extract_openai_cached_tokens(chunk.usage)
 
             if chunk.choices:
                 delta = chunk.choices[0].delta
@@ -191,5 +202,9 @@ class AsyncOpenAIDriver(CostMixin, AsyncDriver):
                     full_text += content
                     yield {"type": "delta", "text": content}
 
-        total_cost = self._calculate_cost("openai", model, prompt_tokens, completion_tokens)
-        yield _build_openai_stream_done(model, full_text, prompt_tokens, completion_tokens, total_cost)
+        total_cost = self._calculate_cost(
+            "openai", model, prompt_tokens, completion_tokens, cached_tokens=cached_prompt_tokens
+        )
+        yield _build_openai_stream_done(
+            model, full_text, prompt_tokens, completion_tokens, total_cost, cached_prompt_tokens
+        )

{prompture-1.0.53.dev2 → prompture-1.0.54.dev1}/prompture/drivers/base.py

@@ -402,6 +402,8 @@ class Driver(ABC):
                 prompt_tokens=meta.get("prompt_tokens", 0),
                 completion_tokens=meta.get("completion_tokens", 0),
                 total_tokens=meta.get("total_tokens", 0),
+                cached_prompt_tokens=meta.get("cached_prompt_tokens", 0),
+                cache_creation_tokens=meta.get("cache_creation_tokens", 0),
                 cost=meta.get("cost", 0.0),
                 elapsed_ms=elapsed_ms,
                 status=status,