92 lines
2.6 KiB
Python
92 lines
2.6 KiB
Python
from dataclasses import dataclass
|
|
|
|
from openai import OpenAI
|
|
from pydantic import BaseModel
|
|
|
|
import config
|
|
from config import SERVOS, HOME, clamp_pose
|
|
from presets import EXAMPLES
|
|
|
|
MAX_STEPS = 16
|
|
MIN_MS, MAX_MS = 80, 6000
|
|
|
|
|
|
class Step(BaseModel):
|
|
s1: int
|
|
s2: int
|
|
s3: int
|
|
s4: int
|
|
ms: int
|
|
|
|
|
|
class Reply(BaseModel):
|
|
summary: str
|
|
steps: list[Step]
|
|
|
|
|
|
@dataclass
|
|
class Plan:
|
|
summary: str
|
|
steps: list # [(pose, ms), ...]
|
|
|
|
|
|
def _fmt(steps):
|
|
return "\n".join(f" {list(pose)} over {ms}ms" for pose, ms in steps)
|
|
|
|
|
|
def _instructions():
|
|
servos = "\n".join(
|
|
f"- {s.key} ({s.name}): {s.role}. range {s.lo}-{s.hi}, home {s.home}"
|
|
for s in SERVOS
|
|
)
|
|
examples = "\n\n".join(f'"{name}":\n{_fmt(steps)}' for name, steps in EXAMPLES.items())
|
|
return f"""You drive a small 4-servo robot arm. Turn the user's spoken command into keyframes.
|
|
|
|
Servos (angles in degrees):
|
|
{servos}
|
|
|
|
Each step is an absolute pose [s1, s2, s3, s4] reached by moving all servos together over `ms` milliseconds.
|
|
Repeat the same pose with a duration to hold still.
|
|
|
|
Rules:
|
|
- Stay inside each servo's range.
|
|
- {MIN_MS}-{MAX_MS} ms per step, at most {MAX_STEPS} steps. ~4 degrees per 30 ms looks natural; go faster only for snappy moves.
|
|
- Unless the user asks to hold a position, finish back at home {list(HOME)}.
|
|
- Relative commands ("a bit more left", "lower") start from the current pose.
|
|
- If the command has nothing to do with moving the arm, return no steps and say why in the summary.
|
|
- summary: a few words describing the motion, in the user's language.
|
|
|
|
Reference moves:
|
|
|
|
{examples}"""
|
|
|
|
|
|
class Planner:
|
|
def __init__(self):
|
|
if not config.OPENAI_API_KEY:
|
|
raise RuntimeError("OPENAI_API_KEY is not set")
|
|
self._client = OpenAI(api_key=config.OPENAI_API_KEY)
|
|
|
|
def plan(self, text, current_pose):
|
|
kwargs = {}
|
|
if config.OPENAI_REASONING_EFFORT:
|
|
kwargs["reasoning"] = {"effort": config.OPENAI_REASONING_EFFORT}
|
|
|
|
resp = self._client.responses.parse(
|
|
model=config.OPENAI_MODEL,
|
|
instructions=_instructions(),
|
|
input=f"Current pose: {list(current_pose)}\nCommand: {text}",
|
|
text_format=Reply,
|
|
timeout=30,
|
|
**kwargs,
|
|
)
|
|
reply = resp.output_parsed
|
|
if reply is None:
|
|
raise RuntimeError("model returned nothing usable")
|
|
|
|
steps = [
|
|
(clamp_pose((s.s1, s.s2, s.s3, s.s4)), max(MIN_MS, min(MAX_MS, s.ms)))
|
|
for s in reply.steps[:MAX_STEPS]
|
|
]
|
|
return Plan(reply.summary.strip(), steps)
|