chat · slides · image · search построены на общем ядре с замороженными
портами: ИИ за интерфейсом AgentDriver, состояние — типизированный RunState,
который течёт в браузер через change-log → SSE. Кликните любую сущность — внутри
«что делает» и пошаговый механизм работы.
01
Pipeline одного запроса
Поток одинаков для всех четырёх агентов — различаются лишь набором
инструментов и формой состояния. Кликайте по шагам.
02
Слои инфраструктуры
Каждая плашка — сущность. Клик открывает разбор: что делает, механизм
работы по шагам, пример кода и связи.
03
Агенты — подробно
Инструменты, режимы и форма RunState каждого продукта.
3·5
Исходники: агенты + рантайм
Буквальный код каждого агента и общего рантайма mlexi_agent (порты:
AgentDriver, RunInput, Tool, StateWriter, fold_changes, …) — прямо из репозитория. Выбери
вкладку и файл; путь указан над кодом, кнопкой copy можно скопировать. Большие файлы
показаны фрагментами.
04
MCP-меш
Агент-продукт сам становится набором инструментов для других. Cookie сессии
форвардится сервис-в-сервис для whoami-авторизации.
Продьюсеры (/mcp)
Кто что потребляет
05
Общие паттерны
На чём держатся все четыре агента.
06
Порты в коде
Реальные фрагменты из packages/python/mlexi_agent и
composition root воркфлоу.
07
Как работает Opik
Opik (Comet, self-hosted, Apache-2.0) — хранилище и UI для трейсов
агентов по стандарту OpenTelemetry gen_ai. Отвечает на вопрос «что именно делал
агент»: какие промпты ушли в модель, что она вернула, сколько токенов, какие инструменты
вызывались, сколько заняло и сколько стоило. У нас трассируются все 4 агента
(chat · image · search · slides). Phoenix полностью удалён.
Конвейер трейса (в общем виде)
Спаны gen_ai — ручная разметка, без авто-инструментации
Трейс = дерево спанов; у каждого gen_ai.operation.name и атрибуты. Виды:
invoke_agent агент ·
chat вызов модели (итерация) ·
execute_tool инструмент.
Спаны эмитит общий PydanticAIDriver вручную (по требованию оператора — никакой
авто-инструментации), одной пост-реконструкцией после прогона из
result.all_messages(). У каждой итерации: точный input (история ДО ответа),
output (ответ), gen_ai.usage.* токены и стоимость mlexi из gateway
(mlexi.cost_credits — int-micros в usage.details, не оценка Opik). Тайминг — из
таймстемпов сообщений; раскладка по проектам через Comet-Workspace+projectName.
chat/execute_tool-спаны с промптами, ответами, токенами, тулзами, стоимостью
codexapi
—
не трассируется: stateless шим, не агент — роундтрип уже виден в спане агента + в gateway
gateway · identity · billing · info
—
не трассируются (не агентные / не LLM-фреймворковые)
По проекту на продукт (chat/image/search/slides); доступ
opik.mlexi.ru только админу (Caddy forward_auth → identity
/internal/v1/check-admin); приём OTLP-HTTP на
opik-frontend-1:5173/api/v1/private/otel/v1/traces (свой compose-стек
/opt/opik); внутренний OTLP самого DBOS отключён (enable_otlp=False).
"""Declared run-state for the chat agent (the frontend contract).
`ChatTurnState` is the live streaming contract for one chat turn. The
generic `WriterEventAdapter` (in mlexi_agent.pydantic_driver) fills it:
finalized assistant text is appended to `text`, and tool-call/tool-return
items are appended to `steps` (a keyed activity log). `status` is the
scalar lifecycle marker the SSE transport watches for terminal values.
Phase B: `chat.turn_changes` (the StateWriter change-log) is the SOLE
durable source for a turn — `chat.message_parts` is gone. The user's
question lives in the change-log too, as the `user_input` SCALAR
(written by the POST handler at seq 1). The static system prompt is NOT
stored — it's passed via RunInput.system_prompt. GET endpoints and
multi-turn history both fold this change-log into a ChatTurnState
(see mlexi_agent.fold + agent/history.py).
"""
from __future__ import annotations
from typing import Any
from pydantic import BaseModel
from mlexi_agent.state import FieldSpec, Kind, RunState
class ChatStep(BaseModel):
"""One activity-log item. `dict` is also accepted by the writer; this
typed shape documents the contract and drives any future TS codegen."""
kind: str # "tool-call" | "tool-return"
tool_name: str
tool_call_id: str
args: dict | None = None # tool-call
content: Any = None # tool-return
class ChatTurnState(RunState):
user_input: str = "" # the user's question (change-log seq 1)
status: str = "streaming"
text: str = "" # assistant text, finalized chunks
steps: list[dict] = []
modes = {
"user_input": FieldSpec(Kind.SCALAR),
"status": FieldSpec(Kind.SCALAR),
"text": FieldSpec(Kind.TEXT),
"steps": FieldSpec(Kind.LIST, key="tool_call_id"),
}
"""Build MCPServerStreamableHTTP instances with cookie auth."""
from __future__ import annotations
import logging
from pydantic_ai.mcp import MCPServerStreamableHTTP
_log = logging.getLogger(__name__)
def build_servers(urls: list[str], *, cookie: str) -> list[MCPServerStreamableHTTP]:
"""One MCPServerStreamableHTTP per URL, with mlexi_session cookie
in the connection-level headers. Cookie stays constant for the
workflow's lifetime."""
servers = []
for url in urls:
try:
srv = MCPServerStreamableHTTP(
url=url,
headers={"Cookie": f"mlexi_session={cookie}"},
# pydantic-ai default is 5s; remote MCPs go through nginx
# + identity-cookie validation, so easily exceed 5s on a
# cold start. 30s gives plenty of margin without making
# users wait forever on a genuinely broken server.
timeout=30,
)
except Exception as exc:
_log.warning("mcp_servers: failed to construct for %s: %s",
url, exc)
continue
servers.append(srv)
return servers
"""DBOS workflow: drive one chat turn via the shared mlexi_agent runtime.
`_run_turn` is the composition root (plain async, testable). It reads the
turn's `user_input` from the change-log (chat.turn_changes — written by the
POST handler at seq 1), builds the neutral history by folding prior completed
turns, wires the MCP toolsets + a StateWriter over chat.turn_changes, and
runs the shared PydanticAIDriver.
Phase B: chat.turn_changes is the SOLE durable source — there is no
message_parts derivation. The live frontend contract, the GET snapshots, and
multi-turn history all run on the folded change-log. Survives chat-service
restart via DBOS recovery; the recovery branch mirrors the terminal `status`
into turn_changes so a reconnecting SSE closes.
"""
from __future__ import annotations
import asyncio
import logging
from dbos import DBOS
from pydantic_ai.exceptions import UsageLimitExceeded
from sqlalchemy import text
from mlexi_agent import LLMConfig, Limits, RunInput, StateWriter
from mlexi_agent.fold import fold_changes
from mlexi_agent.pydantic_driver import PydanticAIDriver
from ..config import get_settings
from ..db import get_session_factory
from .history import state_to_neutral_messages
from .prompt import CHAT_SYSTEM_PROMPT
from .state import ChatTurnState
from .tools.mcp_servers import build_servers
_log = logging.getLogger("chat.workflow")
HEARTBEAT_INTERVAL_SEC = 5
TURN_CHANGES_TABLE = "chat.turn_changes"
_PLACEHOLDER_TITLES = (None, "", "Новый чат", "New chat", "Untitled")
def _channel(turn_id: int) -> str:
return f"chat_turn_{int(turn_id)}"
async def _run_turn(turn_id: int, cookie: str) -> None:
"""Workflow body extracted as a plain async function for testability."""
settings = get_settings()
sf = get_session_factory()
async with sf() as s:
row = (await s.execute(
text("SELECT model_slug, conversation_id FROM chat.turns WHERE id = :id"),
{"id": turn_id},
)).first()
if row is None:
raise RuntimeError(f"turn {turn_id} not found")
model_slug, conv_id = row[0], row[1]
# Phase B: the user's question lives in the change-log (user_input,
# written by the POST handler at seq 1). Fold this turn's change-log to
# read it back — the sole source of truth, no message_parts.
own_rows = (await s.execute(
text("SELECT seq, op, path, value FROM chat.turn_changes "
"WHERE run_id=:id ORDER BY seq"),
{"id": turn_id},
)).all()
max_seq = own_rows[-1].seq if own_rows else 0
# Real OUTPUT = any change-log row whose path is NOT an init write
# (status / user_input). The POST handler writes user_input + status
# (seq 1,2) and the workflow's OWN first write is set("status"), so
# `max_seq > 0` is true the instant the body starts — it does NOT mean
# the model produced anything. A crash AFTER init but BEFORE output
# must RE-RUN, not finalise an empty turn (the search F2 lesson).
has_output = any(
r.path not in ("status", "user_input") for r in own_rows)
# Prior-turn history: fold the change-log of ALL earlier completed
# turns of this conversation, in chronological order, so multi-turn
# chat sees the previous Q/A. Each prior turn's change-log → folded
# ChatTurnState → neutral Messages.
prior_turn_ids = [r.id for r in (await s.execute(text("""
SELECT id FROM chat.turns
WHERE conversation_id = :cid AND id < :tid AND status = 'completed'
ORDER BY started_at ASC
"""), {"cid": conv_id, "tid": turn_id})).all()]
history = []
for pid in prior_turn_ids:
prior_rows = (await s.execute(
text("SELECT seq, op, path, value FROM chat.turn_changes "
"WHERE run_id=:id ORDER BY seq"),
{"id": pid},
)).all()
prior_state = fold_changes(ChatTurnState, prior_rows)
history.extend(state_to_neutral_messages(prior_state))
own_state = fold_changes(ChatTurnState, own_rows)
user_input = own_state.user_input
# DBOS recovery: this turn already produced REAL OUTPUT (text/steps) before
# a crash. Don't re-drive from a partial state — the change-log already
# holds it (the sole source). Finalise as completed AND mirror a terminal
# `status` into turn_changes (the image opus-critical fix), else a
# reconnecting SSE never sees a terminal change and hangs. Init-only
# change-logs (only status / user_input, no output) fall through to drive
# the agent.
if has_output:
_log.warning("turn=%d resume detected (seq=%d, has_output); finalising",
turn_id, max_seq)
recovery_writer = StateWriter(
schema=ChatTurnState, run_id=turn_id, table=TURN_CHANGES_TABLE,
channel=_channel(turn_id), session_factory=sf, initial_seq=max_seq,
)
await _finalise(turn_id, "completed", sf, writer=recovery_writer)
return
# No user_input means the POST handler never wrote it (a recovery where the
# change-log is empty / corrupt) — nothing to drive, finalise as completed.
# Mirror the terminal `status` into turn_changes (Fix 3) — like every other
# terminal branch — else a connected SSE (which closes only on a terminal
# status ROW in the change-log) hangs.
if not user_input:
_log.warning(
"turn=%d resume detected (no user_input in change-log); finalising",
turn_id,
)
recovery_writer = StateWriter(
schema=ChatTurnState, run_id=turn_id, table=TURN_CHANGES_TABLE,
channel=_channel(turn_id), session_factory=sf, initial_seq=max_seq,
)
await _finalise(turn_id, "completed", sf, writer=recovery_writer)
return
# initial_seq=max_seq: on a fresh turn max_seq==2 (POST wrote user_input +
# status); on an init-only recovery it continues the dense log past the
# init rows so the re-driven writes don't collide on (run_id, seq).
writer = StateWriter(
schema=ChatTurnState, run_id=turn_id, table=TURN_CHANGES_TABLE,
channel=_channel(turn_id), session_factory=sf, initial_seq=max_seq,
)
await writer.set("status", "streaming")
toolsets = build_servers(settings.mcp_server_urls, cookie=cookie)
driver = PydanticAIDriver(text_field="text", steps_field="steps")
run_input = RunInput(
prompt=user_input, system_prompt=CHAT_SYSTEM_PROMPT,
history=tuple(history), limits=Limits(requests=10, tool_calls=20),
)
llm = LLMConfig(model_slug=model_slug,
gateway_base_url=settings.gateway_base_url, cookie=cookie)
main_task = asyncio.current_task()
hb = asyncio.create_task(_heartbeat_loop(turn_id, sf, main_task))
try:
await driver.run(
run_input=run_input, tools=[], writer=writer, llm=llm,
toolsets=toolsets,
)
# Phase B: the change-log IS the durable transcript — nothing to
# derive. The driver's WriterEventAdapter already flushed text + steps
# to turn_changes; just mirror the terminal status.
await _finalise(turn_id, "completed", sf, writer=writer)
await _maybe_trigger_title(turn_id, cookie, sf)
except UsageLimitExceeded:
# Partial text/steps already live in turn_changes (the sole source).
await _finalise(turn_id, "failed", sf, writer=writer,
error="usage_limit_exceeded")
# match prior behaviour: usage-limit is not re-raised (turn already failed)
except asyncio.CancelledError:
# Partial assistant activity already lives in turn_changes (the driver
# flushed it pre-cancel); just mirror the terminal status.
async with sf() as s:
st = (await s.execute(text("SELECT status FROM chat.turns WHERE id=:id"),
{"id": turn_id})).scalar_one_or_none()
if st == "streaming":
await _finalise(turn_id, "cancelled", sf, writer=writer)
else:
# The cancel endpoint already flipped chat.turns.status to a
# terminal value before we noticed; _finalise's `if streaming`
# guard would skip it, so the terminal change is never written to
# turn_changes and a live SSE (now listening on the change-log
# channel) never closes. Mirror it through OUR OWN writer (no
# cross-writer race on the (run_id, seq) unique constraint).
await writer.set("status", st)
raise
except Exception as exc:
_log.exception("turn=%d failed", turn_id)
await _finalise(turn_id, "failed", sf, writer=writer, error=str(exc)[:500])
raise
finally:
hb.cancel()
@DBOS.workflow()
async def chat_turn_workflow(turn_id: int, cookie: str) -> None:
"""Drive one chat turn. Survives chat-service restart via DBOS recovery."""
await _run_turn(turn_id, cookie)
async def _maybe_trigger_title(turn_id: int, cookie: str, sf) -> None:
"""If the conversation still has a placeholder title, kick off title-gen.
Runs AFTER the turn is marked completed — the title workflow's first
query asks for the conversation's first `completed` turn, so triggering
it earlier loses a race and the workflow silently returns "no completed
turn"."""
async with sf() as s:
check = (await s.execute(text("""
SELECT c.title, c.id FROM chat.conversations c
JOIN chat.turns t ON t.conversation_id = c.id
WHERE t.id = :tid
"""), {"tid": turn_id})).first()
current_title = check[0] if check else None
conv_id_for_title = check[1] if check else None
if not conv_id_for_title or current_title not in _PLACEHOLDER_TITLES:
return
from dbos import DBOS as _DBOS
from dbos._dbos import _dbos_global_instance
if _dbos_global_instance is not None and _dbos_global_instance._launched:
try:
from .title_gen import generate_title_workflow
await _DBOS.start_workflow_async(generate_title_workflow, conv_id_for_title, cookie)
except Exception as _tg_exc:
_log.warning("title-gen trigger failed for conv=%d: %s",
conv_id_for_title, _tg_exc)
else:
_log.debug("DBOS not launched; skipping title-gen trigger")
async def _finalise(turn_id, status, sf, *, writer: StateWriter | None = None,
error: str | None = None):
async with sf() as s:
await s.execute(
text("UPDATE chat.turns SET status=:st, error=:err, finished_at=now() "
"WHERE id=:id"),
{"st": status, "err": error, "id": turn_id},
)
await s.execute(text(f"NOTIFY chat_turn_status_{int(turn_id)}"))
await s.commit()
if writer is not None:
await writer.set("status", status) # client-facing terminal → closes the SSE
async def _heartbeat_loop(turn_id: int, sf, main_task: asyncio.Task | None = None):
try:
while True:
await asyncio.sleep(HEARTBEAT_INTERVAL_SEC)
async with sf() as s:
st = (await s.execute(
text("SELECT status FROM chat.turns WHERE id = :id"),
{"id": turn_id},
)).scalar_one_or_none()
# Status flipped externally (cancel endpoint or reaper) —
# propagate to the main workflow task and stop heartbeating.
if st is not None and st != "streaming":
if main_task is not None and not main_task.done():
main_task.cancel()
return
await s.execute(
text("UPDATE chat.turns SET last_heartbeat = now() WHERE id = :id"),
{"id": turn_id},
)
await s.commit()
except asyncio.CancelledError:
return
async def _set_status(turn_id: int, status: str, sf, *, error: str | None = None):
async with sf() as s:
await s.execute(
text("UPDATE chat.turns SET status = :st, error = :err, finished_at = now() "
"WHERE id = :id"),
{"st": status, "err": error, "id": turn_id},
)
await s.execute(text(f"NOTIFY chat_turn_status_{int(turn_id)}"))
await s.commit()
"""Declared run-state for the slides agent (the frontend contract).
DeckRunState is **deck-scoped** (spec §5/§6): the change-log `run_id`
column holds the deck id, the channel is `slides_deck_{deck_id}`, and
`seq` is monotonic per deck — one open stream sees all runs of the deck
(this is what kills the old 200 ms poll). Runs are serialized per deck
(one streaming run + a mutation lock), so a per-StateWriter counter
suffices.
Typed fields drive the TS codegen; `modes` drive StateWriter op
validation. The agent fills these progressively:
status/progress/title/global_css/brainstorm — scalars (set)
slides — a keyed growing list (append / merge / remove)
text — agent narration (finalized chunks, append) — kept to
preserve the current trace; driver writes it
steps — the agent activity log (append), written by the driver's
WriterEventAdapter; binary tool-returns sanitized to a marker.
The durable canvas (slides.decks / slides.slides / slide_versions) stays
the product source of truth — the tools keep mutating it. This document
is the LIVE stream layered on top (like chat's message_parts vs
turn_changes).
"""
from __future__ import annotations
from pydantic import BaseModel
from mlexi_agent.state import FieldSpec, Kind, RunState
class Slide(BaseModel):
"""One slide as the frontend renders it (keyed by `id`)."""
id: int
position: int
html: str = ""
notes: str | None = None
class DeckRunState(RunState):
status: str = "streaming" # SCALAR current run lifecycle
progress: dict = {} # SCALAR {done, total}
title: str = "" # SCALAR
global_css: str = "" # SCALAR (large string ok)
slides: list[dict] = [] # LIST key="id" [append, merge, remove]
text: str = "" # TEXT agent narration (trace), append
steps: list[dict] = [] # LIST key="tool_call_id" (agent trace)
brainstorm: dict | None = None # SCALAR pending pause, or null
# Phase B: the user's question. Written to the deck change-log by the
# POST handler before firing the workflow (replacing the dropped
# run_parts user-prompt). The workflow reads it back from the folded
# log; the deck-scoped log carries the LATEST user_input per run.
user_input: str = "" # SCALAR the user's prompt for the run
modes = {
"status": FieldSpec(Kind.SCALAR),
"progress": FieldSpec(Kind.SCALAR),
"title": FieldSpec(Kind.SCALAR),
"global_css": FieldSpec(Kind.SCALAR),
"slides": FieldSpec(Kind.LIST, key="id"),
"text": FieldSpec(Kind.TEXT),
"steps": FieldSpec(Kind.LIST, key="tool_call_id"),
"brainstorm": FieldSpec(Kind.SCALAR),
"user_input": FieldSpec(Kind.SCALAR),
}
"""Host-built slides tools (injected into the shared agent driver).
`build_slides_tools(...)` ports the 14 local `@agent.tool_plain` closures
from `build.py` into `mlexi_agent.Tool` instances. Each tool KEEPS its
durable mutation of the slides tables (decks / slides / slide_versions —
the product source of truth, reusing the `_mutation_lock` /
`_locked_session` pattern from build.py) AND emits the matching
DECK-STATE change through the injected `writer`:
create_slide → insert + writer.append("slides", {id, position, html})
update_slide → version + writer.merge("slides", key=id, {html})
patch_slide → version + writer.merge("slides", key=id, {html})
delete_slide → delete/shift + writer.remove("slides", key=id)
+ re-merge the shifted slides' positions
set_global_css → writer.set("global_css", css)
patch_global_css → writer.set("global_css", new_css)
plan_deck → writer.set("progress", {done, total})
list_slides / read_slide / read_global_css / search_image → read-only
render_slide → returns [header, BinaryContent] to the MODEL
(vision); the driver's adapter records the step
(binary sanitized to a marker); no deck-state write.
ask_user → the durable-pause tool (sets/clears state.brainstorm;
INSERT brainstorm_requests; LISTEN/resolve).
The generic agent trace (every tool call/return) is written by the
driver's WriterEventAdapter — tools only write **product-state** fields.
MCP servers (image-MCP for generate_image, etc.) pass through as
`toolsets=` at the driver, not here.
"""
from __future__ import annotations
import asyncio
import logging
from contextlib import asynccontextmanager
import asyncpg
from pydantic_ai.messages import BinaryContent
from sqlalchemy import select, text
from mlexi_agent import StateWriter, Tool
from . import image_search, marp_renderer
from ..config import get_settings
from ..models import BrainstormRequest, Slide
from ..repo import decks as decks_repo
from ..repo import events as events_repo
from ..repo import slides as slides_repo
from .build import _summarize_render_defects # reuse the rich defect summary
_log = logging.getLogger("slides.agent.tools")
# Default brainstorm wait timeout — same generous 1h budget chat / image
# use. If the user walks away, the tool returns a "no reply" sentinel and
# the agent proceeds rather than hanging the workflow forever.
BRAINSTORM_WAIT_TIMEOUT_SEC = 3600
def _dsn_for_listen(sa_url: str) -> str:
"""Convert SQLAlchemy URL → asyncpg DSN (drop driver suffix)."""
return sa_url.replace("postgresql+asyncpg://", "postgresql://")
def build_slides_tools(
*, writer: StateWriter, sf, cookie: str, deck_id: int, run_id: int,
user_id: int, allow_overdraft: bool,
) -> list[Tool]:
"""Build the flat slide tool set bound to one deck/run.
`sf` is the SQLAlchemy session factory; `writer` is the deck-scoped
StateWriter. `user_id` / `allow_overdraft` are accepted for billing
symmetry with image/chat (image-gen charges via image-MCP's own
accounting today), so they're currently unreferenced.
"""
_ = (user_id, allow_overdraft)
settings = get_settings()
# Deck-mutation lock. PydanticAI executes batched tool calls in
# parallel via asyncio; create/update/patch/delete + global-css all
# mutate the deck. Without serialisation create_slide races on
# slides_deck_pos_uq (two calls pick the same position). Mirrors
# build.py's _mutation_lock.
_mutation_lock = asyncio.Lock()
@asynccontextmanager
async def _locked_session():
async with _mutation_lock:
async with sf() as s:
yield s
@asynccontextmanager
async def _session():
async with sf() as s:
yield s
async def _slide_state(s, position: int) -> dict | None:
"""Snapshot a slide's deck-state shape by position, or None."""
row = (await s.execute(
select(Slide).where(
Slide.deck_id == deck_id, Slide.position == int(position),
)
)).scalar_one_or_none()
if row is None:
return None
return {"id": int(row.id), "position": int(row.position),
"html": row.markdown or ""}
# ── planning / read-only ─────────────────────────────────────────
async def plan_deck(plan: str) -> dict:
"""Record the deck's plan and visual direction as a single text
blob. Call ONCE early so the user can see your outline. Re-call
only if the user explicitly redirects.
Args:
plan: Free-form text — intent summary + per-slide outline
(position, role, headline). One paragraph + a short
list is plenty.
"""
async with _session() as s:
await events_repo.write_event(
s, deck_id=deck_id, run_id=run_id,
kind="deck_planned",
payload={"plan": plan, "length": len(plan)},
)
await s.commit()
done = (await s.execute(
text("SELECT count(*) FROM slides.slides WHERE deck_id=:d"),
{"d": deck_id})).scalar_one()
# Estimate a total slide count from the outline's bullet/line shape
# so the SPA can show a determinate bar; fall back to done.
total = max(int(done or 0), _estimate_total(plan))
await writer.set("progress", {"done": int(done or 0), "total": total})
return {"ok": True, "length": len(plan)}
async def list_slides() -> dict:
"""List all slides in this deck (position + short description).
Cheap — use to orient before edits."""
async with _session() as s:
rows = await slides_repo.list_slides(s, deck_id=deck_id)
return {"slides": rows}
async def read_slide(position: int) -> dict:
"""Get the full Markdown of one slide.
Args:
position: 0-based index of the slide.
"""
try:
async with _session() as s:
return await slides_repo.read_slide(
s, deck_id=deck_id, index=int(position),
)
except slides_repo.SlideNotFound:
return {"error": f"slide {position} not found"}
# ── slide mutation ───────────────────────────────────────────────
async def create_slide(
position: int, markdown: str, layout: str = "default",
) -> dict:
# … (полный файл: 588 строк — путь выше) …
await conn.close()
except Exception:
pass
async def _all_slide_states(s) -> list[dict]:
rows = (await s.execute(
select(Slide).where(Slide.deck_id == deck_id).order_by(Slide.position)
)).scalars().all()
return [{"id": int(r.id), "position": int(r.position),
"html": r.markdown or ""} for r in rows]
return [
Tool(plan_deck),
Tool(list_slides),
Tool(read_slide),
Tool(create_slide),
Tool(update_slide),
Tool(patch_slide),
Tool(delete_slide),
Tool(read_global_css),
Tool(set_global_css),
Tool(patch_global_css),
Tool(search_image),
Tool(render_slide),
# ask_user is the brainstorm durable-pause tool. Named `ask_user`
# per spec §4.5; the SLIDES_SYSTEM_PROMPT references `brainstorm`,
# so expose it under the brainstorm name too (one Tool instance).
Tool(ask_user, name="brainstorm"),
]
"""DBOS workflow: drive one slides authoring run via the shared mlexi_agent runtime.
`_run_deck` is the composition root (plain async, testable). It loads the
run → deck_id/model/user-prompt, wires the injected slide tools + the MCP
toolsets + a DECK-scoped StateWriter over slides.deck_changes, and runs the
shared PydanticAIDriver. The durable canvas (slides.decks / slides.slides /
slide_versions) stays the product source of truth — the tools mutate it;
deck_changes is the LIVE frontend stream layered on top.
Recovery is RUN-scoped (slides.runs.start_seq), not change-log-scoped:
deck_changes is deck-scoped, so a non-empty change-log only means the DECK
has prior changes, not that THIS run drove the agent. On DBOS recovery a
NON-NULL `start_seq` means this run already ran (re-driving would duplicate
slides), so the workflow finalises + mirrors the existing terminal status
instead — without resurrecting a cancelled/failed run to "completed".
start_seq also anchors this run's slice of the deck log (seq > start_seq)
and the user_input read (the latest user_input write at-or-before start_seq).
Phase B: slides.deck_changes is the SOLE durable source — there is no
run_parts derivation. The user's question is the `user_input` SCALAR, written
to the deck log by the POST handler before the workflow fires; the workflow
reads it back by folding the deck log. The live frontend contract and the GET
snapshot's agent trace both run on the folded change-log.
Mirrors services/image/agent/workflow.py + services/chat/agent/workflow.py
+ services/search/agent/workflow.py. Survives slides-service restart via DBOS
recovery. Title-gen unchanged.
"""
from __future__ import annotations
import asyncio
import logging
from dbos import DBOS
from pydantic_ai.exceptions import UsageLimitExceeded
from sqlalchemy import text
from mlexi_agent import LLMConfig, Limits, RunInput, StateWriter
from mlexi_agent.fold import fold_changes
from mlexi_agent.pydantic_driver import PydanticAIDriver
from ..config import get_settings
from ..db import get_session_factory
from .mcp_servers import build_servers
from .state import DeckRunState
from .system_prompt import SLIDES_SYSTEM_PROMPT
from .tools import build_slides_tools
_log = logging.getLogger("slides.workflow")
HEARTBEAT_INTERVAL_SEC = 5
DECK_CHANGES_TABLE = "slides.deck_changes"
def _channel(deck_id: int) -> str:
return f"slides_deck_{int(deck_id)}"
async def _deck_max_seq(sf, deck_id: int) -> int:
async with sf() as s:
return (await s.execute(text(
"SELECT COALESCE(MAX(seq),0) FROM slides.deck_changes WHERE run_id=:d"),
{"d": deck_id})).scalar_one()
async def _fold_deck_user_input(sf, deck_id: int, *, up_to_seq: int) -> str:
"""Read the latest user_input from the deck change-log up to a seq.
Phase B: the user's question is the `user_input` SCALAR the POST handler
wrote to the deck log just before firing this run's workflow. The log is
deck-scoped, so folding it gives the LATEST user_input — which is exactly
this run's prompt, because the POST handler wrote it as the highest seq
at-or-before start_seq (no later run can have written one before this run
records its start_seq under the per-deck single-streaming-run invariant)."""
async with sf() as s:
rows = (await s.execute(
text("SELECT seq, op, path, value FROM slides.deck_changes "
"WHERE run_id=:d AND seq<=:s ORDER BY seq"),
{"d": deck_id, "s": int(up_to_seq)},
)).all()
return fold_changes(DeckRunState, rows).user_input
async def _run_deck(run_id: int, cookie: str) -> None:
"""Workflow body extracted as a plain async function for testability."""
settings = get_settings()
sf = get_session_factory()
async with sf() as s:
row = (await s.execute(
text("SELECT model, deck_id, start_seq "
"FROM slides.runs WHERE id = :id"),
{"id": run_id},
)).first()
if row is None:
raise RuntimeError(f"run {run_id} not found")
model_slug, deck_id = row[0], int(row[1])
start_seq = None if row[2] is None else int(row[2])
# RUN-scoped recovery guard. deck_changes is deck-scoped, so its max_seq
# only tells us the DECK has prior changes — not that THIS run drove the
# agent. start_seq is set once (below) before driving; if we re-enter
# with it already NON-NULL, this is a DBOS recovery of a run that already
# ran. Re-driving would duplicate slides — and the trace already lives in
# deck_changes (the sole source, no run_parts to derive) — so we finalise
# WITHOUT resurrecting a terminal run: mirror the run's CURRENT status into
# deck_changes (only flip to completed if it's still streaming) so a
# reconnecting SSE sees a terminal change and closes.
if start_seq is not None:
max_seq = await _deck_max_seq(sf, deck_id)
_log.warning("run=%d recovery detected (start_seq=%d); finalising",
run_id, start_seq)
recovery_writer = StateWriter(
schema=DeckRunState, run_id=deck_id, table=DECK_CHANGES_TABLE,
channel=_channel(deck_id), session_factory=sf, initial_seq=max_seq,
)
async with sf() as s:
cur = (await s.execute(
text("SELECT status FROM slides.runs WHERE id = :id"),
{"id": run_id})).scalar_one_or_none()
if cur == "streaming":
# _set_status guards WHERE status='streaming', so this only
# commits if the run is genuinely still in-flight.
await _set_status(run_id, "completed", sf)
await _maybe_kick_title_gen(deck_id, cookie, sf)
await recovery_writer.set("status", "completed")
else:
# Already terminal (cancel endpoint / earlier finalise won the
# race). Mirror the existing terminal status into the change-log
# WITHOUT flipping the DB status back to completed.
await recovery_writer.set("status", cur or "completed")
return
# Phase B: the user's question is the `user_input` the POST handler wrote
# to the deck log just before firing this workflow. Read it back from the
# folded deck log (the sole source — no run_parts). The current deck max
# seq is also recorded below as start_seq (this run's slice anchor).
pre_seq = await _deck_max_seq(sf, deck_id)
user_input = await _fold_deck_user_input(sf, deck_id, up_to_seq=pre_seq)
if not user_input:
# No user_input in the deck log (a recovery where the POST handler
# never wrote it / corrupt log) — nothing to drive, finalise.
_log.warning(
"run=%d resume detected (no user_input in deck log); finalising",
run_id,
)
await _set_status(run_id, "completed", sf)
await _maybe_kick_title_gen(deck_id, cookie, sf)
return
# Record start_seq = the deck's current max seq BEFORE driving, in its own
# committed txn. This both (a) marks the run as agent-driven so a crash
# mid-run is recognised as recovery (not re-driven) and (b) anchors this
# run's slice of the deck-scoped log (seq > start_seq).
start_seq = pre_seq
async with sf() as s:
await s.execute(
text("UPDATE slides.runs SET start_seq = :ss WHERE id = :id"),
{"ss": int(start_seq), "id": run_id})
await s.commit()
# Resolve user_id + allow_overdraft via identity.whoami.
from ..deps import get_identity_client
user_id = 0
allow_overdraft = False
try:
whoami = await get_identity_client().whoami(session_cookie=cookie)
user_id = whoami.user_id
allow_overdraft = bool(getattr(whoami, "is_unlimited", False))
except Exception as exc:
_log.warning("run=%d whoami failed (%s); falling back to deck.user_id", run_id, exc)
async with sf() as s:
user_row = (await s.execute(
text("SELECT user_id FROM slides.decks WHERE id = :did"),
{"did": deck_id})).first()
user_id = user_row[0] if user_row else 0
# DECK-scoped StateWriter, seeded at the deck's current max seq (==
# start_seq, just recorded) so a new run's changes continue the deck's
# monotonic log (one stream sees all runs of the deck — what removes the
# old 200 ms poll).
writer = StateWriter(
schema=DeckRunState, run_id=deck_id, table=DECK_CHANGES_TABLE,
channel=_channel(deck_id), session_factory=sf, initial_seq=start_seq,
)
await writer.set("status", "streaming")
tools = build_slides_tools(
writer=writer, sf=sf, cookie=cookie, deck_id=deck_id, run_id=run_id,
user_id=int(user_id), allow_overdraft=allow_overdraft,
)
toolsets = build_servers(settings.mcp_server_urls, cookie=cookie)
driver = PydanticAIDriver(text_field="text", steps_field="steps")
run_input = RunInput(
prompt=user_input, system_prompt=SLIDES_SYSTEM_PROMPT,
history=(), # slides runs are single-shot; deck state read via tools.
# Generous backstop, not a creative leash: a thorough deck legitimately
# runs to dozens of requests / >100 tool calls.
limits=Limits(requests=90, tool_calls=220),
)
llm = LLMConfig(model_slug=model_slug,
gateway_base_url=settings.gateway_base_url, cookie=cookie)
main_task = asyncio.current_task()
hb = asyncio.create_task(_heartbeat_loop(run_id, sf, main_task))
try:
await driver.run(
run_input=run_input, tools=tools, writer=writer, llm=llm,
toolsets=toolsets,
)
# Phase B: deck_changes IS the durable trace — nothing to derive. The
# driver's WriterEventAdapter already flushed text + steps to the deck
# log; just mark terminal + mirror the status.
await _set_status(run_id, "completed", sf)
await writer.set("status", "completed")
# F-M1: final determinate progress so a progress bar reaches done==total.
await _emit_final_progress(deck_id, sf, writer)
await _maybe_kick_title_gen(deck_id, cookie, sf)
except UsageLimitExceeded:
await _finalize_after_usage_limit(
run_id, deck_id, cookie, sf, writer=writer)
except asyncio.CancelledError:
# Phase B: the partial trace already lives in deck_changes (the driver
# flushed it pre-cancel); just mirror the terminal status.
async with sf() as s:
st = (await s.execute(
text("SELECT status FROM slides.runs WHERE id = :id"),
{"id": run_id})).scalar_one_or_none()
if st == "streaming":
await _set_status(run_id, "cancelled", sf)
await writer.set("status", "cancelled")
else:
# The cancel endpoint already flipped slides.runs.status before
# we noticed; _set_status now guards WHERE status='streaming' so
# it would no-op — but the writer's terminal mirror must fire so a
# live SSE on the change-log sees the terminal. Mirror the
# already-terminal status via OUR writer (no cross-writer race on
# the (run_id, seq) unique constraint).
await writer.set("status", st or "cancelled")
raise
except Exception as exc:
_log.exception("run=%d failed", run_id)
# Phase B: partial trace already in deck_changes; mirror the status.
await _set_status(run_id, "failed", sf, error=str(exc)[:500])
await writer.set("status", "failed")
raise
finally:
hb.cancel()
@DBOS.workflow()
async def slides_run_workflow(run_id: int, cookie: str) -> None:
"""Drive one slides run. Survives slides-service restart via DBOS recovery."""
await _run_deck(run_id, cookie)
async def _heartbeat_loop(run_id: int, sf, main_task: asyncio.Task | None = None):
try:
while True:
await asyncio.sleep(HEARTBEAT_INTERVAL_SEC)
async with sf() as s:
status_row = (await s.execute(
text("SELECT status FROM slides.runs WHERE id = :id"),
{"id": run_id})).scalar_one_or_none()
if status_row is not None and status_row != "streaming":
if main_task is not None and not main_task.done():
main_task.cancel()
return
await s.execute(
text("UPDATE slides.runs SET last_heartbeat = now() WHERE id = :id"),
{"id": run_id})
await s.commit()
except asyncio.CancelledError:
return
async def _finalize_after_usage_limit(
run_id: int, deck_id: int, cookie: str, sf, *,
writer: StateWriter | None = None,
) -> None:
"""Disposition a run that exhausted its (generous) usage backstop.
The authoring tools persist slides eagerly, so a run that hit the cap
mid-self-review usually leaves a real, usable deck behind. Treat that
as a soft completion — mark completed and run title-gen — rather than
surfacing an "error" pill. Only a run that produced no slides at all
is a genuine failure worth flagging.
Phase B: the partial trace already lives in deck_changes (the sole
source) — nothing to derive.
"""
async with sf() as s:
n_slides = (await s.execute(
text("SELECT count(*) FROM slides.slides WHERE deck_id = :d"),
{"d": deck_id})).scalar_one()
if int(n_slides or 0) > 0:
_log.info("run=%d hit usage cap with %d slide(s); soft-completing",
run_id, int(n_slides))
await _set_status(run_id, "completed", sf)
if writer is not None:
await writer.set("status", "completed")
await _emit_final_progress(deck_id, sf, writer)
await _maybe_kick_title_gen(deck_id, cookie, sf)
else:
await _set_status(run_id, "failed", sf, error="usage_limit_exceeded")
if writer is not None:
await writer.set("status", "failed")
async def _emit_final_progress(deck_id: int, sf, writer: StateWriter) -> None:
"""F-M1: set a final determinate progress {done, total} on completion.
Best-effort from the deck's slide count so a determinate progress bar
reaches done==total. Swallows errors — progress is a hint, never a
contract, and must not break finalisation.
"""
try:
async with sf() as s:
n = (await s.execute(
text("SELECT count(*) FROM slides.slides WHERE deck_id = :d"),
{"d": deck_id})).scalar_one()
n = int(n or 0)
await writer.set("progress", {"done": n, "total": n})
except Exception as exc:
_log.warning("deck=%d final progress emit failed: %s", deck_id, exc)
async def _maybe_kick_title_gen(deck_id: int, cookie: str, sf) -> None:
"""Fire title-gen workflow if this deck still has a placeholder title.
Fire-and-forget; child workflow lives in DBOS."""
async with sf() as s:
row = (await s.execute(
text("SELECT title FROM slides.decks WHERE id = :id"),
{"id": deck_id})).first()
if row is None:
return
title = row[0]
if title and title not in ("", "Без названия", "Untitled"):
return # user-set title — leave alone
try:
from dbos import SetWorkflowID
from .title_gen import generate_title_workflow
wf_id = f"slides-title-gen-{int(deck_id)}"
with SetWorkflowID(wf_id):
DBOS.start_workflow(generate_title_workflow, deck_id, cookie)
except Exception as exc:
_log.warning("title-gen kick failed for deck=%d: %s", deck_id, exc)
async def _set_status(run_id: int, status: str, sf, *, error: str | None = None):
async with sf() as s:
# F-I2: only transition a run that is still streaming. A cancel
# endpoint (or an earlier finalise) may have already flipped the run
# to a terminal status; this guard means a late completed/failed
# write can never resurrect a cancelled/failed run. RETURNING is
# empty (row is None) when the guard skips the row — the legacy
# lifecycle event below is then correctly NOT re-emitted.
row = (await s.execute(
text("""
UPDATE slides.runs
SET status = :st, error = :err, ended_at = now()
WHERE id = :id AND status = 'streaming'
RETURNING deck_id
"""),
{"st": status, "err": error, "id": run_id},
)).first()
await s.execute(text(f"NOTIFY slides_run_status_{int(run_id)}"))
# Keep the legacy lifecycle event so any not-yet-migrated bootstrap
# path still flips its status pill.
if row is not None:
from ..repo.events import write_event
if status == "completed":
await write_event(
s, deck_id=row.deck_id, run_id=run_id,
kind="run_completed", payload={"run_id": run_id})
else:
await write_event(
s, deck_id=row.deck_id, run_id=run_id,
kind="run_failed",
payload={"run_id": run_id, "status": status,
"error": error or status})
await s.commit()
"""Declared run-state for the image agent (the frontend contract).
Typed fields (drive the TS codegen in Phase 1c) + `modes` (drive
StateWriter op validation). The agent fills these progressively:
status/stage as scalars, images as a keyed growing list, steps as the
activity log written by the driver's WriterEventAdapter.
"""
from __future__ import annotations
from pydantic import BaseModel
from mlexi_agent.state import FieldSpec, Kind, RunState
class ImageItem(BaseModel):
id: int
url: str
prompt: str
aspect: str
width: int | None = None
height: int | None = None
class ImageRunState(RunState):
user_input: str = ""
status: str = "streaming"
stage: str = ""
images: list[ImageItem] = []
steps: list[dict] = []
modes = {
"user_input": FieldSpec(Kind.SCALAR),
"status": FieldSpec(Kind.SCALAR),
"stage": FieldSpec(Kind.SCALAR),
"images": FieldSpec(Kind.LIST, key="id"),
"steps": FieldSpec(Kind.LIST, key="tool_call_id"),
}
"""Host-built image tools (injected into the driver).
Each factory binds db/billing/cookie/user + the StateWriter for this run
and returns a plain async function the driver registers as a tool. Tool
bodies write domain state through the writer (stage + images).
"""
from __future__ import annotations
from uuid import uuid4
from mlexi_agent import StateWriter, Tool
from mlexi_clients import BillingInsufficientCreditsError
from .. import gateway_client, storage
from ..db import get_session_factory
from ..deps import get_billing_client
from ..repo import images as images_repo
from ..repo import settings as settings_repo
def build_image_tools(
*, writer: StateWriter, cookie: str, user_id: int,
allow_overdraft: bool, run_id: int,
) -> list[Tool]:
billing = get_billing_client()
sf = get_session_factory()
async def enhance_prompt(prompt: str) -> dict:
"""Rewrite a casual image request into a detailed prompt.
Args:
prompt: The user's original image description.
"""
await writer.set("stage", "Улучшаю запрос")
async with sf() as s:
cfg = await settings_repo.get_all(s)
rewriter_model = cfg.get("prompt_rewriter_model")
if not rewriter_model:
return {"enhanced": prompt, "note": "prompt_rewriter_model not configured"}
try:
enhanced = await gateway_client.rewrite_prompt_via_chat_completions(
model=rewriter_model, user_text=prompt, cookie=cookie)
except gateway_client.GatewayError as exc:
return {"enhanced": prompt, "error": str(exc)}
return {"enhanced": enhanced}
async def generate_image(prompt: str, aspect: str = "square") -> dict:
"""Generate an image from a text prompt and save it.
Args:
prompt: The image-generation prompt (preferably detailed English).
aspect: 'square' | 'landscape' | 'portrait'.
"""
await writer.set("stage", "Генерирую изображение")
async with sf() as s:
cfg = await settings_repo.get_all(s)
gen_model = cfg.get("gen_model")
if not gen_model:
return {"error": "gen_model not configured"}
per_image = int(cfg.get("credits_per_image") or 10)
key = f"image-tool:{uuid4()}"
try:
res = await billing.reserve(
user_id=user_id, amount=per_image, reason="image-agent",
ref=key, idempotency_key=key, allow_overdraft=allow_overdraft)
except BillingInsufficientCreditsError as exc:
return {"error": f"insufficient_credits: {exc}"}
try:
png = await gateway_client.generate_via_images_endpoint(
model=gen_model, prompt=prompt, aspect=aspect, cookie=cookie)
except gateway_client.GatewayError as exc:
try:
await billing.release(reservation_id=res.reservation_id, used_amount=0)
except Exception:
pass
return {"error": f"generation_failed: {exc}"}
rel, size, w, h = await storage.write_image_bytes(png, "image/png")
async with sf() as s:
row = await images_repo.insert_image(
s, user_id=user_id, parent_id=None, prompt=prompt, aspect=aspect,
model_slug=gen_model, file_path=rel, file_size_bytes=size,
mime_type="image/png", width=w, height=h,
metadata={"run_id": run_id})
await s.commit()
image_id = row.id
await billing.confirm(reservation_id=res.reservation_id, actual_amount=per_image)
url = f"/api/me/images/{image_id}/file"
# Domain state: the finished image becomes a real RunState field.
await writer.append("images", {
"id": image_id, "url": url, "prompt": prompt, "aspect": aspect,
"width": w, "height": h,
})
return {"image_id": image_id, "url": url, "prompt": prompt,
"aspect": aspect, "width": w, "height": h}
return [Tool(enhance_prompt), Tool(generate_image)]
"""DBOS workflow: drive one image run via the shared mlexi_agent runtime."""
from __future__ import annotations
import asyncio
import logging
from dbos import DBOS
from sqlalchemy import text
from mlexi_agent import LLMConfig, Limits, RunInput, StateWriter
from mlexi_agent.pydantic_driver import PydanticAIDriver
from ..config import get_settings
from ..db import get_session_factory
from ..deps import get_identity_client
from ..repo import settings as settings_repo
from .prompt import IMAGE_SYSTEM_PROMPT
from .state import ImageRunState
from .tools import build_image_tools
_log = logging.getLogger("image.workflow")
HEARTBEAT_INTERVAL_SEC = 5
RUN_CHANGES_TABLE = "image.run_changes"
def _channel(run_id: int) -> str:
return f"image_run_{int(run_id)}"
async def _run_image(run_id: int, cookie: str) -> None:
"""Workflow body extracted as a plain async function for testability."""
settings = get_settings()
sf = get_session_factory()
async with sf() as s:
gen = (await s.execute(text(
"SELECT g.user_id, g.original_prompt, g.aspect, g.enhance_pref "
"FROM image.generations g JOIN image.runs r ON r.generation_id=g.id "
"WHERE r.id=:id"), {"id": run_id})).first()
max_seq = (await s.execute(text(
"SELECT COALESCE(MAX(seq),0) FROM image.run_changes WHERE run_id=:id"),
{"id": run_id})).scalar_one()
# Real product OUTPUT = an `images` or `steps` change-log row. The
# workflow's OWN first writes are set("status",...) + set("stage",...),
# so `max_seq > 0` is true the instant the body starts — it does NOT
# mean the agent produced anything. A crash AFTER those init writes but
# BEFORE output must RE-RUN, not finalise an empty run (the search F2
# lesson).
has_output = (await s.execute(text(
"SELECT EXISTS(SELECT 1 FROM image.run_changes "
"WHERE run_id=:id AND path IN ('images','steps'))"),
{"id": run_id})).scalar_one()
if gen is None:
raise RuntimeError(f"run {run_id} not found")
fallback_user_id, prompt, aspect, enhance_pref = gen
# DBOS recovery: this run already produced REAL OUTPUT (images/steps) before
# a crash; finalise rather than re-driving from a partial state. Pass a
# writer (seeded at max_seq) so the terminal `status` is mirrored into
# run_changes too — otherwise a reconnecting SSE never sees a terminal
# change and hangs forever (last run_changes status is still "streaming").
# Init-only change-logs (status/stage but no output) fall through to drive
# the agent.
if has_output:
_log.warning("run=%d resume detected (seq=%d, has_output); finalising",
run_id, max_seq)
recovery_writer = StateWriter(
schema=ImageRunState, run_id=run_id, table=RUN_CHANGES_TABLE,
channel=_channel(run_id), session_factory=sf, initial_seq=max_seq,
)
await _finalise(run_id, "completed", sf, writer=recovery_writer)
return
user_id, allow_overdraft = fallback_user_id, False
try:
whoami = await get_identity_client().whoami(session_cookie=cookie)
user_id = whoami.user_id
allow_overdraft = bool(getattr(whoami, "is_unlimited", False))
except Exception as exc:
_log.warning("run=%d whoami failed (%s); using generation.user_id", run_id, exc)
async with sf() as s:
cfg = await settings_repo.get_all(s)
reasoning_model = (cfg.get("prompt_rewriter_model")
or cfg.get("agent_chat_model") or "codexapi/gpt-5.3-codex")
user_input = prompt
if enhance_pref == "off":
user_input += ("\n\n[Override: пользователь выключил AI-улучшение prompt. "
"НЕ вызывай enhance_prompt — передай prompt в generate_image как есть.]")
elif enhance_pref == "on":
user_input += ("\n\n[Override: пользователь попросил улучшить prompt. "
"ОБЯЗАТЕЛЬНО вызови enhance_prompt первым, затем generate_image.]")
if aspect and aspect != "square":
user_input += f"\n\n[aspect: {aspect}]"
# initial_seq=max_seq: on a fresh run max_seq==0; on an init-only recovery
# (status/stage written, no output) it continues the dense log past the
# init rows so the re-driven writes don't collide on (run_id, seq).
writer = StateWriter(
schema=ImageRunState, run_id=run_id, table=RUN_CHANGES_TABLE,
channel=_channel(run_id), session_factory=sf, initial_seq=max_seq,
)
await writer.set("status", "streaming")
# Carry the user's original question in the change-log (the sole source of
# truth post Phase B). `prompt` is generations.original_prompt; the augmented
# `user_input` (with enhance/aspect override hints) is what we feed the LLM.
await writer.set("user_input", prompt)
await writer.set("stage", "Думаю")
tools = build_image_tools(
writer=writer, cookie=cookie, user_id=user_id,
allow_overdraft=allow_overdraft, run_id=run_id)
driver = PydanticAIDriver(text_field=None, steps_field="steps")
run_input = RunInput(
prompt=user_input, system_prompt=IMAGE_SYSTEM_PROMPT,
limits=Limits(requests=5, tool_calls=5))
llm = LLMConfig(model_slug=reasoning_model,
gateway_base_url=settings.gateway_base_url, cookie=cookie)
main_task = asyncio.current_task()
hb = asyncio.create_task(_heartbeat_loop(run_id, sf, main_task))
try:
await driver.run(run_input=run_input, tools=tools, writer=writer, llm=llm)
await writer.set("stage", "Готово")
await _finalise(run_id, "completed", sf, writer=writer)
except asyncio.CancelledError:
async with sf() as s:
st = (await s.execute(text("SELECT status FROM image.runs WHERE id=:id"),
{"id": run_id})).scalar_one_or_none()
if st == "streaming":
await _finalise(run_id, "cancelled", sf, writer=writer)
raise
except Exception as exc:
_log.exception("run=%d failed", run_id)
await _finalise(run_id, "failed", sf, writer=writer, error=str(exc)[:500])
raise
finally:
hb.cancel()
@DBOS.workflow()
async def image_run_workflow(run_id: int, cookie: str) -> None:
"""Drive one image-generation run. Survives image-service restart via DBOS recovery."""
await _run_image(run_id, cookie)
async def _finalise(run_id, status, sf, *, writer: StateWriter | None = None, error=None):
async with sf() as s:
await s.execute(text(
"UPDATE image.runs SET status=:st, error=:err, finished_at=now() WHERE id=:id"),
{"st": status, "err": error, "id": run_id})
await s.commit()
if writer is not None:
await writer.set("status", status) # client-facing terminal → closes the SSE
async def _heartbeat_loop(run_id, sf, main_task):
try:
while True:
await asyncio.sleep(HEARTBEAT_INTERVAL_SEC)
async with sf() as s:
st = (await s.execute(text("SELECT status FROM image.runs WHERE id=:id"),
{"id": run_id})).scalar_one_or_none()
if st is not None and st != "streaming":
if main_task is not None and not main_task.done():
main_task.cancel()
return
await s.execute(text("UPDATE image.runs SET last_heartbeat=now() WHERE id=:id"),
{"id": run_id})
await s.commit()
except asyncio.CancelledError:
return
"""Declared run-state for the search agent (the frontend contract).
`SearchRunState` is the live streaming contract for one search run, in two
modes:
* **quick** — a single agent with Exa tools. The generic
`WriterEventAdapter` fills `text` (finalized answer narration) and
`steps` (the Exa tool-call/tool-return activity log). `mode="quick"`.
* **deep** — multi-agent fan-out. The DBOS workflow (composition root)
writes the orchestration fields EXPLICITLY through one shared
StateWriter: `plan` (set once), `branches` (each searcher appends a
"running" entry then merges its summary in, keyed by `branch`),
`coverage` (set per round), `answer` (set by the synthesizer). The
browser sub-agent appends `browser` entries carrying the base64
screenshot (transient — nulled by cleanup on completion).
`mode="deep"`.
Phase B: `search.run_changes` (the StateWriter change-log) is the SOLE
durable source for a run — `search.run_parts` is gone. The user's
question lives in the change-log too, as the `user_input` SCALAR
(written by the POST handler at seq 1). The static system prompt is NOT
stored — it's passed via RunInput.system_prompt. GET endpoints and
multi-turn history (`_load_prior_qa`) both fold this change-log into a
SearchRunState (see mlexi_agent.fold).
"""
from __future__ import annotations
from typing import Any
from pydantic import BaseModel
from mlexi_agent.state import FieldSpec, Kind, RunState
class SearchStep(BaseModel):
"""One quick-mode activity-log item (Exa trace). `dict` is also
accepted by the writer; this typed shape documents the contract."""
kind: str # "tool-call" | "tool-return"
tool_name: str
tool_call_id: str
args: dict | None = None # tool-call
content: Any = None # tool-return
class SearchBranch(BaseModel):
"""One deep-mode searcher branch (keyed by `branch`). Appended as
`{branch, status:"running"}`, then merged with its summary on done."""
branch: int
status: str = "running" # running | done | error
summary: str | None = None
key_sources: list[dict] | None = None
error: str | None = None
class BrowserAction(BaseModel):
"""One deep-mode browser screenshot entry (base64-transient)."""
action: str
url: str | None = None
image_b64: str | None = None # nulled by cleanup on completion
media_type: str | None = None
expired: bool = False
class SearchRunState(RunState):
user_input: str = "" # SCALAR the user's question (change-log seq 1)
status: str = "streaming" # SCALAR lifecycle marker
mode: str = "quick" # SCALAR quick | deep
plan: dict | None = None # SCALAR {sub_queries:[...], plan_rationale}
branches: list[dict] = [] # LIST key="branch" [append, merge]
coverage: dict | None = None # SCALAR {round, complete, rationale, ...}
answer: dict | None = None # SCALAR {content, citations[]}
text: str = "" # TEXT (append) — quick-mode answer narration
steps: list[dict] = [] # LIST key="tool_call_id" — quick-mode Exa trace
browser: list[dict] = [] # LIST (append) — {action, url, image_b64}
plan_extensions: list[dict] = [] # LIST (append) — deep follow-up rounds
modes = {
"user_input": FieldSpec(Kind.SCALAR),
"status": FieldSpec(Kind.SCALAR),
"mode": FieldSpec(Kind.SCALAR),
"plan": FieldSpec(Kind.SCALAR),
"coverage": FieldSpec(Kind.SCALAR),
"answer": FieldSpec(Kind.SCALAR),
"text": FieldSpec(Kind.TEXT),
"branches": FieldSpec(Kind.LIST, key="branch"),
"steps": FieldSpec(Kind.LIST, key="tool_call_id"),
"browser": FieldSpec(Kind.LIST),
"plan_extensions": FieldSpec(Kind.LIST),
}
"""Exa tools as injected `mlexi_agent.Tool`s (Phase 4).
Refactored from the old `@agent.tool_plain` decorators on the per-mode
PydanticAI agents into standalone async functions wrapped as
`mlexi_agent.Tool`. The billing closure (reserve → confirm/release per
call, with the platform's idempotency keys) is preserved unchanged; only
the registration mechanism moved (agent-decorator → injected Tool) so the
generic `PydanticAIDriver` can present them to the LLM.
`build_exa_tools(...)` binds the user/billing/credit-config context and
returns the two tools, used by BOTH the quick-mode driver and the deep
searcher driver.
"""
from __future__ import annotations
from uuid import uuid4
from mlexi_agent import Tool
from mlexi_clients import BillingInsufficientCreditsError, BillingServiceClient
from .. import exa_client
def build_exa_tools(
*,
user_id: int,
billing: BillingServiceClient,
allow_overdraft: bool,
per_search: int,
per_content: int,
reason: str = "search-agent",
key_prefix: str = "search-tool",
) -> list[Tool]:
"""Build the Exa search/fetch tools bound to one user's billing context.
`reason`/`key_prefix` differentiate quick-mode ("search-agent") from
deep-searcher ("search-deep") billing rows, matching the prior
per-mode behaviour.
"""
async def exa_search(query: str, top_k: int = 5) -> dict:
"""Search the web for current information. Returns sources with
titles, URLs, snippets and clean page text. Use this for facts
about recent events, statistics, technical details, or any
web-available content.
Args:
query: The search query.
top_k: Number of results to return (1-10, default 5).
"""
amount = per_search + top_k * per_content
key = f"{key_prefix}:{uuid4()}"
try:
res = await billing.reserve(
user_id=user_id, amount=amount,
reason=reason, ref=key, idempotency_key=key,
allow_overdraft=allow_overdraft,
)
except BillingInsufficientCreditsError as exc:
return {"error": f"insufficient_credits: {exc}"}
try:
results = await exa_client.exa_search_with_contents(
query=query, num_results=top_k,
)
except exa_client.ExaError as exc:
try:
await billing.release(reservation_id=res.reservation_id, used_amount=0)
except Exception:
pass
return {"error": f"exa_unavailable: {exc}"}
await billing.confirm(
reservation_id=res.reservation_id,
actual_amount=per_search + len(results) * per_content,
)
return {"sources": [
{
"position": i + 1, "url": r.url, "title": r.title,
"snippet": (r.content_md or "")[:280] if r.content_md else None,
"published_date": (
r.published_date.isoformat() if r.published_date else None
),
"content_md": r.content_md,
}
for i, r in enumerate(results)
]}
async def exa_fetch_url(url: str) -> dict:
"""Fetch the clean text content of a specific URL. Use when you
have a URL and need its full content beyond what search returned.
Args:
url: The URL to fetch (must be http:// or https://).
"""
if not url.startswith(("http://", "https://")):
return {"error": "url must be http(s)://"}
key = f"{key_prefix}:{uuid4()}"
try:
res = await billing.reserve(
user_id=user_id, amount=per_content,
reason=reason, ref=key, idempotency_key=key,
allow_overdraft=allow_overdraft,
)
except BillingInsufficientCreditsError as exc:
return {"error": f"insufficient_credits: {exc}"}
try:
r = await exa_client.exa_contents(url=url)
except exa_client.ExaError as exc:
try:
await billing.release(reservation_id=res.reservation_id, used_amount=0)
except Exception:
pass
return {"error": f"exa_unavailable: {exc}"}
await billing.confirm(
reservation_id=res.reservation_id, actual_amount=per_content,
)
return {
"url": url, "title": r.title, "content_md": r.content_md,
"published_date": (
r.published_date.isoformat() if r.published_date else None
),
}
return [Tool(exa_search), Tool(exa_fetch_url)]
"""DBOS workflow: drive one search run via the shared mlexi_agent runtime.
Quick mode (`_run_quick`, the composition root) reads the run's `user_input`
from the change-log (search.run_changes — written by the POST handler at seq
1), builds the neutral history by folding prior completed runs (final Q/A
only), wires the Exa tools as injected `Tool`s + a StateWriter over
search.run_changes, and runs the shared generic
`PydanticAIDriver(text_field="text", steps_field="steps")`.
Phase B: search.run_changes is the SOLE durable source — there is no
run_parts derivation. The live frontend contract, the GET snapshots, and
multi-turn history all run on the folded change-log. Survives search-service
restart via DBOS recovery; the recovery branch mirrors the terminal `status`
into run_changes so a reconnecting SSE closes.
Deep mode is dispatched to agent.deep.workflow (its own composition root).
"""
from __future__ import annotations
import asyncio
import logging
from dbos import DBOS
from pydantic_ai.exceptions import UsageLimitExceeded
from sqlalchemy import text
from mlexi_agent import LLMConfig, Limits, RunInput, StateWriter
from mlexi_agent.fold import fold_changes
from mlexi_agent.pydantic_driver import PydanticAIDriver
from ..config import get_settings
from ..db import get_session_factory
from ..deps import get_billing_client, get_identity_client
from ..repo import settings as settings_repo
from .deep.workflow import deep_research_workflow # re-export for dispatcher patching
from .exa_tools import build_exa_tools
from .history import state_to_neutral_messages
from .prompt import SEARCH_SYSTEM_PROMPT
from .state import SearchRunState
_log = logging.getLogger("search.workflow")
HEARTBEAT_INTERVAL_SEC = 5
RUN_CHANGES_TABLE = "search.run_changes"
def _channel(run_id: int) -> str:
return f"search_run_{int(run_id)}"
async def _run_quick(run_id: int, cookie: str) -> None:
"""Quick-mode workflow body — plain async, testable composition root."""
settings = get_settings()
sf = get_session_factory()
async with sf() as s:
row = (await s.execute(
text("SELECT model_slug, thread_id FROM search.runs WHERE id = :id"),
{"id": run_id},
)).first()
if row is None:
raise RuntimeError(f"run {run_id} not found")
model_slug, thread_id = row[0], row[1]
# Phase B: the user's question lives in the change-log (user_input,
# written by the POST handler at seq 1). Fold this run's change-log to
# read it back — the sole source of truth, no run_parts.
own_rows = (await s.execute(
text("SELECT seq, op, path, value FROM search.run_changes "
"WHERE run_id=:id ORDER BY seq"),
{"id": run_id},
)).all()
max_seq = own_rows[-1].seq if own_rows else 0
# Real OUTPUT = any change-log row whose path is NOT an init write. The
# POST handler writes user_input + mode + status (seq 1,2,3) and the
# workflow's OWN first write is set("status"), so `max_seq > 0` alone is
# true the instant the body starts — it does NOT mean the model produced
# anything. A crash AFTER init but BEFORE output must RE-RUN, not
# finalise empty (Codex F2).
has_output = any(
r.path not in ("status", "mode", "user_input") for r in own_rows)
# Prior-run history: fold the change-log of ALL earlier completed runs
# in this thread, chronological, into a SearchRunState each → neutral
# Messages (final Q/A only; per-searcher reasoning is NOT replayed).
prior_run_ids = [r.id for r in (await s.execute(text("""
SELECT id FROM search.runs
WHERE thread_id = :tid AND id < :rid AND status = 'completed'
ORDER BY started_at ASC
"""), {"tid": thread_id, "rid": run_id})).all()]
history = []
for pid in prior_run_ids:
prior_rows = (await s.execute(
text("SELECT seq, op, path, value FROM search.run_changes "
"WHERE run_id=:id ORDER BY seq"),
{"id": pid},
)).all()
prior_state = fold_changes(SearchRunState, prior_rows)
history.extend(state_to_neutral_messages(prior_state))
own_state = fold_changes(SearchRunState, own_rows)
user_input = own_state.user_input
# DBOS recovery: this run already produced REAL OUTPUT (text/steps) before a
# crash. Don't re-drive from a partial state — the change-log already holds
# it (the sole source). Finalise as completed AND mirror a terminal status
# into run_changes, else a reconnecting SSE never sees a terminal change and
# hangs. Init-only change-logs (status/mode/user_input, no output) fall
# through to drive the agent.
if has_output:
_log.warning("run=%d resume detected (seq=%d, has_output); finalising",
run_id, max_seq)
recovery_writer = StateWriter(
schema=SearchRunState, run_id=run_id, table=RUN_CHANGES_TABLE,
channel=_channel(run_id), session_factory=sf, initial_seq=max_seq,
)
await _finalise(run_id, "completed", sf, writer=recovery_writer)
return
# No user_input means the POST handler never wrote it (a recovery where the
# change-log is empty / corrupt) — nothing to drive, finalise as completed.
# Mirror a terminal `status` change through a writer — else the new
# change-log SSE never sees a terminal change and a reconnecting client
# hangs (the image/chat lesson). initial_seq=max_seq so any init-only rows
# (status/mode/user_input) don't collide on (run_id, seq).
if not user_input:
_log.warning(
"run=%d resume detected (no user_input in change-log); finalising",
run_id,
)
degenerate_writer = StateWriter(
schema=SearchRunState, run_id=run_id, table=RUN_CHANGES_TABLE,
channel=_channel(run_id), session_factory=sf, initial_seq=max_seq,
)
await _finalise(run_id, "completed", sf, writer=degenerate_writer)
return
user_id, allow_overdraft = await _resolve_user_context(run_id, cookie)
async with sf() as s:
cfg = await settings_repo.get_all(s)
per_search = int(cfg.get("exa_credits_per_search") or 2)
per_content = int(cfg.get("exa_credits_per_contents") or 1)
# initial_seq=max_seq: on a fresh run max_seq==3 (POST wrote user_input +
# mode + status); on an init-only recovery it continues the dense log past
# the init rows so the re-driven writes don't collide on (run_id, seq).
writer = StateWriter(
schema=SearchRunState, run_id=run_id, table=RUN_CHANGES_TABLE,
channel=_channel(run_id), session_factory=sf, initial_seq=max_seq,
)
await writer.set("status", "streaming")
tools = build_exa_tools(
user_id=user_id, billing=get_billing_client(),
allow_overdraft=allow_overdraft,
per_search=per_search, per_content=per_content,
reason="search-agent", key_prefix="search-tool",
)
driver = PydanticAIDriver(text_field="text", steps_field="steps")
run_input = RunInput(
prompt=user_input, system_prompt=SEARCH_SYSTEM_PROMPT,
history=tuple(history), limits=Limits(requests=10, tool_calls=20),
)
llm = LLMConfig(model_slug=model_slug,
gateway_base_url=settings.gateway_base_url, cookie=cookie)
main_task = asyncio.current_task()
hb = asyncio.create_task(_heartbeat_loop(run_id, sf, main_task))
try:
await driver.run(run_input=run_input, tools=tools, writer=writer, llm=llm)
# Phase B: the change-log IS the durable transcript — nothing to derive.
# The driver's WriterEventAdapter already flushed text + steps to
# run_changes; just mirror the terminal status.
await _finalise(run_id, "completed", sf, writer=writer)
await _maybe_kick_title_gen(run_id, thread_id, cookie, sf)
except UsageLimitExceeded:
# Partial text/steps already live in run_changes (the sole source).
await _finalise(run_id, "failed", sf, writer=writer,
error="usage_limit_exceeded")
except asyncio.CancelledError:
# Partial assistant activity already lives in run_changes (the driver
# flushed it pre-cancel); just mirror the terminal status.
async with sf() as s:
st = (await s.execute(text("SELECT status FROM search.runs WHERE id=:id"),
{"id": run_id})).scalar_one_or_none()
if st == "streaming":
await _finalise(run_id, "cancelled", sf, writer=writer)
else:
# The cancel endpoint already flipped search.runs.status to a
# terminal value before we noticed; mirror that exact value through
# OUR OWN writer (no cross-writer seq race) so an open SSE closes.
# `st or "cancelled"` guards the run-row-deleted-mid-flight edge so
# we never write a non-terminal None that leaves the SSE hanging.
await writer.set("status", st or "cancelled")
raise
except Exception as exc:
_log.exception("run=%d failed", run_id)
await _finalise(run_id, "failed", sf, writer=writer, error=str(exc)[:500])
raise
finally:
hb.cancel()
async def _resolve_user_context(run_id: int, cookie: str) -> tuple[int, bool]:
"""Resolve (user_id, allow_overdraft) for Exa-tool billing via whoami,
falling back to thread.user_id. Works on first run AND on recovery
(the cookie is a captured DBOS workflow arg)."""
sf = get_session_factory()
try:
whoami = await get_identity_client().whoami(session_cookie=cookie)
return whoami.user_id, bool(getattr(whoami, "is_unlimited", False))
except Exception as exc:
_log.warning("run=%d whoami failed (%s); using thread.user_id", run_id, exc)
async with sf() as s:
row = (await s.execute(
text("SELECT t.user_id FROM search.threads t "
"JOIN search.runs r ON r.thread_id = t.id WHERE r.id = :rid"),
{"rid": run_id},
)).first()
return (int(row[0]) if row else 0, False)
@DBOS.workflow()
async def quick_run_workflow(run_id: int, cookie: str) -> None:
"""Drive one quick search run. Survives restart via DBOS recovery."""
await _run_quick(run_id, cookie)
async def _heartbeat_loop(run_id: int, sf, main_task: asyncio.Task | None = None):
try:
while True:
await asyncio.sleep(HEARTBEAT_INTERVAL_SEC)
async with sf() as s:
status_row = (await s.execute(
text("SELECT status FROM search.runs WHERE id = :id"),
{"id": run_id},
)).scalar_one_or_none()
# Status flipped externally (cancel endpoint or reaper) —
# propagate to main workflow task and stop heartbeating.
if status_row is not None and status_row != "streaming":
if main_task is not None and not main_task.done():
main_task.cancel()
return
await s.execute(
text("UPDATE search.runs SET last_heartbeat = now() WHERE id = :id"),
{"id": run_id},
)
await s.commit()
except asyncio.CancelledError:
return
async def _maybe_kick_title_gen(
run_id: int, thread_id: int, cookie: str, sf,
) -> None:
"""Fire title-gen workflow if this is the first completed run on a
thread whose title is still null. Fire-and-forget."""
async with sf() as s:
row = (await s.execute(
text("""
SELECT t.title,
(SELECT COUNT(*) FROM search.runs r2
WHERE r2.thread_id = :tid AND r2.status = 'completed') AS done_count
FROM search.threads t WHERE t.id = :tid
"""),
{"tid": thread_id},
)).first()
if row is None:
return
title, done_count = row[0], row[1]
if title or done_count != 1:
return # already titled, or not the first completed run
try:
from dbos import SetWorkflowID
from .title_gen import generate_title_workflow
wf_id = f"search-title-gen-{int(thread_id)}"
with SetWorkflowID(wf_id):
DBOS.start_workflow(generate_title_workflow, thread_id, cookie)
except Exception as exc:
_log.warning("title-gen kick failed for thread=%d: %s", thread_id, exc)
async def _finalise(run_id, status, sf, *, writer: StateWriter | None = None,
error: str | None = None):
async with sf() as s:
await s.execute(
text("UPDATE search.runs SET status=:st, error=:err, finished_at=now() "
"WHERE id=:id"),
{"st": status, "err": error, "id": run_id},
)
await s.execute(text(f"NOTIFY search_run_status_{int(run_id)}"))
await s.commit()
if writer is not None:
await writer.set("status", status) # client-facing terminal → closes the SSE
async def _set_status(run_id: int, status: str, sf, *, error: str | None = None):
async with sf() as s:
await s.execute(
text("UPDATE search.runs SET status = :st, error = :err, finished_at = now() "
"WHERE id = :id"),
{"st": status, "err": error, "id": run_id},
)
await s.execute(text(f"NOTIFY search_run_status_{int(run_id)}"))
await s.commit()
@DBOS.workflow()
async def search_run_workflow(run_id: int, cookie: str) -> None:
"""Mode-aware entry point. Reads Run.mode and dispatches.
Patched in tests via `search.agent.workflow.quick_run_workflow` /
`deep_research_workflow` to verify routing without invoking real
PydanticAI / DB / billing.
"""
sf = get_session_factory()
async with sf() as s:
mode = (await s.execute(
text("SELECT mode FROM search.runs WHERE id = :id"),
{"id": run_id},
)).scalar_one_or_none()
if mode == "deep":
await deep_research_workflow(run_id, cookie)
else:
await quick_run_workflow(run_id, cookie)
"""Deep research workflow — multi-agent orchestration (agent-ports).
Composition root: `_run_deep(run_id, cookie)` drives planner → parallel
searchers → coverage → synthesizer, writing ALL state through ONE shared
`StateWriter` over search.run_changes (channel `search_run_{run_id}`):
set("mode","deep") + set("status","streaming")
planner → set("plan", {sub_queries, plan_rationale})
searchers → asyncio.gather; each: append("branches",{branch,status:"running"})
then run its agent (regular Exa or browser MCP) then
merge("branches",{key, patch:{summary,key_sources,status:"done"}})
(or {status:"error", error} on failure)
coverage → set("coverage", {round, complete, rationale, additional_queries})
+ append("plan_extensions",{round, sub_queries}) for follow-ups
synthesizer → set("answer", {content, citations})
The browser searcher keeps the Playwright MCP toolset; its
`browser_take_screenshot` returns are captured EXPLICITLY into
`browser[]` as base64 (append("browser",{action,url,image_b64,media_type}))
so the SPA shows them live. On completion a cleanup nulls image_b64 +
marks `expired` in the change-log (the sole source).
Phase B: search.run_changes is the SOLE durable source — there is no
run_parts derivation. The user's question lives in the change-log too, as
the `user_input` SCALAR (written by the POST handler at seq 1). GET
snapshots + multi-turn history (`_load_prior_qa`) fold the change-log.
Resume (DBOS recovery): fold search.run_changes → ResumeState
(resume_state_from_changes) to skip completed stages, then mirror the
terminal status so a reconnecting SSE closes.
"""
from __future__ import annotations
import asyncio
import logging
from typing import Any
from pydantic_ai.messages import FunctionToolResultEvent
from sqlalchemy import text
# … (полный файл: 646 строк — путь выше) …
async def deep_research_workflow(run_id: int, cookie: str) -> None:
"""Thin entry — kept named for the dispatcher + existing tests."""
await _run_deep(run_id, cookie)
async def _run_deep(run_id: int, cookie: str) -> None:
sf = get_session_factory()
settings = get_settings()
async with sf() as s:
cfg = await load_deep_settings(s)
row = (await s.execute(text(
"SELECT model_slug, thread_id FROM search.runs WHERE id = :id"
), {"id": run_id})).first()
if row is None:
raise RuntimeError(f"run {run_id} not found")
model_slug, thread_id = row[0], int(row[1])
# Fold the change-log into a ResumeState (the agent-ports source of truth)
# AND a SearchRunState (for user_input — the user's question, written by
# the POST handler at seq 1, the sole source, no run_parts).
change_rows = await _load_changes(sf, run_id)
state = resume_state_from_changes(change_rows)
full_state = fold_changes(SearchRunState, change_rows)
user_question = full_state.user_input
init_seq = await _max_change_seq(sf, run_id)
prior_qa = await _load_prior_qa(run_id, thread_id)
history_text = _format_history(prior_qa)
# ONE shared StateWriter for the whole run — its asyncio.Lock serializes
# concurrent branch writes from the parallel searchers (asyncio.gather).
writer = StateWriter(
schema=SearchRunState, run_id=run_id, table=RUN_CHANGES_TABLE,
channel=_channel(run_id), session_factory=sf, initial_seq=init_seq,
)
if state.has_final_answer:
# Already synthesized before a crash — clean up transient screenshots +
# finalise (the change-log is the durable transcript; nothing to derive).
await _cleanup_safe(run_id)
await _finalise(run_id, "completed", sf, writer=writer)
return
if not full_state.status:
# Defensive: the POST handler writes status + mode (seq 2,3) before
# firing the workflow, so on a fresh run these are already set. Only
# write them if a corrupt/empty change-log somehow lost them.
await writer.set("status", "streaming")
await writer.set("mode", "deep")
main_task = asyncio.current_task()
heartbeat_task = asyncio.create_task(_heartbeat_loop_deep(run_id, sf, main_task))
try:
user_id, allow_overdraft = await _resolve_user_context(run_id, cookie)
gateway_base_url = settings.gateway_base_url
# ──── PLAN ────
if not state.has_plan:
plan = await _run_planner(
run_id=run_id, cookie=cookie, gateway_base_url=gateway_base_url,
model_slug=cfg.planner_model_slug or model_slug,
user_question=user_question, max_sub_queries=cfg.max_sub_queries,
history_text=history_text,
)
if not plan.sub_queries:
_log.warning("planner returned empty plan; aborting deep")
await _finalise(run_id, "failed", sf, writer=writer,
error="planner_empty")
return
await writer.set("plan", plan.model_dump())
else:
plan = state.plan # type: ignore[assignment]
# ──── ROUNDS ────
# all_sub_queries is the SINGLE running list of every branch's SubQuery,
# rebuilt INCREMENTALLY by the per-round in-loop extend below (~end of
# the loop). Do NOT pre-fold state.plan_extensions here: on a resumed
# multi-round run that fold would also be applied per-round, duplicating
# followups as phantom branch ids (Claude F1). start_branch/offset math
# derives from plan.sub_queries + state.plan_extensions directly, so it
# is independent of how all_sub_queries is built.
all_sub_queries: list[SubQuery] = list(plan.sub_queries)
for round_idx in range(cfg.max_rounds):
if round_idx == 0:
start_branch = 1
sqs = plan.sub_queries
else:
ext = state.plan_extensions.get(round_idx, [])
if not ext:
break
start_branch = len(plan.sub_queries) + 1 + sum(
len(state.plan_extensions.get(r, [])) for r in range(1, round_idx))
sqs = [
SubQuery(id=start_branch + i, q=sq["q"], rationale=sq["rationale"],
needs_browser=bool(sq.get("needs_browser", False)))
for i, sq in enumerate(ext)
]
# Launch parallel searchers for branches not yet completed.
pending: list[tuple[int, SubQuery]] = []
for i, sq in enumerate(sqs):
branch = start_branch + i
if branch in state.summary_branches():
continue
pending.append((branch, sq))
if pending:
# Cap browser searchers per round: demote excess to regular.
browser_indices = [i for i, (_, sq) in enumerate(pending)
if sq.needs_browser]
if len(browser_indices) > cfg.max_browser_per_round:
for idx in browser_indices[cfg.max_browser_per_round:]:
branch, sq = pending[idx]
pending[idx] = (branch, sq.model_copy(
update={"needs_browser": False}))
tasks = [
_dispatch_one_searcher(
run_id=run_id, cookie=cookie,
gateway_base_url=gateway_base_url, model_slug=model_slug,
cfg=cfg, writer=writer, branch=branch, sub_query=sq,
plan_rationale=plan.plan_rationale,
user_id=user_id, allow_overdraft=allow_overdraft,
)
for (branch, sq) in pending
]
await asyncio.gather(*tasks, return_exceptions=True)
state = resume_state_from_changes(await _load_changes(sf, run_id))
if round_idx == cfg.max_rounds - 1:
break
# ──── COVERAGE ────
if round_idx in state.coverage_rounds():
coverage = state.coverage(round_idx)
else:
coverage = await _run_coverage(
run_id=run_id, cookie=cookie, gateway_base_url=gateway_base_url,
model_slug=cfg.synthesizer_model_slug or model_slug,
plan=plan, all_sub_queries=all_sub_queries,
summaries={b: state.summary_for_branch(b)
for b in state.summary_branches()},
round_idx=round_idx, user_question=user_question,
)
payload = coverage.model_dump()
payload["round"] = round_idx
await writer.set("coverage", payload)
state = resume_state_from_changes(await _load_changes(sf, run_id))
if coverage.complete or not coverage.additional_queries:
break
next_round = round_idx + 1
if next_round in state.plan_extensions:
ext_payload_sqs = state.plan_extensions[next_round]
else:
ext_payload_sqs = [
{"q": fq.q, "rationale": fq.rationale,
"needs_browser": fq.needs_browser}
for fq in coverage.additional_queries
]
await writer.append("plan_extensions",
{"round": next_round, "sub_queries": ext_payload_sqs})
"""mlexi-agent: shared agent state-contract core + frozen agent ports.
State contract: RunState schema + field kinds, the change wire model,
StateWriter (output port), StateProjection SSE transport.
Agent ports: RunInput/Message (input), Tool, LLMConfig (LLM access),
AgentDriver (the swap boundary). The concrete PydanticAIDriver lives in
`mlexi_agent.pydantic_driver` (import on demand — needs the `pydanticai`
extra) so this core stays importable without pydantic-ai installed.
"""
from .changes import Change, Op
from .driver import AgentDriver
from .dsn import asyncpg_dsn
from .fold import fold_changes
from .llm import LLMConfig, gateway_openai_client
from .messages import Limits, Message, RunInput
from .projection import stream_changes
from .state import FieldSpec, Kind, RunState, UndeclaredField
from .tools import Tool
from .writer import IllegalOp, StateWriter
__all__ = [
# state contract
"Kind", "FieldSpec", "RunState", "UndeclaredField",
"Op", "Change",
"StateWriter", "IllegalOp",
"fold_changes",
"stream_changes", "asyncpg_dsn",
# agent ports
"Message", "Limits", "RunInput",
"Tool",
"LLMConfig", "gateway_openai_client",
"AgentDriver",
]
"""AgentDriver port — the swap boundary.
A driver implements exactly this: build the loop from the injected
capabilities and run it, writing state through the injected StateWriter.
PydanticAIDriver is the first impl; a LangGraphDriver would implement the
same Protocol. The host (composition root) injects tools/writer/llm.
"""
from __future__ import annotations
from typing import TYPE_CHECKING, Any, Protocol, Sequence
if TYPE_CHECKING:
from .messages import RunInput
from .llm import LLMConfig
from .tools import Tool
from .writer import StateWriter
class AgentDriver(Protocol):
async def run(
self, *, run_input: "RunInput", tools: "Sequence[Tool]",
writer: "StateWriter", llm: "LLMConfig",
toolsets: "Sequence[Any] | None" = None,
) -> None:
...
"""Neutral run input + history — the driver's INPUT port.
Deliberately framework-agnostic: never expose PydanticAI message types
here. Each driver converts these into its own message representation.
"""
from __future__ import annotations
from dataclasses import dataclass
@dataclass(frozen=True)
class Message:
"""One prior conversation turn, neutral shape.
Plain text turn:
Message(role="user", content="hello") # kind=None or "text"
Message(role="assistant", content="hi there")
Tool-call step (assistant emitted a tool invocation):
Message(role="assistant", content="", kind="tool-call",
tool_name="search", tool_call_id="call_1",
args={"query": "cats"})
Tool-return step (result fed back to the model):
Message(role="user", content='["cat1"]', kind="tool-return",
tool_name="search", tool_call_id="call_1")
All new fields are optional so existing Message(role, content) callers
are unaffected.
"""
role: str # "user" | "assistant" | "system"
content: str
kind: str | None = None # None/"text" = plain text; "tool-call"; "tool-return"
tool_name: str | None = None # required for tool-call and tool-return kinds
tool_call_id: str | None = None # links a tool-call to its matching tool-return
args: dict | None = None # arguments dict, for tool-call kind only
@dataclass(frozen=True)
class Limits:
"""Neutral usage limits; drivers map to their framework's type."""
requests: int | None = None
tool_calls: int | None = None
@dataclass(frozen=True)
class RunInput:
"""Everything the AI core needs to run one turn."""
prompt: str
system_prompt: str | None = None
history: tuple[Message, ...] = ()
limits: Limits | None = None
"""Tool port — a host-built, scope-bound, typed async callable.
The function's type-hinted signature IS the args schema (drivers read it
to present the tool to the LLM). name/description default to the
function's __name__/__doc__. `ask_user` is just one instance of this port.
"""
from __future__ import annotations
import inspect
from dataclasses import dataclass
from typing import Any, Awaitable, Callable
@dataclass(frozen=True)
class Tool:
fn: Callable[..., Awaitable[Any]]
name: str | None = None
description: str | None = None
def resolved_name(self) -> str:
return self.name or self.fn.__name__
def resolved_description(self) -> str:
return self.description or inspect.getdoc(self.fn) or ""
"""LLM access port — gateway connection params + the shared OpenAI-shape
client factory. All platform LLM traffic goes through the gateway
(billing/auth/proxy); drivers must build their model from an LLMConfig,
never talk to a provider directly.
`gateway_openai_client` lazy-imports openai/httpx so the core package
imports without the `pydanticai` extra installed.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Any
@dataclass(frozen=True)
class LLMConfig:
model_slug: str
gateway_base_url: str
cookie: str
read_timeout: float = 90.0
def _normalise_base_url(url: str) -> str:
base = url.rstrip("/")
if not base.endswith("/v1"):
base = base + "/v1"
return base
def gateway_openai_client(
*, cookie: str, gateway_base_url: str, read_timeout: float = 90.0,
) -> Any:
"""Build an AsyncOpenAI client pointed at the MLexi gateway.
Cookie auth per-request; any inbound Authorization header is stripped
(the gateway authenticates by the session cookie). Mirrors the factory
previously duplicated in services/{image,chat}/.../_gateway_client.py.
"""
import httpx
from openai import AsyncOpenAI
async def _strip_authorization(request):
request.headers.pop("Authorization", None)
http = httpx.AsyncClient(
headers={"Cookie": f"mlexi_session={cookie}"},
event_hooks={"request": [_strip_authorization]},
timeout=httpx.Timeout(connect=10.0, read=read_timeout, write=10.0, pool=5.0),
)
return AsyncOpenAI(
api_key="unused",
base_url=_normalise_base_url(gateway_base_url),
http_client=http,
)
"""RunState schema base + per-field kind declaration.
Each field declares a Kind (and, for lists, the element key used by
merge/remove). The StateWriter validates ops against these kinds; the
frontend reducer applies the same op set. See the design spec §2.1.
"""
from __future__ import annotations
from dataclasses import dataclass
from enum import Enum
from typing import ClassVar
from pydantic import BaseModel
class Kind(str, Enum):
"""Field kind → which ops the StateWriter allows on it."""
SCALAR = "scalar" # set
TEXT = "text" # set | append (string concat)
LIST = "list" # append | merge | remove (merge/remove need a key)
@dataclass(frozen=True)
class FieldSpec:
kind: Kind
key: str | None = None # element-id field name; required for merge/remove
class UndeclaredField(KeyError):
"""Raised when a path is not in the schema's `modes` map."""
class RunState(BaseModel):
"""Base for per-service typed run-state documents.
Subclasses declare `modes`: field-name → FieldSpec. The set of paths
in `modes` IS the declared contract; writing any other path raises.
"""
modes: ClassVar[dict[str, FieldSpec]] = {}
@classmethod
def spec(cls, path: str) -> FieldSpec:
try:
return cls.modes[path]
except KeyError:
raise UndeclaredField(path) from None
"""Wire model for one change in the run_changes log.
A change is { seq, path, op, value }. For merge/remove the value carries
the key so the frontend reducer stays schema-free:
merge → {"key_field": "id", "key": 3, "patch": {...}}
remove → {"key_field": "id", "key": 3}
"""
from __future__ import annotations
from dataclasses import dataclass
from enum import Enum
from typing import Any
class Op(str, Enum):
SET = "set"
APPEND = "append"
MERGE = "merge"
REMOVE = "remove"
@dataclass(frozen=True)
class Change:
seq: int
path: str
op: Op
value: Any
"""StateWriter — the backend output port.
Validates set/append/merge/remove against a RunState schema, appends one
row per change to <table>, and NOTIFYs <channel>. Each change is its own
short transaction so the NOTIFY fires post-commit (subscribers never read
ahead of MVCC). seq is monotonic per writer instance (one run = one
workflow = single async task; a lock guards concurrent tool/adapter
writes). Mirrors today's persistence.save_part + the scattered NOTIFY.
"""
from __future__ import annotations
import asyncio
import json
import re
from typing import Any, Callable
from sqlalchemy import text
from .changes import Op
from .state import Kind, RunState
_IDENT = re.compile(r"^[a-z_][a-z0-9_]*$")
class IllegalOp(Exception):
"""Raised when an op is not allowed for a field's kind."""
def _strip_nul(obj: Any) -> Any:
"""Drop NUL (U+0000) — Postgres JSONB/text cannot store it. Mirrors
services/slides/.../persistence.py::_strip_nul."""
if isinstance(obj, str):
return obj.replace(chr(0), "")
if isinstance(obj, dict):
return {_strip_nul(k): _strip_nul(v) for k, v in obj.items()}
if isinstance(obj, (list, tuple)):
return [_strip_nul(v) for v in obj]
return obj
class StateWriter:
def __init__(
self, *, schema: type[RunState], run_id: int, table: str,
channel: str, session_factory: Callable, initial_seq: int = 0,
):
if not _IDENT.match(channel):
raise ValueError(f"channel is not a safe identifier: {channel!r}")
self._schema = schema
self._run_id = run_id
self._table = table
self._channel = channel
self._sf = session_factory
self._seq = initial_seq
self._lock = asyncio.Lock()
async def set(self, path: str, value: Any) -> int:
spec = self._schema.spec(path)
if spec.kind not in (Kind.SCALAR, Kind.TEXT):
raise IllegalOp(f"set not allowed on {spec.kind.value} field {path!r}")
return await self._emit(path, Op.SET, value)
async def append(self, path: str, item: Any) -> int:
spec = self._schema.spec(path)
if spec.kind not in (Kind.TEXT, Kind.LIST):
raise IllegalOp(f"append not allowed on {spec.kind.value} field {path!r}")
return await self._emit(path, Op.APPEND, item)
async def merge(self, path: str, *, key: Any, patch: dict) -> int:
spec = self._schema.spec(path)
if spec.kind is not Kind.LIST or spec.key is None:
raise IllegalOp(f"merge needs a keyed list field, got {path!r}")
return await self._emit(
path, Op.MERGE, {"key_field": spec.key, "key": key, "patch": patch})
async def remove(self, path: str, *, key: Any) -> int:
spec = self._schema.spec(path)
if spec.kind is not Kind.LIST or spec.key is None:
raise IllegalOp(f"remove needs a keyed list field, got {path!r}")
return await self._emit(
path, Op.REMOVE, {"key_field": spec.key, "key": key})
async def _emit(self, path: str, op: Op, value: Any) -> int:
async with self._lock:
# Advance the counter only AFTER a successful commit. If the
# INSERT/commit raises, self._seq is unchanged so the next
# write re-uses this number (dense log, no wasted seq).
next_seq = self._seq + 1
async with self._sf() as s:
await s.execute(
text(
f"INSERT INTO {self._table} "
"(run_id, seq, op, path, value, created_at) "
"VALUES (:r, :seq, :op, :path, CAST(:val AS JSONB), now())"
),
{"r": self._run_id, "seq": next_seq, "op": op.value, "path": path,
"val": json.dumps(_strip_nul(value), ensure_ascii=False)},
)
await s.execute(text(f"NOTIFY {self._channel}"))
await s.commit()
self._seq = next_seq
return next_seq
"""fold_changes — replay a run's change-log into a RunState.
The change-log (the `*_changes` table written by StateWriter) is the single
source of truth for a run (Phase B). This helper folds those rows, in `seq`
order, onto a fresh `schema()` instance — the backend mirror of the frontend
reducer in ``packages/ui/src/agent/reducer.js``. The two MUST agree op-for-op:
set → state[path] = value (replace; for TEXT too)
append → TEXT field: string concat; LIST field: push value
merge → LIST keyed by FieldSpec.key: shallow-merge patch into the item
whose [key_field] == key; no match → no-op (NOT an append)
remove → LIST keyed: drop the item whose [key_field] == key; no match → no-op
unknown op → ignored (reducer.js `default:` returns state unchanged)
Differences from the schema-free JS reducer, by design:
* The reducer decides string-vs-list `append` by runtime type
(``typeof cur === "string"``). Here we dispatch on the field's declared
``Kind`` (TEXT → concat, LIST → push). For a well-formed log seeded from
the same schema (TEXT starts ``""``, LIST starts ``[]``) the result is
identical.
* An ``append``/``merge``/``remove`` on a SCALAR (or ``set`` on a LIST) is a
*declared-contract* violation and raises ``IllegalOp`` — the JS reducer
can't see the schema, so it can't catch this; the backend should.
* A path the current schema does not declare is *ignored* (forward-compat
for rolling deploys where a newer writer emits a field this reader does
not yet know), rather than raising.
Dependency-light: pure state + changes, no pydantic_ai.
"""
from __future__ import annotations
import json
from typing import Any, Iterable, Mapping
from .changes import Op
from .state import Kind, RunState, UndeclaredField
from .writer import IllegalOp
def _normalize_row(row: Any) -> tuple[int, str, str, Any]:
"""Coerce a change row into (seq, op, path, value).
Accepts: a Mapping ({"seq","op","path","value"}), a SQLAlchemy Row
(attribute access), or a plain (seq, op, path, value) tuple/sequence.
"""
if isinstance(row, Mapping):
return (row["seq"], row["op"], row["path"], row["value"])
# SQLAlchemy Row and similar named-tuple-like objects expose attributes.
if hasattr(row, "seq") and hasattr(row, "op") and hasattr(row, "path"):
return (row.seq, row.op, row.path, row.value)
seq, op, path, value = row # plain tuple / list
return (seq, op, path, value)
def _decode_value(value: Any, kind: "Kind | None" = None) -> Any:
"""Coerce a change-log value, kind-aware.
asyncpg/SQLAlchemy hand back JSONB already DECODED to native Python
(a JSONB ``"4"`` comes back as the str ``"4"``, a JSONB object as a
``dict``). So for a TEXT or SCALAR field the value is used verbatim —
re-running ``json.loads`` would corrupt a string that happens to be
valid JSON (``"4"`` → ``4``, ``"true"`` → ``True``), the chat title-gen
bug. Only for a LIST field (whose items are objects) do we keep the
raw-JSON-string fallback, for transports that hand the value back as an
un-decoded JSON string. ``bytes`` is always decoded (never a native
JSONB result, so it must be raw)."""
if isinstance(value, bytes):
try:
return json.loads(value)
except (ValueError, TypeError):
return value
if isinstance(value, str) and kind is Kind.LIST:
try:
return json.loads(value)
except (ValueError, TypeError):
return value
return value
def _op_value(op: Any) -> str:
return op.value if isinstance(op, Op) else str(op)
def fold_changes(schema: type[RunState], rows: Iterable[Mapping | tuple]) -> RunState:
"""Fold change-log ``rows`` into a fresh ``schema()`` instance.
Rows may be dicts, SQLAlchemy Rows, or ``(seq, op, path, value)`` tuples,
in any order — they are sorted by ``seq`` defensively before applying.
``value`` may be a decoded object or a JSON string.
"""
state = schema()
normalized = [_normalize_row(r) for r in rows]
normalized.sort(key=lambda r: r[0])
for _seq, raw_op, path, raw_value in normalized:
op = _op_value(raw_op)
# Forward-compat: a path this schema doesn't declare is skipped, not fatal.
try:
spec = schema.spec(path)
except UndeclaredField:
continue
# Decode kind-aware: TEXT/SCALAR values are used verbatim (asyncpg
# already decoded the JSONB), LIST values may be raw JSON strings.
value = _decode_value(raw_value, spec.kind)
if op == Op.SET.value:
if spec.kind not in (Kind.SCALAR, Kind.TEXT):
raise IllegalOp(f"set not allowed on {spec.kind.value} field {path!r}")
setattr(state, path, value)
elif op == Op.APPEND.value:
if spec.kind is Kind.TEXT:
setattr(state, path, (getattr(state, path) or "") + value)
elif spec.kind is Kind.LIST:
cur = list(getattr(state, path) or [])
cur.append(value)
setattr(state, path, cur)
else:
raise IllegalOp(
f"append not allowed on {spec.kind.value} field {path!r}")
elif op == Op.MERGE.value:
if spec.kind is not Kind.LIST:
raise IllegalOp(f"merge not allowed on {spec.kind.value} field {path!r}")
key_field = value["key_field"]
key = value["key"]
patch = value.get("patch", {})
cur = list(getattr(state, path) or [])
# mirror reducer.js: map, replacing only the matching element;
# absent key → no element changes (no append).
setattr(state, path, [
{**e, **patch} if e.get(key_field) == key else e
for e in cur
])
elif op == Op.REMOVE.value:
if spec.kind is not Kind.LIST:
raise IllegalOp(f"remove not allowed on {spec.kind.value} field {path!r}")
key_field = value["key_field"]
key = value["key"]
cur = list(getattr(state, path) or [])
setattr(state, path, [e for e in cur if e.get(key_field) != key])
# Unknown op → ignore (reducer.js `default:` returns state unchanged).
return state
"""StateProjection — the shared NOTIFY-driven SSE transport.
Replay (changes after `after_seq`) then live-tail via Postgres
LISTEN/NOTIFY. One mechanism for all services (replaces the per-service
projections and the slides 200ms poll). The generator returns once a
terminal `status` change has been emitted, so the HTTP layer's SSE
response ends cleanly. Frame: `id: <seq>\\nevent: change\\ndata: <json>`.
"""
from __future__ import annotations
import asyncio
import json
from typing import Any, AsyncIterator, Awaitable, Callable
import asyncpg
from sqlalchemy import text
KEEPALIVE_INTERVAL_SEC = 20
TERMINAL_STATUSES = {"completed", "failed", "cancelled", "interrupted"}
def _frame(seq: int, op: str, path: str, value: Any) -> bytes:
body = json.dumps(
{"seq": seq, "op": op, "path": path, "value": value},
ensure_ascii=False,
)
return f"id: {seq}\nevent: change\ndata: {body}\n\n".encode("utf-8")
def _is_terminal(op: str, path: str, value: Any, status_field: str) -> bool:
return op == "set" and path == status_field and value in TERMINAL_STATUSES
async def _read_after(session_factory, table, run_id, after_seq):
async with session_factory() as s:
return (await s.execute(
text(f"SELECT seq, op, path, value FROM {table} "
"WHERE run_id = :r AND seq > :a ORDER BY seq"),
{"r": run_id, "a": after_seq},
)).all()
async def stream_changes(
*, table: str, run_id: int, channel: str, after_seq: int,
session_factory: Callable, listen_dsn: str,
status_field: str = "status",
is_disconnected: Callable[[], Awaitable[bool]] | None = None,
terminate_on_terminal: bool = True,
) -> AsyncIterator[bytes]:
"""Replay (changes after `after_seq`) then live-tail via LISTEN/NOTIFY.
`terminate_on_terminal` (default True) is the run-scoped transport used by
image/chat: a terminal `status` change ends the stream so the SSE response
closes cleanly. Set it False for a DECK-scoped stream (slides), where the
change-log `run_id` is the deck id and one open stream sees MANY runs: a
terminal status only means the CURRENT run finished, not that the deck is
done. With False, both the connect-time already-terminal guard AND the
mid-tail terminal return are skipped — the generator only ends on client
disconnect (`is_disconnected`) or error. (Without it, once run #1 sets
`status=completed`, every later run on that deck would stream zero frames.)
"""
last = after_seq
# 1. Bootstrap replay. Replay EVERY row first — do NOT close on a terminal
# status seen mid-replay. A RESUMED run's change-log holds a prior terminal
# status followed by new streaming + changes; closing at the intermediate
# terminal would drop the resumed rows. Termination is decided after the
# full replay based on the LATEST status (step 1b), not a mid-replay one.
for r in await _read_after(session_factory, table, run_id, last):
yield _frame(r.seq, r.op, r.path, r.value)
last = r.seq
# 1b. Latest-status terminal guard. After replaying up to the current max
# seq, close ONLY if the LATEST status (the status value at/by max seq) is
# terminal — no further NOTIFY will fire, so the live tail would hang.
# This also covers the reconnect case where `after_seq` is at/after the
# terminal change (replay empty). For a resumed run whose latest status is
# streaming, we fall through to the live tail. Skipped for a deck-scoped
# stream (terminate_on_terminal=False): a later run on the deck WILL fire
# more NOTIFYs, so we must enter the live tail.
if terminate_on_terminal:
async with session_factory() as s:
latest_status = (await s.execute(
text(f"SELECT value FROM {table} "
"WHERE run_id = :r AND path = :p AND op = 'set' "
"ORDER BY seq DESC LIMIT 1"),
{"r": run_id, "p": status_field},
)).scalar_one_or_none()
if latest_status in TERMINAL_STATUSES:
return
# 2. Live tail.
conn = await asyncpg.connect(dsn=listen_dsn)
try:
q: asyncio.Queue = asyncio.Queue()
await conn.add_listener(channel, lambda *a: q.put_nowait(1))
while True:
# Drain anything written since `last` (covers a NOTIFY that
# fired between the bootstrap SELECT and add_listener).
for r in await _read_after(session_factory, table, run_id, last):
yield _frame(r.seq, r.op, r.path, r.value)
last = r.seq
if terminate_on_terminal and _is_terminal(
r.op, r.path, r.value, status_field,
):
return
if is_disconnected is not None:
try:
if await is_disconnected():
return
except Exception:
pass
try:
await asyncio.wait_for(q.get(), timeout=KEEPALIVE_INTERVAL_SEC)
except asyncio.TimeoutError:
yield b": keepalive\n\n"
finally:
await conn.close()
"""PydanticAIDriver — the first AgentDriver implementation.
Builds a PydanticAI Agent from the injected tools + LLMConfig, runs it,
and translates the framework's stream events into StateWriter calls via
WriterEventAdapter. The event-handling mirrors the proven per-service
ProgressPersister (services/*/agent/events.py), generalised to write
through StateWriter instead of save_part.
Not imported by mlexi_agent.__init__ (keeps the core importable without
the `pydanticai` extra). Import directly:
from mlexi_agent.pydantic_driver import PydanticAIDriver
"""
from __future__ import annotations
from typing import Any, Callable, Sequence
from pydantic_ai import Agent, Tool as PydAITool
from pydantic_ai.messages import (
ModelRequest, ModelResponse,
SystemPromptPart, UserPromptPart, TextPart, TextPartDelta,
ToolCallPart, ToolReturnPart,
PartStartEvent, PartDeltaEvent,
FunctionToolCallEvent, FunctionToolResultEvent,
)
from pydantic_ai.models.openai import OpenAIChatModel
from pydantic_ai.providers.openai import OpenAIProvider
from pydantic_ai.usage import UsageLimits
from .llm import LLMConfig, gateway_openai_client
from .messages import Limits, Message, RunInput
from .tools import Tool
from .writer import StateWriter
def _default_model_factory(llm: LLMConfig) -> Any:
client = gateway_openai_client(
cookie=llm.cookie, gateway_base_url=llm.gateway_base_url,
read_timeout=llm.read_timeout,
)
return OpenAIChatModel(llm.model_slug, provider=OpenAIProvider(openai_client=client))
def _to_usage_limits(limits: Limits | None) -> UsageLimits | None:
if limits is None:
return None
return UsageLimits(
request_limit=limits.requests, tool_calls_limit=limits.tool_calls,
)
_JSON_SCALARS = (str, int, float, bool, type(None))
def _sanitize_step_content(content: Any) -> Any:
"""Recursively make a tool-return value safe for a JSONB INSERT.
A tool may return content the model consumes directly but the
change-log cannot store — e.g. slides' `render_slide` returns
`[header, BinaryContent(png)]` for a vision self-check. Bytes must
never reach the JSONB INSERT in StateWriter.append. Replace any
`pydantic_ai.messages.BinaryContent` (and any other
non-(str/int/float/bool/None/dict/list) object) with a marker
`{"_binary": <media_type or "binary">}`; recurse into dict/list so a
binary nested in `[caption, png]` or `{"img": png}` is also collapsed.
Plain str/dict/list/scalars pass through byte-identical, so chat/image
tool returns (already JSON-native) are unaffected.
"""
# bool is a subclass of int — both are JSON scalars, kept as-is.
if isinstance(content, _JSON_SCALARS):
return content
if isinstance(content, dict):
return {k: _sanitize_step_content(v) for k, v in content.items()}
if isinstance(content, (list, tuple)):
return [_sanitize_step_content(v) for v in content]
# Anything else is non-JSON-native. Prefer a media_type if it carries
# one (BinaryContent does); otherwise a generic binary marker.
media = getattr(content, "media_type", None)
return {"_binary": media if isinstance(media, str) and media else "binary"}
def _to_message_history(history: Sequence[Message]) -> list:
"""Convert neutral Message history to PydanticAI model-message objects.
Five cases:
1. role="user", kind=None/"text" → ModelRequest[UserPromptPart]
2. role="system", kind=None/"text" → ModelRequest[SystemPromptPart]
3. role="assistant", kind=None/"text" → ModelResponse[TextPart]
4. role="assistant", kind="tool-call" → ModelResponse[ToolCallPart]
5. role="user", kind="tool-return" → ModelRequest[ToolReturnPart]
"""
out: list = []
for m in history:
if m.kind == "tool-return":
# case 5: user carries the tool result back to the model
out.append(ModelRequest(parts=[
ToolReturnPart(
tool_name=m.tool_name or "",
content=m.content,
tool_call_id=m.tool_call_id or "",
)
]))
elif m.kind == "tool-call":
# case 4: assistant emitted a tool invocation
out.append(ModelResponse(parts=[
ToolCallPart(
tool_name=m.tool_name or "",
args=m.args,
tool_call_id=m.tool_call_id or "",
)
]))
elif m.role == "user":
# case 1: plain user text
out.append(ModelRequest(parts=[UserPromptPart(content=m.content)]))
elif m.role == "system":
# case 2: system prompt carried in history
out.append(ModelRequest(parts=[SystemPromptPart(content=m.content)]))
else:
# case 3: assistant plain text (kind=None or "text")
out.append(ModelResponse(parts=[TextPart(content=m.content)]))
return out
class WriterEventAdapter:
"""PydanticAI event_stream_handler → StateWriter.
Appends finalized assistant text to `text_field` (if set) and a
tool-call/tool-return activity log to `steps_field` (if set). Either
field may be None when the product doesn't surface that stream.
Token deltas are buffered, never written (finalized chunks only).
"""
def __init__(
self, writer: StateWriter, *,
text_field: str | None = "text", steps_field: str | None = "steps",
):
self._w = writer
self._text_field = text_field
self._steps_field = steps_field
self._text_buf: dict[int, str] = {}
async def __call__(self, ctx, stream):
async for ev in stream:
await self.on_event(ev)
# Flush at the end of EACH round/invocation. PartStartEvent reuses
# index 0 every round, so text accumulated this round must be
# finalized before the next round overwrites the buffer — otherwise
# earlier assistant text is lost (multi-round tool-using chats).
await self.flush_pending_text()
async def on_event(self, ev: Any) -> None:
if isinstance(ev, PartStartEvent):
if isinstance(ev.part, TextPart):
self._text_buf[ev.index] = ev.part.content or ""
elif isinstance(ev, PartDeltaEvent):
if isinstance(ev.delta, TextPartDelta):
self._text_buf[ev.index] = (
self._text_buf.get(ev.index, "") + ev.delta.content_delta
)
elif isinstance(ev, FunctionToolCallEvent):
# Finalize any preceding text BEFORE logging the tool call, so
# text emitted before a tool call keeps its place in the change
# log (text → tool-call order, not reordered to the end).
await self.flush_pending_text()
if self._steps_field is not None:
await self._w.append(self._steps_field, {
"kind": "tool-call", "tool_name": ev.part.tool_name,
"args": ev.part.args, "tool_call_id": ev.part.tool_call_id,
})
elif isinstance(ev, FunctionToolResultEvent):
r = getattr(ev, "part", None) or getattr(ev, "result", None)
if self._steps_field is not None:
await self._w.append(self._steps_field, {
"kind": "tool-return", "tool_name": r.tool_name,
# Sanitize so a binary tool-return (e.g. render_slide's
# BinaryContent png) can't break the JSONB INSERT.
"content": _sanitize_step_content(r.content),
"tool_call_id": r.tool_call_id,
})
async def flush_pending_text(self) -> None:
for index in sorted(self._text_buf.keys()):
content = self._text_buf.pop(index)
if content and self._text_field is not None:
await self._w.append(self._text_field, content)
class PydanticAIDriver:
def __init__(
self, *, text_field: str | None = "text", steps_field: str | None = "steps",
model_factory: Callable[[LLMConfig], Any] = _default_model_factory,
):
self._text_field = text_field
self._steps_field = steps_field
self._model_factory = model_factory
async def run(
self, *, run_input: RunInput, tools: Sequence[Tool],
writer: StateWriter, llm: LLMConfig,
toolsets: Sequence[Any] | None = None,
) -> None:
model = self._model_factory(llm)
pyd_tools = [
PydAITool(t.fn, takes_ctx=False,
name=t.resolved_name(), description=t.resolved_description())
for t in tools
]
agent = Agent(
model=model,
system_prompt=run_input.system_prompt or "",
tools=pyd_tools,
toolsets=list(toolsets or []),
)
adapter = WriterEventAdapter(
writer, text_field=self._text_field, steps_field=self._steps_field,
)
await agent.run(
run_input.prompt,
message_history=_to_message_history(run_input.history),
event_stream_handler=adapter,
usage_limits=_to_usage_limits(run_input.limits),
)
await adapter.flush_pending_text()