234 lines
8.2 KiB
Python
234 lines
8.2 KiB
Python
"""Interactive chat mode: talk to the model, it drives the capsule.
|
|
|
|
Type natural-language commands (any language) and the LLM translates them
|
|
into AICC tool calls: 'go to the beacon', 'turn left', 'look around',
|
|
'activate the beacon'. Every tool call and result is printed as a transcript,
|
|
the current vision frame is saved to ``chat_frame.png`` and the sensor-built
|
|
top-down map to ``chat_map.png`` after each turn. Works great together with
|
|
the live viewer (scripts/run_live.sh) — watch the capsule in your browser
|
|
while you talk to it.
|
|
|
|
Commands:
|
|
/state print current position/heading (proprioception)
|
|
/look take a vision frame and save it
|
|
/map save the top-down map (chat_map.png)
|
|
/models list models on the endpoint
|
|
/model X switch the model mid-session
|
|
/steps N auto-continue budget for 'go to X' style requests
|
|
/help this text
|
|
/exit quit
|
|
|
|
Usage:
|
|
python -m testbed.chat [--model gemma4:e2b] [--model gemma4:12b]
|
|
python -m testbed.chat --base-url https://api.openai.com/v1 --model gpt-4o-mini
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import asyncio
|
|
import sys
|
|
|
|
from aicc.client import AICCClient
|
|
from aicc.transport.websocket import WebSocketClientTransport
|
|
|
|
from testbed.llm_agent import CHAT_MISSION, LLMController, ollama_models
|
|
|
|
DEFAULT_URL = "ws://127.0.0.1:8765"
|
|
DEFAULT_BASE_URL = "http://localhost:11434/v1"
|
|
DEFAULT_MODEL = "gemma4:e2b"
|
|
|
|
HELP = """\
|
|
You are talking to the capsule's brain (LLM). Type what you want, e.g.:
|
|
'иди к маяку' / 'go to the beacon' — the model walks there step by step
|
|
'повернись налево' / 'turn left' — quick turn
|
|
'осмотрись' / 'look around' — vision + description
|
|
'активируй маяк' / 'activate it' — interact (must be close)
|
|
Commands: /state /look /map /models /model X /steps N /help /exit"""
|
|
|
|
|
|
async def chat_loop(client: AICCClient, manifest, args: argparse.Namespace) -> int:
|
|
model = args.model
|
|
controller = LLMController(
|
|
client,
|
|
manifest,
|
|
base_url=args.base_url,
|
|
api_key=args.api_key,
|
|
model=model,
|
|
system_prompt=CHAT_MISSION,
|
|
log=lambda role, msg: print(f" [{role}] {msg}"),
|
|
)
|
|
auto_steps = args.auto_steps
|
|
print(f"[chat] model: {model} (endpoint {args.base_url})")
|
|
print("[chat] type your commands; /help for the command list; /exit to quit\n")
|
|
|
|
async def cmd_state() -> None:
|
|
out = (await client.call_tool("proprioception", {})).output
|
|
print(
|
|
f" pos ({out['position']['x']:.2f}, {out['position']['z']:.2f}) "
|
|
f"heading {out['rotation']['yaw_deg']:.1f} deg health {out['health']}"
|
|
)
|
|
|
|
async def cmd_look() -> None:
|
|
out = (await client.call_tool("vision", {})).output
|
|
import base64
|
|
|
|
raw = base64.b64decode(out["png_b64"])
|
|
|
|
def _write() -> None:
|
|
with open("chat_frame.png", "wb") as fh:
|
|
fh.write(raw)
|
|
|
|
await asyncio.to_thread(_write)
|
|
print(f" frame saved to chat_frame.png ({out['width']}x{out['height']})")
|
|
|
|
def save_map() -> None:
|
|
if controller.pos is None:
|
|
print(" (no position yet — ask the capsule to move first)")
|
|
return
|
|
from testbed.room.mapview import draw_sensor_map
|
|
|
|
layout = {"obstacles": [], "beacon": {"x": 12.5, "z": 12.5}}
|
|
if controller.beacon_pos is not None:
|
|
layout["beacon"] = {
|
|
"x": controller.beacon_pos[0],
|
|
"z": controller.beacon_pos[1],
|
|
}
|
|
img = draw_sensor_map(layout, controller.pos, controller.yaw, controller.path)
|
|
img.save("chat_map.png")
|
|
print(
|
|
f" map saved to chat_map.png (pos {controller.pos[0]:.1f}, {controller.pos[1]:.1f})"
|
|
)
|
|
|
|
async def cmd_models() -> None:
|
|
print(" fetching models…")
|
|
for m in ollama_models(args.base_url, args.api_key):
|
|
print(f" {m}")
|
|
|
|
while True:
|
|
try:
|
|
text = input("you> ").strip()
|
|
except (EOFError, KeyboardInterrupt):
|
|
print("\n[chat] bye")
|
|
return 0
|
|
if not text:
|
|
continue
|
|
low = text.lower()
|
|
if low in ("/exit", "exit", "quit", "выход"):
|
|
print("[chat] bye")
|
|
return 0
|
|
if low == "/help":
|
|
print(HELP)
|
|
continue
|
|
if low == "/state":
|
|
await cmd_state()
|
|
continue
|
|
if low == "/look":
|
|
await cmd_look()
|
|
continue
|
|
if low == "/map":
|
|
save_map()
|
|
continue
|
|
if low == "/models":
|
|
await cmd_models()
|
|
continue
|
|
if low.startswith("/model "):
|
|
model = low.split(maxsplit=1)[1]
|
|
controller.model = model
|
|
print(f"[chat] switching to {model}")
|
|
continue
|
|
if low.startswith("/steps "):
|
|
try:
|
|
auto_steps = max(0, int(low.split(maxsplit=1)[1]))
|
|
print(f"[chat] auto-continue budget: {auto_steps} steps")
|
|
except ValueError:
|
|
print(" usage: /steps N")
|
|
continue
|
|
|
|
controller.messages.append({"role": "user", "content": text})
|
|
steps_used = 0
|
|
while True:
|
|
try:
|
|
turn = await controller.invoke()
|
|
except RuntimeError as exc:
|
|
print(f" [error] {exc}")
|
|
break
|
|
steps_used += 1
|
|
if turn.text:
|
|
print(f" model> {turn.text}")
|
|
if turn.interacted:
|
|
print(f"\n[chat] BEACON ACTIVATED: {turn.message}")
|
|
save_map()
|
|
return 0
|
|
if not turn.calls:
|
|
if not turn.text and auto_steps and steps_used <= auto_steps:
|
|
# The model answered with nothing: nudge it to actually act.
|
|
controller.messages.append(
|
|
{
|
|
"role": "user",
|
|
"content": (
|
|
"Your last turn contained no action and no answer. "
|
|
"Carry out the user's request: call a tool now "
|
|
"(check the CURRENT STATE note)."
|
|
),
|
|
}
|
|
)
|
|
print(" [chat] (nudge: model gave an empty turn)")
|
|
continue
|
|
break # the model answered in words; wait for the user
|
|
if auto_steps and steps_used >= auto_steps:
|
|
print(
|
|
f" [chat] (auto-continue budget of {auto_steps} reached — say 'continue' to keep going)"
|
|
)
|
|
break
|
|
if not auto_steps:
|
|
break
|
|
# The model acted without commenting; let it keep going for 'go to X'
|
|
if controller.pos is not None:
|
|
save_map()
|
|
|
|
return 0
|
|
|
|
|
|
async def run(args: argparse.Namespace) -> int:
|
|
transport = WebSocketClientTransport(args.url)
|
|
async with AICCClient(transport) as client:
|
|
manifest = await client.handshake()
|
|
print(f"[handshake] session {manifest.session_id} world {manifest.world.name}")
|
|
print(f"[handshake] tools: {[t.id for t in manifest.tools]}")
|
|
return await chat_loop(client, manifest, args)
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(
|
|
description="Interactive chat: talk to the LLM, it drives the capsule."
|
|
)
|
|
parser.add_argument(
|
|
"--url",
|
|
default=DEFAULT_URL,
|
|
help=f"bridge WebSocket URL (default {DEFAULT_URL})",
|
|
)
|
|
parser.add_argument(
|
|
"--base-url", default=DEFAULT_BASE_URL, help="OpenAI-compatible endpoint"
|
|
)
|
|
parser.add_argument(
|
|
"--model", default=DEFAULT_MODEL, help="model id on the endpoint"
|
|
)
|
|
parser.add_argument("--api-key", default="ollama", help="API key for the endpoint")
|
|
parser.add_argument(
|
|
"--auto-steps",
|
|
type=int,
|
|
default=8,
|
|
help="how many tool steps the model may chain per request (0 = one action per turn)",
|
|
)
|
|
args = parser.parse_args()
|
|
try:
|
|
return asyncio.run(run(args))
|
|
except KeyboardInterrupt:
|
|
print("\n[chat] bye")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|