464 lines
17 KiB
Python
464 lines
17 KiB
Python
"""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
|