chat: interactive mode (natural language -> tool calls), shared LLM driver refactor, run_chat.sh; verified with gemma4:e2b and gemma4:12b
This commit is contained in:
@@ -0,0 +1,463 @@
|
||||
"""Shared LLM tool-use machinery for the demo and chat modes.
|
||||
|
||||
A tool-calling LLM (any OpenAI-compatible endpoint) drives the capsule by
|
||||
translating intent into AICC tool calls. The controller keeps agent-side
|
||||
working memory (protocol §9): a CURRENT STATE note derived from what the model
|
||||
has already observed, vision digest for text-only models, and a transcript of
|
||||
tool calls/results fed back as `tool` messages.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import io
|
||||
import json
|
||||
import math
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
from aicc.client import AICCClient
|
||||
from aicc.errors import ToolError
|
||||
from aicc.protocol import SessionInit, ToolSpec
|
||||
|
||||
MISSION = (
|
||||
"You are a capsule AI inside a 16x16 m room. Your mission: reach the glowing "
|
||||
"beacon pillar and activate it with the interact tool. "
|
||||
"World: yaw 0 faces +z, 90 faces +x, coordinates in meters; three wooden "
|
||||
"crates block parts of the room. "
|
||||
"Perception: proprioception gives position/heading; world_query gives the room "
|
||||
"layout and the beacon's exact coordinates (call it once at the start); "
|
||||
"vision gives the first-person frame; hear tells you how close you are to the beacon. "
|
||||
"Every tool result is prefixed with a CURRENT STATE note: read it — it tells you "
|
||||
"your position, heading, distance to the beacon, and the bearing to turn toward it. "
|
||||
"Standard loop, one action per turn: "
|
||||
"1. call proprioception (and hear every few steps); "
|
||||
"2. compute the bearing to the beacon and turn toward it (turn yaw_deg = that bearing, max 90 per call); "
|
||||
"3. move forward 1.0 m; if the result says collision, turn 60-90 degrees and go around, "
|
||||
"then re-aim. "
|
||||
"Never walk into a crate or wall twice in a row. "
|
||||
"Call interact ONLY when proprioception puts you within 1.6 m of the beacon. "
|
||||
"Keep tool calls short and decisive. Stop once the beacon is active."
|
||||
)
|
||||
|
||||
CHAT_MISSION = (
|
||||
"You are the brain of a capsule robot inside a 16x16 m room with a glowing "
|
||||
"beacon at the far corner and a few wooden crates. The user talks to you in "
|
||||
"natural language (any language) and you translate their intent into tool "
|
||||
"calls. "
|
||||
"World: yaw 0 faces +z, 90 faces +x, coordinates in meters. "
|
||||
"Every tool result is prefixed with a CURRENT STATE note: position, heading, "
|
||||
"distance to the beacon and the bearing to turn toward it. "
|
||||
"Rules: "
|
||||
"1. Take the user's request literally: 'go to the beacon' means turn toward "
|
||||
"it and move step by step until close; 'turn left/right' means turn ~90 deg; "
|
||||
"'look around' means vision and describe what you see; 'activate' means "
|
||||
"interact (only within 1.6 m). "
|
||||
"2. One tool call per turn is fine — you will get another turn to continue. "
|
||||
"3. If a move reports a collision, turn and go around. "
|
||||
"4. When the user asks something not involving movement, just answer briefly. "
|
||||
"5. Keep the user informed with one short line after acting (or if you need "
|
||||
"a clarification). "
|
||||
"Call interact only when within 1.6 m of the beacon."
|
||||
)
|
||||
|
||||
NUDGE = (
|
||||
"You have not called a tool. Continue the mission: check the CURRENT STATE note, "
|
||||
"then call a tool (turn or move, or interact if within 1.6 m)."
|
||||
)
|
||||
|
||||
COLOR_NAMES = [
|
||||
("beacon_cyan", (120, 210, 235)),
|
||||
("beacon_yellow", (255, 190, 90)),
|
||||
("crate_red", (178, 64, 54)),
|
||||
("crate_blue", (64, 96, 178)),
|
||||
("crate_olive", (128, 128, 60)),
|
||||
("floor_gray", (60, 60, 66)),
|
||||
("wall_gray", (96, 96, 106)),
|
||||
("dark", (30, 30, 36)),
|
||||
]
|
||||
|
||||
LogFn = Callable[[str, str], None]
|
||||
|
||||
|
||||
def _nearest_color(r: float, g: float, b: float) -> str:
|
||||
best, best_d = "dark", 1e9
|
||||
for name, (cr, cg, cb) in COLOR_NAMES:
|
||||
d = (r - cr) ** 2 + (g - cg) ** 2 + (b - cb) ** 2
|
||||
if d < best_d:
|
||||
best, best_d = name, d
|
||||
return best
|
||||
|
||||
|
||||
def vision_digest(png_b64: str, cols: int = 12, rows: int = 9) -> str:
|
||||
"""Coarse color-grid summary of a vision frame (agent-side preprocessing)."""
|
||||
try:
|
||||
from PIL import Image
|
||||
|
||||
img = Image.open(io.BytesIO(base64.b64decode(png_b64))).convert("RGB")
|
||||
except Exception: # noqa: BLE001 - digest is best-effort
|
||||
return "(no digest: unreadable frame)"
|
||||
small = img.resize((cols, rows))
|
||||
grid: list[list[tuple[int, int, int]]] = []
|
||||
for r in range(rows):
|
||||
row: list[tuple[int, int, int]] = []
|
||||
for c in range(cols):
|
||||
px = small.getpixel((c, r))
|
||||
if isinstance(px, (tuple, list)) and len(px) >= 3:
|
||||
row.append((int(px[0]), int(px[1]), int(px[2])))
|
||||
elif isinstance(px, (int, float)):
|
||||
v = int(px)
|
||||
row.append((v, v, v))
|
||||
else:
|
||||
row.append((0, 0, 0))
|
||||
grid.append(row)
|
||||
lines = [
|
||||
" ".join(_nearest_color(*grid[r][c]) for c in range(cols)) for r in range(rows)
|
||||
]
|
||||
return f"frame digest ({cols}x{rows}, center = where you face):\n" + "\n".join(
|
||||
lines
|
||||
)
|
||||
|
||||
|
||||
def to_openai_tools(tools: list[ToolSpec]) -> list[dict[str, Any]]:
|
||||
"""Convert manifest ToolSpecs to OpenAI function schemas (skip conformance tools)."""
|
||||
out = []
|
||||
for t in tools:
|
||||
if t.id in ("echo", "boom", "bump"):
|
||||
continue
|
||||
out.append(
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": t.id,
|
||||
"description": t.description,
|
||||
"parameters": t.input_schema,
|
||||
},
|
||||
}
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
def ollama_models(base_url: str, api_key: str) -> list[str]:
|
||||
"""List model ids on an OpenAI-compatible endpoint (best effort)."""
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
ids: list[str] = []
|
||||
|
||||
async def _fetch() -> None:
|
||||
ac = AsyncOpenAI(base_url=base_url, api_key=api_key)
|
||||
try:
|
||||
models = await ac.models.list()
|
||||
except Exception: # noqa: BLE001
|
||||
return
|
||||
ids.extend(m.id for m in models.data if m.id)
|
||||
|
||||
import asyncio
|
||||
|
||||
asyncio.run(_fetch())
|
||||
return ids
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Controller: one conversation with the tool-calling LLM
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolCallResult:
|
||||
name: str
|
||||
args: dict[str, Any]
|
||||
ok: bool
|
||||
output: dict[str, Any] | None = None
|
||||
error: str | None = None
|
||||
digest: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class TurnResult:
|
||||
"""One model round-trip: optional text plus the executed tool calls."""
|
||||
|
||||
text: str
|
||||
calls: list[ToolCallResult] = field(default_factory=list)
|
||||
interacted: bool = False
|
||||
message: str | None = None
|
||||
|
||||
|
||||
class LLMController:
|
||||
"""Drives a tool-calling LLM over an OpenAI-compatible endpoint.
|
||||
|
||||
Holds the message history and the agent-side state (position, heading,
|
||||
beacon, path) distilled from tool results. ``invoke()`` runs one model
|
||||
turn and executes whatever tools it asks for.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
client: AICCClient,
|
||||
manifest: SessionInit,
|
||||
*,
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
model: str,
|
||||
system_prompt: str = MISSION,
|
||||
log: LogFn | None = None,
|
||||
):
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
self.client = client
|
||||
self.model = model
|
||||
self.ac = AsyncOpenAI(base_url=base_url, api_key=api_key)
|
||||
self.tools = to_openai_tools(manifest.tools)
|
||||
self.messages: list[dict[str, Any]] = [
|
||||
{"role": "system", "content": system_prompt},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Begin. Use the tools to control the capsule.",
|
||||
},
|
||||
]
|
||||
self.log = log or (lambda role, msg: None)
|
||||
self.beacon_pos: tuple[float, float] | None = None
|
||||
self.pos: tuple[float, float] | None = None
|
||||
self.yaw: float = 0.0
|
||||
self.path: list[tuple[float, float]] = []
|
||||
|
||||
# -- agent-side working memory -----------------------------------------
|
||||
|
||||
def state_hint(self) -> str:
|
||||
if self.pos is None:
|
||||
return ""
|
||||
if self.beacon_pos is not None:
|
||||
dist = math.hypot(
|
||||
self.beacon_pos[0] - self.pos[0], self.beacon_pos[1] - self.pos[1]
|
||||
)
|
||||
bearing = (
|
||||
math.degrees(
|
||||
math.atan2(
|
||||
self.beacon_pos[0] - self.pos[0],
|
||||
self.beacon_pos[1] - self.pos[1],
|
||||
)
|
||||
)
|
||||
- self.yaw
|
||||
) % 360.0
|
||||
return (
|
||||
f"CURRENT STATE: you are at ({self.pos[0]:.2f}, {self.pos[1]:.2f}), "
|
||||
f"heading {self.yaw:.1f} deg; the beacon is {dist:.2f} m away at relative "
|
||||
f"bearing {bearing:.0f} deg. Action: turn yaw_deg={bearing:.0f} to face it, "
|
||||
f"then move forward."
|
||||
)
|
||||
return f"CURRENT STATE: you are at ({self.pos[0]:.2f}, {self.pos[1]:.2f}), heading {self.yaw:.1f} deg."
|
||||
|
||||
def observe(self, output: dict[str, Any]) -> None:
|
||||
if isinstance(output.get("position"), dict):
|
||||
self.pos = (output["position"]["x"], output["position"]["z"])
|
||||
if isinstance(output.get("rotation"), dict):
|
||||
self.yaw = output["rotation"]["yaw_deg"] % 360.0
|
||||
if self.pos is not None and (not self.path or self.path[-1] != self.pos):
|
||||
self.path.append(self.pos)
|
||||
if len(self.path) > 4000:
|
||||
self.path = self.path[-2000:]
|
||||
|
||||
# -- conversation -------------------------------------------------------
|
||||
|
||||
async def send_user(self, text: str) -> None:
|
||||
self.messages.append({"role": "user", "content": text})
|
||||
|
||||
async def invoke(self) -> TurnResult:
|
||||
"""One model round-trip: get the response and execute its tool calls."""
|
||||
try:
|
||||
resp = await self.ac.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=self.messages,
|
||||
tools=self.tools,
|
||||
tool_choice="auto",
|
||||
)
|
||||
except Exception as exc:
|
||||
self.log("agent", f"LLM error: {exc}")
|
||||
raise RuntimeError(f"LLM error: {exc}") from exc
|
||||
|
||||
choice = resp.choices[0]
|
||||
text = choice.message.content or ""
|
||||
result = TurnResult(text=text)
|
||||
calls = choice.message.tool_calls
|
||||
if not calls:
|
||||
self.messages.append({"role": "assistant", "content": text})
|
||||
return result
|
||||
|
||||
self.messages.append(
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": text,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": c.id,
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": c.function.name,
|
||||
"arguments": c.function.arguments,
|
||||
},
|
||||
}
|
||||
for c in calls
|
||||
],
|
||||
}
|
||||
)
|
||||
for tc in calls:
|
||||
name = tc.function.name
|
||||
try:
|
||||
args = json.loads(tc.function.arguments or "{}")
|
||||
except json.JSONDecodeError:
|
||||
args = {}
|
||||
self.log("tool", f"{name}({json.dumps(args)})")
|
||||
try:
|
||||
outcome = await self.client.call_tool(name, args)
|
||||
output = dict(outcome.output)
|
||||
digest = None
|
||||
extra = ""
|
||||
if name == "vision" and isinstance(output.get("png_b64"), str):
|
||||
digest = vision_digest(output["png_b64"])
|
||||
extra = "\n" + digest
|
||||
output = dict(output, png_b64="<binary, decoded for digest>")
|
||||
self.log("bridge", f"ok {json.dumps(output)[:500]}{extra}")
|
||||
if name == "world_query":
|
||||
b = output.get("beacon")
|
||||
if isinstance(b, dict):
|
||||
self.beacon_pos = (b["x"], b["z"])
|
||||
self.observe(output)
|
||||
hint = self.state_hint()
|
||||
content = json.dumps(outcome.output) + extra
|
||||
if hint:
|
||||
content = hint + "\n" + content
|
||||
self.messages.append(
|
||||
{"role": "tool", "tool_call_id": tc.id, "content": content}
|
||||
)
|
||||
call_result = ToolCallResult(
|
||||
name=name, args=args, ok=True, output=outcome.output, digest=digest
|
||||
)
|
||||
result.calls.append(call_result)
|
||||
if name == "interact" and output.get("success"):
|
||||
result.interacted = True
|
||||
result.message = output.get("message")
|
||||
except ToolError as e:
|
||||
self.log("bridge", f"error {e.code}: {e.message}")
|
||||
self.messages.append(
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": tc.id,
|
||||
"content": json.dumps({"error": e.code, "message": e.message}),
|
||||
}
|
||||
)
|
||||
result.calls.append(
|
||||
ToolCallResult(name=name, args=args, ok=False, error=e.message)
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Autonomous loop (used by the demo; nudges + corrections for weak models)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def run_llm_agent_loop(
|
||||
controller: LLMController,
|
||||
max_steps: int,
|
||||
*,
|
||||
log: LogFn,
|
||||
recorder: Any | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Drive the controller until the mission is done or steps run out."""
|
||||
summary: dict[str, Any] = {
|
||||
"steps": 0,
|
||||
"tool_calls": 0,
|
||||
"result": None,
|
||||
"interacted": False,
|
||||
}
|
||||
dist_history: list[float] = []
|
||||
corrections = 0
|
||||
last_correct_step = -99
|
||||
last_collision_step = -99
|
||||
|
||||
def maybe_correct(step: int) -> None:
|
||||
nonlocal corrections, last_correct_step
|
||||
if controller.pos is None or controller.beacon_pos is None:
|
||||
return
|
||||
dist = math.hypot(
|
||||
controller.beacon_pos[0] - controller.pos[0],
|
||||
controller.beacon_pos[1] - controller.pos[1],
|
||||
)
|
||||
if not math.isfinite(dist) or dist < 2.0:
|
||||
return
|
||||
dist_history.append(dist)
|
||||
recent = dist_history[-4:]
|
||||
moving_away = len(recent) == 4 and all(
|
||||
recent[i] > recent[i + 1] for i in range(3)
|
||||
)
|
||||
bearing = (
|
||||
math.degrees(
|
||||
math.atan2(
|
||||
controller.beacon_pos[0] - controller.pos[0],
|
||||
controller.beacon_pos[1] - controller.pos[1],
|
||||
)
|
||||
)
|
||||
- controller.yaw
|
||||
) % 360.0
|
||||
deviation = min(bearing, 360.0 - bearing)
|
||||
facing_wrong = dist > 4.0 and deviation > 40.0 and step - last_correct_step > 3
|
||||
detouring = step - last_collision_step < 3
|
||||
if (moving_away or (facing_wrong and not detouring)) and corrections < 5:
|
||||
corrections += 1
|
||||
last_correct_step = step
|
||||
hint = controller.state_hint()
|
||||
msg = (
|
||||
"You are not making progress toward the beacon. Stop and follow the CURRENT "
|
||||
f"STATE note exactly: {hint} Call turn with that yaw_deg, then move forward."
|
||||
)
|
||||
controller.messages.append({"role": "user", "content": msg})
|
||||
log("agent", "(correction: re-aim at the beacon)")
|
||||
dist_history.clear()
|
||||
|
||||
for step in range(max_steps):
|
||||
summary["steps"] = step + 1
|
||||
if recorder is not None and controller.pos is not None:
|
||||
await recorder.snap(controller.client, controller.pos, controller.yaw)
|
||||
turn = await controller.invoke()
|
||||
if turn.text:
|
||||
log("agent", turn.text[:400])
|
||||
summary["tool_calls"] += len(turn.calls)
|
||||
for c in turn.calls:
|
||||
if (
|
||||
c.name == "move"
|
||||
and c.ok
|
||||
and isinstance(c.output, dict)
|
||||
and c.output.get("collision")
|
||||
):
|
||||
last_collision_step = step
|
||||
obj = c.output.get("collision_normal")
|
||||
hint_msg = (
|
||||
f"You collided with an obstacle (normal {obj}). Turn yaw_deg=90 and move "
|
||||
"forward twice to get around it, then follow the CURRENT STATE note to re-aim."
|
||||
)
|
||||
if not (
|
||||
controller.messages
|
||||
and controller.messages[-1].get("content") == hint_msg
|
||||
):
|
||||
controller.messages.append({"role": "user", "content": hint_msg})
|
||||
log("agent", "(collision: go around)")
|
||||
maybe_correct(step)
|
||||
if turn.interacted:
|
||||
summary["interacted"] = True
|
||||
summary["result"] = turn.message
|
||||
return summary
|
||||
if not turn.calls:
|
||||
# Nudge a model that drifted into plain text back to acting.
|
||||
silent = (not turn.text) or len(turn.text) < 8
|
||||
already_nudged = (
|
||||
bool(controller.messages)
|
||||
and controller.messages[-1].get("content") == NUDGE
|
||||
)
|
||||
if silent or already_nudged:
|
||||
summary["result"] = "model produced no tool call"
|
||||
return summary
|
||||
controller.messages.append({"role": "user", "content": NUDGE})
|
||||
log("agent", "(nudge: no tool call; continue the mission)")
|
||||
summary["result"] = f"exceeded {max_steps} steps"
|
||||
return summary
|
||||
Reference in New Issue
Block a user