Files
mnote/wolai-backend/app/services/ai_document_agent.py
T
lix-2026 41e958769e feat: land page aggregate and phase7 document ai mainline
- 收口 page aggregate 读取、本地状态与命令客户端\n- 接入 phase7 document ai sidecar 与前端编排入口\n- 更新 architecture 与 design 状态迁移
2026-04-23 07:38:34 +08:00

1620 lines
53 KiB
Python

from __future__ import annotations
import asyncio
import contextlib
import json
import os
import time
from collections import deque
from dataclasses import dataclass, field
from functools import lru_cache
from pathlib import Path
from urllib.parse import urlsplit, urlunsplit
from typing import Any, AsyncGenerator, Callable, Literal, Mapping
from uuid import uuid4
import httpx
from pydantic import BaseModel, Field
from app.config import settings
try:
from agents import (
Agent,
ItemHelpers,
ModelResponse,
MultiProvider,
OpenAIProvider,
OpenAIResponsesModel,
RunConfig,
RunContextWrapper,
Runner,
SQLiteSession,
Usage,
function_tool,
trace,
)
from agents.models.interface import ModelProvider
from agents.models.multi_provider import MultiProviderMap
from agents.models.openai_responses import Converter as ResponsesConverter
from agents.util._json import _to_dump_compatible
except ImportError: # pragma: no cover - 依赖缺失时走降级分支
Agent = None
ItemHelpers = None
ModelResponse = None
MultiProvider = None
OpenAIProvider = None
OpenAIResponsesModel = None
RunConfig = None
RunContextWrapper = Any
Runner = None
SQLiteSession = None
Usage = None
function_tool = None
trace = None
MultiProviderMap = None
ModelProvider = object
ResponsesConverter = None
_to_dump_compatible = None
def sdk_available() -> bool:
return all(
dependency is not None
for dependency in (
Agent,
ModelResponse,
OpenAIProvider,
OpenAIResponsesModel,
ResponsesConverter,
RunConfig,
Runner,
SQLiteSession,
Usage,
function_tool,
trace,
)
)
def has_configured_openai_api_key() -> bool:
value = get_effective_openai_api_key()
if not value:
return False
return "please-set-your-key" not in value
def _is_placeholder_openai_key(value: str | None) -> bool:
return "please-set-your-key" in str(value or "").strip()
@lru_cache(maxsize=1)
def load_codex_auth_openai_key() -> str:
auth_path = Path.home() / ".codex" / "auth.json"
try:
payload = json.loads(auth_path.read_text("utf-8"))
except Exception:
return ""
value = str(payload.get("OPENAI_API_KEY", "")).strip() if isinstance(payload, dict) else ""
if not value or _is_placeholder_openai_key(value):
return ""
return value
def get_effective_openai_api_key(env: Mapping[str, str] | None = None) -> str:
env_map = env if env is not None else os.environ
env_value = str(env_map.get("OPENAI_API_KEY", "")).strip()
settings_value = settings.openai_api_key.strip()
auth_value = load_codex_auth_openai_key()
if env_value and not _is_placeholder_openai_key(env_value):
return env_value
if settings_value and not _is_placeholder_openai_key(settings_value):
return settings_value
if auth_value and not _is_placeholder_openai_key(auth_value):
return auth_value
return env_value or settings_value or auth_value
@lru_cache(maxsize=8)
def load_available_openai_models(
base_url: str,
api_key: str,
) -> tuple[str, ...]:
normalized_base_url = base_url.rstrip("/")
if not normalized_base_url or not api_key.strip():
return ()
try:
with httpx.Client(timeout=10.0) as client:
response = client.get(
f"{normalized_base_url}/models",
headers={"Authorization": f"Bearer {api_key.strip()}"},
)
response.raise_for_status()
payload = response.json()
except Exception:
return ()
data = payload.get("data") if isinstance(payload, dict) else None
if not isinstance(data, list):
return ()
ids: list[str] = []
for item in data:
if not isinstance(item, dict):
continue
model_id = item.get("id")
if isinstance(model_id, str) and model_id.strip():
ids.append(model_id.strip())
return tuple(ids)
@dataclass(frozen=True)
class DocumentAiModelSpec:
key: str
title: str
description: str
gateway_model: str
resolved_combo: str | None = None
resolved_runtime_model: str | None = None
status: str = "active"
is_default: bool = False
@dataclass(frozen=True)
class DocumentAiProfileSpec:
id: str
title: str
description: str
instructions: str
session_mode: Literal["off", "page", "workspace"] = "page"
tool_names: tuple[str, ...] = ()
is_default: bool = False
status: str = "active"
@dataclass(frozen=True)
class DocumentAiToolSpec:
name: str
title: str
description: str
scope: Literal["document", "tree", "workspace"] = "document"
mode: Literal["read", "write"] = "read"
status: str = "active"
version: str = "v1"
DEFAULT_OMNIROUTE_BASE_URL = "http://localhost:20128/v1"
DOCUMENT_AI_MODEL_SPECS: tuple[DocumentAiModelSpec, ...] = (
DocumentAiModelSpec(
key="gpt-5.4",
title="GPT-5.4",
description="文档页默认主写模型键,经由 OmniRoute 二次路由。",
gateway_model="gpt-5.4",
resolved_combo="slow",
resolved_runtime_model="gpt-5.4",
is_default=True,
),
DocumentAiModelSpec(
key="gpt-5.4-mini",
title="GPT-5.4 Mini",
description="更偏快响应的轻量模型键,经由 OmniRoute 二次路由。",
gateway_model="gpt-5.4-mini",
resolved_runtime_model="gpt-5.4-mini",
),
DocumentAiModelSpec(
key="gpt-5.3-codex",
title="GPT-5.3 Codex",
description="偏代码与结构化执行的模型键,经由 OmniRoute 二次路由。",
gateway_model="gpt-5.3-codex",
resolved_combo="codex",
resolved_runtime_model="gpt-5.3-codex",
),
DocumentAiModelSpec(
key="gemini-3.1-pro-preview",
title="Gemini 3.1 Pro Preview",
description="偏长上下文与深推理的模型键,经由 OmniRoute 二次路由。",
gateway_model="gemini-3.1-pro-preview",
resolved_runtime_model="gemini-3.1-pro-preview",
),
DocumentAiModelSpec(
key="third",
title="Third",
description="第三路由保底模型键,经由 OmniRoute combo 直通。",
gateway_model="third",
resolved_combo="third",
resolved_runtime_model="third",
),
)
DOCUMENT_AI_TOOL_SPECS: tuple[DocumentAiToolSpec, ...] = (
DocumentAiToolSpec(
name="doc_get",
title="读取当前页摘要",
description="读取当前页块摘要,适合先理解页面结构。",
mode="read",
),
DocumentAiToolSpec(
name="doc_find",
title="按块查找",
description="按文本查找当前页块,适合定位待改写内容。",
mode="read",
),
DocumentAiToolSpec(
name="doc_insert_blocks",
title="插入块",
description="向当前页插入新块,适合新增标题和段落。",
mode="write",
),
DocumentAiToolSpec(
name="doc_replace_range",
title="改写块文本",
description="改写当前页指定块文本。",
mode="write",
),
DocumentAiToolSpec(
name="slash_run",
title="斜杠命令",
description="执行当前页允许的斜杠命令;当前仅限重命名当前页。",
mode="write",
),
)
DOCUMENT_AI_PROFILE_SPECS: tuple[DocumentAiProfileSpec, ...] = (
DocumentAiProfileSpec(
id="page_writer_ai_first",
title="AI 主写",
description="AI 优先直接进入主编辑区起草,人负责审核和继续编辑。",
instructions=(
"当前 profile:AI 优先直接进入主编辑区起草,人负责审核和继续编辑。"
" 如果用户目标明确,优先输出可直接落入当前页的修改结果,而不是长篇空谈。"
),
tool_names=tuple(spec.name for spec in DOCUMENT_AI_TOOL_SPECS),
is_default=True,
),
DocumentAiProfileSpec(
id="page_writer_strict",
title="严格审慎",
description="更保守地读取、定位、改写,优先减少误写。",
instructions=(
"当前 profile:严格审慎。写入前优先读取和定位目标块;若信息不足,先解释缺口,不要冒进写入。"
),
tool_names=tuple(spec.name for spec in DOCUMENT_AI_TOOL_SPECS),
),
DocumentAiProfileSpec(
id="page_writer_polish",
title="润色整理",
description="更偏改写、收敛、结构整理,不主动扩写超出当前页语境的内容。",
instructions=(
"当前 profile:润色整理。优先压缩、润色、重组当前页已有内容,不主动扩写超出当前页证据范围的新事实。"
),
tool_names=tuple(spec.name for spec in DOCUMENT_AI_TOOL_SPECS),
),
)
def _document_ai_model_spec_by_key(model_key: str | None) -> DocumentAiModelSpec | None:
value = str(model_key or "").strip()
if not value:
return None
return next((spec for spec in DOCUMENT_AI_MODEL_SPECS if spec.key == value), None)
def _document_ai_model_spec_by_gateway_model(gateway_model: str | None) -> DocumentAiModelSpec | None:
value = str(gateway_model or "").strip()
if not value:
return None
return next((spec for spec in DOCUMENT_AI_MODEL_SPECS if spec.gateway_model == value), None)
def _resolve_default_document_ai_model_key(
env: Mapping[str, str] | None = None,
) -> str:
env_map = env if env is not None else os.environ
for value in _iter_candidate_model_names(
request_model=None,
configured_default_model=settings.mnote_ai_default_model,
env=env_map,
):
spec = _document_ai_model_spec_by_key(value) or _document_ai_model_spec_by_gateway_model(value)
if spec is not None:
return spec.key
default_spec = next((spec for spec in DOCUMENT_AI_MODEL_SPECS if spec.is_default), DOCUMENT_AI_MODEL_SPECS[0])
return default_spec.key
def _resolve_default_document_ai_profile_id() -> str:
default_spec = next((spec for spec in DOCUMENT_AI_PROFILE_SPECS if spec.is_default), DOCUMENT_AI_PROFILE_SPECS[0])
return default_spec.id
def resolve_document_ai_profile(profile_id: str | None) -> DocumentAiProfileSpec:
value = str(profile_id or "").strip()
if value:
picked = next((spec for spec in DOCUMENT_AI_PROFILE_SPECS if spec.id == value), None)
if picked is not None:
return picked
default_id = _resolve_default_document_ai_profile_id()
return next(spec for spec in DOCUMENT_AI_PROFILE_SPECS if spec.id == default_id)
def build_document_ai_config_payload(
env: Mapping[str, str] | None = None,
) -> dict[str, Any]:
env_map = env if env is not None else os.environ
default_model_key = _resolve_default_document_ai_model_key(env_map)
default_profile_id = _resolve_default_document_ai_profile_id()
raw_base_url = str(env_map.get("OPENAI_BASE_URL", "")).strip() or DEFAULT_OMNIROUTE_BASE_URL
return {
"provider": "online",
"transport": "openai_responses",
"baseUrl": normalize_loopback_base_url(raw_base_url),
"sessionEnabled": True,
"defaultModelKey": default_model_key,
"defaultProfileId": default_profile_id,
"models": [
{
"key": spec.key,
"title": spec.title,
"description": spec.description,
"gatewayModel": spec.gateway_model,
"resolvedCombo": spec.resolved_combo,
"resolvedRuntimeModel": spec.resolved_runtime_model,
"status": spec.status,
"default": spec.key == default_model_key,
}
for spec in DOCUMENT_AI_MODEL_SPECS
],
"profiles": [
{
"id": spec.id,
"title": spec.title,
"description": spec.description,
"sessionMode": spec.session_mode,
"toolNames": list(spec.tool_names),
"status": spec.status,
"default": spec.id == default_profile_id,
}
for spec in DOCUMENT_AI_PROFILE_SPECS
],
"tools": [
{
"name": spec.name,
"title": spec.title,
"description": spec.description,
"scope": spec.scope,
"mode": spec.mode,
"status": spec.status,
"version": spec.version,
}
for spec in DOCUMENT_AI_TOOL_SPECS
],
}
def normalize_loopback_base_url(base_url: str) -> str:
value = str(base_url or "").strip()
if not value:
return DEFAULT_OMNIROUTE_BASE_URL
parsed = urlsplit(value)
if parsed.hostname != "127.0.0.1":
return value
netloc = parsed.netloc.replace("127.0.0.1", "localhost", 1)
return urlunsplit((parsed.scheme, netloc, parsed.path, parsed.query, parsed.fragment))
def _iter_candidate_model_names(
*,
request_model: str | None,
configured_default_model: str | None,
env: Mapping[str, str],
) -> tuple[str, ...]:
values = (
(request_model or "").strip(),
(configured_default_model or "").strip(),
str(env.get("ONLINE_AI_MODEL", "")).strip(),
str(env.get("OPENAI_MODEL", "")).strip(),
str(env.get("OPENAI_DEFAULT_MODEL", "")).strip(),
str(env.get("CLAW_DEFAULT_MODEL", "")).strip(),
)
return tuple(value for value in values if value)
def _extract_model_prefix(model_name: str) -> str | None:
value = str(model_name or "").strip()
if "/" not in value:
return None
prefix, _ = value.split("/", 1)
return prefix.strip() or None
class GatewayPrefixedOpenAIProvider(ModelProvider):
def __init__(
self,
prefix: str,
*,
api_key: str | None,
base_url: str | None,
use_responses: bool = True,
) -> None:
self._prefix = prefix
self._api_key = api_key or ""
self._base_url = base_url or ""
self._use_responses = use_responses
def get_model(self, model_name: str | None):
full_model_name = f"{self._prefix}/{model_name}" if model_name else self._prefix
return CompatOpenAIResponsesModel(
model=full_model_name,
openai_client=None,
model_is_explicit=True,
base_url=self._base_url,
api_key=self._api_key,
use_responses=self._use_responses,
)
class DocumentAiModelProvider(ModelProvider):
def __init__(
self,
*,
provider_map: Mapping[str, ModelProvider],
api_key: str | None,
base_url: str | None,
use_responses: bool = True,
) -> None:
self._provider_map = dict(provider_map)
self._api_key = api_key or ""
self._base_url = base_url or ""
self._use_responses = use_responses
def get_model(self, model_name: str | None):
value = str(model_name or "").strip()
if "/" in value:
prefix, suffix = value.split("/", 1)
provider = self._provider_map.get(prefix)
if provider is None:
provider = GatewayPrefixedOpenAIProvider(
prefix,
api_key=self._api_key,
base_url=self._base_url,
use_responses=self._use_responses,
)
return provider.get_model(suffix)
return CompatOpenAIResponsesModel(
model=value or None,
openai_client=None,
model_is_explicit=bool(value),
base_url=self._base_url,
api_key=self._api_key,
use_responses=self._use_responses,
)
def build_agent_usage_from_response_usage(response_usage: Any):
if response_usage is None:
return Usage()
def pick(name: str, default: Any = None) -> Any:
if isinstance(response_usage, Mapping):
return response_usage.get(name, default)
return getattr(response_usage, name, default)
input_tokens = int(pick("input_tokens", 0) or 0)
output_tokens = int(pick("output_tokens", 0) or 0)
total_tokens_raw = pick("total_tokens", None)
total_tokens = input_tokens + output_tokens
if total_tokens_raw not in (None, ""):
try:
total_tokens = int(total_tokens_raw)
except (TypeError, ValueError):
total_tokens = input_tokens + output_tokens
return Usage(
requests=1,
input_tokens=input_tokens,
output_tokens=output_tokens,
total_tokens=total_tokens,
input_tokens_details=pick("input_tokens_details", None),
output_tokens_details=pick("output_tokens_details", None),
)
def _is_openai_omit(value: Any) -> bool:
return value is None or value.__class__.__name__ == "Omit"
def parse_gateway_raw_response_text(raw_text: str) -> dict[str, Any]:
text = raw_text.strip()
if not text:
raise RuntimeError("网关未返回有效响应体")
if text.startswith("{"):
payload = json.loads(text)
if not isinstance(payload, dict):
raise RuntimeError("网关响应不是对象")
return payload
parsed_payload: dict[str, Any] | None = None
for block in text.split("\n\n"):
lines = [
line[len("data:") :].strip()
for line in block.splitlines()
if line.startswith("data:")
]
if not lines:
continue
try:
data = json.loads("\n".join(lines))
except json.JSONDecodeError:
continue
if isinstance(data, dict) and isinstance(data.get("response"), dict):
parsed_payload = data["response"]
continue
if isinstance(data, dict):
parsed_payload = data
if parsed_payload is None:
raise RuntimeError(f"无法解析网关 Responses 返回:{text[:200]}")
return parsed_payload
class CompatOpenAIResponsesModel(OpenAIResponsesModel):
def __init__(
self,
model,
openai_client,
*,
model_is_explicit: bool = True,
base_url: str,
api_key: str,
use_responses: bool = True,
) -> None:
super().__init__(model=model, openai_client=openai_client, model_is_explicit=model_is_explicit)
self._gateway_base_url = base_url.rstrip("/")
self._gateway_api_key = api_key
self._gateway_use_responses = use_responses
async def get_response(
self,
system_instructions,
input,
model_settings,
tools,
output_schema,
handoffs,
tracing,
previous_response_id: str | None = None,
conversation_id: str | None = None,
prompt=None,
):
list_input = ItemHelpers.input_to_new_input_list(input)
list_input = _to_dump_compatible(list_input)
list_input = self._remove_openai_responses_api_incompatible_fields(list_input)
if model_settings.parallel_tool_calls and tools:
parallel_tool_calls = True
elif model_settings.parallel_tool_calls is False:
parallel_tool_calls = False
else:
parallel_tool_calls = None
tool_choice = ResponsesConverter.convert_tool_choice(model_settings.tool_choice)
converted_tools = ResponsesConverter.convert_tools(tools, handoffs)
converted_tools_payload = _to_dump_compatible(converted_tools.tools)
response_format = ResponsesConverter.get_response_format(output_schema)
should_omit_model = prompt is not None and not self._model_is_explicit
should_omit_tools = prompt is not None and len(converted_tools_payload) == 0
include_set: set[str] = set(converted_tools.includes)
if model_settings.response_include is not None:
include_set.update(model_settings.response_include)
if model_settings.top_logprobs is not None:
include_set.add("message.output_text.logprobs")
include = list(include_set)
payload: dict[str, Any] = {
"previous_response_id": previous_response_id,
"conversation": conversation_id,
"instructions": system_instructions,
"model": None if should_omit_model else self.model,
"input": list_input,
"include": include or None,
"tools": None if should_omit_tools else converted_tools_payload,
"prompt": prompt,
"temperature": model_settings.temperature,
"top_p": model_settings.top_p,
"truncation": model_settings.truncation,
"max_output_tokens": model_settings.max_tokens,
"tool_choice": tool_choice,
"parallel_tool_calls": parallel_tool_calls,
"text": response_format,
"store": model_settings.store,
"prompt_cache_retention": model_settings.prompt_cache_retention,
"reasoning": model_settings.reasoning,
"metadata": model_settings.metadata,
}
if model_settings.top_logprobs is not None:
payload["top_logprobs"] = model_settings.top_logprobs
if model_settings.verbosity is not None:
if isinstance(payload.get("text"), dict):
payload["text"]["verbosity"] = model_settings.verbosity
else:
payload["text"] = {"verbosity": model_settings.verbosity}
if model_settings.extra_body:
payload.update(dict(model_settings.extra_body))
payload = {
key: value
for key, value in payload.items()
if not _is_openai_omit(value)
}
headers = {
"Authorization": f"Bearer {self._gateway_api_key}",
"Content-Type": "application/json",
}
if isinstance(model_settings.extra_headers, Mapping):
headers.update({str(k): str(v) for k, v in model_settings.extra_headers.items()})
async with httpx.AsyncClient(timeout=60.0) as client:
raw_response = await client.post(
f"{self._gateway_base_url}/responses",
headers=headers,
json=payload,
params=model_settings.extra_query,
)
raw_response.raise_for_status()
parsed_payload = parse_gateway_raw_response_text(raw_response.text)
if isinstance(parsed_payload.get("error"), dict):
raise RuntimeError(
str(parsed_payload["error"].get("message") or "Responses API 返回错误")
)
return ModelResponse(
output=parsed_payload.get("output") or [],
usage=build_agent_usage_from_response_usage(parsed_payload.get("usage")),
response_id=str(parsed_payload.get("id") or "") or None,
)
def wrap_openai_model_for_gateway_compat(model: Any):
if OpenAIResponsesModel is not None and isinstance(model, OpenAIResponsesModel):
return CompatOpenAIResponsesModel(
model=model.model,
openai_client=model._client,
model_is_explicit=getattr(model, "_model_is_explicit", True),
base_url=getattr(getattr(model, "_client", None), "base_url", ""),
api_key="",
)
return model
def build_document_model_provider(
*,
request_model: str | None,
configured_default_model: str | None,
env: Mapping[str, str] | None = None,
available_models_loader: Callable[[str, str], list[str] | tuple[str, ...]] | None = None,
):
if not sdk_available():
raise RuntimeError("openai-agents 未安装,无法构建文档页 model provider")
env_map = env if env is not None else os.environ
base_url = str(env_map.get("OPENAI_BASE_URL", "")).strip()
api_key = get_effective_openai_api_key(env_map)
loader = available_models_loader or load_available_openai_models
available_models = list(loader(base_url, api_key)) if base_url and api_key else []
prefixes: set[str] = set()
for model_name in _iter_candidate_model_names(
request_model=request_model,
configured_default_model=configured_default_model,
env=env_map,
):
prefix = _extract_model_prefix(model_name)
if prefix:
prefixes.add(prefix)
for model_name in available_models:
prefix = _extract_model_prefix(str(model_name))
if prefix:
prefixes.add(prefix)
provider_map: dict[str, ModelProvider] = {}
for prefix in sorted(prefixes):
if prefix in {"openai", "litellm"}:
continue
provider_map[prefix] = GatewayPrefixedOpenAIProvider(
prefix,
api_key=api_key,
base_url=base_url,
use_responses=True,
)
return DocumentAiModelProvider(
provider_map=provider_map,
api_key=api_key,
base_url=base_url,
use_responses=True,
)
def resolve_document_agent_model(
*,
request_model: str | None,
configured_default_model: str | None,
env: Mapping[str, str] | None = None,
available_models_loader: Callable[[str, str], list[str] | tuple[str, ...]] | None = None,
) -> str:
env_map = env if env is not None else os.environ
candidates = _iter_candidate_model_names(
request_model=request_model,
configured_default_model=configured_default_model,
env=env_map,
)
raw_model = candidates[0] if candidates else ""
if not raw_model:
raise RuntimeError(
"未配置文档页 AI model;请设置 MNOTE_AI_DEFAULT_MODEL 或在请求中显式传入 provider/model"
)
if "/" in raw_model:
return raw_model
base_url = str(env_map.get("OPENAI_BASE_URL", "")).strip()
api_key = get_effective_openai_api_key(env_map)
if not base_url or not api_key:
raise RuntimeError(
f"文档页 AI model '{raw_model}' 缺少 provider 前缀,且当前无法查询可用模型;请显式设置 provider/model"
)
loader = available_models_loader or load_available_openai_models
available_models = [model.strip() for model in loader(base_url, api_key) if str(model).strip()]
if not available_models:
raise RuntimeError(
f"文档页 AI model '{raw_model}' 缺少 provider 前缀,且当前无法查询可用模型;请显式设置 provider/model"
)
# 当前 sidecar 走 Responses API,优先选择网关上可用的 responses provider 前缀。
responses_model = f"ju/{raw_model}"
if responses_model in available_models:
return responses_model
suffix_matches = sorted(
{
model
for model in available_models
if model.endswith(f"/{raw_model}")
}
)
if len(suffix_matches) == 1:
return suffix_matches[0]
if len(suffix_matches) > 1:
raise RuntimeError(
f"文档页 AI model '{raw_model}' 存在多个候选,请使用 provider/model 显式指定"
)
raise RuntimeError(
f"未找到可用的文档页 AI model '{raw_model}';请设置 MNOTE_AI_DEFAULT_MODEL 或在请求中显式传入 provider/model"
)
def resolve_document_agent_request_model(
*,
request_model: str | None,
request_model_key: str | None,
configured_default_model: str | None,
env: Mapping[str, str] | None = None,
available_models_loader: Callable[[str, str], list[str] | tuple[str, ...]] | None = None,
) -> str:
direct_model_key = str(request_model_key or "").strip()
if direct_model_key:
spec = _document_ai_model_spec_by_key(direct_model_key)
if spec is not None:
return spec.gateway_model
return direct_model_key
direct_request_model = str(request_model or "").strip()
if direct_request_model:
spec = _document_ai_model_spec_by_key(direct_request_model)
if spec is not None:
return spec.gateway_model
return resolve_document_agent_model(
request_model=request_model,
configured_default_model=configured_default_model,
env=env,
available_models_loader=available_models_loader,
)
def make_id(prefix: str) -> str:
return f"{prefix}_{uuid4().hex}"
def to_sse_frame(event: str, data: Any) -> str:
return f"event: {event}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n"
class DocumentAiMessage(BaseModel):
role: Literal["user", "assistant"]
content: str
class DocumentAiContextPayload(BaseModel):
document_id: str | None = Field(default=None, alias="documentId")
document_blocks: Any = Field(default=None, alias="documentBlocks")
node: Any = None
subtree: Any = None
outline: Any = None
evidence: Any = None
page_options: Any = Field(default=None, alias="pageOptions")
editor_runtime_page_options: Any = Field(
default=None,
alias="editorRuntimePageOptions",
)
model_config = {
"populate_by_name": True,
}
class DocumentAiRunRequest(BaseModel):
user_id: str = Field(alias="userId")
session_id: str | None = Field(default=None, alias="sessionId")
model: str | None = None
model_key: str | None = Field(default=None, alias="modelKey")
profile_id: str | None = Field(default=None, alias="profileId")
max_steps: int = Field(default=10, alias="maxSteps")
messages: list[DocumentAiMessage]
context: DocumentAiContextPayload
model_config = {
"populate_by_name": True,
}
class InsertBlockSpec(BaseModel):
block_type: Literal["paragraph", "heading"] = Field(
default="paragraph",
alias="type",
)
text: str = ""
level: int = 2
model_config = {
"populate_by_name": True,
}
@dataclass
class DocumentAiEmitter:
queue: asyncio.Queue[tuple[str, Any] | None] = field(
default_factory=asyncio.Queue,
)
tool_count: int = 0
_tool_sequence: int = 0
async def emit(self, event: str, data: Any) -> None:
await self.queue.put((event, data))
async def close(self) -> None:
await self.queue.put(None)
def next_tool_id(self, tool: str) -> str:
self._tool_sequence += 1
return f"agents_{tool}_{self._tool_sequence}"
@dataclass
class DocumentAiRunContext:
user_id: str
request_id: str
trace_id: str
session_id: str | None
workspace_id: str | None
document_id: str | None
document_blocks: Any
node: Any
subtree: Any
outline: Any
evidence: Any
page_options: Any
editor_runtime_page_options: Any
forward_headers: dict[str, str]
emitter: DocumentAiEmitter
bridge: "LocalBridgeRuntime"
@dataclass(frozen=True)
class BridgeRuntimeInvocation:
command: str
args: tuple[str, ...]
cwd: str
@lru_cache(maxsize=1)
def resolve_bridge_repo_root() -> Path:
repo_root = Path(__file__).resolve().parents[3]
manifest = repo_root / "rust" / "Cargo.toml"
if not manifest.exists():
raise RuntimeError(f"未找到 Rust bridge manifest: {manifest}")
return repo_root
@lru_cache(maxsize=1)
def resolve_bridge_runtime_invocation() -> BridgeRuntimeInvocation:
repo_root = resolve_bridge_repo_root()
explicit_bin = str(os.environ.get("MNOTE_RUST_BRIDGE_BIN", "")).strip()
if explicit_bin:
return BridgeRuntimeInvocation(
command=explicit_bin,
args=(),
cwd=str(repo_root),
)
for candidate in (
repo_root / "rust" / "target" / "debug" / "bridge-runtime",
repo_root / "rust" / "target" / "release" / "bridge-runtime",
):
if candidate.exists():
return BridgeRuntimeInvocation(
command=str(candidate),
args=(),
cwd=str(repo_root),
)
return BridgeRuntimeInvocation(
command="cargo",
args=(
"run",
"--quiet",
"--manifest-path",
str(repo_root / "rust" / "Cargo.toml"),
"-p",
"bridge-runtime",
"--",
),
cwd=str(repo_root),
)
async def run_bridge_runtime(input_payload: dict[str, Any]) -> dict[str, Any]:
invocation = resolve_bridge_runtime_invocation()
payload = json.dumps(input_payload, ensure_ascii=False).encode("utf-8")
try:
process = await asyncio.create_subprocess_exec(
invocation.command,
*invocation.args,
stdin=asyncio.subprocess.PIPE,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
cwd=invocation.cwd,
env={
**os.environ,
"CARGO_TERM_COLOR": "never",
},
)
except Exception as error:
raise RuntimeError(f"bridge-runtime 启动失败: {error}") from error
try:
stdout, stderr = await asyncio.wait_for(
process.communicate(payload),
timeout=30.0,
)
except asyncio.TimeoutError as error:
process.kill()
with contextlib.suppress(ProcessLookupError):
await process.wait()
raise RuntimeError("bridge-runtime 执行超时") from error
stdout_text = stdout.decode("utf-8", errors="replace").strip()
stderr_text = stderr.decode("utf-8", errors="replace").strip()
if not stdout_text:
raise RuntimeError(stderr_text or "bridge-runtime 未返回任何结果")
try:
response = json.loads(stdout_text)
except json.JSONDecodeError as error:
raise RuntimeError(f"bridge-runtime 返回非法 JSON: {error}") from error
if not isinstance(response, dict):
raise RuntimeError("bridge-runtime 返回格式非法")
if process.returncode != 0 or not bool(response.get("ok")):
error_payload = response.get("error")
if isinstance(error_payload, dict):
message = str(error_payload.get("message", "")).strip()
if message:
raise RuntimeError(message)
raise RuntimeError(
stderr_text or f"bridge-runtime 执行失败(code={process.returncode})"
)
return response
def infer_tool_invocation_kind(tool: str) -> Literal["command", "query", "job"]:
if tool in {"doc_get", "doc_find"}:
return "query"
return "command"
class LocalBridgeRuntime:
async def execute_tool(
self,
*,
tool: str,
context: DocumentAiRunContext,
invocation_kind: Literal["command", "query", "job"] | None = None,
args_json: dict[str, Any],
data: Any = None,
target: dict[str, Any] | None = None,
reason: str | None = None,
refs: list[str] | None = None,
) -> Any:
payload = {
"kind": "tool",
"context": {
"deploymentId": None,
"projectId": None,
"workspaceId": context.workspace_id,
"requestId": context.request_id,
"traceId": context.trace_id,
"actor": {
"actorType": "user",
"actorId": context.user_id,
"sessionId": context.session_id,
},
"source": {
"channel": "ai_orchestrator",
"client": "wolai-backend",
},
"tenantId": None,
"authToken": None,
"idempotencyKey": None,
"validateOnly": False,
"dryRun": False,
},
"tool": {
"tool": tool,
"kind": invocation_kind or infer_tool_invocation_kind(tool),
"mode": "result",
"argsJson": args_json,
"target": {
"workspaceId": context.workspace_id,
"pageId": target.get("pageId") if target else None,
"blockId": target.get("blockId") if target else None,
}
if target
else None,
"reason": reason,
"refs": refs or [],
},
}
if data is not None:
payload["data"] = data
response = await run_bridge_runtime(payload)
if "result" not in response:
raise RuntimeError("bridge-runtime 未返回 result")
return response["result"]
def _safe_json_preview(label: str, value: Any, limit: int) -> str | None:
if value is None:
return None
try:
text = json.dumps(value, ensure_ascii=False)
except TypeError:
return f"{label}=provided"
return f"{label}={text[:limit]}"
def _extract_inline_text(block: Any) -> str:
if not isinstance(block, dict):
return ""
content = block.get("content")
if not isinstance(content, list):
return ""
parts: list[str] = []
for node in content:
if isinstance(node, dict) and isinstance(node.get("text"), str):
parts.append(node["text"])
return "".join(parts).strip()
def _normalize_blocks(value: Any) -> list[Any]:
if isinstance(value, list):
return value
if isinstance(value, dict) and isinstance(value.get("blocks"), list):
return value["blocks"]
return []
def _summarize_blocks(value: Any, limit: int = 24) -> list[dict[str, Any]]:
root_blocks = _normalize_blocks(value)
queue: deque[tuple[Any, int]] = deque((block, 0) for block in root_blocks)
summaries: list[dict[str, Any]] = []
while queue and len(summaries) < limit:
block, depth = queue.popleft()
if not isinstance(block, dict):
continue
block_id = str(block.get("id") or "").strip()
block_type = str(block.get("type") or "unknown").strip()
children = block.get("children")
child_list = children if isinstance(children, list) else []
if block_id:
summaries.append(
{
"id": block_id,
"type": block_type,
"depth": depth,
"text": _extract_inline_text(block)[:120],
"childCount": len(child_list),
}
)
for child in child_list:
queue.append((child, depth + 1))
return summaries
def build_document_agent_instructions(
request: DocumentAiRunRequest,
) -> str:
profile = resolve_document_ai_profile(request.profile_id)
lines = [
"你是 mnote 的文档页 AI 编写代理。请始终使用简体中文。",
profile.instructions,
"你的工作对象只限当前文档页,不要假设你能跨页批量操作,也不要创建新的产品契约。",
"需要理解当前页内容时,先用 doc_get 或 doc_find;需要写入时,只用 doc_insert_blocks、doc_replace_range、slash_run。",
"只有在用户明确要求改标题时,才允许用 slash_run,并且只允许对当前 documentId 生成 /rename 命令。",
"如果工具没有返回成功结果,不要声称页面已经写入完成。",
]
latest_user_message = next(
(
message.content.strip()
for message in reversed(request.messages)
if message.role == "user" and message.content.strip()
),
"",
)
normalized_latest_user_message = latest_user_message.lower()
hard_rules: list[str] = []
if request.context.document_id and (
"标题" in latest_user_message or "rename" in normalized_latest_user_message
):
hard_rules.extend(
[
"当前请求包含标题修改意图;你必须调用一次 slash_run,不允许只返回文字说明。",
f"slash_run 的 text 必须严格使用这个格式:/rename {request.context.document_id} <新标题>。",
"只有在 slash_run 返回成功后,才能用一句简短中文确认标题已修改。",
]
)
if "blockid" in normalized_latest_user_message and (
"改写" in latest_user_message or "替换" in latest_user_message
):
hard_rules.extend(
[
"当前请求要求按 blockId 改写正文;你必须直接调用 doc_replace_range,不要只返回文字说明。",
"如果用户消息里已经给了 blockId 和目标文本,就不要先调用 doc_find。",
"只有在 doc_replace_range 返回成功后,才能用一句简短中文确认正文已修改。",
]
)
if ("插入" in latest_user_message or "新增" in latest_user_message) and (
"正文" in latest_user_message or "段" in latest_user_message
):
hard_rules.extend(
[
"当前请求要求新增正文;你必须调用 doc_insert_blocks,不允许只返回文字说明。",
"doc_insert_blocks 至少要插入一个 paragraph 或 heading block;除非用户明确要求,不要顺手改标题。",
"只有在 doc_insert_blocks 返回成功后,才能用一句简短中文确认正文已写入。",
]
)
if hard_rules:
lines.append("当前请求的强制执行规则:\n- " + "\n- ".join(hard_rules))
context = request.context
lines.append(f"profileId={profile.id}")
if context.document_id:
lines.append(f"documentId={context.document_id}")
lines.append(
"documentBlockSummaries="
+ json.dumps(
_summarize_blocks(context.document_blocks),
ensure_ascii=False,
)
)
for label, value, limit in (
("pageOptions", context.page_options, 1600),
("editorRuntimePageOptions", context.editor_runtime_page_options, 1200),
("node", context.node, 1600),
("subtree", context.subtree, 2400),
("outline", context.outline, 2400),
("evidence", context.evidence, 2400),
):
preview = _safe_json_preview(label, value, limit)
if preview:
lines.append(preview)
return "\n\n".join(lines).strip()
def _trim_messages(messages: list[DocumentAiMessage]) -> list[dict[str, str]]:
return [
{"role": message.role, "content": message.content}
for message in messages[-24:]
if message.content.strip()
]
async def _emit_tool_event(
ctx: DocumentAiRunContext,
tool: str,
args_json: dict[str, Any],
execute: Callable[[], Any],
) -> Any:
tool_id = ctx.emitter.next_tool_id(tool)
await ctx.emitter.emit(
"tool_call",
{
"id": tool_id,
"tool": tool,
"args": args_json,
},
)
started = time.perf_counter()
try:
result = await execute()
except Exception as error:
elapsed_ms = max(0, round((time.perf_counter() - started) * 1000))
await ctx.emitter.emit(
"tool_result",
{
"id": tool_id,
"tool": tool,
"ok": False,
"ms": elapsed_ms,
"result": {"error": str(error)},
},
)
raise
ctx.emitter.tool_count += 1
elapsed_ms = max(0, round((time.perf_counter() - started) * 1000))
await ctx.emitter.emit(
"tool_result",
{
"id": tool_id,
"tool": tool,
"ok": True,
"ms": elapsed_ms,
"result": result,
},
)
return result
def build_document_tools():
if not sdk_available():
raise RuntimeError("openai-agents 未安装,无法构建文档页 tools")
@function_tool
async def doc_get(
wrapper: RunContextWrapper[DocumentAiRunContext],
maxBlocks: int = 80,
) -> Any:
"""读取当前页的块摘要。需要先理解页面结构时调用。"""
ctx = wrapper.context
args_json = {"maxBlocks": max(10, min(int(maxBlocks), 240))}
return await _emit_tool_event(
ctx,
"doc_get",
args_json,
lambda: ctx.bridge.execute_tool(
tool="doc_get",
context=ctx,
invocation_kind="query",
args_json=args_json,
data=ctx.document_blocks,
target={"pageId": ctx.document_id},
),
)
@function_tool
async def doc_find(
wrapper: RunContextWrapper[DocumentAiRunContext],
query: str,
maxResults: int = 8,
) -> Any:
"""按文本查找当前页的块。准备改写某段内容前先调用。"""
ctx = wrapper.context
query_text = str(query).strip()
if not query_text:
raise RuntimeError("doc_find 缺少 query")
args_json = {
"query": query_text,
"maxResults": max(1, min(int(maxResults), 30)),
}
return await _emit_tool_event(
ctx,
"doc_find",
args_json,
lambda: ctx.bridge.execute_tool(
tool="doc_find",
context=ctx,
invocation_kind="query",
args_json=args_json,
data=ctx.document_blocks,
target={"pageId": ctx.document_id},
),
)
@function_tool
async def doc_insert_blocks(
wrapper: RunContextWrapper[DocumentAiRunContext],
blocks: list[InsertBlockSpec],
afterBlockId: str | None = None,
beforeBlockId: str | None = None,
) -> Any:
"""在当前页插入新块。新增标题或段落时调用。"""
ctx = wrapper.context
if not blocks:
raise RuntimeError("doc_insert_blocks 缺少 blocks")
args_json = {
"blocks": [
{
"type": block.block_type,
"text": block.text,
"level": block.level,
}
for block in blocks[:20]
],
}
if afterBlockId:
args_json["afterBlockId"] = afterBlockId
if beforeBlockId:
args_json["beforeBlockId"] = beforeBlockId
return await _emit_tool_event(
ctx,
"doc_insert_blocks",
args_json,
lambda: ctx.bridge.execute_tool(
tool="doc_insert_blocks",
context=ctx,
invocation_kind="command",
args_json=args_json,
data=ctx.document_blocks,
target={"pageId": ctx.document_id},
),
)
@function_tool
async def doc_replace_range(
wrapper: RunContextWrapper[DocumentAiRunContext],
blockId: str,
text: str,
mode: Literal["replace", "append", "prepend"] = "replace",
) -> Any:
"""改写当前页某个块的文本。"""
ctx = wrapper.context
block_id = str(blockId).strip()
content = str(text).strip()
if not block_id or not content:
raise RuntimeError("doc_replace_range 缺少 blockId 或 text")
safe_mode = mode if mode in {"replace", "append", "prepend"} else "replace"
args_json = {
"blockId": block_id,
"text": content,
"mode": safe_mode,
}
return await _emit_tool_event(
ctx,
"doc_replace_range",
args_json,
lambda: ctx.bridge.execute_tool(
tool="doc_replace_range",
context=ctx,
invocation_kind="command",
args_json=args_json,
data=ctx.document_blocks,
target={"pageId": ctx.document_id, "blockId": block_id},
),
)
@function_tool
async def slash_run(
wrapper: RunContextWrapper[DocumentAiRunContext],
text: str,
) -> Any:
"""执行当前页允许的斜杠命令。当前只允许改当前页标题。"""
ctx = wrapper.context
command_text = str(text).strip()
if not command_text.startswith("/rename "):
raise RuntimeError("当前只允许 /rename 当前页标题")
if ctx.document_id and f"/rename {ctx.document_id} " not in command_text:
raise RuntimeError("slash_run 只允许改当前文档页标题")
args_json = {"text": command_text}
return await _emit_tool_event(
ctx,
"slash_run",
args_json,
lambda: ctx.bridge.execute_tool(
tool="slash_run",
context=ctx,
invocation_kind="command",
args_json=args_json,
data={"source": "ai-orchestrator"},
target={"pageId": ctx.document_id},
),
)
return [
doc_get,
doc_find,
doc_insert_blocks,
doc_replace_range,
slash_run,
]
def build_document_agent(
request: DocumentAiRunRequest,
):
if not sdk_available():
raise RuntimeError("openai-agents 未安装,无法构建文档页 agent")
kwargs: dict[str, Any] = {
"name": "mnote-document-page-writer",
"instructions": build_document_agent_instructions(request),
"tools": build_document_tools(),
}
kwargs["model"] = resolve_document_agent_request_model(
request_model=request.model,
request_model_key=request.model_key,
configured_default_model=settings.mnote_ai_default_model,
)
return Agent[DocumentAiRunContext](**kwargs)
def build_forward_headers(source_headers: dict[str, str]) -> dict[str, str]:
allowed = {
"authorization",
"cookie",
"x-request-id",
"x-trace-id",
"x-session-id",
"x-mnote-workspace-id",
"x-mnote-source-channel",
"x-mnote-source-client",
"x-mnote-actor-id",
"x-mnote-actor-type",
"user-agent",
}
return {
key: value
for key, value in source_headers.items()
if key.lower() in allowed and value
}
def build_session(session_id: str | None):
if not sdk_available() or not session_id:
return None
raw_path = settings.openai_agents_session_db_path.strip()
db_path = (
Path(raw_path)
if raw_path
else Path(__file__).resolve().parents[2] / ".data" / "openai-agents-sessions.sqlite3"
)
db_path.parent.mkdir(parents=True, exist_ok=True)
return SQLiteSession(session_id, str(db_path))
async def run_document_agent_stream(
request: DocumentAiRunRequest,
*,
source_headers: dict[str, str],
) -> AsyncGenerator[str, None]:
if not sdk_available():
raise RuntimeError("openai-agents 未安装")
if not has_configured_openai_api_key():
raise RuntimeError("后端未配置 OPENAI_API_KEY")
effective_api_key = get_effective_openai_api_key()
if effective_api_key:
os.environ.setdefault("OPENAI_API_KEY", effective_api_key)
emitter = DocumentAiEmitter()
run_context = DocumentAiRunContext(
user_id=request.user_id,
request_id=make_id("req"),
trace_id=make_id("trace"),
session_id=request.session_id,
workspace_id=source_headers.get("x-mnote-workspace-id") or None,
document_id=request.context.document_id,
document_blocks=request.context.document_blocks,
node=request.context.node,
subtree=request.context.subtree,
outline=request.context.outline,
evidence=request.context.evidence,
page_options=request.context.page_options,
editor_runtime_page_options=request.context.editor_runtime_page_options,
forward_headers=build_forward_headers(source_headers),
emitter=emitter,
bridge=LocalBridgeRuntime(),
)
async def worker() -> None:
try:
agent = build_document_agent(request)
run_config = RunConfig(
model_provider=build_document_model_provider(
request_model=request.model,
configured_default_model=settings.mnote_ai_default_model,
)
)
session = build_session(request.session_id)
run_input = _trim_messages(request.messages)
if not run_input:
raise RuntimeError("缺少 messages")
max_turns = max(1, min(int(request.max_steps), 12))
with trace("mnote-document-ai-run", group_id=run_context.trace_id):
result = await Runner.run(
agent,
input=run_input,
context=run_context,
session=session,
max_turns=max_turns,
run_config=run_config,
)
text = str(result.final_output or "").strip() or "(无输出)"
await emitter.emit("assistant_message", {"text": text})
await emitter.emit(
"completion",
{
"ok": True,
"text": text,
"steps": max(1, emitter.tool_count or 1),
},
)
except Exception as error:
await emitter.emit(
"error",
{
"ok": False,
"message": str(error),
},
)
finally:
await emitter.close()
task = asyncio.create_task(worker())
try:
while True:
item = await emitter.queue.get()
if item is None:
break
event, data = item
yield to_sse_frame(event, data)
finally:
if not task.done():
task.cancel()
with contextlib.suppress(asyncio.CancelledError):
await task