run_turn takes its tools, system prompt and tool caller; the chat passes the MCP tools as before. The assistant node offers one tool per wired target, reads its other inputs as context, streams its reply a sentence at a time, warms the provider on `listening` and says so when the provider is not there. Completions share the per-server HTTP client. wren replies when nothing reads `unhandled`, stt shows a note when nothing was heard. /ai/control lists one conversation per entry node with the messages to watch.
357 lines
12 KiB
Python
357 lines
12 KiB
Python
"""The assistant: its provider, and the conversation with it.
|
|
|
|
The provider is instance-wide and superuser-only to set, because the key it
|
|
holds is spent by everyone who chats. The conversation is per person and runs
|
|
as them: every tool call carries their own token, so the assistant can do
|
|
nothing they could not do in the browser.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
import uuid
|
|
from collections.abc import AsyncIterator
|
|
from typing import Any, Literal
|
|
|
|
import httpx
|
|
from fastapi import APIRouter, Depends, HTTPException, Request
|
|
from fastapi.responses import StreamingResponse
|
|
from pydantic import BaseModel, Field
|
|
|
|
from fluksio.ai import agent, client
|
|
from fluksio.ai import config as ai_config
|
|
from fluksio.api.deps import (
|
|
CurrentUser,
|
|
FlowControllerDep,
|
|
TokenDep,
|
|
get_current_active_superuser,
|
|
)
|
|
from fluksio.core.config import settings
|
|
from fluksio.flow.messages import qualify
|
|
from fluksio.flow.schemas import NAME_PATTERN, NodeDef
|
|
from fluksio.mcp import server as tools
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
router = APIRouter(prefix="/ai", tags=["ai"])
|
|
|
|
#: Whose turn is running: a second message while one is answered is refused
|
|
#: rather than interleaved into the same transcript.
|
|
_busy: set[uuid.UUID] = set()
|
|
|
|
|
|
class AIConfigPublic(BaseModel):
|
|
api_base: str
|
|
model: str
|
|
max_rounds: int
|
|
api_key_set: bool
|
|
api_key_masked: str | None
|
|
|
|
|
|
class AIConfigUpdate(BaseModel):
|
|
api_base: str = Field(min_length=1, max_length=512)
|
|
#: Blank or absent keeps the stored key.
|
|
api_key: str | None = Field(default=None, max_length=512)
|
|
model: str = Field(default="", max_length=128)
|
|
max_rounds: int = Field(default=12, ge=1, le=50)
|
|
|
|
|
|
class AIModelsBody(BaseModel):
|
|
api_base: str = Field(min_length=1, max_length=512)
|
|
#: Absent tries the stored key, so a saved provider can be re-listed.
|
|
api_key: str | None = Field(default=None, max_length=512)
|
|
|
|
|
|
class AIStatus(BaseModel):
|
|
available: bool
|
|
model: str | None
|
|
|
|
|
|
class ChatItem(BaseModel):
|
|
role: str
|
|
content: str | None = None
|
|
tool_calls: list[dict[str, Any]] | None = None
|
|
#: On a tool's answer: whether it worked, and the error if not.
|
|
ok: bool | None = None
|
|
error: str | None = None
|
|
|
|
|
|
type WatchKind = Literal["partial", "text", "reply", "actions"]
|
|
|
|
|
|
class ControlWatch(BaseModel):
|
|
"""A message the conversation shows, as one kind of line."""
|
|
|
|
msg: str
|
|
kind: WatchKind
|
|
#: The node whose values on it count. Empty counts every value, as the
|
|
#: prompt does: a command typed here or spoken to a device.
|
|
node: str = ""
|
|
|
|
|
|
class ControlNode(BaseModel):
|
|
"""A conversation the chat can hold: where a command goes, and what to show."""
|
|
|
|
#: Qualified, ``flow.node``.
|
|
id: str
|
|
flow: str
|
|
title: str
|
|
prompt_msg: str
|
|
watch: list[ControlWatch]
|
|
|
|
|
|
class TurnBody(BaseModel):
|
|
message: str = Field(min_length=1, max_length=20_000)
|
|
#: The flow open in the editor, when the chat is opened from one.
|
|
flow: str | None = Field(default=None, pattern=NAME_PATTERN.pattern)
|
|
|
|
|
|
def _public(config: ai_config.AIConfig | None) -> AIConfigPublic:
|
|
if config is None:
|
|
return AIConfigPublic(
|
|
api_base="",
|
|
model="",
|
|
max_rounds=12,
|
|
api_key_set=False,
|
|
api_key_masked=None,
|
|
)
|
|
return AIConfigPublic(**config.public())
|
|
|
|
|
|
@router.get(
|
|
"/config",
|
|
dependencies=[Depends(get_current_active_superuser)],
|
|
response_model=AIConfigPublic,
|
|
)
|
|
def read_ai_config() -> Any:
|
|
return _public(ai_config.load())
|
|
|
|
|
|
@router.put(
|
|
"/config",
|
|
dependencies=[Depends(get_current_active_superuser)],
|
|
response_model=AIConfigPublic,
|
|
)
|
|
def update_ai_config(body: AIConfigUpdate) -> Any:
|
|
try:
|
|
client.check_base(body.api_base)
|
|
except ValueError as exc:
|
|
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
|
current = ai_config.load()
|
|
key = body.api_key or (current.api_key if current else "")
|
|
config = ai_config.AIConfig(
|
|
api_base=body.api_base.rstrip("/"),
|
|
api_key=key,
|
|
model=body.model.strip(),
|
|
max_rounds=body.max_rounds,
|
|
)
|
|
ai_config.save(config)
|
|
return _public(config)
|
|
|
|
|
|
@router.delete("/config", dependencies=[Depends(get_current_active_superuser)])
|
|
def delete_ai_config() -> Any:
|
|
ai_config.delete()
|
|
return _public(None)
|
|
|
|
|
|
@router.post(
|
|
"/config/models",
|
|
dependencies=[Depends(get_current_active_superuser)],
|
|
response_model=list[str],
|
|
)
|
|
async def list_ai_models(body: AIModelsBody) -> Any:
|
|
"""What the provider offers, tried with the key given or the one stored."""
|
|
try:
|
|
client.check_base(body.api_base)
|
|
except ValueError as exc:
|
|
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
|
current = ai_config.load()
|
|
key = body.api_key or (current.api_key if current else "")
|
|
try:
|
|
return await client.list_models(body.api_base.rstrip("/"), key)
|
|
except (client.ProviderError, httpx.HTTPError) as exc:
|
|
raise HTTPException(
|
|
status_code=400, detail=f"The provider answered {exc}"
|
|
) from exc
|
|
|
|
|
|
@router.get("/status", response_model=AIStatus)
|
|
def read_ai_status(_: CurrentUser) -> Any:
|
|
"""Whether there is an assistant to show. Readable by anyone signed in."""
|
|
config = ai_config.load()
|
|
if config is None or not config.ready:
|
|
return AIStatus(available=False, model=None)
|
|
return AIStatus(available=True, model=config.model)
|
|
|
|
|
|
#: The node types a command can be typed to.
|
|
_TALKERS = ("wren", "assistant")
|
|
|
|
|
|
def _wired(flow: str, node: NodeDef, port: str) -> str:
|
|
"""The message a node's port is wired to, qualified, or empty."""
|
|
for spec in [*node.requires, *node.provides]:
|
|
if spec.port == port and spec.name:
|
|
return qualify(flow, spec.name)
|
|
return ""
|
|
|
|
|
|
@router.get("/control", response_model=list[ControlNode])
|
|
def read_control_nodes(_: CurrentUser, controller: FlowControllerDep) -> Any:
|
|
"""The conversations of the running live flows, for the chat's Control tab.
|
|
|
|
One per wren or assistant node a command reaches first. A node whose
|
|
prompt is another's ``unhandled`` answers in that one's conversation,
|
|
after it. The conversation shows the prompt, the words heard so far (the
|
|
``partial`` of a ``stt`` node writing the prompt), and each node's
|
|
``reply`` and ``actions``, so typed and spoken commands are one.
|
|
"""
|
|
found = []
|
|
for flow in controller.store.read_all():
|
|
if flow.mode != "live" or flow.name in controller.disabled:
|
|
continue
|
|
talkers = [
|
|
node
|
|
for node in flow.nodes
|
|
if node.type in _TALKERS and _wired(flow.name, node, "prompt")
|
|
]
|
|
passed = {_wired(flow.name, node, "unhandled") for node in talkers}
|
|
for node in talkers:
|
|
prompt = _wired(flow.name, node, "prompt")
|
|
if prompt in passed:
|
|
continue
|
|
watch = [ControlWatch(msg=prompt, kind="text")]
|
|
watch += [
|
|
ControlWatch(msg=partial, kind="partial", node=f"{flow.name}.{stt.id}")
|
|
for stt in flow.nodes
|
|
if stt.type == "stt"
|
|
and _wired(flow.name, stt, "text") == prompt
|
|
and (partial := _wired(flow.name, stt, "partial"))
|
|
]
|
|
chain = [node]
|
|
for talker in chain:
|
|
unhandled = _wired(flow.name, talker, "unhandled")
|
|
chain += [
|
|
other
|
|
for other in talkers
|
|
if unhandled
|
|
and _wired(flow.name, other, "prompt") == unhandled
|
|
and other not in chain
|
|
]
|
|
answers: tuple[WatchKind, ...] = ("reply", "actions")
|
|
watch += [
|
|
ControlWatch(msg=msg, kind=kind, node=f"{flow.name}.{talker.id}")
|
|
for talker in chain
|
|
for kind in answers
|
|
if (msg := _wired(flow.name, talker, kind))
|
|
]
|
|
if any(w.kind in answers for w in watch):
|
|
found.append(
|
|
ControlNode(
|
|
id=f"{flow.name}.{node.id}",
|
|
flow=flow.name,
|
|
title=f"{flow.title or flow.name} · {node.title or node.id}",
|
|
prompt_msg=prompt,
|
|
watch=watch,
|
|
)
|
|
)
|
|
return found
|
|
|
|
|
|
def _chat_file(user_id: uuid.UUID) -> Any:
|
|
return settings.CHATS_DIR / f"{user_id}.json"
|
|
|
|
|
|
def _load_chat(user_id: uuid.UUID) -> list[dict[str, Any]]:
|
|
try:
|
|
messages: list[dict[str, Any]] = json.loads(_chat_file(user_id).read_text())
|
|
except (OSError, ValueError):
|
|
return []
|
|
return messages
|
|
|
|
|
|
def _save_chat(user_id: uuid.UUID, messages: list[dict[str, Any]]) -> None:
|
|
path = _chat_file(user_id)
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
path.write_text(json.dumps(messages))
|
|
|
|
|
|
@router.get("/chat", response_model=list[ChatItem])
|
|
def read_chat(user: CurrentUser) -> Any:
|
|
"""The transcript, with each tool's answer reduced to whether it worked."""
|
|
items = []
|
|
for message in _load_chat(user.id):
|
|
if message.get("role") == "tool":
|
|
items.append(
|
|
ChatItem(role="tool", **agent.summarize(message.get("content")))
|
|
)
|
|
else:
|
|
items.append(
|
|
ChatItem(
|
|
role=message["role"],
|
|
content=message.get("content"),
|
|
tool_calls=message.get("tool_calls"),
|
|
)
|
|
)
|
|
return items
|
|
|
|
|
|
@router.delete("/chat", response_model=list[ChatItem])
|
|
def clear_chat(user: CurrentUser) -> Any:
|
|
_chat_file(user.id).unlink(missing_ok=True)
|
|
return []
|
|
|
|
|
|
def _sse(event: dict[str, Any]) -> str:
|
|
return f"data: {json.dumps(event)}\n\n"
|
|
|
|
|
|
@router.post("/chat")
|
|
def chat(
|
|
body: TurnBody, request: Request, user: CurrentUser, token: TokenDep
|
|
) -> StreamingResponse:
|
|
"""One turn, streamed as server-sent events.
|
|
|
|
``delta`` as the model writes, ``message`` with a round's whole text,
|
|
``tool_call`` and ``tool_result`` as it works, ``error`` when something
|
|
stops it, and ``done`` last either way.
|
|
"""
|
|
config = ai_config.load()
|
|
if config is None or not config.ready:
|
|
# 404 rather than 403: an unconfigured instance says nothing about it.
|
|
raise HTTPException(status_code=404, detail="The assistant is not set up")
|
|
if user.id in _busy:
|
|
raise HTTPException(status_code=409, detail="Still answering the last one")
|
|
tools.connect(request.app)
|
|
messages = _load_chat(user.id)
|
|
|
|
async def generate() -> AsyncIterator[str]:
|
|
_busy.add(user.id)
|
|
try:
|
|
yield _sse({"type": "start"})
|
|
async for event in agent.run_turn(
|
|
config,
|
|
messages,
|
|
body.message,
|
|
tools=await agent.tool_specs(),
|
|
system=agent.system_prompt(body.flow),
|
|
call_tool=agent.mcp_caller(f"Bearer {token}"),
|
|
):
|
|
yield _sse(event)
|
|
except Exception as exc: # noqa: BLE001 - the stream is the error channel
|
|
logger.exception("The assistant's turn failed")
|
|
yield _sse({"type": "error", "detail": f"{type(exc).__name__}: {exc}"})
|
|
finally:
|
|
# What happened so far, even if the browser left part-way.
|
|
_save_chat(user.id, messages)
|
|
_busy.discard(user.id)
|
|
yield _sse({"type": "done"})
|
|
|
|
return StreamingResponse(
|
|
generate(),
|
|
media_type="text/event-stream",
|
|
headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
|
|
)
|