File size: 3,801 Bytes
c1f5d12 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 | """Shaktimaan runtime: route a task with the 0.5B router, then call the chosen agent(s) via OpenAI-compatible APIs.
Usage:
SHAKTIMAAN_ENDPOINTS='{"glm-5.3-flash": "http://host:8000/v1", ...}' \
python orchestrate.py [--image chart.png] [--code] "task"
"""
import argparse, asyncio, base64, json, mimetypes, os, sys
import httpx
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from common import (CODE_SYSTEM, POOL, ROUTER_SYSTEM, SOLVER_SYSTEM, STRONGEST, capable, endpoints, est_tokens,
max_out, router_view, strip_think)
ROUTER = os.environ.get("SHAKTIMAAN_ROUTER", "podstack/Shaktimaan")
ENDPOINTS = endpoints()
def route(task, has_image):
tok = AutoTokenizer.from_pretrained(ROUTER)
model = AutoModelForCausalLM.from_pretrained(ROUTER, dtype=torch.bfloat16)
enc = tok.apply_chat_template([{"role": "system", "content": ROUTER_SYSTEM},
{"role": "user", "content": router_view(task, has_image)}],
add_generation_prompt=True, return_tensors="pt", return_dict=True)
out = model.generate(**enc, max_new_tokens=96, do_sample=False)
try:
plan = json.loads(tok.decode(out[0, enc["input_ids"].shape[1]:], skip_special_tokens=True))
except Exception: # noqa: BLE001
plan = {"strategy": "single", "agents": []}
# keep only agents that are reachable and can actually take this input
ok = [a for a in plan.get("agents", []) if a in ENDPOINTS and capable(a, task, has_image)]
plan["agents"] = ok or [a for a in [STRONGEST, "kimi-vl"] if a in ENDPOINTS and capable(a, task, has_image)][:1]
return plan
async def ask(client, agent, messages, n_tokens):
r = await client.post(f"{ENDPOINTS[agent]}/chat/completions",
json={"model": POOL[agent]["model"], "messages": messages,
"max_tokens": max_out(agent, n_tokens)})
r.raise_for_status()
return strip_think(r.json()["choices"][0]["message"]["content"])
async def run(task, image=None, code=False):
has_image = image is not None
plan = route(task, has_image)
print("route:", json.dumps(plan), file=sys.stderr)
system = CODE_SYSTEM if code else SOLVER_SYSTEM
user = [{"type": "image_url", "image_url": {"url": image}}, {"type": "text", "text": task}] if has_image else task
msgs = [{"role": "system", "content": system}, {"role": "user", "content": user}]
n = est_tokens(task) + 1500 * has_image
async with httpx.AsyncClient(timeout=1800) as client:
if plan.get("strategy") != "ensemble" or len(plan["agents"]) == 1:
return await ask(client, plan["agents"][0], msgs, n)
drafts = await asyncio.gather(*[ask(client, a, msgs, n) for a in plan["agents"]])
joined = "\n\n".join(f"### Candidate {i + 1}\n{d}" for i, d in enumerate(drafts))
agg = f"{task}\n\nIndependent candidate solutions:\n{joined}\n\n" \
"Check them, resolve disagreements, and give the best final answer."
return await ask(client, plan["agents"][-1], [{"role": "system", "content": system},
{"role": "user", "content": agg}], est_tokens(agg))
if __name__ == "__main__":
p = argparse.ArgumentParser()
p.add_argument("task", nargs="+")
p.add_argument("--image", help="path to an image to attach")
p.add_argument("--code", action="store_true", help="ask agents for a Python code solution")
args = p.parse_args()
uri = None
if args.image:
mime = mimetypes.guess_type(args.image)[0] or "image/png"
uri = f"data:{mime};base64," + base64.b64encode(open(args.image, "rb").read()).decode()
print(asyncio.run(run(" ".join(args.task), uri, args.code)))
|