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