app/backend/fluksio/ai/agent.py

"""The turn: the model, the tools, and the rounds between them.

The tools are 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 assistant the moment it exists, and the
assistant can do nothing an external agent could not.

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
from typing import Any

import httpx
from mcp.types import 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.inputSchema,
            },
        }
        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 {}


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 isinstance(result, dict):
        text = json.dumps(result)
    else:
        text = "".join(c.text for c in result 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,
    *,
    flow: str | None,
    authorization: str,
) -> 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.
    """
    messages.append({"role": "user", "content": text})
    tools = await tool_specs()
    system = {"role": "system", "content": system_prompt(flow)}
    token = caller_authorization.set(authorization)
    try:
        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=[system, *messages],
                    tools=tools,
                ):
                    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)
        yield {
            "type": "error",
            "detail": f"Stopped after {config.max_rounds} rounds of tool calls.",
        }
    finally:
        caller_authorization.reset(token)