70 lines
4.0 KiB
Python
70 lines
4.0 KiB
Python
"""Measure one prepared Gemini request without executing research tools."""
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
from pathlib import Path
|
|
from jsonschema import Draft202012Validator
|
|
from .llm import GatewayConfig, LLMError
|
|
from .llm.gateway import PRIVACY
|
|
from .minimal_prompt_test import measure
|
|
from .telemetry.safety import safe_code
|
|
|
|
DEFAULT = Path('service-data/benchmarks/compact-investigation-20261017/request.json')
|
|
|
|
|
|
def prepare(provider,request):
|
|
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'])
|
|
generation={'maxOutputTokens':provider.max_output_tokens,'responseMimeType':'application/json',
|
|
'responseJsonSchema':request['output_schema']}
|
|
if provider.reasoning_effort is not None:generation['thinkingConfig']={'thinkingLevel':provider.reasoning_effort}
|
|
payload={'systemInstruction':{'parts':[{'text':messages[0]['content']}]},
|
|
'contents':[{'role':'user','parts':[{'text':messages[1]['content']}]}],'generationConfig':generation}
|
|
return {'id':request['prompt_version'],'prompt_characters':sum(len(m['content']) for m in messages),
|
|
'payload':payload}
|
|
|
|
|
|
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('--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)
|
|
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}
|
|
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(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 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())
|