Add optional Gemini context caching and cache usage accounting

src/viano_bug_analysis/llm/gateway.py

@@ -2,7 +2,6 @@ import asyncio
 import copy
 import logging
 import time
-import math
 import re
 from dataclasses import replace
 
@@ -16,6 +15,8 @@ from .errors import GatewayUnavailable, LLMError, ProviderUnavailable
 from .openai_compatible import OpenAICompatibleProvider
 from .provider_base import ProviderResponse
 from .diagnostics import sanitize, count
+from .gemini import GeminiProvider
+from .usage import usage_metadata
 
 LOGGER = logging.getLogger("viano_bug_analysis.llm")
 PRIVACY = {"local_only": {"local"}, "self_hosted": {"local", "self_hosted"}, "cloud_allowed": {"local", "self_hosted", "cloud"}}
@@ -36,7 +37,7 @@ def _schema_check(schema):
 class LLMGateway:
     def __init__(self, config, adapters=None):
         self.config = config
-        self.adapters = {"openai_compatible": OpenAICompatibleProvider()} | (adapters or {})
+        self.adapters = {"openai_compatible": OpenAICompatibleProvider(), "gemini": GeminiProvider()} | (adapters or {})
         if any(p.enabled and p.adapter not in self.adapters for p in config.providers):
             raise LLMError("LLM_ADAPTER")
 
@@ -84,8 +85,8 @@ class LLMGateway:
                 response = await asyncio.wait_for(self.adapters[provider.adapter].generate(provider, copy.deepcopy(request), max_response_bytes=self.config.max_response_bytes), timeout=provider.timeout_seconds)
                 if not isinstance(response.model, str) or not response.model.strip() or len(response.model) > 256 or any(ord(c) < 32 for c in response.model):
                     raise LLMError("LLM_RESPONSE")
-                for count in (response.input_tokens, response.output_tokens):
-                    if count is not None and (isinstance(count, bool) or not isinstance(count, int) or not 0 <= count <= 10**12):
+                for tokens in (response.input_tokens, response.output_tokens, response.cached_input_tokens, response.reasoning_tokens):
+                    if tokens is not None and (isinstance(tokens, bool) or not isinstance(tokens, int) or not 0 <= tokens <= 10**12):
                         raise LLMError("LLM_RESPONSE")
                 encoded = json_bytes(response.document)
                 if len(encoded) > self.config.max_response_bytes:
@@ -93,13 +94,16 @@ class LLMGateway:
                 if not isinstance(response.document, dict):
                     raise LLMError("LLM_OUTPUT_SCHEMA")
                 Draft202012Validator(request["output_schema"], format_checker=FORMAT_CHECKER).validate(response.document)
-                cost = None
-                if all(value is not None for value in (provider.input_cost_per_million, provider.output_cost_per_million, response.input_tokens, response.output_tokens)):
-                    cost = (response.input_tokens * provider.input_cost_per_million + response.output_tokens * provider.output_cost_per_million) / 1_000_000
+                usage = usage_metadata(provider, response)
+                cost = usage['estimated_cost']
                 metadata = {"provider": provider.id, "model": response.model, "prompt_version": request["prompt_version"],
                             "input_tokens": response.input_tokens, "output_tokens": response.output_tokens,
                             "estimated_cost": cost, "elapsed_seconds": round(time.monotonic() - started, 6), "attempts": attempts}
                 metadata["cost_currency"] = provider.cost_currency if cost is not None else None
+                metadata.update(usage)
+                if any(usage[k] is not None for k in ('estimated_input_cost', 'estimated_cached_input_cost',
+                                                      'estimated_output_cost', 'estimated_cache_storage_cost')):
+                    metadata['cost_currency'] = provider.cost_currency
                 LOGGER.info("LLM request succeeded provider=%s", provider.id)
                 return {"document": parse_manifest(encoded), "metadata": metadata}
             except (ProviderUnavailable, TimeoutError) as error:
@@ -108,6 +112,11 @@ class LLMGateway:
                 http_status = getattr(error, "http_status", None)
                 if type(http_status) is int and 100 <= http_status <= 599:
                     attempts[-1]["http_status"] = http_status
+                cache_details = getattr(error, 'diagnostics', {})
+                if 'cache_created_tokens' in cache_details and cache_details['cache_created_tokens'] != 0:
+                    # Allocation can have succeeded even if generation was unavailable.
+                    attempts[-1]['cache_usage'] = usage_metadata(provider, diagnostics=error.diagnostics)
+                    attempts[-1]['cache_usage']['cost_currency'] = provider.cost_currency
                 LOGGER.warning("LLM endpoint unavailable provider=%s code=%s", provider.id, code)
             except ValidationError as invalid:
                 error = LLMError("LLM_OUTPUT_SCHEMA",diagnostics={'validation_keyword':invalid.validator})
@@ -131,12 +140,10 @@ class LLMGateway:
         details['max_output_tokens'] = count(provider.max_output_tokens)
         error.diagnostics = details
         input_tokens, output_tokens = details.get('input_tokens'),details.get('output_tokens')
-        cost = None
-        if all(v is not None for v in (input_tokens,output_tokens,provider.input_cost_per_million,provider.output_cost_per_million)):
-            value = (input_tokens*provider.input_cost_per_million+output_tokens*provider.output_cost_per_million)/1_000_000
-            if math.isfinite(value):
-                cost = value
         error.metadata = {'provider':provider.id,'model':provider.model,'prompt_version':request['prompt_version'],
-                          'input_tokens':input_tokens,'output_tokens':output_tokens,'estimated_cost':cost,
-                          'cost_currency':provider.cost_currency if cost is not None else None,
+                          'input_tokens':input_tokens,'output_tokens':output_tokens,'cost_currency':None,
                           'elapsed_seconds':round(time.monotonic()-started,6),'attempts':copy.deepcopy(attempts)}
+        error.metadata.update(usage_metadata(provider, response, details))
+        if any(error.metadata[k] is not None for k in ('estimated_cost','estimated_input_cost',
+                'estimated_cached_input_cost','estimated_output_cost','estimated_cache_storage_cost')):
+            error.metadata['cost_currency'] = provider.cost_currency