stroblme
1b5a16f636
ai/client.py streams the completion and puts the chunks back together (Assembler), covering OpenAI's piecewise tool-call arguments, Ollama's whole calls, OpenRouter's comment lines, unnumbered calls and in-stream errors. The turn emits a delta event per piece and a message per round as before, and the chat grows the last answer as deltas arrive.
301 lines
9.3 KiB
Python
301 lines
9.3 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
|
|
|
|
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
|
|
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
|
|
|
|
|
|
class ControlNode(BaseModel):
|
|
"""A wren node the chat can talk to: where it sends and where it reads."""
|
|
|
|
#: Qualified, ``flow.node``.
|
|
id: str
|
|
flow: str
|
|
title: str
|
|
prompt_msg: str
|
|
actions_msg: str
|
|
|
|
|
|
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)
|
|
|
|
|
|
@router.get("/control", response_model=list[ControlNode])
|
|
def read_control_nodes(_: CurrentUser, controller: FlowControllerDep) -> Any:
|
|
"""The wren nodes of the running live flows, for the chat's Control tab.
|
|
|
|
Only those with both ends wired: the chat publishes the command on the
|
|
prompt's message and reads the answer off the actions'.
|
|
"""
|
|
found = []
|
|
for flow in controller.store.read_all():
|
|
if flow.mode != "live" or flow.name in controller.disabled:
|
|
continue
|
|
for node in flow.nodes:
|
|
if node.type != "wren":
|
|
continue
|
|
ends = {
|
|
spec.port: qualify(flow.name, spec.name)
|
|
for spec in [*node.requires, *node.provides]
|
|
if spec.port in ("prompt", "actions") and spec.name
|
|
}
|
|
if len(ends) == 2:
|
|
found.append(
|
|
ControlNode(
|
|
id=f"{flow.name}.{node.id}",
|
|
flow=flow.name,
|
|
title=node.title or node.id,
|
|
prompt_msg=ends["prompt"],
|
|
actions_msg=ends["actions"],
|
|
)
|
|
)
|
|
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,
|
|
flow=body.flow,
|
|
authorization=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"},
|
|
)
|