| """ |
| clankerDiffusion — tool-using agent (ReAct loop) with mid-response RAG. |
| |
| The model emits <tool name="...">args</tool> blocks. We parse + execute |
| them locally (tools.py) and feed the <result> back, looping until the |
| model stops calling tools. Each turn alternates AR <-> diffusion, so the |
| agent is literally "sometimes normal, sometimes diffusion". |
| |
| Knowledge injection happens in TWO ways, both mid-response: |
| * model-driven : the model calls <tool name="retrieve">q</tool>; we run the |
| retriever and feed <result> back, then keep generating. |
| * controller-driven : the RAG controller watches the partial response, builds |
| a query from the user question + what was just said, and if it finds |
| relevant passages it injects <context>...</context> straight into the |
| stream and the model continues -- no need for the model to ask. |
| """ |
| import argparse |
| import infer |
| from tools import execute_tool, parse_tool_calls |
| from rag import default_kb |
|
|
|
|
| SYSTEM = ("You are clanker, a helpful assistant that can THINK, USE TOOLS, and " |
| "READ CONTEXT. When you need to compute, inspect files, or look up " |
| "knowledge, emit a tool call like " |
| "<tool name=\"calc\">17 * 23</tool>, " |
| "<tool name=\"python\">print(2**10)</tool>, " |
| "<tool name=\"read_file\">notes.txt</tool>, " |
| "<tool name=\"list_dir\">.</tool>, or " |
| "<tool name=\"retrieve\">the capital of France</tool>. " |
| "Available tools: calc(expr), python(code), read_file(path), " |
| "list_dir(path), retrieve(query). If <context>...</context> is " |
| "provided, use it to answer. After a <result>...</result> appears, " |
| "continue and give the final answer. You may reason in <think>...</think>.") |
|
|
|
|
| def run_agent(prompt, max_turns=8, mode_cycle=("diff", "ar"), |
| max_new=200, gen_len=160, steps=24, |
| rag=True, rag_k=3, rag_budget=2): |
| model, tok = infer._MODEL, infer._TOK |
| kb = default_kb() |
| injected = set() |
|
|
| ctx = [tok.bos_id] + tok.encode( |
| f"<system>{SYSTEM}</system><user>{prompt}</user><assistant>") |
|
|
| print(f"\n=== USER ===\n{prompt}\n") |
| for turn in range(max_turns): |
| mode = mode_cycle[turn % len(mode_cycle)] |
| prev_len = len(ctx) |
| print(f"--- turn {turn} [{mode}] ---", flush=True) |
| if mode == "ar": |
| ids = infer.generate_ar(model, tok, ctx, max_new=max_new) |
| else: |
| ids = infer.generate_diff(model, tok, ctx, gen_len=gen_len, steps=steps) |
| new_ids = ids[prev_len:] |
| ctx = ids |
| text_new = tok.decode(new_ids) |
| print(text_new, flush=True) |
|
|
| |
| calls = parse_tool_calls(text_new) |
| if calls: |
| for name, arg in calls: |
| print(f" [tool] {name}({arg[:120]})", flush=True) |
| res = execute_tool(name, arg) |
| print(f" [result] {res[:300]}", flush=True) |
| ctx = ctx + tok.encode(f"<result>{res}</result>") |
| continue |
|
|
| |
| if rag and len(injected) < rag_budget * rag_k: |
| q = prompt + " " + text_new[-300:] |
| fresh = [p for p in kb.retrieve(q, k=rag_k) if p not in injected] |
| if fresh: |
| injected.update(fresh) |
| block = "<context>\n" + "\n---\n".join(fresh) + "\n</context>" |
| ctx = ctx + tok.encode(block) |
| print(f" [rag] injected {len(fresh)} passage(s) mid-response", |
| flush=True) |
| continue |
|
|
| break |
|
|
| final = tok.decode(ctx) |
| i = final.rfind("<assistant>") |
| print("\n=== clanker (final) ===\n" + (final[i + 10:] if i >= 0 else final)) |
| return final |
|
|
|
|
| if __name__ == "__main__": |
| ap = argparse.ArgumentParser() |
| ap.add_argument("prompt", nargs="?", default="What is 17 * 23, and what files are in the current folder?") |
| ap.add_argument("--ckpt", default=None) |
| ap.add_argument("--turns", type=int, default=8) |
| ap.add_argument("--no-rag", action="store_true") |
| a = ap.parse_args() |
| infer.init(ckpt_path=a.ckpt) |
| run_agent(a.prompt, max_turns=a.turns, rag=not a.no_rag) |
|
|
