space/runtime.py

"""Run the real Lua loop, with explicit provenance and metric denominators."""
from datetime import datetime, timezone
import hashlib
import importlib.metadata
import json
from pathlib import Path
import threading
import time

from lupa.lua54 import LuaRuntime

ROOT = Path(__file__).resolve().parents[1]
LOCK = threading.RLock()
EXPECTED = {
    'calc-single': (['calc'], '1018039; must use calc'),
    'read-then-calc': (['read_file'], '12; must read the fixture and use calc'),
    'recover-parse-error': (['calc'], '4; scripted run must recover from a parse error'),
    'recover-missing-file': (['read_file', 'list_dir'], 'alpha; scripted run must recover from a missing file'),
    'loop-bait': (['read_file'], 'Report no secret; scripted run must trigger repetition detection'),
    'oversized-observation': (['read_file'], 'omega; scripted run must exercise context eviction'),
    'approved-write': (['write_file'], 'Write 42 and read it back after approval'),
    'denied-write': (['write_file'], 'Do not write; report the refusal'),
}


def new_runtime():
    lua = LuaRuntime(unpack_returned_tuples=True)
    lua.globals().lab_root = ROOT.as_posix()
    lua.execute("package.path = lab_root .. '/?.lua;' .. lab_root .. '/space/vendor/lua_llama/?.lua;' .. package.path")
    lab = lua.eval("require('space.lab')")[0]
    return lua, lab


def tasks():
    _, lab = new_runtime()
    return json.loads(lab.tasks_json())


def fixtures():
    return {p.relative_to(ROOT).as_posix(): p.read_text(encoding='utf-8')
            for p in (ROOT / 'eval' / 'fixtures').glob('*.txt')}


def source_paths():
    names = ['agent.lua', 'grammar.lua', 'protocol.lua', 'context.lua', 'tools.lua',
             'trace.lua', 'bpe.lua', 'backend/local.lua', 'eval/tasks/tasks.lua',
             'tools/llama2c.py', 'space/runtime.py', 'space/models.py', 'space/lab.lua',
             'space/revisions.py', 'space/requirements.txt',
             'space/vendor/lua_llama/tokenizer.lua', 'space/vendor/lua_llama/utils.lua']
    return sorted([ROOT / name for name in names] + list((ROOT / 'eval' / 'fixtures').glob('*.txt')))


def source_digest(root=ROOT):
    digest = hashlib.sha256()
    for path in source_paths():
        relative = path.relative_to(ROOT)
        digest.update(relative.as_posix().encode())
        digest.update((root / relative).read_bytes().replace(b'\r\n', b'\n'))
    return digest.hexdigest()


def metrics(result, expected_tools=None):
    steps = result.get('steps') or []
    generated = [s for s in steps if 'raw' in s]
    parsed = [s for s in generated if s.get('call')]
    first = generated[0].get('call', {}).get('tool') if generated else None
    return {
        'generated_turns': len(generated), 'valid_calls': len(parsed),
        'tool_call_validity': 100 * len(parsed) / len(generated) if generated else None,
        'correct_first_tool': first in expected_tools if expected_tools and generated else None,
        'first_tool': first, 'steps': len(steps),
        'malformed_calls': sum(s.get('status') == 'parse_error' for s in steps),
        'loop_detected': any(s.get('status') == 'repeat' for s in steps),
        'context_evicted': any(s.get('evicted', False) for s in steps),
        'forced_finishes': sum(bool(s.get('forced_finish')) for s in steps),
    }


class Experiment:
    """A batch may reuse kernels and grammar caches, never files or agent history."""
    def __init__(self):
        self.lua, self.lab = new_runtime()
        self.backends = {}

    def factory(self, registry, config_json):
        from space.models import tiny_backend, qwen_backend
        config = json.loads(config_json)
        names = tuple(registry.names(registry).values())
        key = (config['model'], config['grammar'], names)
        if key not in self.backends:
            builder = tiny_backend if config['model'] == 'tiny' else qwen_backend
            self.backends[key] = builder(self.lua, registry, config)
        backend, tokenizer = self.backends[key]
        if backend.reset:
            backend.reset(backend)
        return backend, tokenizer

    def run(self, model='tiny', task_id='calc-single', prompt=None, grammar=True,
            context_limit=None, approval_gate=True, approve_write=None,
            loop_detector=True, error_recovery=True, max_steps=8,
            policy='elide_observations', **extra):
        if model not in {'tiny', 'qwen', 'scripted'}:
            raise ValueError('Unknown model')
        task = next((t for t in tasks() if t['id'] == task_id), None)
        prompt = prompt if prompt is not None else (task['prompt'] if task else '')
        if not prompt.strip() or len(prompt.encode('utf-8')) > 2000:
            raise ValueError('Enter a task of 1–2000 UTF-8 bytes.')
        if not 1 <= int(max_steps) <= 10:
            raise ValueError('Step limit must be between 1 and 10.')
        matched = task is not None and task['prompt'] == prompt
        if model == 'scripted' and not matched:
            raise ValueError('Scripted replay uses the preset prompt. Choose a live model for a custom task.')
        context_limit = int(context_limit or (160 if model == 'tiny' else (task['budget'] if task else 2048)))
        maximum = 160 if model == 'tiny' else 3072 if model == 'qwen' else 100000
        if not 64 <= context_limit <= maximum:
            raise ValueError(f'Context budget must be 64–{maximum} for this model.')
        expected_tools, expected_outcome = EXPECTED.get(task_id, (None, None)) if matched else (None, None)
        if 'expected_tools' in extra:
            expected_tools = extra.pop('expected_tools')
        config = dict(model=model, task_id=task_id if matched else None, prompt=prompt,
                      grammar=bool(grammar) if model != 'scripted' else False,
                      context_limit=context_limit, approval_gate=bool(approval_gate),
                      approve_write=bool(approve_write if approve_write is not None else task and task['approve']),
                      loop_detector=bool(loop_detector), error_recovery=bool(error_recovery),
                      max_steps=int(max_steps), policy=policy, **extra)
        with LOCK:
            start = time.perf_counter()
            result = json.loads(self.lab.run(json.dumps(config), json.dumps(fixtures()), self.factory))
        result.update(model=model, task_id=config['task_id'], prompt=prompt, settings=config,
                      expected_tools=expected_tools, expected_outcome=expected_outcome,
                      elapsed_seconds=round(time.perf_counter() - start, 3),
                      run_at=datetime.now(timezone.utc).isoformat(), source_sha256=source_digest(),
                      inference_kind='scripted' if model == 'scripted' else 'live',
                      decoder='none (scripted)' if model == 'scripted' else
                              ('grammar.lua + PyTorch kernels' if model == 'tiny' else 'llama.cpp JSON schema') if grammar else 'free',
                      token_accounting='estimated bytes / 4' if model == 'scripted' else
                                       'TinyStories BPE' if model == 'tiny' else 'Qwen tokens; chat framing excluded')
        result['metrics'] = metrics(result, expected_tools)
        package = 'torch' if model == 'tiny' else 'llama-cpp-python' if model == 'qwen' else 'lupa'
        result['inference_library'] = {package: importlib.metadata.version(package)}
        return result


def run_experiment(**kwargs):
    return Experiment().run(**kwargs)