ares/runtime.py

"""Local runtime for separately trained Ares and Xiphos checkpoints.
Neither role executes a tool. Xiphos returns a reviewable proposal only.
"""
from dataclasses import dataclass
from pathlib import Path
import re, torch
from .config import AresConfig
from .model import AresTransformer
from .tokenizer import ByteBPETokenizer

@dataclass
class LoadedModel:
    model: AresTransformer
    tokenizer: ByteBPETokenizer
    device: str
    trained: bool
    checkpoint: str

def load_model(checkpoint, tokenizer_path, device=None, expected_role=None):
    device=device or ('cuda' if torch.cuda.is_available() else 'cpu')
    tok=ByteBPETokenizer.load(tokenizer_path)
    ckpt=torch.load(checkpoint,map_location=device,weights_only=False)
    if expected_role and ckpt.get('role') not in (None, expected_role):
        raise ValueError(f'Checkpoint role {ckpt.get("role")!r} cannot be loaded as {expected_role!r}')
    model=AresTransformer(AresConfig(**ckpt['config'])).to(device)
    model.load_state_dict(ckpt['model']); model.eval()
    return LoadedModel(model,tok,device,bool(ckpt.get('training_complete',False)),str(checkpoint))

def respond(role, loaded, message, max_new=160):
    if not loaded.trained:
        return f'{role} is installed but its checkpoint is not trained yet. Train and evaluate it before using its output.'
    # Role headers are data boundaries, not a security boundary; server never executes generated text.
    prompt=f'<system>You are {role}. ' + ('Answer the user directly and helpfully.' if role=='Ares' else 'Draft a concise, non-executable implementation plan for user review.') + f'</system><user>{message}</user><assistant>'
    ids=torch.tensor([loaded.tokenizer.encode(prompt,add_bos=True)],device=loaded.device)
    if ids.size(1)>loaded.model.config.max_seq_len-1: ids=ids[:,-loaded.model.config.max_seq_len+1:]
    out=loaded.model.generate(ids,max_new=max_new,temperature=.7)[0].tolist()
    answer=loaded.tokenizer.decode(out[len(ids[0]):])
    # Strip control tokens if a model emits them; never parse a response as commands.
    return re.sub(r'<[^>]+>','',answer).strip()

def proposed_plan(xiphos, request):
    """Return model text in a plan envelope. Approval/execution is deliberately separate."""
    return {'status':'proposed','executable':False,'request':request,'proposal':respond('Xiphos',xiphos,request),'required_review':['scope and files','data licences/privacy','resource estimate','tests and rollback','exact allowlisted commands']}