diff --git a/README.md b/README.md index 7cf0d42..816871b 100644 --- a/README.md +++ b/README.md @@ -184,10 +184,12 @@ python -m flow schedule ## Agentic Editing (new in 0.3) Beyond the headless pipeline, Flow ships an **in-app video agent** — an LLM that -operates your project through 41 tools (plan scenes, generate/regenerate, reorder, -trim, keyframes, captions, color, cast characters, narrate, batch-generate…). The -ordered scenes *are* the timeline; ffmpeg assembles them with narration/caption -tracks. +operates your project through 42 tools (plan scenes, generate/regenerate, reorder, +trim, keyframes, captions, color, cast/create characters, narrate, batch-generate…). +The ordered scenes *are* the timeline; ffmpeg assembles them with narration/caption +tracks. The agent **streams its reply token-by-token** over SSE (`token` events), +with `tool_start`/`tool_result` events as each tool runs — so a UI can type the +answer out live and show tool activity as it happens. ```bash # Run the agent API (default model: kimi via NVIDIA build) diff --git a/src/flow/agent/loop.py b/src/flow/agent/loop.py index bfb29d0..978d70b 100644 --- a/src/flow/agent/loop.py +++ b/src/flow/agent/loop.py @@ -51,7 +51,11 @@ def __init__(self, client: NvidiaClient, ctx: ToolContext, def run_stream(self, user_message: str, history: list[dict] | None = None): """Run one turn, yielding events as they happen (for SSE). - Events: {type: assistant|tool_call|reply, ...}. Final event is 'reply'. + Events: + - ``{type: token, content}`` — a streamed reply text delta + - ``{type: tool_start, tool, args}`` — a tool is about to run + - ``{type: tool_result, tool, result}`` — that tool's result + - ``{type: reply, content}`` — the final assistant answer (terminal) """ messages: list[dict] = [{"role": "system", "content": build_system_prompt(self.ctx)}] messages += history or [] @@ -59,13 +63,16 @@ def run_stream(self, user_message: str, history: list[dict] | None = None): tools = openai_schemas() for _ in range(self.max_iterations): - resp = self.client.chat(messages, tools=tools) - msg = resp["choices"][0]["message"] + final_msg: dict | None = None + for ev in self.client.chat_stream(messages, tools=tools): + if ev["type"] == "content": + yield {"type": "token", "content": ev["text"]} + elif ev["type"] == "final": + final_msg = ev["message"] + msg = final_msg or {"role": "assistant", "content": ""} messages.append(msg) tool_calls = msg.get("tool_calls") or [] - if msg.get("content") and tool_calls: - yield {"type": "assistant", "content": msg["content"]} if not tool_calls: yield {"type": "reply", "content": msg.get("content", "")} return @@ -76,8 +83,9 @@ def run_stream(self, user_message: str, history: list[dict] | None = None): args = json.loads(tc["function"]["arguments"] or "{}") except json.JSONDecodeError: args = {} + yield {"type": "tool_start", "tool": name, "args": args} res = dispatch(self.ctx, name, args) - yield {"type": "tool_call", "tool": name, "args": args, "result": res} + yield {"type": "tool_result", "tool": name, "args": args, "result": res} messages.append({"role": "tool", "tool_call_id": tc["id"], "content": json.dumps(res)}) @@ -86,10 +94,15 @@ def run_stream(self, user_message: str, history: list[dict] | None = None): def run(self, user_message: str, history: list[dict] | None = None) -> dict: """Run one user turn to completion (collects the stream).""" calls_made: list[dict] = [] - reply = "" + parts: list[str] = [] + final_reply = "" for ev in self.run_stream(user_message, history): - if ev["type"] == "tool_call": + kind = ev["type"] + if kind == "tool_result": calls_made.append(ev) - elif ev["type"] == "reply": - reply = ev["content"] + elif kind == "token": + parts.append(ev["content"]) + elif kind == "reply": + final_reply = ev["content"] + reply = "".join(parts) or final_reply return {"reply": reply, "tool_calls": calls_made} diff --git a/src/flow/agent/nvidia.py b/src/flow/agent/nvidia.py index 10b7476..deb40bb 100644 --- a/src/flow/agent/nvidia.py +++ b/src/flow/agent/nvidia.py @@ -6,6 +6,7 @@ from __future__ import annotations +import json import os import httpx @@ -49,3 +50,55 @@ def chat(self, messages: list[dict], tools: list[dict] | None = None, json=payload) r.raise_for_status() return r.json() + + def chat_stream(self, messages: list[dict], tools: list[dict] | None = None, + temperature: float = 0.6, max_tokens: int = 4096): + """Stream a chat-completions turn. Yields ``{"type":"content","text":...}`` + for each token delta, then one ``{"type":"final","message":{...}}`` with the + fully assembled assistant message (content + any tool_calls). Lets callers + type the reply out live instead of waiting for the whole response.""" + payload: dict = { + "model": self.model, "messages": messages, + "temperature": temperature, "max_tokens": max_tokens, "stream": True, + } + if tools: + payload["tools"] = tools + payload["tool_choice"] = "auto" + + parts: list[str] = [] + tool_calls: dict[int, dict] = {} + with httpx.Client(timeout=self.timeout) as c: + with c.stream("POST", f"{self.base_url}/chat/completions", + headers=self._headers(), json=payload) as r: + r.raise_for_status() + for line in r.iter_lines(): + if not line.startswith("data:"): + continue + data = line[5:].strip() + if data == "[DONE]": + break + try: + choice = json.loads(data)["choices"][0] + except (json.JSONDecodeError, KeyError, IndexError): + continue + delta = choice.get("delta") or {} + if delta.get("content"): + parts.append(delta["content"]) + yield {"type": "content", "text": delta["content"]} + for tc in (delta.get("tool_calls") or []): + idx = tc.get("index", 0) + slot = tool_calls.setdefault( + idx, {"id": "", "type": "function", + "function": {"name": "", "arguments": ""}}) + if tc.get("id"): + slot["id"] = tc["id"] + fn = tc.get("function") or {} + if fn.get("name"): + slot["function"]["name"] += fn["name"] + if fn.get("arguments"): + slot["function"]["arguments"] += fn["arguments"] + + message: dict = {"role": "assistant", "content": "".join(parts)} + if tool_calls: + message["tool_calls"] = [tool_calls[i] for i in sorted(tool_calls)] + yield {"type": "final", "message": message} diff --git a/src/flow/tools/flow_native.py b/src/flow/tools/flow_native.py index f341b23..27ae049 100644 --- a/src/flow/tools/flow_native.py +++ b/src/flow/tools/flow_native.py @@ -8,6 +8,7 @@ from __future__ import annotations +from flow.schemas import Character from flow.store.frames import seconds_to_frames from flow.store.media import GenerationStatus, MediaAsset, MediaType from flow.store.models import Clip, ClipStatus @@ -16,6 +17,25 @@ from flow.tools.registry import tool +@tool("create_character", "Add a reusable character to the project's cast for " + "cross-scene subject consistency. Optionally pin a reference image; cast it " + "into scenes with attach_character_to_scene.", + {"project_id": "string?", "name": "string", "description": "string", + "reference_image": {"type": "string", "optional": True}}, + mutating=True) +def create_character(ctx: ToolContext, args: dict) -> dict: + p = ctx.project + name = args["name"] + if name in p.characters: + return result.error("exists", f"character {name!r} is already in the cast", + hint="pick another name or attach_character_to_scene") + p.characters[name] = Character( + description=args.get("description", ""), + reference_image=args.get("reference_image"), + ) + return result.ok(summary=f"created character {name}", character=name) + + @tool("attach_character_to_scene", "Cast a character (from the project's reusable " "cast) into a scene for subject consistency. Set regenerate to re-roll now.", {"project_id": "string?", "scene_id": "string", "character_name": "string", diff --git a/tests/test_agent.py b/tests/test_agent.py new file mode 100644 index 0000000..e40003d --- /dev/null +++ b/tests/test_agent.py @@ -0,0 +1,114 @@ +"""Tests for the in-app agent: token-streaming loop + the create_character tool.""" + +import json + +import flow.tools.all # noqa: F401 — register all tools +from flow.agent.loop import Agent +from flow.store.project import Project +from flow.store.store import ProjectStore +from flow.tools.context import ToolContext, dispatch + + +def make_ctx() -> ToolContext: + store = ProjectStore(url="sqlite://") # in-memory + project = Project(title="Test") + store.save(project) + return ToolContext(project=project, store=store) + + +class FakeClient: + """Scripted streaming client. Each turn is (content_chunks, tool_calls).""" + + def __init__(self, turns: list[tuple[list[str], list[dict] | None]]) -> None: + self.turns = list(turns) + self.calls = 0 + + def chat_stream(self, messages, tools=None, **kw): + chunks, tool_calls = self.turns[self.calls] + self.calls += 1 + for ch in chunks: + yield {"type": "content", "text": ch} + msg: dict = {"role": "assistant", "content": "".join(chunks)} + if tool_calls: + msg["tool_calls"] = tool_calls + yield {"type": "final", "message": msg} + + +def _tool_call(name: str, args: dict, cid: str = "call_1") -> dict: + return {"id": cid, "type": "function", + "function": {"name": name, "arguments": json.dumps(args)}} + + +# --- create_character tool --- + +def test_create_character_adds_to_cast(): + ctx = make_ctx() + res = dispatch(ctx, "create_character", {"name": "hero", "description": "a knight"}) + assert res["ok"] is True + assert "hero" in ctx.project.characters + assert ctx.project.characters["hero"].description == "a knight" + + +def test_create_character_rejects_duplicate(): + ctx = make_ctx() + dispatch(ctx, "create_character", {"name": "hero", "description": "a knight"}) + res = dispatch(ctx, "create_character", {"name": "hero", "description": "other"}) + assert res["ok"] is False + assert res["error"]["code"] == "exists" + + +def test_create_then_attach_character(): + ctx = make_ctx() + dispatch(ctx, "create_character", {"name": "hero", "description": "a knight"}) + dispatch(ctx, "plan_video", {"scenes": [{"prompt": "opening shot"}]}) + scene_id = ctx.project.ordered_clips()[0].clip_id + res = dispatch(ctx, "attach_character_to_scene", + {"scene_id": scene_id, "character_name": "hero"}) + assert res["ok"] is True + assert "hero" in ctx.project.get_clip(scene_id).characters + + +# --- streaming agent loop --- + +def test_run_stream_emits_tokens_then_reply(): + ctx = make_ctx() + agent = Agent(FakeClient([(["Hel", "lo!"], None)]), ctx) + events = list(agent.run_stream("hi")) + assert [e["type"] for e in events] == ["token", "token", "reply"] + assert events[0]["content"] == "Hel" + assert events[-1]["content"] == "Hello!" + + +def test_run_stream_executes_tool_then_replies(): + ctx = make_ctx() + turns = [ + ([], [_tool_call("create_character", {"name": "hero", "description": "knight"})]), + (["done"], None), + ] + agent = Agent(FakeClient(turns), ctx) + events = list(agent.run_stream("make a hero")) + types = [e["type"] for e in events] + + assert "tool_start" in types and "tool_result" in types + assert types.index("tool_start") < types.index("tool_result") + assert types[-1] == "reply" + + start = next(e for e in events if e["type"] == "tool_start") + assert start["tool"] == "create_character" + res = next(e for e in events if e["type"] == "tool_result") + assert res["result"]["ok"] is True + # the tool really ran against the project + assert "hero" in ctx.project.characters + assert events[-1]["content"] == "done" + + +def test_run_collects_stream(): + ctx = make_ctx() + turns = [ + ([], [_tool_call("create_character", {"name": "hero", "description": "knight"})]), + (["all ", "set"], None), + ] + out = Agent(FakeClient(turns), ctx).run("make a hero") + assert out["reply"] == "all set" + assert len(out["tool_calls"]) == 1 + assert out["tool_calls"][0]["tool"] == "create_character"