155 lines
8.8 KiB
Python
155 lines
8.8 KiB
Python
"""Measure one prepared Gemini request without executing research tools."""
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import re
|
|
from pathlib import Path
|
|
from jsonschema import Draft202012Validator
|
|
from .llm import GatewayConfig, LLMError
|
|
from .llm.gateway import PRIVACY
|
|
from .llm.gemini import GeminiProvider
|
|
from .minimal_prompt_test import measure
|
|
from .telemetry.safety import safe_code
|
|
|
|
DEFAULT = Path('service-data/benchmarks/compact-investigation-20261017/request.json')
|
|
|
|
|
|
def gemini_schema(schema):
|
|
"""Provider projection only; full original schema remains the local validator.
|
|
|
|
Gemini's responseJsonSchema subset omits const and several validation
|
|
keywords. In particular, a const-only kind has no usable type after const
|
|
is ignored. Translate it rather than handing the internal contract through.
|
|
"""
|
|
allowed = {'type','format','title','description','enum','items','prefixItems',
|
|
'minItems','maxItems','minimum','maximum','anyOf','oneOf',
|
|
'properties','additionalProperties','required'}
|
|
if not isinstance(schema,dict):
|
|
raise LLMError('PROMPT_PROBE_SCHEMA_UNSUPPORTED')
|
|
if any(k in schema for k in ('$ref','$defs','allOf','not','if','then','else')):
|
|
raise LLMError('PROMPT_PROBE_SCHEMA_UNSUPPORTED')
|
|
result = {}
|
|
for key,value in schema.items():
|
|
if key not in allowed: continue
|
|
if key == 'properties': result[key] = {k:gemini_schema(v) for k,v in value.items()}
|
|
elif key in {'items','additionalProperties'} and isinstance(value,dict): result[key] = gemini_schema(value)
|
|
elif key in {'anyOf','oneOf','prefixItems'}: result['anyOf' if key=='oneOf' else key] = [gemini_schema(v) for v in value]
|
|
else: result[key] = value
|
|
if 'const' in schema:
|
|
value=schema['const']
|
|
if not isinstance(value,str): raise LLMError('PROMPT_PROBE_SCHEMA_UNSUPPORTED')
|
|
result['type']='string';result['enum']=[value]
|
|
if isinstance(result.get('type'),list):
|
|
types=result.pop('type')
|
|
result['anyOf']=[{'type':value} for value in types]
|
|
return result
|
|
|
|
|
|
def prepare(provider,request, *, schema_mode="native"):
|
|
if (set(request)!={'prompt_version','messages','output_schema'} or
|
|
not isinstance(request['prompt_version'],str) or not request['prompt_version'].strip()):
|
|
raise LLMError('PROMPT_PROBE_REQUEST')
|
|
messages=request['messages']
|
|
if (not isinstance(messages,list) or [m.get('role') for m in messages if isinstance(m,dict)]!=['system','user'] or
|
|
any(set(m)!={'role','content'} or not isinstance(m['content'],str) or not m['content'].strip() for m in messages)):
|
|
raise LLMError('PROMPT_PROBE_REQUEST')
|
|
Draft202012Validator.check_schema(request['output_schema'])
|
|
if schema_mode not in {'native','prompt'}:raise LLMError('PROMPT_PROBE_SCHEMA_MODE')
|
|
generation={'maxOutputTokens':provider.max_output_tokens,'responseMimeType':'application/json'}
|
|
system_text=messages[0]['content']
|
|
if schema_mode=='native':
|
|
generation['responseJsonSchema']=gemini_schema(request['output_schema'])
|
|
else:
|
|
system_text += '\nOutput JSON contract (validated locally):\n'+json.dumps(request['output_schema'],ensure_ascii=False,separators=(',',':'))
|
|
if provider.reasoning_effort is not None:generation['thinkingConfig']={'thinkingLevel':provider.reasoning_effort}
|
|
payload={'systemInstruction':{'parts':[{'text':system_text}]},
|
|
'contents':[{'role':'user','parts':[{'text':messages[1]['content']}]}],'generationConfig':generation}
|
|
return {'id':request['prompt_version'],'prompt_characters':len(system_text)+len(messages[1]['content']),
|
|
'payload':payload,'validation_schema':request['output_schema']}
|
|
|
|
|
|
def provider_diagnostic(body, provider):
|
|
"""Bounded local probe feedback, excluding raw response bodies and secrets."""
|
|
error=body.get('error',{}) if isinstance(body,dict) else {}
|
|
if not isinstance(error,dict): return {}
|
|
def clean(text):
|
|
if not isinstance(text,str): return None
|
|
secret=os.environ.get(provider.api_key_env) if provider.api_key_env else None
|
|
if secret: text=text.replace(secret,'[redacted]')
|
|
text=re.sub(r'https?://[^\s]+','[url]',text)
|
|
# Echoed literal values can contain prompt data; field paths remain useful.
|
|
text=re.sub(r'"[^"\n]*"|\'[^\'\n]*\'', '[quoted value]', text)
|
|
return ''.join(c if c.isprintable() else ' ' for c in text)[:1000]
|
|
result={}
|
|
if error.get('status') in {'INVALID_ARGUMENT','FAILED_PRECONDITION','RESOURCE_EXHAUSTED','PERMISSION_DENIED','UNAUTHENTICATED','NOT_FOUND','INTERNAL','UNAVAILABLE'}:
|
|
result['status']=error['status']
|
|
if isinstance(error.get('message'),str): result['message']=clean(error['message'])
|
|
violations=[]
|
|
details=error.get('details',[])
|
|
if isinstance(details,list):
|
|
for detail in details[:8]:
|
|
if not isinstance(detail,dict):continue
|
|
values=detail.get('fieldViolations',[])
|
|
if not isinstance(values,list):continue
|
|
for value in values[:8]:
|
|
if not isinstance(value,dict):continue
|
|
violations.append({k:clean(value[k]) for k in ('field','description') if isinstance(value.get(k),str)})
|
|
if violations:result['field_violations']=violations[:8]
|
|
return result
|
|
|
|
|
|
def measure_diagnostic(provider,case,limit):
|
|
diagnostic={}
|
|
def post(url,payload,headers,response_limit,deadline):
|
|
return GeminiProvider._post(url,payload,headers,response_limit,deadline,
|
|
error_capture=lambda body:diagnostic.update(provider_diagnostic(body,provider)))
|
|
result=measure(provider,case,limit,post=post)
|
|
if diagnostic:result['provider_error']=diagnostic
|
|
return result
|
|
|
|
|
|
def main(argv=None):
|
|
p=argparse.ArgumentParser(description=__doc__)
|
|
p.add_argument('--request',type=Path,default=DEFAULT)
|
|
p.add_argument('--llm-config',default='config/llm.working-brief.local.toml')
|
|
p.add_argument('--provider')
|
|
p.add_argument('--schema-mode',choices=['native','prompt'],default='native',help='Native schema or JSON mode with the full contract once in the prompt')
|
|
p.add_argument('--output',type=Path)
|
|
p.add_argument('--run',action='store_true',help='Make exactly one potentially billable API call')
|
|
args=p.parse_args(argv)
|
|
try:
|
|
config=GatewayConfig.load(args.llm_config)
|
|
providers=[v for v in config.providers if v.enabled and v.privacy in PRIVACY[config.privacy_profile]]
|
|
provider=next((v for v in providers if v.id==args.provider),None) if args.provider else next(iter(providers),None)
|
|
if provider is None or provider.adapter!='gemini':raise LLMError('PROMPT_PROBE_PROVIDER')
|
|
raw=args.request.read_bytes()
|
|
if len(raw)>config.max_request_bytes:raise LLMError('LLM_REQUEST_LIMIT')
|
|
request=json.loads(raw);case=prepare(provider,request,schema_mode=args.schema_mode)
|
|
if len(json.dumps(case['payload'],ensure_ascii=False).encode('utf-8'))>config.max_request_bytes:raise LLMError('LLM_REQUEST_LIMIT')
|
|
metadata={'request_sha256':hashlib.sha256(raw).hexdigest(),'provider':provider.id,'model':provider.model,
|
|
'reasoning_effort':provider.reasoning_effort,'max_output_tokens':provider.max_output_tokens,
|
|
'timeout_seconds':provider.timeout_seconds,'explicit_cache':False,'research_queries_executed':0,
|
|
'wire_schema_sha256':hashlib.sha256(json.dumps(case['payload']['generationConfig'].get('responseJsonSchema'),sort_keys=True,separators=(',',':')).encode()).hexdigest(),
|
|
'schema_transport':'gemini_subset_v1_local_full_validation' if args.schema_mode=='native' else 'json_mode_schema_once_in_prompt_local_full_validation'}
|
|
if not args.run:
|
|
print(json.dumps({'status':'PREPARED','network_used':False,**metadata,
|
|
'message_characters':case['prompt_characters']}));return 0
|
|
output=args.output or args.request.with_name('result.json')
|
|
# Exclusive reservation prevents accidental duplicate paid execution.
|
|
with output.open('x',encoding='utf-8') as f:json.dump({'status':'RUNNING',**metadata},f)
|
|
result={**measure_diagnostic(provider,case,config.max_response_bytes),**metadata}
|
|
output.write_text(json.dumps(result,ensure_ascii=False,indent=2)+'\n',encoding='utf-8')
|
|
print(json.dumps(result,ensure_ascii=True))
|
|
return 0 if result['status']=='SUCCESS' else 2
|
|
except FileExistsError:
|
|
print(json.dumps({'status':'ERROR','code':'PROMPT_PROBE_OUTPUT_EXISTS',
|
|
'guidance':'Use --output with a new filename. No API request was sent.'}));return 2
|
|
except KeyboardInterrupt:
|
|
print(json.dumps({'status':'INTERRUPTED'}));return 130
|
|
except Exception as error:
|
|
print(json.dumps({'status':'ERROR','code':safe_code(getattr(error,'code',None))}));return 2
|
|
|
|
|
|
if __name__=='__main__':raise SystemExit(main())
|