lfmputer_runtime/analytics/engine.py

"""Model-side input/output codec and registered analytic circuits.

The codec parses JSON and tokenizes text. It does not call Python solvers,
sympy, regex validators, SAT packages or external programs to compute answers.
"""
import json
import torch
from torch import nn
from .circuits import (SchemaCircuit,WorkflowCircuit,IndexCircuit,SyntaxScopeCircuit,
                       PolynomialCircuit,ConstraintCircuit,UnitCircuit,SATCircuit,
                       ClosureCircuit,BinaryCircuit,UNIT_NAMES)

FEATURES=('schema','workflow','index','syntax_scope','polynomial','constraints','units','sat','closure','binary')


class AnalyticEngine(nn.Module):
    def __init__(self,spec=None):
        super().__init__()
        defaults=dict(fields=16,key_bytes=32,value_bytes=64,enums=8,nodes=32,
            docs=64,doc_tokens=8192,query_tokens=64,text_bytes=65536,stack=32,scopes=8,symbols=64,
            degree=8,coefficient_bound=1024,sat_variables=12,sat_clauses=32,bytes=512)
        if spec and set(spec)-set(defaults):raise ValueError('Unknown analytic capacity')
        self.spec={**defaults,**(spec or {})}
        limits={**defaults,'fields':32,'key_bytes':128,'value_bytes':256,'enums':32,
                'text_bytes':1048576,'stack':64,'scopes':16,'symbols':256,'bytes':1024}
        for key,value in self.spec.items():
            if type(value) is not int or not 1<=value<=limits[key]:raise ValueError('Unsupported analytic capacity: '+key)
        s=self.spec
        self.layers=nn.ModuleDict(dict(schema=SchemaCircuit(),workflow=WorkflowCircuit(),index=IndexCircuit(),
            syntax_scope=SyntaxScopeCircuit(),polynomial=PolynomialCircuit(s['degree']),
            constraints=ConstraintCircuit(),units=UnitCircuit(),sat=SATCircuit(s['sat_variables']),
            closure=ClosureCircuit(),binary=BinaryCircuit()))
        self.register_buffer('enabled',torch.ones(len(FEATURES),dtype=torch.bool))
        self.register_buffer('calls',torch.zeros(len(FEATURES),dtype=torch.long),persistent=False)

    @property
    def device(self):return self.enabled.device

    def prepare(self):
        torch.set_float32_matmul_precision('highest')
        if self.device.type=='cuda':torch.backends.cuda.matmul.allow_tf32=False
        for name,parameter in self.named_parameters():
            double=(name.startswith('layers.polynomial.') or name.startswith('layers.units.units.') or name.startswith('layers.units.combine.'))
            expected=torch.float64 if double else torch.float32
            if parameter.dtype!=expected or parameter.requires_grad:
                raise TypeError('Analytic coefficients require their compiled FP32/FP64 dtype and frozen gradients: '+name)

    def set_enabled(self,names):
        if set(names)-set(FEATURES):raise ValueError('Unknown analytic feature')
        self.enabled.copy_(torch.tensor([name in names for name in FEATURES],device=self.device))

    def forward(self,feature,**kwargs):
        self.prepare()
        i=FEATURES.index(feature)
        if not bool(self.enabled[i]):raise RuntimeError(f'Feature disabled: {feature}')
        self.calls[i].add_(1)
        return self.layers[feature](**kwargs)

    def tensor(self,data,dtype=torch.float32):return torch.tensor(data,device=self.device,dtype=dtype)

    def text(self,text,limit=None):
        if not isinstance(text,str):raise TypeError('Expected text')
        encoded=text.encode('utf-8')
        if len(encoded)>(limit or self.spec['text_bytes']):raise ValueError('Text capacity exceeded')
        return self.tensor(list(encoded),torch.long)

    @staticmethod
    def integer(value,lo,hi):
        if type(value) is not int or not lo<=value<=hi:raise ValueError(f'Integer must be in [{lo},{hi}]')
        return value

    @torch.inference_mode()
    def run(self,feature,payload,tokenizer=None):
        self.prepare()
        if feature not in FEATURES:raise ValueError('Unknown feature')
        if not isinstance(payload,dict):raise TypeError('Payload must be an object')
        if not bool(self.enabled[FEATURES.index(feature)]):raise RuntimeError('Feature disabled: '+feature)
        from .interfaces import validate_payload
        validate_payload(feature,payload)
        s=self.spec
        def primitive(value):
            if isinstance(value,torch.Tensor):
                return value.item() if value.numel()==1 else value.cpu().tolist()
            return value
        def result(values):return {k:primitive(v) for k,v in values.items()}

        if feature=='schema':return self.schema(payload)
        if feature=='workflow':
            n=self.integer(payload['nodes'],1,s['nodes']);deps=torch.zeros(n,n,device=self.device)
            for before,after in payload.get('dependencies',[]):
                deps[self.integer(after,0,n-1),self.integer(before,0,n-1)]=1
            complete=torch.zeros(n,device=self.device)
            for node in payload.get('completed',[]):complete[self.integer(node,0,n-1)]=1
            out=self(feature,dependencies=deps,completed=complete)
            return dict(available=out['available'].nonzero().flatten().cpu().tolist(),cyclic=bool(out['cyclic']))
        if feature=='index':
            records=payload['documents']
            if not 0<=len(records)<=s['docs']:raise ValueError('Document capacity exceeded')
            def tokens(text,limit,empty=False):
                ids=list(text.encode('utf-8'))
                if not int(not empty)<=len(ids)<=limit:raise ValueError('Index byte capacity exceeded')
                return ids
            docs=torch.zeros(len(records),s['doc_tokens'],device=self.device,dtype=torch.long)
            mask=torch.zeros_like(docs,dtype=torch.float32)
            keys=[];versions=[]
            for i,doc in enumerate(records):
                ids=tokens(doc['text'],s['doc_tokens'],empty=True);docs[i,:len(ids)]=self.tensor(ids,torch.long);mask[i,:len(ids)]=1
                keys.append(self.integer(doc.get('key',doc['id']),0,2**32-1))
                versions.append(self.integer(doc.get('version',1),1,2**31-1))
            query=tokens(payload['query'],s['query_tokens'])
            out=self(feature,documents=docs,doc_mask=mask,query=self.tensor(query,torch.long),
                versions=self.tensor(versions,torch.long),keys=self.tensor(keys,torch.long))
            found=bool(out['found']);index=int(out['index'])
            return dict(found=found,document_id=records[index]['id'] if found else None,
                        byte_offset=int(out['positions'][index]) if found else None,scores=out['scores'].cpu().tolist())
        if feature=='syntax_scope':
            out={}
            if 'text' in payload:
                parsed=self(feature,text=self.text(payload['text']),max_depth=s['stack'])
                out.update(balanced=bool(parsed['valid']),remaining_depth=int(parsed['depth']))
            if 'events' in payload:
                events=[]
                for event in payload['events']:
                    op=self.integer(event[0],0,3);identifier=self.integer(event[1],0,s['symbols']-1)
                    events.append([op,identifier])
                if len(events)>s['text_bytes']:raise ValueError('Scope event capacity exceeded')
                # Same registered syntax/scope module, explicit finite symbol IR.
                self.calls[FEATURES.index(feature)].add_(1)
                if not bool(self.enabled[FEATURES.index(feature)]):raise RuntimeError('Feature disabled')
                checked=self.layers[feature].scope(self.tensor(events,torch.long).reshape(-1,2),s['scopes'],s['symbols'])
                out['scope_valid']=bool(checked['valid'])
            if not out:raise ValueError('Provide text or explicit scope events')
            return out
        if feature=='polynomial':
            def coefficients(values):
                if not 1<=len(values)<=s['degree']+1:raise ValueError('Polynomial degree capacity exceeded')
                for v in values:self.integer(v,-s['coefficient_bound'],s['coefficient_bound'])
                return self.tensor(values+[0]*(s['degree']+1-len(values)),torch.float64)
            left=coefficients(payload['left']);right=coefficients(payload.get('right',[0]))
            value=self(feature,left=left,right=right,operation=payload['operation']).round().long().cpu().tolist()
            while len(value)>1 and value[-1]==0:value.pop()
            return {'coefficients':value}
        if feature=='constraints':
            words=self.integer(payload.get('words',-1),-1,10000)
            paragraphs=self.integer(payload.get('paragraphs',-1),-1,10000)
            forbidden=self.text(payload.get('forbidden',''),64)
            out=self(feature,text=self.text(payload['text']),words=words,paragraphs=paragraphs,forbidden=forbidden)
            return dict(valid=bool(out['valid']),words=int(out['words']),paragraphs=int(out['paragraphs']),forbidden=bool(out['forbidden']))
        if feature=='units':
            op=payload.get('operation','compatible')
            if op not in ('compatible','convert','add','multiply','divide'):raise ValueError('Unsupported unit operation')
            if payload['left'] not in UNIT_NAMES or payload['right'] not in UNIT_NAMES:raise ValueError('Unit outside registry')
            out=self(feature,left=self.tensor(UNIT_NAMES.index(payload['left']),torch.long),
                right=self.tensor(UNIT_NAMES.index(payload['right']),torch.long),operation=op)
            if op in ('convert','add') and not bool(out['compatible']):raise ValueError('Incompatible unit dimensions for '+op)
            return dict(compatible=bool(out['compatible']),dimensions=out['dimensions'].long().cpu().tolist(),
                        ratio_num=int(out['scale_num']),ratio_den=int(out['scale_den']))
        if feature=='sat':
            n=self.integer(payload['variables'],1,s['sat_variables']);clauses=payload['clauses']
            if not 1<=len(clauses)<=s['sat_clauses']:raise ValueError('Clause capacity exceeded')
            padded=[]
            for clause in clauses:
                if not 1<=len(clause)<=3:raise ValueError('Only clauses of one to three literals supported')
                for literal in clause:
                    self.integer(literal,-n,n)
                    if literal==0:raise ValueError('Literal zero is invalid')
                padded.append(clause+[clause[-1]]*(3-len(clause)))
            out=self(feature,clauses=self.tensor(padded,torch.long),n_variables=n)
            return dict(satisfiable=bool(out['satisfiable']),assignment=out['assignment'].long().cpu().tolist() if bool(out['satisfiable']) else None)
        if feature=='closure':
            n=self.integer(payload['nodes'],1,s['nodes']);adj=torch.zeros(n,n,device=self.device)
            for a,b in payload.get('edges',[]):adj[self.integer(a,0,n-1),self.integer(b,0,n-1)]=1
            closure=self(feature,adjacency=adj)
            if 'query' in payload:
                a,b=payload['query'];return {'reachable':bool(closure[self.integer(a,0,n-1),self.integer(b,0,n-1)])}
            return {'adjacency':closure.long().cpu().tolist()}
        if feature=='binary':
            op=payload['operation'];module=self.layers[feature]
            if not bool(self.enabled[FEATURES.index(feature)]):raise RuntimeError('Feature disabled')
            if op in ('hex_encode','base64_encode','base64_decode'):
                if op=='base64_decode':data=self.text(payload['text'],s['bytes']*2)
                else:
                    values=payload['bytes']
                    if len(values)>s['bytes']:raise ValueError('Byte capacity exceeded')
                    data=self.tensor([self.integer(v,0,255) for v in values],torch.long)
                self.calls[FEATURES.index(feature)].add_(1)
                if op=='hex_encode':chars=module.hex_encode(data)
                elif op=='base64_encode':chars=module.base64_encode(data)
                else:
                    if data.numel()==0:return {'bytes':[]}
                    if data.numel()%4:raise ValueError('Base64 length must be a multiple of four')
                    decoded,valid,pad=module.base64_decode(data)
                    if not bool(valid):raise ValueError('Invalid or noncanonical Base64')
                    if decoded.numel()-int(pad)>s['bytes']:raise ValueError('Decoded byte capacity exceeded')
                    return {'bytes':decoded[:decoded.numel()-int(pad)].cpu().tolist()}
                return {'text':bytes(chars.cpu().tolist()).decode('ascii')}
            left=payload.get('left',[]);right=payload.get('right',[0]*len(left))
            if not 1<=len(left)<=s['bytes'] or len(left)!=len(right):raise ValueError('Invalid SIMD vector dimensions')
            for value in left+right:self.integer(value,0,2**32-1)
            rotation=self.integer(payload.get('rotation',0),0,31)
            out=self(feature,left=self.tensor(left,torch.long),right=self.tensor(right,torch.long),operation=op,rotation=rotation)
            return {'values':out.cpu().tolist()}
        raise AssertionError(feature)

    def schema(self,payload):
        s=self.spec;schema=payload['schema'];value=payload['value']
        if not isinstance(schema,dict) or not isinstance(value,dict):raise ValueError('Flat object schemas only')
        unsupported=set(schema)-{'type','properties','required','additionalProperties','description','title','$schema'}
        if unsupported or schema.get('type','object')!='object':raise ValueError('Unsupported schema keyword/type')
        properties=schema.get('properties',{})
        if not isinstance(properties,dict) or type(schema.get('additionalProperties',True)) is not bool:
            raise ValueError('Only boolean additionalProperties supported')
        if not isinstance(schema.get('required',[]),list) or any(not isinstance(k,str) for k in schema.get('required',[])):
            raise ValueError('required must be a string list')
        if len(properties)>s['fields'] or len(value)>s['fields']:raise ValueError('Field capacity exceeded')
        F,K,V,E=s['fields'],s['key_bytes'],s['value_bytes'],s['enums']
        keys_s=torch.zeros(F,K,device=self.device);keys_a=torch.zeros_like(keys_s)
        values_a=torch.zeros(F,V,device=self.device);enums=torch.zeros(F,E,V,device=self.device)
        ps=torch.zeros(F,device=self.device);pa=ps.clone();required=ps.clone();ep=torch.zeros(F,E,device=self.device)
        types=torch.zeros(F,7,device=self.device);actual=torch.zeros_like(types)
        numeric=torch.zeros(F,device=self.device,dtype=torch.float64)
        enum_numeric=torch.zeros(F,E,device=self.device,dtype=torch.float64)
        enum_is_numeric=torch.zeros(F,E,device=self.device)
        type_names=['null','boolean','integer','number','string','array','object']
        def encoded(data,limit):
            if not isinstance(data,str):raise ValueError('Schema keys must be strings')
            b=data.encode('utf-8')
            if not b or len(b)>limit or b'\0' in b:raise ValueError('Schema string capacity exceeded or NUL')
            return self.tensor(list(b)+[0]*(limit-len(b)))
        def canonical(v):return json.dumps(v,ensure_ascii=False,sort_keys=True,separators=(',',':'),allow_nan=False)
        def numeric_value(v):
            if type(v) in (int,float):
                if not -2**40<=v<=2**40:raise ValueError('Schema numeric domain exceeds +/-2^40')
                return float(v),True
            return 0.,False
        for i,(name,prop) in enumerate(properties.items()):
            if set(prop)-{'type','enum','description','title','default','examples'}:raise ValueError('Unsupported property keyword')
            keys_s[i]=encoded(name,K);ps[i]=1;required[i]=name in schema.get('required',[])
            allowed=prop.get('type',type_names)
            allowed=[allowed] if isinstance(allowed,str) else allowed
            for t in allowed:
                if t not in type_names:raise ValueError('Unsupported type')
                types[i,type_names.index(t)]=1
                if t=='number':types[i,2]=1
            choices=prop.get('enum',[])
            if len(choices)>E or ('enum' in prop and not choices):raise ValueError('Invalid enum capacity')
            for j,choice in enumerate(choices):
                if isinstance(choice,(list,dict)):raise ValueError('Only scalar enum members supported')
                enums[i,j]=encoded(canonical(choice),V);ep[i,j]=1
                number,is_number=numeric_value(choice);enum_numeric[i,j]=number;enum_is_numeric[i,j]=is_number
        if any(name not in properties for name in schema.get('required',[])):raise ValueError('Required property absent in schema')
        for i,(name,v) in enumerate(value.items()):
            keys_a[i]=encoded(name,K);pa[i]=1
            wire=canonical(v)
            # Values longer than every supported enum literal cannot equal an
            # enum. Keep their type; an all-zero code matches no JSON scalar.
            if len(wire.encode('utf-8'))<=V:values_a[i]=encoded(wire,V)
            typ={type(None):0,bool:1,int:2,float:3,str:4,list:5,dict:6}.get(type(v))
            if typ is None:raise ValueError('Non-JSON value')
            actual[i,typ]=1
            number,is_number=numeric_value(v)
            if is_number:numeric[i]=number
        out=self('schema',key_schema=keys_s,key_actual=keys_a,present_schema=ps,present_actual=pa,
            required=required,allowed_types=types,actual_types=actual,enum_values=enums,enum_present=ep,
            actual_values=values_a,numeric=numeric,enum_numeric=enum_numeric,enum_is_numeric=enum_is_numeric,
            additional=bool(schema.get('additionalProperties',True)))
        return {k:(bool(v) if k=='valid' else int(v)) for k,v in out.items()}