"""In-model generation/control driver. Numerical answers come from registered layers."""
import hashlib
import json
import re
import time
import uuid
import torch
try:
from lfmputer_runtime.computer.chat import TOOLS as BASE_TOOLS,parse_call,execute as base_execute,fulfilled
except ModuleNotFoundError:
from lfmcomputer.chat import TOOLS as BASE_TOOLS,parse_call,execute as base_execute,fulfilled
from .engine import FEATURES
from .interfaces import parameters,normalize,implicit_call,answer_contract
DESCRIPTIONS={
'schema':'Validate a bounded flat JSON object schema: type, required, scalar enum, additionalProperties. payload={"schema":...,"value":...}.',
'workflow':'Compute available task IDs and detect cycles. payload={"nodes":N,"dependencies":[[prerequisite,task],...],"completed":[ids]}.',
'index':'Exact UTF-8 source-span lookup, latest version per key. payload={"documents":[{"id":1,"key":1,"version":1,"text":"..."}],"query":"..."}. Omit documents to search RAM files.',
'syntax_scope':'Check quote/comment-aware ()[]{} balance; NOT full Python syntax. payload={"text":"..."}. Optional events: [op,id], op 0=enter,1=leave,2=declare,3=use.',
'polynomial':'Exact bounded integer polynomial add/multiply/derivative. Coefficients are ascending powers, max degree 8 and abs coefficient 1024. payload={"operation":"multiply","left":[1,1],"right":[1,-1]}.',
'constraints':'Check ASCII-whitespace word count, blank-line paragraphs and literal forbidden substring. payload={"text":"...","words":4,"paragraphs":2,"forbidden":"..."}. Omit unused constraints.',
'units':'SI dimensions, compatibility and rational scale conversion; no affine temperature conversion. payload={"operation":"compatible","left":"N","right":"J"}. Units include m,cm,mm,km,s,ms,min,h,kg,g,A,K,mol,cd,N,J,W,Pa,Hz,C,V,ohm,m/s,m/s^2,m2,m3.',
'sat':'Exact bounded 3-CNF satisfiability, at most 12 variables and 32 clauses. Literals are nonzero signed variable IDs. payload={"variables":2,"clauses":[[1,2],[-1]]}.',
'closure':'Directed graph transitive reachability only. payload={"nodes":N,"edges":[[from,to],...],"query":[from,to]}. Does not infer arbitrary relationship semantics.',
'binary':'Exact uint32 xor/and/or/not/rol/ror or hex_encode/base64_encode/base64_decode. payload={"operation":"xor","left":[1,2],"right":[3,4]}; encoders use bytes:[0..255], decoder uses text.',
}
def tools(model):
result=list(BASE_TOOLS) if getattr(model,'include_computer_tools',True) else []
for i,name in enumerate(FEATURES):
if bool(model.analytics.enabled[i]):
result.append({'type':'function','function':{'name':'analytic_'+name,'description':DESCRIPTIONS[name],
'parameters':parameters(name)}})
return result+list(getattr(model,'external_tools',[]))
def decode_payload(text):
if len(text)>200000:raise ValueError('Payload capacity exceeded')
def unique(pairs):
value={}
for key,item in pairs:
if key in value:raise ValueError('Duplicate JSON key')
value[key]=item
return value
def bad_constant(x):raise ValueError('Non-finite JSON number')
return json.loads(text,object_pairs_hook=unique,parse_constant=bad_constant)
def dispatch(model,session,call):
name=call['name'];args=call['arguments']
declared={t['function']['name']:t['function']['parameters'] for t in tools(model)}
if name not in declared:raise ValueError('Function is not available in this configuration')
if name.startswith('analytic_') and name[len('analytic_'):] in FEATURES:
feature=name[len('analytic_'):]
if isinstance(args,dict) and set(args)=={'payload'}:
args=decode_payload(args['payload']) if isinstance(args['payload'],str) else args['payload']
if not isinstance(args,dict):raise ValueError('Arguments must be an object')
args=normalize(feature,args);call['arguments']=args
if bool(model.analytics.enabled[FEATURES.index('schema')]):
checked=model(mode='analytic',feature='schema',payload={'schema':declared[name],'value':args})
if not checked['valid']:raise ValueError('Tool-call schema violation: '+json.dumps(checked))
if name in getattr(model,'external_dispatch',{}):return model.external_dispatch[name](args)
if not name.startswith('analytic_'):return base_execute(model,session,call)
feature=name[len('analytic_'):]
if feature not in FEATURES:raise ValueError('Unknown analytic feature')
payload=args
if feature=='index' and 'documents' not in payload:
records=model.files(session.memory,'list')
payload={**payload,'documents':[{'id':r['name'],'key':i,'version':r['revision'],
'text':model.files(session.memory,'get',name=r['name'])['content']} for i,r in enumerate(records)]}
return model(mode='analytic',feature=feature,payload=payload)
@torch.inference_mode()
def generate(model,session,max_tokens,available,stop_after_tool_call=True):
from transformers import StoppingCriteria,StoppingCriteriaList
tok=model.tokenizer;close=tok.encode('',add_special_tokens=False)
class Stop(StoppingCriteria):
def __call__(self,input_ids,scores,**kwargs):
return stop_after_tool_call and input_ids.shape[1]>=len(close) and input_ids[0,-len(close):].tolist()==close
encoded=tok.apply_chat_template(session.messages,tools=available,add_generation_prompt=True,
enable_thinking=False,tokenize=True,return_dict=True,return_tensors='pt')
prompt_tokens=encoded['input_ids'].shape[1]
if prompt_tokens>getattr(model,'context_limit',16384):raise RuntimeError('Configured context budget exceeded')
remaining_context=getattr(model,'context_limit',16384)-prompt_tokens
if remaining_context<=0:raise RuntimeError('Configured context budget exhausted')
max_tokens=min(max_tokens,remaining_context)
encoded=encoded.to(next(model.language_model.parameters()).device)
out=model.language_model.generate(**encoded,max_new_tokens=max_tokens,do_sample=False,repetition_penalty=1.1,
stopping_criteria=StoppingCriteriaList([Stop()]),
pad_token_id=tok.pad_token_id if tok.pad_token_id is not None else tok.eos_token_id)
suffix=out[0,prompt_tokens:]
eos=model.language_model.generation_config.eos_token_id or tok.eos_token_id
eos=[eos] if isinstance(eos,int) else eos
stopped=bool(suffix.numel() and (int(suffix[-1]) in eos or suffix[-len(close):].tolist()==close))
session.last_generation_finished=stopped
return tok.decode(suffix,skip_special_tokens=False),prompt_tokens,suffix.numel()
def sequence_ok(model,events,requirements,available):
if not requirements:return True
if not bool(model.analytics.enabled[FEATURES.index('workflow')]):return fulfilled(events,requirements)
names=[t['function']['name'] for t in available]
if any(name not in names for name in requirements):return False
codes={name:i+1 for i,name in enumerate(names)}
ids=model.analytics.tensor([codes.get(e['call']['name'],0) for e in events],torch.long)
successes=model.analytics.tensor([int(e['response']['ok']) for e in events])
expected=model.analytics.tensor([codes[name] for name in requirements],torch.long)
model.analytics.calls[FEATURES.index('workflow')].add_(1)
return bool(model.analytics.layers['workflow'].check_sequence(ids,successes,expected))
def chat(model,session,prompt,required_tools=(),required_files=None,final_constraints=None,
max_calls=8,max_new_tokens=512,total_new_tokens=2048,verbose=True,response_schema=None):
if model.tokenizer is None:raise RuntimeError('Tokenizer missing')
if response_schema is None:response_schema=answer_contract(prompt)
session.messages.append({'role':'user','content':prompt})
available=tools(model);events=[];started=time.perf_counter();used=0;prompt_tokens=0
for attempt in range(max_calls+4):
remaining=total_new_tokens-used
if remaining<=0:raise RuntimeError('Total generated-token budget exhausted')
raw,prefill,count=generate(model,session,min(max_new_tokens,remaining),available)
used+=count;prompt_tokens+=prefill
session.generation_stats={'generated_tokens':used,'prompt_tokens':prompt_tokens}
try:call=parse_call(raw)
except (ValueError,KeyError,TypeError) as error:
session.messages.append({'role':'assistant','content':raw})
session.messages.append({'role':'user','content':'Malformed tool call. Correct its XML/JSON syntax. '+str(error)})
continue
if call is None:
final_raw=raw.rsplit('',1)[-1] if '' in raw else raw
answer=re.sub(r'.*?','',final_raw,flags=re.S)
if '' in answer:raise RuntimeError('Incomplete reasoning')
answer=re.sub(r'<\|(?:im_end|im_start|endoftext|startoftext)\|>','',answer).strip()
try:unfinished=implicit_call(json.loads(answer))
except (ValueError,TypeError):unfinished=None
if unfinished and any(t['function']['name']==unfinished['name'] for t in available):
# The model supplied the entire formal input but forgot XML framing.
# Route only the unique public contract, never derive an input from gold.
call=unfinished
if call is not None:
raw=''+json.dumps(call)+''
if call is None:
errors=[]
if response_schema is not None and bool(model.analytics.enabled[FEATURES.index('schema')]):
try:
parsed=json.loads(answer)
checked=model(mode='analytic',feature='schema',payload={'schema':response_schema,'value':parsed})
if not checked['valid']:errors.append('Final JSON violates the requested response schema: '+json.dumps(response_schema))
except (ValueError,TypeError):errors.append('Return valid JSON matching '+json.dumps(response_schema)+'; do not add prose or code fences.')
if not sequence_ok(model,events,required_tools,available):errors.append('Execute the required tools in order: '+', '.join(required_tools))
for name,text in (required_files or {}).items():
try:
if text not in model.files(session.memory,'get',name=name)['content']:errors.append('Required content is missing in '+name)
except KeyError:errors.append('Required file is absent: '+name)
if final_constraints and bool(model.analytics.enabled[FEATURES.index('constraints')]):
try:
checked=model(mode='analytic',feature='constraints',payload={'text':answer,**final_constraints})
if not checked['valid']:errors.append('Output violates these user constraints: '+json.dumps(final_constraints)+'; observed '+json.dumps(checked))
except ValueError as error:
errors.append('Output cannot be checked: '+str(error)+'. Produce a shorter answer satisfying the original constraints.')
if errors:
session.messages.append({'role':'assistant','content':raw})
session.messages.append({'role':'user','content':'Complete/correct the following before answering: '+'; '.join(errors)})
continue
session.messages.append({'role':'assistant','content':answer})
return dict(answer=answer,events=events,seconds=time.perf_counter()-started,
generated_tokens=used,prompt_tokens=prompt_tokens,verified=bool(required_tools or required_files or final_constraints),
truncated=not getattr(session,'last_generation_finished',True))
if len(events)>=max_calls:raise RuntimeError('Tool-call budget exhausted')
identifier='call_'+uuid.uuid4().hex
content=re.sub(r'.*?','',raw,flags=re.S).strip()
session.messages.append({'role':'assistant','content':content,
'tool_calls':[{'id':identifier,'type':'function','function':call}]})
try:response={'ok':True,'result':dispatch(model,session,call)}
except (ValueError,KeyError,TypeError,RuntimeError,MemoryError,OverflowError) as error:
response={'ok':False,'error':f'{type(error).__name__}: {error}'}
event={'call':call,'response':response};events.append(event)
previous=session.audit[-1]['hash'] if session.audit else '0'*64
digest=hashlib.sha256((previous+json.dumps(event,sort_keys=True)).encode()).hexdigest()
session.audit.append({**event,'previous':previous,'hash':digest})
if verbose:print('→',call['name'],json.dumps(response,ensure_ascii=False)[:500],flush=True)
session.messages.append({'role':'tool','tool_call_id':identifier,'name':call['name'],'content':json.dumps(response,ensure_ascii=False)})
raise RuntimeError('Generation/repair budget exhausted')