integration/pi_tool_router.py

"""Route a short tool intent to Pi's live tool schemas with Needle."""

import json
import sys
from pathlib import Path

from needle import SimpleAttentionNetwork, generate, get_tokenizer, load_checkpoint


MODEL_ROOT = Path(__file__).resolve().parents[1]
params, config = load_checkpoint(MODEL_ROOT / "needle-pi-coding-agent.pkl")
model = SimpleAttentionNetwork(config)
tokenizer = get_tokenizer()


def main() -> None:
    request = json.load(sys.stdin)
    result = generate(
        model,
        params,
        tokenizer,
        query=request["intent"],
        tools=json.dumps(request["tools"], separators=(",", ":")),
        stream=False,
    )
    calls = json.loads(result)
    if not isinstance(calls, list) or not calls:
        raise ValueError("Needle returned no tool calls")
    # Exa plans one step at a time. Running only Needle's first call keeps Pi's
    # real tool result in the loop before either model chooses another action.
    json.dump(calls[:1], sys.stdout, separators=(",", ":"))


if __name__ == "__main__":
    main()