Files
app/backend/fluksio/ai/agent.py
T

204 lines
7.4 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:
- 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.
- save_flow writes 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 `save_flow` takes, 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 named {flow!r}."
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.",
}