File size: 4,523 Bytes
df43f42 | 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 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | """
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:] # ONLY the newly generated tokens
ctx = ids # full updated context (no dup)
text_new = tok.decode(new_ids)
print(text_new, flush=True)
# 1) model-driven tool calls (incl. retrieve) -> mid-response injection
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 # keep going within the same answer
# 2) controller-driven RAG: inject context mid-response, then 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 # nothing to do -> answer is final
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)
|