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