260 lines
17 KiB
Python
260 lines
17 KiB
Python
import asyncio
|
|
import copy
|
|
import logging
|
|
import time
|
|
import re
|
|
from dataclasses import replace
|
|
|
|
from jsonschema import Draft202012Validator
|
|
from jsonschema.exceptions import SchemaError, ValidationError
|
|
|
|
from viano_bug_analysis.lifecycle.storage import json_bytes
|
|
from viano_bug_analysis.package.manifest import FORMAT_CHECKER, parse_manifest
|
|
|
|
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
|
|
from .response_recovery import REPAIR_PROMPT
|
|
|
|
LOGGER = logging.getLogger("viano_bug_analysis.llm")
|
|
PRIVACY = {"local_only": {"local"}, "self_hosted": {"local", "self_hosted"}, "cloud_allowed": {"local", "self_hosted", "cloud"}}
|
|
|
|
|
|
def _schema_check(schema):
|
|
# No external resolution/network access during validation, including nested refs.
|
|
if isinstance(schema, dict):
|
|
for key, value in schema.items():
|
|
if key in {"$ref", "$dynamicRef"} and (not isinstance(value, str) or not value.startswith("#")):
|
|
raise LLMError("LLM_SCHEMA_REFERENCE")
|
|
_schema_check(value)
|
|
elif isinstance(schema, list):
|
|
for value in schema:
|
|
_schema_check(value)
|
|
|
|
|
|
def is_investigation_review(request):
|
|
"""Route the controller's versioned narrow review, never arbitrary research text."""
|
|
if request['prompt_version'] not in {'investigation_project_v16','investigation_project_v17','investigation_project_v18','investigation_project_v19','investigation_project_v20'}:
|
|
return False
|
|
try:
|
|
message = next(m for m in reversed(request['messages']) if m['role'] == 'user')
|
|
document = parse_manifest(message['content'].encode('utf-8'))
|
|
action_schema = request['output_schema']['properties']['next']
|
|
if request['output_schema'].get('x-viano-investigation-window') == 1 and 'anyOf' in action_schema:
|
|
action_schema = action_schema['anyOf'][0]
|
|
if request['output_schema'].get('x-viano-investigation-window') == 1 and 'oneOf' in action_schema:
|
|
action_schema = action_schema['oneOf'][0]
|
|
return (document.get('task') == 'assess_pending_result_for_early_review'
|
|
and document['input']['early_review']['request']['status'] == 'preparing'
|
|
and set(action_schema['properties']['kind']['enum'])
|
|
== {'defer','reject'})
|
|
except (ValueError, TypeError, KeyError, AttributeError, StopIteration, RecursionError):
|
|
return False
|
|
|
|
|
|
def review_metadata(provider, request):
|
|
if provider.investigation_review_reasoning_effort is not None and is_investigation_review(request):
|
|
return {'reasoning_profile':'investigation_review', 'reasoning_effort':provider.reasoning_effort}
|
|
return {}
|
|
|
|
|
|
class LLMGateway:
|
|
def __init__(self, config, adapters=None):
|
|
self.config = config
|
|
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")
|
|
|
|
def validate_request(self, request):
|
|
try:
|
|
encoded = json_bytes(request)
|
|
if len(encoded) > self.config.max_request_bytes:
|
|
raise LLMError("LLM_REQUEST_LIMIT")
|
|
request = parse_manifest(encoded)
|
|
if not isinstance(request, dict) or set(request) != {"prompt_version", "privacy", "messages", "output_schema"}:
|
|
raise LLMError("LLM_REQUEST")
|
|
if request["privacy"] == "configured":
|
|
request["privacy"] = self.config.privacy_profile
|
|
if request["privacy"] not in PRIVACY or not isinstance(request["prompt_version"], str) or not request["prompt_version"].strip():
|
|
raise LLMError("LLM_REQUEST")
|
|
if not isinstance(request["messages"], list) or not request["messages"]:
|
|
raise LLMError("LLM_REQUEST")
|
|
for message in request["messages"]:
|
|
if not isinstance(message, dict) or set(message) != {"role", "content"} or message["role"] not in {"system", "user", "assistant"} or not isinstance(message["content"], str):
|
|
raise LLMError("LLM_REQUEST")
|
|
if not isinstance(request["output_schema"], dict) or request["output_schema"].get("type") != "object":
|
|
raise LLMError("LLM_REQUEST")
|
|
_schema_check(request["output_schema"])
|
|
Draft202012Validator.check_schema(request["output_schema"])
|
|
from .investigation_size import is_investigation, check, TRANSPORT_RESERVE
|
|
if is_investigation(request):
|
|
check(request, reserve=TRANSPORT_RESERVE)
|
|
eligible = PRIVACY[request["privacy"]] & PRIVACY[self.config.privacy_profile]
|
|
if not any(p.enabled and p.privacy in eligible for p in self.config.providers):
|
|
raise LLMError("LLM_PRIVACY_NO_PROVIDER")
|
|
return request
|
|
except (ValueError, TypeError, RecursionError, SchemaError):
|
|
raise LLMError("LLM_REQUEST") from None
|
|
|
|
async def complete_round(self, request: dict) -> dict:
|
|
"""Try eligible providers in order exactly once; never sleep or retry internally."""
|
|
request = self.validate_request(request)
|
|
attempts = []
|
|
for provider in self.config.providers:
|
|
if not provider.enabled or provider.privacy not in (PRIVACY[request["privacy"]] & PRIVACY[self.config.privacy_profile]):
|
|
continue
|
|
if (provider.entity_resolution_reasoning_effort is not None
|
|
and re.fullmatch(r"entity_resolution_(plan|assessment)_v[1-9][0-9]*", request["prompt_version"])):
|
|
provider = replace(provider, reasoning_effort=provider.entity_resolution_reasoning_effort)
|
|
if provider.investigation_review_reasoning_effort is not None and is_investigation_review(request):
|
|
provider = replace(provider, reasoning_effort=provider.investigation_review_reasoning_effort)
|
|
started = time.monotonic()
|
|
response = None
|
|
try:
|
|
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 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")
|
|
if provider.cache_required and (provider.cache_required_prompt_prefix is None
|
|
or request['prompt_version'].startswith(provider.cache_required_prompt_prefix)):
|
|
if (response.cached_input_tokens is None or response.cached_input_tokens <= 0
|
|
or response.input_tokens is None or response.cached_input_tokens > response.input_tokens):
|
|
raise LLMError('LLM_CACHE_USAGE_UNCONFIRMED')
|
|
from .response_recovery import enabled,normalize,apply_repair
|
|
if enabled(request):
|
|
document,audit=normalize(response.document)
|
|
if response.normalization:
|
|
original=copy.deepcopy(response.normalization)
|
|
if audit:original['rules'].extend(audit['rules'])
|
|
audit=original
|
|
if audit and len(audit['rules'])>16:raise LLMError('LLM_RESPONSE_LIMIT')
|
|
response=replace(response,document=document,normalization=audit)
|
|
from .answer_recovery import normalize_schema
|
|
recovered,rules=normalize_schema(response.document,request['output_schema'])
|
|
if rules and len((response.normalization or {}).get('rules',[]))+len(rules)<=16:
|
|
audit=copy.deepcopy(response.normalization) if response.normalization else {'version':1,'original_document':copy.deepcopy(response.document),'rules':[]}
|
|
audit['rules'].extend(rules)
|
|
response=replace(response,document=recovered,normalization=audit)
|
|
encoded = json_bytes(response.document)
|
|
if len(encoded) > self.config.max_response_bytes:
|
|
raise LLMError("LLM_RESPONSE_LIMIT")
|
|
if response.normalization and len(json_bytes({'document':response.document,
|
|
'normalization':response.normalization}))>self.config.max_response_bytes:
|
|
raise LLMError('LLM_RESPONSE_LIMIT')
|
|
if not isinstance(response.document, dict):
|
|
raise LLMError("LLM_OUTPUT_SCHEMA")
|
|
Draft202012Validator(request["output_schema"], format_checker=FORMAT_CHECKER).validate(response.document)
|
|
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)
|
|
metadata.update(review_metadata(provider, request))
|
|
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
|
|
result={"document": parse_manifest(encoded), "metadata": metadata}
|
|
if response.normalization:result['normalization']=self._audit_budget(response.normalization)
|
|
if response.response_capture:result['response_capture']=self._capture_budget(response.response_capture)
|
|
if request['prompt_version']==REPAIR_PROMPT:
|
|
result['document']=apply_repair(request,response.document)
|
|
if len(json_bytes(result['document']))>self.config.max_response_bytes:
|
|
raise LLMError('LLM_RESPONSE_LIMIT')
|
|
result['repair']={'version':1,'patch':copy.deepcopy(response.document)}
|
|
if len(json_bytes({'document':result['document'],'repair':result['repair']}))>self.config.max_response_bytes:
|
|
raise LLMError('LLM_RESPONSE_LIMIT')
|
|
LOGGER.info("LLM request succeeded provider=%s", provider.id)
|
|
return result
|
|
except (ProviderUnavailable, TimeoutError) as error:
|
|
if provider.cache_required and (provider.cache_required_prompt_prefix is None
|
|
or request['prompt_version'].startswith(provider.cache_required_prompt_prefix)):
|
|
required_error = LLMError('LLM_CACHE_REQUIRED',
|
|
http_status=getattr(error, 'http_status', None), diagnostics=getattr(error, 'diagnostics', None))
|
|
self._failure_metadata(required_error, provider, request, started, response, attempts)
|
|
raise required_error from None
|
|
code = error.code if isinstance(error, ProviderUnavailable) else "LLM_TIMEOUT"
|
|
attempts.append({"provider": provider.id, "code": code})
|
|
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:
|
|
info={'validation_keyword':invalid.validator}
|
|
from .diagnostics import schema_feedback
|
|
info.update(schema_feedback(invalid,include_keyword=True))
|
|
error = LLMError("LLM_OUTPUT_SCHEMA",diagnostics=info)
|
|
from .response_recovery import enabled,LIMIT
|
|
if isinstance(response.document,dict) and len(json_bytes(response.document))<=LIMIT:
|
|
error.rejected_document=copy.deepcopy(response.document)
|
|
if response.normalization:error.normalization=copy.deepcopy(response.normalization)
|
|
self._retain_rejected_repair(error,request,response)
|
|
self._failure_metadata(error,provider,request,started,response,attempts)
|
|
raise error from None
|
|
except LLMError as error:
|
|
self._retain_rejected_repair(error,request,response)
|
|
self._failure_metadata(error,provider,request,started,response,attempts)
|
|
raise
|
|
except Exception:
|
|
error = LLMError("LLM_PROVIDER_ERROR")
|
|
self._failure_metadata(error,provider,request,started,response,attempts)
|
|
raise error from None
|
|
raise GatewayUnavailable(attempts)
|
|
|
|
def _retain_rejected_repair(self,error,request,response):
|
|
from .response_recovery import REPAIR_PROMPT,LIMIT,sha
|
|
if request['prompt_version']!=REPAIR_PROMPT:return
|
|
if (request['prompt_version']==REPAIR_PROMPT and isinstance(response,ProviderResponse)
|
|
and isinstance(response.document,dict)
|
|
and len(json_bytes(response.document))<=min(LIMIT,self.config.max_response_bytes)):
|
|
error.rejected_repair={'version':1,'reason':error.code,
|
|
'patch':copy.deepcopy(response.document),'sha256':sha(response.document)}
|
|
elif isinstance(getattr(error,'repair_response_text',None),str):
|
|
text=error.repair_response_text
|
|
if len(text.encode('utf-8'))<=min(LIMIT,self.config.max_response_bytes):
|
|
error.rejected_repair={'version':1,'reason':error.code,'original_text':text,'sha256':sha(text)}
|
|
|
|
def _audit_budget(self,value):
|
|
value=copy.deepcopy(value)
|
|
raw=value.get('original_text')
|
|
if isinstance(raw,str) and len(raw.encode('utf-8'))>min(self.config.response_capture_bytes,self.config.max_response_bytes):value.pop('original_text')
|
|
return value
|
|
|
|
def _capture_budget(self,value):
|
|
value=copy.deepcopy(value)
|
|
if value['bytes']>min(self.config.response_capture_bytes,self.config.max_response_bytes):
|
|
value.pop('text',None);value['retained']=False
|
|
return value
|
|
|
|
def _failure_metadata(self,error,provider,request,started,response,attempts):
|
|
if isinstance(response,ProviderResponse) and response.response_capture:
|
|
error.response_capture=response.response_capture
|
|
if hasattr(error,'response_capture'):error.response_capture=self._capture_budget(error.response_capture)
|
|
if hasattr(error,'normalization'):error.normalization=self._audit_budget(error.normalization)
|
|
details = sanitize(error.diagnostics)
|
|
if isinstance(response,ProviderResponse):
|
|
details = sanitize(response.diagnostics) | details
|
|
details.update(input_tokens=count(response.input_tokens),output_tokens=count(response.output_tokens))
|
|
details['max_output_tokens'] = count(provider.max_output_tokens)
|
|
error.diagnostics = details
|
|
input_tokens, output_tokens = details.get('input_tokens'),details.get('output_tokens')
|
|
error.metadata = {'provider':provider.id,'model':provider.model,'prompt_version':request['prompt_version'],
|
|
'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))
|
|
error.metadata.update(review_metadata(provider, request))
|
|
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
|