191 lines
10 KiB
Python
191 lines
10 KiB
Python
"""Optional native Gemini transport with disposable instruction/schema caching.
|
|
|
|
No conversation state or research evidence is stored exclusively at the provider.
|
|
Cache handles and credential fingerprints never enter the portable case.
|
|
"""
|
|
import asyncio
|
|
import hashlib
|
|
import http.client
|
|
import os
|
|
import re
|
|
import threading
|
|
import time
|
|
import urllib.error
|
|
import urllib.request
|
|
|
|
from viano_bug_analysis.lifecycle.storage import json_bytes
|
|
from viano_bug_analysis.package.manifest import parse_manifest
|
|
from .errors import LLMError, ProviderUnavailable
|
|
from .openai_compatible import _NoRedirect
|
|
from .provider_base import ProviderResponse
|
|
from .diagnostics import count, sanitize
|
|
|
|
|
|
class GeminiProvider:
|
|
def __init__(self):
|
|
self._cache = {}
|
|
self._lock = threading.Lock()
|
|
|
|
async def generate(self, config, request, *, max_response_bytes):
|
|
return await asyncio.to_thread(self._generate, config, request, max_response_bytes)
|
|
|
|
@staticmethod
|
|
def _post(url, payload, headers, limit, deadline):
|
|
remaining = deadline - time.monotonic()
|
|
if remaining <= 0:
|
|
raise ProviderUnavailable('LLM_TIMEOUT')
|
|
opener = urllib.request.build_opener(urllib.request.ProxyHandler({}), _NoRedirect())
|
|
try:
|
|
req = urllib.request.Request(url, data=json_bytes(payload), headers=headers)
|
|
with opener.open(req, timeout=remaining) as response:
|
|
chunks, size = [], 0
|
|
while True:
|
|
if time.monotonic() >= deadline:
|
|
raise ProviderUnavailable('LLM_TIMEOUT')
|
|
chunk = response.read(min(65536, limit + 1 - size))
|
|
if not chunk:
|
|
break
|
|
size += len(chunk)
|
|
if size > limit:
|
|
raise LLMError('LLM_RESPONSE_LIMIT')
|
|
chunks.append(chunk)
|
|
return parse_manifest(b''.join(chunks))
|
|
except urllib.error.HTTPError as error:
|
|
error.close()
|
|
if error.code in {408, 429} or 500 <= error.code < 600:
|
|
raise ProviderUnavailable('LLM_HTTP_TEMPORARY', http_status=error.code) from None
|
|
raise LLMError('LLM_HTTP_CONFIGURATION', http_status=error.code) from None
|
|
except (urllib.error.URLError, TimeoutError, OSError, http.client.HTTPException):
|
|
raise ProviderUnavailable('LLM_CONNECTION') from None
|
|
except (ValueError, TypeError, RecursionError):
|
|
raise LLMError('LLM_RESPONSE') from None
|
|
|
|
@staticmethod
|
|
def _messages(request):
|
|
system, contents = [], []
|
|
for message in request['messages']:
|
|
if message['role'] == 'system':
|
|
system.append({'text': message['content']})
|
|
continue
|
|
text = message['content']
|
|
# Move only an exact duplicate of the authoritative schema. Everything
|
|
# else, including every research result, stays in the live input.
|
|
try:
|
|
value = parse_manifest(text.encode('utf-8'))
|
|
if isinstance(value, dict) and value.get('output_schema') == request['output_schema']:
|
|
value.pop('output_schema')
|
|
text = json_bytes(value).decode('utf-8')
|
|
except (ValueError, TypeError, RecursionError):
|
|
pass
|
|
contents.append({'role': 'model' if message['role'] == 'assistant' else 'user',
|
|
'parts': [{'text': text}]})
|
|
prefix = {'role': 'user', 'parts': [{'text': json_bytes({
|
|
'output_schema': request['output_schema']}).decode('utf-8')}]}
|
|
return {'parts': system}, [prefix], contents
|
|
|
|
def _generate(self, config, request, limit):
|
|
required = config.cache_required and (config.cache_required_prompt_prefix is None
|
|
or request['prompt_version'].startswith(config.cache_required_prompt_prefix))
|
|
secret = os.environ.get(config.api_key_env) if config.api_key_env else None
|
|
if config.api_key_env and (not secret or any(ord(c) < 32 or ord(c) == 127 for c in secret)):
|
|
raise LLMError('LLM_CREDENTIALS')
|
|
headers = {'Content-Type': 'application/json', 'Accept': 'application/json'}
|
|
if secret:
|
|
headers['x-goog-api-key'] = secret
|
|
base = config.base_url.rstrip('/')
|
|
deadline = time.monotonic() + config.timeout_seconds
|
|
system, prefix, contents = self._messages(request)
|
|
generation = {'maxOutputTokens': config.max_output_tokens}
|
|
if config.reasoning_effort is not None:
|
|
generation['thinkingConfig'] = {'thinkingLevel': config.reasoning_effort}
|
|
if config.response_format != 'none':
|
|
generation['responseMimeType'] = 'application/json'
|
|
if config.response_format == 'json_schema':
|
|
generation['responseJsonSchema'] = request['output_schema']
|
|
full = {'contents': prefix + contents, 'generationConfig': generation}
|
|
if system['parts']:
|
|
full['systemInstruction'] = system
|
|
status, created, name = 'disabled', 0, None
|
|
cache_http_status = None
|
|
fingerprint = hashlib.sha256(json_bytes([base, config.id, config.model, secret,
|
|
system, prefix, config.cache_ttl_seconds])).hexdigest()
|
|
if config.use_context_cache:
|
|
status = 'below_minimum'
|
|
if len(json_bytes([system, prefix]).decode('utf-8')) >= config.cache_min_prefix_chars:
|
|
with self._lock:
|
|
now = time.monotonic()
|
|
self._cache = {k: v for k, v in self._cache.items() if v[1] > now}
|
|
entry = self._cache.get(fingerprint)
|
|
if entry:
|
|
name, _ = entry
|
|
status = 'reused' if name else 'fallback'
|
|
elif len(self._cache) < 8:
|
|
status = 'fallback'
|
|
try:
|
|
cache_payload = {'model': 'models/' + config.model, 'contents': prefix,
|
|
'ttl': str(config.cache_ttl_seconds) + 's'}
|
|
if system['parts']:
|
|
cache_payload['systemInstruction'] = system
|
|
# A lost creation response may still have allocated
|
|
# remote storage; do not report a fabricated zero.
|
|
created = None
|
|
cached = self._post(base + '/cachedContents', cache_payload, headers, limit,
|
|
min(deadline, now + (config.timeout_seconds / 2 if required
|
|
else min(5, config.timeout_seconds / 4))))
|
|
if not isinstance(cached, dict) or not isinstance(cached.get('name'), str) or not re.fullmatch(
|
|
r'cachedContents/[A-Za-z0-9_-]{1,256}', cached['name']):
|
|
raise LLMError('LLM_RESPONSE')
|
|
name = cached['name']
|
|
created = count((cached.get('usageMetadata') or {}).get('totalTokenCount'))
|
|
status = 'created'
|
|
except (LLMError, TypeError, AttributeError) as error:
|
|
name = None
|
|
cache_http_status = getattr(error, 'http_status', None)
|
|
# Negative caching prevents repeated failed creation requests.
|
|
self._cache[fingerprint] = (name, now + (config.cache_ttl_seconds if name else 60))
|
|
else:
|
|
status = 'fallback'
|
|
payload = {'contents': contents, 'generationConfig': generation, 'cachedContent': name} if name else full
|
|
diagnostics = {'cache_status': status, 'cache_created_tokens': created,
|
|
'cache_ttl_seconds': config.cache_ttl_seconds, 'max_output_tokens': config.max_output_tokens}
|
|
try:
|
|
if required and not name:
|
|
raise LLMError('LLM_CACHE_REQUIRED', http_status=cache_http_status)
|
|
try:
|
|
envelope = self._post(base + '/models/' + config.model + ':generateContent', payload, headers, limit, deadline)
|
|
except LLMError as error:
|
|
if required:
|
|
raise LLMError('LLM_CACHE_REQUIRED', http_status=error.http_status) from None
|
|
if not name or error.http_status not in {400, 403, 404}:
|
|
raise
|
|
with self._lock:
|
|
self._cache[fingerprint] = (None, time.monotonic() + 60)
|
|
diagnostics['cache_status'] = 'fallback'
|
|
envelope = self._post(base + '/models/' + config.model + ':generateContent', full, headers, limit, deadline)
|
|
usage = envelope.get('usageMetadata') or {}
|
|
prompt = count(usage.get('promptTokenCount'))
|
|
cached = count(usage.get('cachedContentTokenCount'))
|
|
thoughts = count(usage.get('thoughtsTokenCount'))
|
|
candidates = count(usage.get('candidatesTokenCount'))
|
|
output = candidates + (thoughts or 0) if candidates is not None else None
|
|
diagnostics.update(input_tokens=prompt, cached_input_tokens=cached,
|
|
output_tokens=output, reasoning_tokens=thoughts)
|
|
if required and (cached is None or cached <= 0 or prompt is None or cached > prompt):
|
|
raise LLMError('LLM_CACHE_USAGE_UNCONFIRMED')
|
|
candidate = envelope['candidates'][0]
|
|
diagnostics['finish_reason'] = candidate.get('finishReason')
|
|
if candidate.get('finishReason') != 'STOP':
|
|
raise LLMError('LLM_INCOMPLETE_OUTPUT')
|
|
parts = candidate['content']['parts']
|
|
if any('functionCall' in part for part in parts):
|
|
raise LLMError('LLM_INCOMPLETE_OUTPUT')
|
|
text = ''.join(part['text'] for part in parts if 'text' in part and not part.get('thought'))
|
|
document = parse_manifest(text.encode('utf-8'))
|
|
return ProviderResponse(document, envelope.get('modelVersion', config.model), prompt, output,
|
|
sanitize(diagnostics), cached, thoughts)
|
|
except LLMError as error:
|
|
error.diagnostics = sanitize(diagnostics) | sanitize(error.diagnostics)
|
|
raise
|
|
except (ValueError, TypeError, KeyError, IndexError, AttributeError, RecursionError):
|
|
raise LLMError('LLM_RESPONSE', diagnostics=sanitize(diagnostics)) from None
|