| import hashlib |
| import json |
| import os |
| import re |
|
|
| for _key in ("OMP_NUM_THREADS", "MKL_NUM_THREADS", "OPENBLAS_NUM_THREADS"): |
| try: |
| _threads = int(os.environ.get(_key, "4")) |
| except ValueError: |
| _threads = 4 |
| os.environ[_key] = str(max(1, min(4, _threads))) |
| import numpy as np |
|
|
| ROLES = ["none", "tool", "arg", "value", "field", "limit"] |
| MODES = ["single", "sequential", "parallel"] |
| OUTPUTS = ["all", "last"] |
|
|
|
|
| class Diagnostic(ValueError): |
| pass |
|
|
|
|
| NUMBER = re.compile(r"-?(?:0|[1-9]\d*)(?:\.\d+)?(?:[eE][+-]?\d+)?\Z") |
| TOKEN = re.compile(r'"(?:\\.|[^"\\])*"|[A-Za-z_$][\w$]*|[+-]?\d[\w.+-]*|[^\s]') |
|
|
|
|
| def number_ok(v): |
| import math |
|
|
| if type(v) is int: |
| try: |
| fv = float(v) |
| except OverflowError as error: |
| raise Diagnostic("number: JS overflow") from error |
| if not math.isfinite(fv) or int(fv) != v: |
| raise Diagnostic("number: integer not exactly JS-representable") |
| elif type(v) is float: |
| if not math.isfinite(v): |
| raise Diagnostic("number: nonfinite") |
| else: |
| raise Diagnostic("number: primitive required") |
|
|
|
|
| def catalog_check(tools): |
| if not isinstance(tools, list) or not 1 <= len(tools) <= 8: |
| raise Diagnostic("catalog: expected 1-8 tools") |
| names = set() |
| for t in tools: |
| if not isinstance(t, dict) or set(t) != { |
| "name", |
| "arguments", |
| "required", |
| "returns", |
| }: |
| raise Diagnostic("catalog: tool shape") |
| name = t["name"] |
| a = t["arguments"] |
| req = t["required"] |
| ret = t["returns"] |
| if ( |
| not isinstance(name, str) |
| or not re.fullmatch(r"[A-Za-z_$][A-Za-z0-9_$]*", name) |
| or name in names |
| ): |
| raise Diagnostic("catalog: unsafe/duplicate tool") |
| names.add(name) |
| if ( |
| not isinstance(a, dict) |
| or len(a) > 4 |
| or any( |
| not isinstance(k, str) |
| or not k |
| or v not in ("string", "number", "boolean") |
| for k, v in a.items() |
| ) |
| ): |
| raise Diagnostic("catalog: primitive argument schema") |
| if ( |
| not isinstance(req, list) |
| or any(not isinstance(k, str) for k in req) |
| or len(set(req)) != len(req) |
| or set(req) - set(a) |
| ): |
| raise Diagnostic("catalog: required keys") |
| if ( |
| not isinstance(ret, dict) |
| or set(ret) - {"kind", "properties"} |
| or ret.get("kind") not in ("text", "object", "array") |
| ): |
| raise Diagnostic("catalog: return schema") |
| props = ret.get("properties", []) |
| if ( |
| not isinstance(props, list) |
| or any(not isinstance(k, str) or not k for k in props) |
| or len(set(props)) != len(props) |
| ): |
| raise Diagnostic("catalog: properties") |
| return tools |
|
|
|
|
| def frontend(item): |
| if not isinstance(item, dict) or not isinstance(item.get("instruction"), str): |
| raise Diagnostic("input: instruction string required") |
| tools = catalog_check(item.get("tools")) |
| s = item["instruction"] |
| toks = TOKEN.findall(s) |
| if not toks: |
| raise Diagnostic("input: empty instruction") |
| if len(toks) > 128: |
| raise Diagnostic("token bound") |
| names = {t["name"] for t in tools} |
| args = {a for t in tools for a in t["arguments"]} |
| fields = {a for t in tools for a in t["returns"].get("properties", [])} |
| ids = [] |
| features = [] |
| values = [] |
| for x in toks: |
| literal = False |
| v = x |
| if x.startswith('"'): |
| try: |
| v = json.loads(x) |
| literal = True |
| except Exception as error: |
| raise Diagnostic("malformed literal") from error |
| elif x in ("true", "false"): |
| v = x == "true" |
| literal = True |
| elif x[0].isdigit() or (x[0] in "+-" and len(x) > 1 and x[1].isdigit()): |
| if not NUMBER.fullmatch(x): |
| raise Diagnostic("number: malformed JSON syntax") |
| try: |
| v = json.loads(x) |
| number_ok(v) |
| except (ValueError, OverflowError) as e: |
| raise Diagnostic(str(e)) from e |
| from decimal import Decimal |
|
|
| decimal = Decimal(x) |
| if v == 0 and decimal != 0: |
| raise Diagnostic("number: JS underflow") |
| if decimal == decimal.to_integral_value() and int(v) != int(decimal): |
| raise Diagnostic("number: integer not exactly JS-representable") |
| if ( |
| type(v) is float |
| and v.is_integer() |
| and decimal != decimal.to_integral_value() |
| ): |
| raise Diagnostic("number: fractional part lost in JS") |
| literal = True |
| f = [ |
| float(x in names and not literal), |
| float(x in args and not literal), |
| float(x in fields and not literal), |
| float(literal and isinstance(v, str)), |
| float(literal and type(v) in (int, float)), |
| float(literal and type(v) is bool), |
| ] |
| f += [ |
| float(not literal and any(t["arguments"].get(x) == typ for t in tools)) |
| for typ in ("string", "number", "boolean") |
| ] |
| f += [ |
| float( |
| not literal |
| and any(t["name"] == x and t["returns"]["kind"] == kind for t in tools) |
| ) |
| for kind in ("text", "object", "array") |
| ] |
| norm = "<schema>" if any(f[:3]) else "<literal>" if literal else x.lower() |
| ids.append( |
| int.from_bytes(hashlib.sha256(norm.encode()).digest()[:4], "little") % 1024 |
| ) |
| features.append(f) |
| values.append(v) |
| for i, x in enumerate(toks): |
| if x == "=" and (i + 1 >= len(toks) or not any(features[i + 1][3:6])): |
| raise Diagnostic("literal: assignment requires JSON string/number/boolean") |
| return ( |
| toks, |
| np.array(ids, dtype=np.int64), |
| np.array(features, dtype=np.float32), |
| values, |
| ) |
|
|
|
|
| def validate(plan, tools): |
| catalog_check(tools) |
| if not isinstance(plan, dict) or set(plan) != {"mode", "calls", "output"}: |
| raise Diagnostic("plan: shape") |
| catalog = {t["name"]: t for t in tools} |
| calls = plan["calls"] |
| out = plan["output"] |
| if not isinstance(calls, list) or not 1 <= len(calls) <= 3: |
| raise Diagnostic("call bound") |
| if ( |
| not isinstance(plan["mode"], str) |
| or plan["mode"] not in MODES |
| or (len(calls) == 1) != (plan["mode"] == "single") |
| ): |
| raise Diagnostic("mode mismatch") |
| if ( |
| not isinstance(out, dict) |
| or set(out) - {"kind", "fields", "limit"} |
| or out.get("kind") not in OUTPUTS |
| ): |
| raise Diagnostic("output kind/shape") |
| for c in calls: |
| if ( |
| not isinstance(c, dict) |
| or set(c) != {"tool", "args"} |
| or not isinstance(c["tool"], str) |
| or c["tool"] not in catalog |
| ): |
| raise Diagnostic("call: unknown tool/shape") |
| t = catalog[c["tool"]] |
| a = c["args"] |
| if ( |
| not isinstance(a, dict) |
| or len(a) > 4 |
| or any(not isinstance(k, str) for k in a) |
| or set(a) - set(t["arguments"]) |
| or set(t["required"]) - set(a) |
| ): |
| raise Diagnostic("unknown/missing arguments") |
| for k, v in a.items(): |
| typ = ( |
| "boolean" |
| if type(v) is bool |
| else "number" |
| if type(v) in (int, float) |
| else "string" |
| if type(v) is str |
| else "invalid" |
| ) |
| if typ != t["arguments"][k]: |
| raise Diagnostic("argument type") |
| if typ == "number": |
| number_ok(v) |
| if "fields" in out or "limit" in out: |
| t = catalog[calls[0]["tool"]] |
| if len(calls) != 1 or t["returns"]["kind"] != "array": |
| raise Diagnostic("shaping scope") |
| if "fields" in out: |
| fs = out["fields"] |
| if ( |
| not isinstance(fs, list) |
| or not 1 <= len(fs) <= 3 |
| or any(not isinstance(k, str) for k in fs) |
| or len(set(fs)) != len(fs) |
| or set(fs) - set(t["returns"].get("properties", [])) |
| ): |
| raise Diagnostic("projection fields") |
| if "limit" in out and ( |
| type(out["limit"]) is not int or not 0 <= out["limit"] <= 100 |
| ): |
| raise Diagnostic("limit range") |
| return plan |
|
|
|
|
| def assemble(item, roles, mode, output): |
| toks, _, _, vals = frontend(item) |
| if ( |
| not hasattr(roles, "__len__") |
| or len(roles) != len(toks) |
| or any( |
| not isinstance(r, (int, np.integer)) |
| or isinstance(r, (bool, np.bool_)) |
| or r not in range(6) |
| for r in roles |
| ) |
| ): |
| raise Diagnostic("roles: length/value mismatch") |
| if ( |
| not isinstance(mode, (int, np.integer)) |
| or not isinstance(output, (int, np.integer)) |
| or isinstance(mode, (bool, np.bool_)) |
| or isinstance(output, (bool, np.bool_)) |
| or mode not in range(3) |
| or output not in range(2) |
| ): |
| raise Diagnostic("heads: invalid class") |
| calls = [] |
| pending = None |
| fields = [] |
| limit = None |
| for x, v, r in zip(toks, vals, roles, strict=False): |
| if r == 1: |
| if pending is not None: |
| raise Diagnostic("unbound argument") |
| calls.append({"tool": x, "args": {}}) |
| elif r == 2: |
| if not calls or pending is not None: |
| raise Diagnostic("argument structure") |
| pending = x |
| elif r == 3: |
| if pending is None or pending in calls[-1]["args"]: |
| raise Diagnostic("value structure/duplicate") |
| calls[-1]["args"][pending] = v |
| pending = None |
| elif r == 4: |
| fields.append(x) |
| elif r == 5: |
| if limit is not None: |
| raise Diagnostic("multiple limits") |
| limit = v |
| if pending is not None: |
| raise Diagnostic("missing value") |
| out = {"kind": OUTPUTS[output]} |
| if fields: |
| out["fields"] = fields |
| if limit is not None: |
| out["limit"] = limit |
| return validate({"mode": MODES[mode], "calls": calls, "output": out}, item["tools"]) |
|
|
|
|
| def compile_js(plan, tools): |
| validate(plan, tools) |
|
|
| def j(x): |
| return json.dumps(x, ensure_ascii=True, separators=(",", ":")) |
|
|
| inv = [ |
| f"tools[{j(c['tool'])}](JSON.parse({j(j(c['args']))}))" for c in plan["calls"] |
| ] |
| pre = "const err=e=>String(e && e.message || e); const invoke=async f=>{try{return {ok:true,value:await f()}}catch(e){return {ok:false,error:err(e)}}};\n" |
| if plan["mode"] == "parallel": |
| pre += ( |
| "const settled=await Promise.allSettled([" |
| + ",".join("invoke(()=>" + x + ")" for x in inv) |
| + ']); const r=settled.map(s=>s.status==="fulfilled"?s.value:{ok:false,error:err(s.reason)});\n' |
| ) |
| else: |
| pre += "const r=[];\n" + "".join( |
| "r.push(await invoke(()=>" + x + "));\n" for x in inv |
| ) |
| out = plan["output"] |
| if "fields" in out or "limit" in out: |
| pre += 'if(r[0].ok){try{let v=r[0].value; if(!Array.isArray(v))throw Error("result-schema: expected array");' |
| if "fields" in out: |
| pre += ( |
| "const fields=" |
| + j(out["fields"]) |
| + '; v=v.map(row=>{if(row===null||typeof row!=="object"||Array.isArray(row)||fields.some(k=>!Object.hasOwn(row,k)))throw Error("result-schema: missing field"); return Object.fromEntries(fields.map(k=>[k,row[k]]));});' |
| ) |
| if "limit" in out: |
| pre += f"v=v.slice(0,{out['limit']});" |
| pre += "r[0]={ok:true,value:v};}catch(e){r[0]={ok:false,error:err(e)}}}\n" |
| return pre + ("return r;" if out["kind"] == "all" else "return r[r.length-1];") |
|
|
