"""Ollama /api/generate streaming client; no third-party dependencies.
Usage: python3 llm_benchmark.py --model MODEL --prompt-file prompt.txt --runs 5
Each output line is JSON. TTFC is first nonempty response content, not exact
model-token latency. Thinking output is timed separately. Memory is not measured.
"""
import argparse,json,time,urllib.request

def measure(url,payload):
    start=time.monotonic();first=None;thinking=None;final=None
    request=urllib.request.Request(url,data=json.dumps(payload).encode(),headers={'Content-Type':'application/json'})
    with urllib.request.urlopen(request,timeout=600) as response:
        for line in response:
            if not line.strip():continue
            event=json.loads(line);now=time.monotonic()
            if event.get('error'):raise RuntimeError(event['error'])
            if event.get('thinking') and thinking is None:thinking=now-start
            if event.get('response') and first is None:first=now-start
            if event.get('done'):final=event;break
    if final is None:raise RuntimeError('Stream ended without done=true')
    duration=final.get('eval_duration',0)
    return {'ttfc_s':first,'first_thinking_s':thinking,'client_total_s':time.monotonic()-start,
            'generation_tokens_s':final.get('eval_count',0)*1e9/duration if duration>0 else None,
            **{key:final.get(key) for key in ('model','total_duration','load_duration','prompt_eval_count','prompt_eval_duration','eval_count','eval_duration','done_reason')}}
if __name__=='__main__':
    p=argparse.ArgumentParser();p.add_argument('--model',required=True);p.add_argument('--prompt-file',required=True);p.add_argument('--url',default='http://127.0.0.1:11434/api/generate');p.add_argument('--runs',type=int,default=5);a=p.parse_args()
    if a.runs<1:p.error('--runs must be positive')
    with open(a.prompt_file,encoding='utf8') as f:prompt=f.read()
    payload={'model':a.model,'prompt':prompt,'stream':True,'keep_alive':'5m','options':{'temperature':0,'seed':42,'num_predict':128,'num_ctx':4096}}
    for run in range(a.runs):
        try:result={'run':run+1,'status':'ok',**measure(a.url,payload)}
        except Exception as exc:result={'run':run+1,'status':'error','error':str(exc)}
        print(json.dumps(result),flush=True)