211 lines
7.7 KiB
Python
211 lines
7.7 KiB
Python
"""The turn: the model, the tools, and the rounds between them.
|
|
|
|
The turn takes its tools, its system prompt and the function that calls a
|
|
tool from whoever runs it. The chat passes the MCP server's own, called
|
|
in-process rather than over the endpoint: same names, same schemas, same REST
|
|
calls underneath. A tool added for external agents is there for the chat the
|
|
moment it exists, and it can do nothing an external agent could not. The
|
|
assistant node passes the targets wired to it, and nothing else.
|
|
|
|
The transcript is a list of chat-completion messages, mutated in place by
|
|
:func:`run_turn` and persisted by the caller. It is also what the browser
|
|
renders, so nothing is stored twice.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
from collections.abc import AsyncIterator, Awaitable, Callable
|
|
from typing import Any
|
|
|
|
import httpx
|
|
from mcp.types import CallToolResult, TextContent
|
|
|
|
from fluksio.ai import client
|
|
from fluksio.ai.config import AIConfig
|
|
from fluksio.flow.schemas import FlowDef
|
|
from fluksio.mcp.server import caller_authorization, mcp
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_SYSTEM = """You are the assistant inside a Fluksio instance. You read, build, edit and \
|
|
operate flows through the tools; they are the same API the editor uses, so \
|
|
anything you save appears in the editor at once.
|
|
|
|
How to work:
|
|
- A flow has an id and a key, its `name`, which also prefixes its messages. \
|
|
Tools take either; write id:<x> or key:<x> when a tool says the name could \
|
|
mean two flows.
|
|
- Read before you write: get_flow before save_flow, and pass back the version \
|
|
you read. On a version conflict, read again and reapply the change.
|
|
- create_flow makes a new flow and save_flow changes one; both write a draft, \
|
|
and validate_flow tells you whether it is sound. \
|
|
Publish only when the user asks for the change to go live; otherwise leave \
|
|
the draft for review in the editor and say so.
|
|
- list_node_types gives the node types and their parameters; get_node_source \
|
|
shows how a Python node is written before you write one.
|
|
- Keep answers short and name what you changed. Answers render as markdown: \
|
|
use lists, `inline code` and fenced code blocks where they help, and no \
|
|
tables, images or HTML.
|
|
|
|
The flow document `create_flow` and `save_flow` take, as a JSON schema:
|
|
"""
|
|
|
|
# ponytail: a tool's answer is replayed on every later round of the
|
|
# conversation, so one big one is paid for again and again; cut it here rather
|
|
# than summarise it. A per-tool digest is the upgrade if this bites.
|
|
_RESULT_LIMIT = 20_000
|
|
|
|
|
|
def system_prompt(flow: str | None) -> str:
|
|
prompt = _SYSTEM + json.dumps(FlowDef.model_json_schema())
|
|
if flow:
|
|
prompt += (
|
|
f"\n\nThe user is looking at the flow with the id {flow!r}. Name it "
|
|
f"to the tools as id:{flow}."
|
|
)
|
|
return prompt
|
|
|
|
|
|
async def tool_specs() -> list[dict[str, Any]]:
|
|
return [
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": tool.name,
|
|
"description": tool.description or "",
|
|
"parameters": tool.input_schema,
|
|
},
|
|
}
|
|
for tool in await mcp.list_tools()
|
|
]
|
|
|
|
|
|
def _arguments(raw: Any) -> dict[str, Any]:
|
|
if isinstance(raw, dict):
|
|
return raw
|
|
try:
|
|
parsed = json.loads(raw or "{}")
|
|
except ValueError:
|
|
return {}
|
|
return parsed if isinstance(parsed, dict) else {}
|
|
|
|
|
|
#: Runs one tool call and answers with what the model reads back.
|
|
type ToolCaller = Callable[[str, dict[str, Any]], Awaitable[str]]
|
|
|
|
|
|
def mcp_caller(authorization: str) -> ToolCaller:
|
|
"""The MCP tools, each call made as the person chatting."""
|
|
|
|
async def call(name: str, arguments: dict[str, Any]) -> str:
|
|
token = caller_authorization.set(authorization)
|
|
try:
|
|
return await _call_tool(name, arguments)
|
|
finally:
|
|
caller_authorization.reset(token)
|
|
|
|
return call
|
|
|
|
|
|
async def _call_tool(name: str, arguments: dict[str, Any]) -> str:
|
|
"""Run one tool; whatever goes wrong is text the model can read."""
|
|
try:
|
|
result = await mcp.call_tool(name, arguments)
|
|
except Exception as exc: # noqa: BLE001 - an unknown tool, or refused arguments
|
|
return json.dumps({"error": str(exc)})
|
|
if not isinstance(result, CallToolResult):
|
|
return json.dumps({"error": "The tool asked for more input than a chat gives."})
|
|
text = "".join(c.text for c in result.content if isinstance(c, TextContent))
|
|
if len(text) > _RESULT_LIMIT:
|
|
text = text[:_RESULT_LIMIT] + "… (truncated)"
|
|
return text
|
|
|
|
|
|
def summarize(content: str | None) -> dict[str, Any]:
|
|
"""What the browser shows for a tool's answer: whether it worked."""
|
|
try:
|
|
result = json.loads(content or "")
|
|
except ValueError:
|
|
return {"ok": True}
|
|
if isinstance(result, dict) and "error" in result:
|
|
return {"ok": False, "error": str(result["error"])[:300]}
|
|
return {"ok": True}
|
|
|
|
|
|
async def run_turn(
|
|
config: AIConfig,
|
|
messages: list[dict[str, Any]],
|
|
text: str,
|
|
*,
|
|
tools: list[dict[str, Any]],
|
|
system: str,
|
|
call_tool: ToolCaller,
|
|
extra: dict[str, Any] | None = None,
|
|
timeout: float | httpx.Timeout = 120.0,
|
|
answer_tools: bool = True,
|
|
) -> AsyncIterator[dict[str, Any]]:
|
|
"""One user turn, yielding progress events as it goes.
|
|
|
|
A round's text arrives as ``delta`` events while the model writes it, and
|
|
once more whole as ``message`` when the round is done. A round is appended
|
|
to *messages* only once every tool in it has answered, so a turn cut off
|
|
part-way leaves a transcript the provider still accepts. With
|
|
*answer_tools* false the turn ends once its tools have answered, for a
|
|
caller that words what they did itself.
|
|
"""
|
|
messages.append({"role": "user", "content": text})
|
|
head = {"role": "system", "content": system}
|
|
for _ in range(config.max_rounds):
|
|
choice: dict[str, Any] = {}
|
|
try:
|
|
async for part in client.stream_completion(
|
|
api_base=config.api_base,
|
|
api_key=config.api_key,
|
|
model=config.model,
|
|
messages=[head, *messages],
|
|
tools=tools,
|
|
timeout=timeout,
|
|
extra=extra,
|
|
):
|
|
if isinstance(part, str):
|
|
yield {"type": "delta", "content": part}
|
|
else:
|
|
choice = part
|
|
except (client.ProviderError, httpx.HTTPError) as exc:
|
|
yield {"type": "error", "detail": f"The provider answered {exc}"}
|
|
return
|
|
content = choice.get("content") or ""
|
|
calls = choice.get("tool_calls") or []
|
|
assistant: dict[str, Any] = {"role": "assistant", "content": content}
|
|
if content:
|
|
yield {"type": "message", "content": content}
|
|
if not calls:
|
|
messages.append(assistant)
|
|
return
|
|
assistant["tool_calls"] = calls
|
|
answers: list[dict[str, Any]] = []
|
|
for call in calls:
|
|
function = call.get("function") or {}
|
|
name = str(function.get("name", ""))
|
|
arguments = _arguments(function.get("arguments"))
|
|
yield {"type": "tool_call", "name": name, "arguments": arguments}
|
|
result = await call_tool(name, arguments)
|
|
answers.append(
|
|
{
|
|
"role": "tool",
|
|
"tool_call_id": call.get("id", ""),
|
|
"content": result,
|
|
}
|
|
)
|
|
yield {"type": "tool_result", "name": name, **summarize(result)}
|
|
messages.append(assistant)
|
|
messages.extend(answers)
|
|
if not answer_tools:
|
|
return
|
|
yield {
|
|
"type": "error",
|
|
"detail": f"Stopped after {config.max_rounds} rounds of tool calls.",
|
|
}
|