M4b：ai_provider 删除手术①-⑥——神识全面接管模型出口
- ① 后端删除 app/ai_provider/（38 文件，含 anthropic/openai 格式、billing、audio、embedding）、services/ai_provider_admin.py、core/config/model_config.py、tests/ai_provider/（21 文件） - ② 连带退役 application/agent/（旧 Agent 循环/steer/follow-up/确认注册表）、application/skill_executor.py、app/skills/（python skills 孤儿）、endpoints/audio.py、application/imports.py（ProviderManager 消费方）；settings 端点收敛为生成偏好 - ③ runs.py 修剪：AgentRunRegistry/steer/follow-up/confirmations 端点退役，保留 list/get/events/cancel/stream；safety 确认改为直接落库裁决（pending→approved/denied） - ④ 前端删除：models 页、ModelGatewayView、gateway.ts、GatewayPanel（玄关）、ModelGatewayPanel/ModelStrategyForm、audio.ts/useVoiceInput、网关类型与 mock；设置页收敛为外观/语言/通知 - ⑤ 配置收敛：config/models|providers|pricing.example.yaml 归档 shenshi/config/；.env.example 删 provider keys/MODEL_ENDPOINT_POOL/MODEL_CONFIG_PATH，收敛 SHENSHI_* - ⑥ 验收：rg 审计灵台 Provider SDK import 与 ai_provider 引用归零；后端 528 passed（2 个预存失败经 HEAD 复跑验证与本手术无关）、前端 type-check+201 passed、神识 go vet+test、BFF 4 passed
This commit is contained in:
chenxin committed
1 parent
0d383725b5
commit
462d570751
129 files changed
+81
-15867
No files matched your search
+7
-48
@@ -88,7 +88,7 @@ LOG_ROTATION_SIZE=10 MB
|
||||
# 本地单用户范围标识，当前阶段不实现多用户、组织、租户、OAuth 或 SSO
|
||||
DEFAULT_OWNER_SCOPE=local-user
|
||||
|
||||
# 默认逻辑模型名。必须指向 config/models.yaml 中真实远程 provider 的逻辑模型
|
||||
# 默认逻辑模型名（提示值；实际路由由神识端点目录决定）
|
||||
DEFAULT_MODEL_HINT=deepseek-v4-flash
|
||||
|
||||
# --- 管理端点和认证 ---
|
||||
@@ -207,7 +207,7 @@ BIG_CYCLE_MEMORY_THRESHOLD=20
|
||||
# 两次大周天最小间隔（秒），避免频繁触发
|
||||
BIG_CYCLE_MIN_INTERVAL_SECONDS=300
|
||||
|
||||
# 大周天候选生成所用模型（config/models.yaml 的逻辑模型名）
|
||||
# 大周天候选生成所用模型（逻辑模型名；路由由神识决定）
|
||||
BIG_CYCLE_MODEL=deepseek-v4-flash
|
||||
|
||||
# 单次大周天生成的演化候选上限
|
||||
@@ -238,7 +238,7 @@ MCP_STDIO_CALL_TIMEOUT_SECONDS=30
|
||||
# MCP stdio SIGTERM → SIGKILL 等待秒数
|
||||
MCP_STDIO_STOP_TIMEOUT_SECONDS=5
|
||||
|
||||
# Embedding 模型逻辑名。指向 config/models.yaml 中已注册的 embedding 模型
|
||||
# Embedding 模型逻辑名（路由由神识决定）
|
||||
EMBEDDING_MODEL_HINT=bge-m3
|
||||
|
||||
# --- 感知与导入 ---
|
||||
@@ -259,53 +259,12 @@ IMPORT_ALLOWED_AUDIO_MIME_TYPES=audio/wav,audio/mpeg,audio/mp4,audio/flac
|
||||
|
||||
|
||||
# =====================================================================
|
||||
# 模型网关 · 神识（LLM Provider API Keys）
|
||||
# 模型网关 · 神识（Go 网关统一接管，灵台不经手任何 Provider 密钥）
|
||||
# =====================================================================
|
||||
# 模型名称、provider 映射、能力标签、超时、限流、价格等，统一放在 config/models.yaml 中。
|
||||
# .env 只保存各 provider 的 URL 与 API key，避免模型策略散落在环境变量中。
|
||||
# 模型端点目录（含 base_url / api_key / 上游模型 / 优先级）统一经神识环境变量 SHENSHI_ENDPOINTS 配置，
|
||||
# 详见下方「神识 · 模型网关」段；旧 config/models|providers|pricing.yaml 与灵台侧各 provider
|
||||
# API Key 已随 ai_provider 退役（归档于 shenshi/config/ 仅供参考）。
|
||||
|
||||
# --- DeepSeek ---
|
||||
DEEPSEEK_API_KEY=
|
||||
DEEPSEEK_BASE_URL=https://api.deepseek.com
|
||||
|
||||
# --- 火山方舟 / 豆包 ---
|
||||
VOLCENGINE_ARK_API_KEY=
|
||||
VOLCENGINE_ARK_BASE_URL=https://ark.cn-beijing.volces.com/api/v3
|
||||
|
||||
# --- 小米 MiMo ---
|
||||
XIAOMI_MIMO_API_KEY=
|
||||
XIAOMI_MIMO_BASE_URL=https://token-plan-cn.xiaomimimo.com/v1
|
||||
|
||||
# --- OpenAI ---
|
||||
OPENAI_API_KEY=
|
||||
OPENAI_BASE_URL=https://api.openai.com/v1
|
||||
|
||||
# --- Anthropic ---
|
||||
ANTHROPIC_API_KEY=
|
||||
ANTHROPIC_BASE_URL=https://api.anthropic.com
|
||||
|
||||
# --- 元炁（本地推理引擎）---
|
||||
# 作为 OpenAI 兼容 provider 接入，供 bge-m3 embedding 等模型调用。
|
||||
YUANQI_BASE_URL=http://localhost:11434
|
||||
YUANQI_API_KEY=
|
||||
|
||||
# --- 其他兼容网关预留 ---
|
||||
CUSTOM_LLM_API_KEY=
|
||||
CUSTOM_LLM_BASE_URL=
|
||||
|
||||
# --- 模型 endpoint 池路由（高级调优项，保持默认即可正常工作）---
|
||||
# 物理 endpoint 池的 failover 与重试策略
|
||||
MODEL_ENDPOINT_POOL_ENABLED=true
|
||||
MODEL_ENDPOINT_POOL_DRY_RUN_SELECTION=false
|
||||
MODEL_ENDPOINT_POOL_EXPOSE_DEBUG_HEADERS=false
|
||||
MODEL_ENDPOINT_POOL_AGENT_FAILOVER_ENABLED=false
|
||||
MODEL_ENDPOINT_POOL_PER_ENDPOINT_MAX_RETRIES=1
|
||||
MODEL_ENDPOINT_POOL_RETRY_BASE_DELAY_SECONDS=0.25
|
||||
MODEL_ENDPOINT_POOL_RETRY_MAX_DELAY_SECONDS=2.0
|
||||
|
||||
# 模型配置文件路径
|
||||
# 默认读取项目根目录 config/models.yaml；不存在时 fallback 到 config/models.example.yaml
|
||||
MODEL_CONFIG_PATH=../config/models.yaml
|
||||
|
||||
|
||||
# =====================================================================
|
||||
|
||||
@@ -244,7 +244,7 @@ make test # 全部测试（backend / shenshi / bff / fronten
|
||||
```bash
|
||||
# 先准备配置（只需一次）
|
||||
cp .env.example .env # 填写 POSTGRES_PASSWORD 与各 *_API_KEY
|
||||
cp config/models.example.yaml config/models.yaml # 按需填写真实模型 endpoint
|
||||
# 模型端点目录经 SHENSHI_ENDPOINTS 配置（见 .env.example 神识段），无需复制 yaml
|
||||
|
||||
# 1) 启动基础设施（首次创建 hunjing-net 网络；Milvus 会自动带起 hunjing-etcd / hunjing-minio）
|
||||
docker compose -f docker/docker-compose.infra.yml up -d
|
||||
|
||||
@@ -1,63 +0,0 @@
|
||||
from app.ai_provider.config.loader import load_models_config, load_providers_config
|
||||
from app.ai_provider.config.models import ModelConfig, ProviderConfig
|
||||
|
||||
# 嵌入 / 重排序
|
||||
from app.ai_provider.embedding import (
|
||||
EmbeddingProtocol,
|
||||
EmbeddingRequest,
|
||||
EmbeddingResponse,
|
||||
EmbeddingResult,
|
||||
EmbeddingService,
|
||||
EmbeddingServiceError,
|
||||
)
|
||||
from app.ai_provider.manager import ProviderManager, ProviderManagerFactory
|
||||
from app.ai_provider.reranker import (
|
||||
RerankerProtocol,
|
||||
RerankerService,
|
||||
RerankerServiceError,
|
||||
RerankRequest,
|
||||
RerankResponse,
|
||||
RerankResult,
|
||||
)
|
||||
from app.ai_provider.types.chat import ChatChunk, ChatResult, TokenUsage
|
||||
from app.ai_provider.types.errors import AIProviderError
|
||||
from app.ai_provider.types.messages import (
|
||||
AIMessage,
|
||||
BaseMessage,
|
||||
HumanMessage,
|
||||
SystemMessage,
|
||||
ToolMessage,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
# LLM
|
||||
"ProviderManager",
|
||||
"ProviderManagerFactory",
|
||||
"BaseMessage",
|
||||
"SystemMessage",
|
||||
"HumanMessage",
|
||||
"AIMessage",
|
||||
"ToolMessage",
|
||||
"ChatChunk",
|
||||
"ChatResult",
|
||||
"TokenUsage",
|
||||
"AIProviderError",
|
||||
"ModelConfig",
|
||||
"ProviderConfig",
|
||||
"load_models_config",
|
||||
"load_providers_config",
|
||||
# Embedding
|
||||
"EmbeddingProtocol",
|
||||
"EmbeddingRequest",
|
||||
"EmbeddingResponse",
|
||||
"EmbeddingResult",
|
||||
"EmbeddingService",
|
||||
"EmbeddingServiceError",
|
||||
# Reranker
|
||||
"RerankerProtocol",
|
||||
"RerankRequest",
|
||||
"RerankResponse",
|
||||
"RerankResult",
|
||||
"RerankerService",
|
||||
"RerankerServiceError",
|
||||
]
|
||||
@@ -1,197 +0,0 @@
|
||||
"""音频服务 — ASR / TTS 适配器外观。
|
||||
|
||||
用法::
|
||||
|
||||
from app.ai_provider.audio import AudioService
|
||||
|
||||
service = AudioService(factory)
|
||||
result = await service.transcribe(audio_bytes, format="wav")
|
||||
result = await service.synthesize("你好，魂镜", voice="default")
|
||||
|
||||
``AudioService`` 通过注册表中带有 ``asr`` / ``tts`` 能力标签的模型来解析
|
||||
使用哪个模型，然后将创建委托给 ``ModelFactory.create_audio()``。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from app.ai_provider.audio.base import (
|
||||
ASRProtocol,
|
||||
ASRRequest,
|
||||
ASRResponse,
|
||||
AudioServiceError,
|
||||
TTSProtocol,
|
||||
TTSRequest,
|
||||
TTSResponse,
|
||||
)
|
||||
|
||||
|
||||
class AudioService:
|
||||
"""外观类，负责解析 ASR / TTS 模型并委托给 ProviderManager。
|
||||
|
||||
接受 ProviderManager 或实现了以下接口的适配器：
|
||||
- ``.registry`` → 具有 ``list_capability_models(capability: str)`` 方法
|
||||
- ``.create_audio(name: str)`` → 返回 ASRProtocol / TTSProtocol 实例
|
||||
"""
|
||||
|
||||
def __init__(self, manager) -> None:
|
||||
self._manager = manager
|
||||
self._asr_adapter: ASRProtocol | None = None
|
||||
self._tts_adapter: TTSProtocol | None = None
|
||||
|
||||
# ── Registry 兼容适配 ─────────────────────────────────────────
|
||||
|
||||
@property
|
||||
def _registry(self):
|
||||
"""按需获取 registry，兼容 adapter 和裸 ProviderManager/ModelRegistry。"""
|
||||
mgr = self._manager
|
||||
if hasattr(mgr, "registry"):
|
||||
return mgr.registry
|
||||
if hasattr(mgr, "model_registry"):
|
||||
return mgr.model_registry
|
||||
return mgr
|
||||
|
||||
def _list_capability_models(self, tag: str):
|
||||
reg = self._registry
|
||||
# 适配器接口（list_capability_models）vs ModelRegistry 接口（list_by_capability）
|
||||
if hasattr(reg, "list_capability_models"):
|
||||
return reg.list_capability_models(tag)
|
||||
if hasattr(reg, "list_by_capability"):
|
||||
return reg.list_by_capability(tag)
|
||||
return []
|
||||
|
||||
@property
|
||||
def _factory(self):
|
||||
mgr = self._manager
|
||||
if hasattr(mgr, "create_audio"):
|
||||
return mgr
|
||||
return None
|
||||
|
||||
# ── ASR ──────────────────────────────────────────────────────────────
|
||||
|
||||
@property
|
||||
def asr(self) -> ASRProtocol:
|
||||
"""解析并缓存 ASR 适配器（取注册表中第一个活跃 ASR 模型）。"""
|
||||
if self._asr_adapter is None:
|
||||
self._asr_adapter = self._resolve_asr()
|
||||
return self._asr_adapter
|
||||
|
||||
def asr_for(self, model_name: str) -> ASRProtocol:
|
||||
"""解析并缓存指定模型名的 ASR 适配器（覆盖默认第一个）。"""
|
||||
self._asr_adapter = self._factory.create_audio(model_name)
|
||||
return self._asr_adapter
|
||||
|
||||
async def transcribe(
|
||||
self,
|
||||
audio_data: bytes,
|
||||
format: str = "wav",
|
||||
language: str | None = None,
|
||||
model_name: str | None = None,
|
||||
) -> ASRResponse:
|
||||
"""快捷方法：构造请求并调用 ASR 适配器。
|
||||
|
||||
Args:
|
||||
model_name: 可选，指定使用哪个 ASR 模型（按注册表逻辑名）。
|
||||
省略时取注册表中第一个活跃的 ASR 模型。
|
||||
"""
|
||||
if model_name is not None:
|
||||
adapter = self.asr_for(model_name)
|
||||
else:
|
||||
adapter = self.asr
|
||||
request = ASRRequest(
|
||||
audio_data=audio_data,
|
||||
format=format,
|
||||
language=language,
|
||||
)
|
||||
return await adapter.transcribe(request)
|
||||
|
||||
# ── TTS ──────────────────────────────────────────────────────────────
|
||||
|
||||
@property
|
||||
def tts(self) -> TTSProtocol:
|
||||
"""解析并缓存 TTS 适配器。"""
|
||||
if self._tts_adapter is None:
|
||||
self._tts_adapter = self._resolve_tts()
|
||||
return self._tts_adapter
|
||||
|
||||
async def synthesize(
|
||||
self,
|
||||
text: str,
|
||||
voice: str = "default",
|
||||
response_format: str = "wav",
|
||||
extra: dict | None = None,
|
||||
) -> TTSResponse:
|
||||
"""快捷方法：构造请求并调用 TTS 适配器。"""
|
||||
request = TTSRequest(
|
||||
text=text,
|
||||
voice=voice,
|
||||
response_format=response_format,
|
||||
extra=extra,
|
||||
)
|
||||
return await self.tts.synthesize(request)
|
||||
|
||||
# ── 解析 ───────────────────────────────────────────────────────
|
||||
|
||||
def _resolve_asr(self) -> ASRProtocol:
|
||||
"""查找第一个活跃的 ASR 模型并通过工厂创建。"""
|
||||
models = self._list_capability_models("asr")
|
||||
if not models:
|
||||
raise AudioServiceError(
|
||||
status_code=503,
|
||||
code="audio_asr_no_model",
|
||||
message="No ASR model configured in the model registry.",
|
||||
)
|
||||
factory = self._factory
|
||||
if factory is None:
|
||||
raise AudioServiceError(
|
||||
status_code=501,
|
||||
code="audio_asr_no_factory",
|
||||
message="Audio adapter factory not available (provider manager missing create_audio).",
|
||||
)
|
||||
return factory.create_audio(models[0].name)
|
||||
|
||||
def _resolve_tts(self) -> TTSProtocol:
|
||||
"""查找第一个活跃的 TTS 模型并通过工厂创建。"""
|
||||
models = self._list_capability_models("tts")
|
||||
if not models:
|
||||
raise AudioServiceError(
|
||||
status_code=503,
|
||||
code="audio_tts_no_model",
|
||||
message="No TTS model configured in the model registry.",
|
||||
)
|
||||
factory = self._factory
|
||||
if factory is None:
|
||||
raise AudioServiceError(
|
||||
status_code=501,
|
||||
code="audio_tts_no_factory",
|
||||
message="Audio adapter factory not available (provider manager missing create_audio).",
|
||||
)
|
||||
return factory.create_audio(models[0].name)
|
||||
|
||||
# ── 生命周期 ────────────────────────────────────────────────────────
|
||||
|
||||
async def close(self) -> None:
|
||||
"""释放适配器持有的 HTTP 连接。"""
|
||||
for adapter in (self._asr_adapter, self._tts_adapter):
|
||||
if adapter is None:
|
||||
continue
|
||||
close = getattr(adapter, "close", None)
|
||||
if close is not None:
|
||||
await close()
|
||||
|
||||
async def __aenter__(self) -> AudioService:
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *exc_info: object) -> None:
|
||||
await self.close()
|
||||
|
||||
|
||||
__all__ = [
|
||||
"ASRProtocol",
|
||||
"ASRRequest",
|
||||
"ASRResponse",
|
||||
"AudioService",
|
||||
"AudioServiceError",
|
||||
"TTSProtocol",
|
||||
"TTSRequest",
|
||||
"TTSResponse",
|
||||
]
|
||||
@@ -1,263 +0,0 @@
|
||||
"""MiMo ASR adapter — ``POST /v1/chat/completions`` with ``input_audio``.
|
||||
|
||||
MiMo 的 ASR 不走独立音频端点，而是通过聊天补全接口扩展实现：
|
||||
|
||||
POST {base_url}/chat/completions
|
||||
Content-Type: application/json
|
||||
Authorization: Bearer {api_key}
|
||||
|
||||
{
|
||||
"model": "mimo-v2.5-asr",
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "input_audio",
|
||||
"input_audio": {
|
||||
"data": "data:audio/wav;base64,<base64_encoded_audio>"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
],
|
||||
"asr_options": {
|
||||
"language": "auto"
|
||||
}
|
||||
}
|
||||
|
||||
响应：:
|
||||
|
||||
{
|
||||
"choices": [{
|
||||
"message": {
|
||||
"content": "转写后的文本"
|
||||
}
|
||||
}]
|
||||
}
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import time
|
||||
|
||||
import httpx
|
||||
|
||||
from app.ai_provider.audio.base import (
|
||||
ASRProtocol,
|
||||
ASRRequest,
|
||||
ASRResponse,
|
||||
AudioAdapterBase,
|
||||
AudioServiceError,
|
||||
)
|
||||
from app.core.logging import logger
|
||||
|
||||
_MIME_MAP = {
|
||||
"wav": "audio/wav",
|
||||
"mp3": "audio/mpeg",
|
||||
"m4a": "audio/mp4",
|
||||
"flac": "audio/flac",
|
||||
}
|
||||
|
||||
# ASR 缓存：默认最多 128 条，TTL 1 小时
|
||||
_ASR_CACHE_MAX_SIZE = 128
|
||||
_ASR_CACHE_TTL_SECONDS = 3600
|
||||
|
||||
|
||||
class _ASRCacheEntry:
|
||||
__slots__ = ("text", "stored_at")
|
||||
|
||||
def __init__(self, text: str) -> None:
|
||||
self.text = text
|
||||
self.stored_at = time.monotonic()
|
||||
|
||||
|
||||
_asr_cache: dict[str, _ASRCacheEntry] = {}
|
||||
|
||||
|
||||
def _asr_cache_key(model: str, audio_data: bytes, language: str, fmt: str) -> str:
|
||||
"""生成 ASR 缓存键：SHA256(model + audio_bytes + language + format)。"""
|
||||
raw = f"{model}\0{language}\0{fmt}\0"
|
||||
raw_bytes = raw.encode() + audio_data
|
||||
return hashlib.sha256(raw_bytes).hexdigest()
|
||||
|
||||
|
||||
def _asr_cache_get(key: str) -> str | None:
|
||||
entry = _asr_cache.get(key)
|
||||
if entry is None:
|
||||
return None
|
||||
if time.monotonic() - entry.stored_at > _ASR_CACHE_TTL_SECONDS:
|
||||
del _asr_cache[key]
|
||||
return None
|
||||
return entry.text
|
||||
|
||||
|
||||
def _asr_cache_set(key: str, text: str) -> None:
|
||||
if len(_asr_cache) >= _ASR_CACHE_MAX_SIZE:
|
||||
sorted_keys = sorted(_asr_cache.keys(), key=lambda k: _asr_cache[k].stored_at)
|
||||
for old_key in sorted_keys[: _ASR_CACHE_MAX_SIZE // 2]:
|
||||
del _asr_cache[old_key]
|
||||
_asr_cache[key] = _ASRCacheEntry(text)
|
||||
|
||||
|
||||
class HunjingASROpenAI(AudioAdapterBase, ASRProtocol):
|
||||
"""MiMo ASR 适配器，通过 ``POST /chat/completions`` 携带 ``input_audio`` 实现。
|
||||
|
||||
将音频字节包装为 chat 消息中的 ``input_audio`` 内容类型的 data URL，
|
||||
然后从响应中提取转写文本。
|
||||
"""
|
||||
|
||||
_adapter_type = "asr"
|
||||
_default_timeout = 300
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
model: str = "mimo-v2.5-asr",
|
||||
request_timeout_seconds: int = 300,
|
||||
) -> None:
|
||||
self._base_url = base_url.rstrip("/")
|
||||
self._api_key = api_key
|
||||
self._model = model
|
||||
self._timeout = request_timeout_seconds
|
||||
self._client: httpx.AsyncClient | None = None
|
||||
|
||||
# ── 客户端 ───────────────────────────────────────────────────────────
|
||||
|
||||
@property
|
||||
def _http(self) -> httpx.AsyncClient:
|
||||
if self._client is None or self._client.is_closed:
|
||||
self._client = httpx.AsyncClient(timeout=httpx.Timeout(self._timeout))
|
||||
return self._client
|
||||
|
||||
async def close(self) -> None:
|
||||
if self._client is not None and not self._client.is_closed:
|
||||
await self._client.aclose()
|
||||
|
||||
# ── ASRProtocol ──────────────────────────────────────────────────────
|
||||
|
||||
async def transcribe(self, request: ASRRequest) -> ASRResponse:
|
||||
"""通过 ``POST /chat/completions`` 转写音频。
|
||||
|
||||
将音频编码为 data URL 并以 ``input_audio`` 内容块的形式发送到 chat 消息中。
|
||||
"""
|
||||
# ── 缓存查找 ──────────────────────────────────────────────────
|
||||
cache_key = _asr_cache_key(
|
||||
self._model, request.audio_data, request.language or "auto", request.format
|
||||
)
|
||||
cached = _asr_cache_get(cache_key)
|
||||
if cached is not None:
|
||||
logger.info("audio_asr.cache_hit", model=self._model, text_length=len(cached))
|
||||
return ASRResponse(text=cached)
|
||||
|
||||
# ── 构造请求 ──────────────────────────────────────────────────
|
||||
mime = _MIME_MAP.get(request.format, "audio/wav")
|
||||
b64_audio = base64.b64encode(request.audio_data).decode()
|
||||
data_url = f"data:{mime};base64,{b64_audio}"
|
||||
|
||||
payload: dict = {
|
||||
"model": self._model,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "input_audio",
|
||||
"input_audio": {
|
||||
"data": data_url,
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
"asr_options": {
|
||||
"language": request.language or "auto",
|
||||
},
|
||||
}
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {self._api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
url = f"{self._base_url}/chat/completions"
|
||||
|
||||
logger.info(
|
||||
"audio_asr.request",
|
||||
model=self._model,
|
||||
format=request.format,
|
||||
language=request.language,
|
||||
url=url,
|
||||
size_bytes=len(request.audio_data),
|
||||
)
|
||||
|
||||
# ── API 调用 ────────────────────────────────────────────────
|
||||
try:
|
||||
response = await self._http.post(url, json=payload, headers=headers)
|
||||
except httpx.TimeoutException as exc:
|
||||
raise AudioServiceError(
|
||||
status_code=504,
|
||||
code="audio_asr_timeout",
|
||||
message="ASR request timed out calling chat completions.",
|
||||
) from exc
|
||||
except httpx.RequestError as exc:
|
||||
raise AudioServiceError(
|
||||
status_code=503,
|
||||
code="audio_asr_connection_error",
|
||||
message=f"ASR connection error: {exc}",
|
||||
) from exc
|
||||
|
||||
if response.status_code != 200:
|
||||
body = response.text[:500]
|
||||
raise AudioServiceError(
|
||||
status_code=response.status_code,
|
||||
code="audio_asr_api_error",
|
||||
message=f"ASR API returned {response.status_code}: {body}",
|
||||
)
|
||||
|
||||
# 解析响应，提取转写文本
|
||||
text = self._extract_text(response.json())
|
||||
if not text:
|
||||
raise AudioServiceError(
|
||||
status_code=500,
|
||||
code="audio_asr_empty_response",
|
||||
message="ASR API returned empty transcription.",
|
||||
)
|
||||
|
||||
# ── 写入缓存 ────────────────────────────────────────────────
|
||||
_asr_cache_set(cache_key, text)
|
||||
|
||||
logger.info(
|
||||
"audio_asr.success",
|
||||
model=self._model,
|
||||
text_length=len(text),
|
||||
cache_size=len(_asr_cache),
|
||||
)
|
||||
return ASRResponse(text=text)
|
||||
|
||||
# ── 辅助方法 ──────────────────────────────────────────────────────────
|
||||
|
||||
@staticmethod
|
||||
def _extract_text(body: dict) -> str:
|
||||
"""从 chat completion 响应中提取转写文本。"""
|
||||
try:
|
||||
choices = body.get("choices", [])
|
||||
if not choices:
|
||||
raise AudioServiceError(
|
||||
status_code=500,
|
||||
code="audio_asr_no_choices",
|
||||
message="ASR response has no choices.",
|
||||
)
|
||||
return (choices[0].get("message") or {}).get("content") or ""
|
||||
except AudioServiceError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
raise AudioServiceError(
|
||||
status_code=500,
|
||||
code="audio_asr_response_parse_error",
|
||||
message=f"Failed to parse ASR response: {exc}",
|
||||
) from exc
|
||||
@@ -1,245 +0,0 @@
|
||||
"""音频服务类型 — 请求/响应模型和协议定义。
|
||||
|
||||
ASR（自动语音识别）和 TTS（文本转语音）适配器**不是** LangChain 的
|
||||
``BaseChatModel`` 子类 —— 它们是通过 HTTP 调用的服务，经由 ``AudioService``
|
||||
或直接从 ``ModelFactory`` 消费。
|
||||
|
||||
用法::
|
||||
|
||||
service = AudioService(factory)
|
||||
result = await service.transcribe(audio_data)
|
||||
result = await service.synthesize("你好")
|
||||
|
||||
# 或通过工厂直接调用:
|
||||
tts = factory.create_audio("mimo-v2.5-tts")
|
||||
result = await tts.synthesize(TTSRequest(text="你好"))
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
|
||||
|
||||
from app.ai_provider.types.errors import AIProviderError
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from app.ai_provider.config.models import ModelConfig, ModelEndpointConfig
|
||||
|
||||
|
||||
class AudioServiceError(AIProviderError):
|
||||
"""ASR / TTS 服务失败的基础错误。"""
|
||||
|
||||
|
||||
# ── ASR ───────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class ASRRequest:
|
||||
"""将音频转写为文本的请求。
|
||||
|
||||
属性:
|
||||
audio_data: 原始音频字节。
|
||||
format: 音频容器格式（如 ``"wav"``、``"mp3"``、``"m4a"``）。
|
||||
language: 可选的 ISO-639-1 语言提示（如 ``"zh"``、``"en"``）。
|
||||
"""
|
||||
|
||||
__slots__ = ("audio_data", "format", "language")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
audio_data: bytes,
|
||||
format: str = "wav",
|
||||
language: str | None = None,
|
||||
) -> None:
|
||||
self.audio_data = audio_data
|
||||
self.format = format
|
||||
self.language = language
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ASRResponse:
|
||||
"""成功转写的结果。"""
|
||||
|
||||
text: str
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class ASRProtocol(Protocol):
|
||||
"""ASR 适配器的协议。"""
|
||||
|
||||
async def transcribe(self, request: ASRRequest) -> ASRResponse:
|
||||
"""将音频字节转写为文本。
|
||||
|
||||
异常:
|
||||
AudioServiceError: API 调用失败或响应异常时抛出。
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
# ── TTS ───────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TTSRequest:
|
||||
"""将文本合成为语音的请求。
|
||||
|
||||
属性:
|
||||
text: 待朗读的文本。
|
||||
voice: 音色标识符（模型决定可用音色）。
|
||||
response_format: 期望的音频格式（``"wav"``、``"mp3"``、``"opus"``）。
|
||||
extra: 厂商特定的额外参数（如音色克隆/设计参数）。
|
||||
"""
|
||||
|
||||
__slots__ = ("text", "voice", "response_format", "extra")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
text: str,
|
||||
voice: str = "default",
|
||||
response_format: str = "wav",
|
||||
extra: dict | None = None,
|
||||
) -> None:
|
||||
self.text = text
|
||||
self.voice = voice
|
||||
self.response_format = response_format
|
||||
self.extra = extra or {}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TTSResponse:
|
||||
"""成功合成的结果。"""
|
||||
|
||||
audio_data: bytes
|
||||
format: str
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class TTSProtocol(Protocol):
|
||||
"""TTS 适配器的协议。"""
|
||||
|
||||
async def synthesize(self, request: TTSRequest) -> TTSResponse:
|
||||
"""将文本合成为音频字节。
|
||||
|
||||
异常:
|
||||
AudioServiceError: API 调用失败或响应异常时抛出。
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
# ── 共享工厂基类 ──────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class AudioAdapterBase:
|
||||
"""ASR/TTS 适配器的基类，提供共享的工厂方法。
|
||||
|
||||
子类必须设置 ``_adapter_type``（如 ``"asr"``、``"tts"``），并在
|
||||
``__init__`` 中接受 ``base_url``、``api_key``、``model``、
|
||||
``request_timeout_seconds`` 参数。
|
||||
"""
|
||||
|
||||
_adapter_type: str = "audio"
|
||||
_default_timeout: int = 120
|
||||
|
||||
@classmethod
|
||||
def from_endpoint(
|
||||
cls,
|
||||
endpoint: ModelEndpointConfig,
|
||||
*,
|
||||
provider: ProviderConfig | None = None,
|
||||
timeout_seconds: int | None = None,
|
||||
) -> Any:
|
||||
"""从已解析的端点配置构建适配器。
|
||||
|
||||
凭证解析优先级：
|
||||
1. endpoint 直填的 ``base_url`` / ``api_key``
|
||||
2. endpoint 的 ``base_url_env`` / ``api_key_env`` 指向的 settings 字段
|
||||
3. provider 级凭证（调用方按需传入 ProviderConfig）
|
||||
"""
|
||||
base_url = _resolve_base_url(endpoint, provider)
|
||||
api_key = _resolve_api_key(endpoint, provider)
|
||||
if not base_url:
|
||||
raise AudioServiceError(
|
||||
status_code=500,
|
||||
code=f"audio_{cls._adapter_type}_endpoint_invalid",
|
||||
message=(
|
||||
f"{cls._adapter_type.upper()} endpoint '{endpoint.name}' "
|
||||
"has no resolved base_url."
|
||||
),
|
||||
)
|
||||
return cls(
|
||||
base_url=base_url,
|
||||
api_key=api_key,
|
||||
model=endpoint.model,
|
||||
request_timeout_seconds=timeout_seconds
|
||||
or endpoint.request_timeout_seconds
|
||||
or cls._default_timeout,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_model_config(
|
||||
cls,
|
||||
config: ModelConfig,
|
||||
*,
|
||||
provider: ProviderConfig | None = None,
|
||||
timeout_seconds: int | None = None,
|
||||
) -> Any:
|
||||
"""从逻辑模型配置构建适配器（使用第一个活跃端点）。"""
|
||||
for ep in config.endpoints:
|
||||
if ep.status == "active":
|
||||
return cls.from_endpoint(
|
||||
ep, provider=provider, timeout_seconds=timeout_seconds
|
||||
)
|
||||
raise AudioServiceError(
|
||||
status_code=500,
|
||||
code=f"audio_{cls._adapter_type}_no_endpoint",
|
||||
message=f"{cls._adapter_type.upper()} model '{config.name}' has no active endpoint.",
|
||||
)
|
||||
|
||||
|
||||
def _resolve_base_url(
|
||||
endpoint: ModelEndpointConfig,
|
||||
provider: ProviderConfig | None,
|
||||
) -> str:
|
||||
"""端点 base_url → endpoint.base_url_env 对应 settings → provider.base_url(_env)。"""
|
||||
if endpoint.base_url:
|
||||
return endpoint.base_url
|
||||
if endpoint.base_url_env:
|
||||
val = _settings_get(endpoint.base_url_env)
|
||||
if val:
|
||||
return val
|
||||
if provider is not None:
|
||||
if provider.base_url:
|
||||
return provider.base_url
|
||||
if provider.base_url_env:
|
||||
val = _settings_get(provider.base_url_env)
|
||||
if val:
|
||||
return val
|
||||
return ""
|
||||
|
||||
|
||||
def _resolve_api_key(
|
||||
endpoint: ModelEndpointConfig,
|
||||
provider: ProviderConfig | None,
|
||||
) -> str:
|
||||
"""端点 api_key → endpoint.api_key_env 对应 settings → provider 级凭证。"""
|
||||
if endpoint.api_key:
|
||||
return endpoint.api_key
|
||||
if endpoint.api_key_env:
|
||||
val = _settings_get(endpoint.api_key_env)
|
||||
if val:
|
||||
return val
|
||||
if provider is not None:
|
||||
# ProviderConfig 的 api_key_env / base_url_env 也走 settings 解析
|
||||
val = _settings_get(provider.api_key_env)
|
||||
if val:
|
||||
return val
|
||||
return ""
|
||||
|
||||
|
||||
def _settings_get(env_name: str) -> str | None:
|
||||
"""按大写 env 名查 settings 上的小写字段（与现有 _resolve_credentials 约定一致）。"""
|
||||
from app.core.config import settings
|
||||
|
||||
key = env_name.lower()
|
||||
if hasattr(settings, key):
|
||||
value = getattr(settings, key)
|
||||
return value if isinstance(value, str) else None
|
||||
return None
|
||||
@@ -1,325 +0,0 @@
|
||||
"""MiMo TTS adapter — ``POST /v1/chat/completions`` with ``audio`` parameter.
|
||||
|
||||
MiMo 的 TTS 不走独立音频端点，而是通过聊天补全接口扩展实现：
|
||||
|
||||
POST {base_url}/chat/completions
|
||||
Content-Type: application/json
|
||||
Authorization: Bearer {api_key}
|
||||
|
||||
{
|
||||
"model": "mimo-v2.5-tts",
|
||||
"messages": [
|
||||
{"role": "user", "content": "（可选）音色风格描述，voicedesign 必填"},
|
||||
{"role": "assistant", "content": "要朗读的文本"}
|
||||
],
|
||||
"audio": {
|
||||
"format": "wav",
|
||||
"voice": "mimo_default"
|
||||
}
|
||||
}
|
||||
|
||||
响应：:
|
||||
|
||||
{
|
||||
"choices": [{
|
||||
"message": {
|
||||
"audio": {
|
||||
"data": "<base64 编码的音频数据>",
|
||||
"id": "...",
|
||||
"expires_at": null,
|
||||
"transcript": null
|
||||
}
|
||||
}
|
||||
}]
|
||||
}
|
||||
|
||||
模型差异：
|
||||
|
||||
- ``mimo-v2.5-tts`` — ``voice`` 可选，默认 ``mimo_default``，仅支持预置音色
|
||||
- ``mimo-v2.5-tts-voiceclone`` — ``voice`` 必填，传入音频样本的 base64（mp3/wav）
|
||||
- ``mimo-v2.5-tts-voicedesign`` — 不支持 ``voice``，\n ``user`` 消息必填（风格描述），支持 ``optimize_text_preview``
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import time
|
||||
|
||||
import httpx
|
||||
|
||||
from app.ai_provider.audio.base import (
|
||||
AudioAdapterBase,
|
||||
AudioServiceError,
|
||||
TTSProtocol,
|
||||
TTSRequest,
|
||||
TTSResponse,
|
||||
)
|
||||
from app.core.logging import logger
|
||||
|
||||
# voicedesign 模型的默认风格描述（用户未提供时使用）
|
||||
_DEFAULT_VOICE_DESIGN_PROMPT = "自然、清晰的朗读风格，语气平和，语速适中，适合日常对话场景。"
|
||||
|
||||
# TTS 音频缓存：默认最多 256 条，TTL 1 小时
|
||||
_TTS_CACHE_MAX_SIZE = 256
|
||||
_TTS_CACHE_TTL_SECONDS = 3600
|
||||
|
||||
|
||||
class _TTSCacheEntry:
|
||||
__slots__ = ("data", "format", "stored_at")
|
||||
|
||||
def __init__(self, data: bytes, format: str) -> None:
|
||||
self.data = data
|
||||
self.format = format
|
||||
self.stored_at = time.monotonic()
|
||||
|
||||
|
||||
_tts_cache: dict[str, _TTSCacheEntry] = {}
|
||||
|
||||
|
||||
def _tts_cache_key(model: str, voice: str, fmt: str, text: str, extra_json: str) -> str:
|
||||
"""生成 TTS 缓存键：SHA256(model + voice + format + text + extra)。"""
|
||||
raw = f"{model}\0{voice}\0{fmt}\0{text}\0{extra_json}"
|
||||
return hashlib.sha256(raw.encode()).hexdigest()
|
||||
|
||||
|
||||
def _tts_cache_get(key: str) -> _TTSCacheEntry | None:
|
||||
entry = _tts_cache.get(key)
|
||||
if entry is None:
|
||||
return None
|
||||
if time.monotonic() - entry.stored_at > _TTS_CACHE_TTL_SECONDS:
|
||||
del _tts_cache[key]
|
||||
return None
|
||||
return entry
|
||||
|
||||
|
||||
def _tts_cache_set(key: str, entry: _TTSCacheEntry) -> None:
|
||||
if len(_tts_cache) >= _TTS_CACHE_MAX_SIZE:
|
||||
# 淘汰最旧的一半
|
||||
sorted_keys = sorted(_tts_cache.keys(), key=lambda k: _tts_cache[k].stored_at)
|
||||
for old_key in sorted_keys[: _TTS_CACHE_MAX_SIZE // 2]:
|
||||
del _tts_cache[old_key]
|
||||
_tts_cache[key] = entry
|
||||
|
||||
|
||||
class HunjingTTSOpenAI(AudioAdapterBase, TTSProtocol):
|
||||
"""MiMo TTS 适配器，通过 ``POST /chat/completions`` 携带 ``audio`` 参数实现。
|
||||
|
||||
调用兼容 OpenAI 格式的 chat completions 端点，在请求体中添加 ``audio`` 字段，
|
||||
然后从响应中提取并解码 base64 音频数据。
|
||||
"""
|
||||
|
||||
_adapter_type = "tts"
|
||||
_default_timeout = 120
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
model: str = "mimo-v2.5-tts",
|
||||
request_timeout_seconds: int = 120,
|
||||
) -> None:
|
||||
self._base_url = base_url.rstrip("/")
|
||||
self._api_key = api_key
|
||||
self._model = model
|
||||
self._timeout = request_timeout_seconds
|
||||
self._client: httpx.AsyncClient | None = None
|
||||
|
||||
# ── 客户端 ───────────────────────────────────────────────────────────
|
||||
|
||||
@property
|
||||
def _http(self) -> httpx.AsyncClient:
|
||||
if self._client is None or self._client.is_closed:
|
||||
self._client = httpx.AsyncClient(timeout=httpx.Timeout(self._timeout))
|
||||
return self._client
|
||||
|
||||
async def close(self) -> None:
|
||||
if self._client is not None and not self._client.is_closed:
|
||||
await self._client.aclose()
|
||||
|
||||
# ── TTSProtocol ──────────────────────────────────────────────────────
|
||||
|
||||
async def synthesize(self, request: TTSRequest) -> TTSResponse:
|
||||
"""通过 ``POST /chat/completions`` 将文本合成为音频。
|
||||
|
||||
构造 chat 消息和 ``audio`` 参数，匹配 MiMo 兼容 OpenAI 格式的 TTS，
|
||||
然后从响应中解码 base64 音频。
|
||||
|
||||
结果按 ``(model, voice, format, text)`` 在进程内缓存，避免重复 API 调用。
|
||||
"""
|
||||
# ── 缓存查找 ──────────────────────────────────────────────────────
|
||||
cache_key = _tts_cache_key(
|
||||
model=self._model,
|
||||
voice=request.voice,
|
||||
fmt=request.response_format,
|
||||
text=request.text,
|
||||
extra_json=str(sorted((request.extra or {}).items())),
|
||||
)
|
||||
cached = _tts_cache_get(cache_key)
|
||||
if cached is not None:
|
||||
logger.info("audio_tts.cache_hit", model=self._model, text_length=len(request.text))
|
||||
return TTSResponse(audio_data=cached.data, format=cached.format)
|
||||
|
||||
# ── API 调用 ──────────────────────────────────────────────────────
|
||||
url = f"{self._base_url}/chat/completions"
|
||||
|
||||
# 构造 messages
|
||||
messages = self._build_messages(request)
|
||||
|
||||
# 构造 audio 参数
|
||||
audio_params = self._build_audio_params(request)
|
||||
|
||||
payload: dict = {
|
||||
"model": self._model,
|
||||
"messages": messages,
|
||||
"audio": audio_params,
|
||||
"max_completion_tokens": 8192,
|
||||
}
|
||||
# 合并 extra（如 frequency_penalty, temperature 等）
|
||||
if request.extra:
|
||||
payload.update(request.extra)
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {self._api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
logger.info(
|
||||
"audio_tts.request",
|
||||
model=self._model,
|
||||
voice=request.voice,
|
||||
audio_format=request.response_format,
|
||||
text_length=len(request.text),
|
||||
url=url,
|
||||
)
|
||||
|
||||
try:
|
||||
response = await self._http.post(url, json=payload, headers=headers)
|
||||
except httpx.TimeoutException as exc:
|
||||
raise AudioServiceError(
|
||||
status_code=504,
|
||||
code="audio_tts_timeout",
|
||||
message="TTS request timed out calling chat completions.",
|
||||
) from exc
|
||||
except httpx.RequestError as exc:
|
||||
raise AudioServiceError(
|
||||
status_code=503,
|
||||
code="audio_tts_connection_error",
|
||||
message=f"TTS connection error: {exc}",
|
||||
) from exc
|
||||
|
||||
if response.status_code != 200:
|
||||
body = response.text[:500]
|
||||
raise AudioServiceError(
|
||||
status_code=response.status_code,
|
||||
code="audio_tts_api_error",
|
||||
message=f"TTS API returned {response.status_code}: {body}",
|
||||
)
|
||||
|
||||
# 解析响应，提取 base64 音频
|
||||
audio_data = self._extract_audio(response.json())
|
||||
|
||||
# ── 写入缓存 ──────────────────────────────────────────────────────
|
||||
_tts_cache_set(cache_key, _TTSCacheEntry(audio_data, request.response_format))
|
||||
|
||||
logger.info(
|
||||
"audio_tts.success",
|
||||
model=self._model,
|
||||
audio_bytes=len(audio_data),
|
||||
cache_size=len(_tts_cache),
|
||||
)
|
||||
return TTSResponse(audio_data=audio_data, format=request.response_format)
|
||||
|
||||
# ── 辅助方法 ──────────────────────────────────────────────────────────
|
||||
|
||||
def _build_messages(self, request: TTSRequest) -> list[dict]:
|
||||
"""构造 chat messages 数组。
|
||||
|
||||
- voicedesign 模型：user 消息必填（风格描述）
|
||||
- 其他模型：user 消息可选，assistant 消息存放待朗读文本
|
||||
"""
|
||||
is_voicedesign = "voicedesign" in self._model
|
||||
|
||||
if is_voicedesign:
|
||||
# voicedesign：user 消息为风格提示，assistant 消息为朗读文本
|
||||
style_prompt = (request.extra or {}).pop("style_prompt", _DEFAULT_VOICE_DESIGN_PROMPT)
|
||||
return [
|
||||
{"role": "user", "content": style_prompt},
|
||||
{"role": "assistant", "content": request.text},
|
||||
]
|
||||
# 基础 TTS / voiceclone：只传 assistant 消息
|
||||
return [
|
||||
{"role": "assistant", "content": request.text},
|
||||
]
|
||||
|
||||
def _build_audio_params(self, request: TTSRequest) -> dict:
|
||||
"""构造 ``audio`` 参数字段，按模型处理 voice/format/optimize_text_preview 等。"""
|
||||
params: dict = {
|
||||
"format": _resolve_api_format(request.response_format),
|
||||
}
|
||||
|
||||
is_voicedesign = "voicedesign" in self._model
|
||||
|
||||
if is_voicedesign:
|
||||
# voicedesign 不支持 voice 字段
|
||||
optimize = (request.extra or {}).pop("optimize_text_preview", None)
|
||||
if optimize is not None:
|
||||
params["optimize_text_preview"] = optimize
|
||||
else:
|
||||
# 基础 TTS / voiceclone：按需传 voice
|
||||
if request.voice and request.voice != "default":
|
||||
params["voice"] = request.voice
|
||||
else:
|
||||
params["voice"] = "mimo_default"
|
||||
|
||||
return params
|
||||
|
||||
@staticmethod
|
||||
def _extract_audio(body: dict) -> bytes:
|
||||
"""从 chat completion 响应中提取并解码 base64 音频数据。"""
|
||||
try:
|
||||
choices = body.get("choices", [])
|
||||
if not choices:
|
||||
raise AudioServiceError(
|
||||
status_code=500,
|
||||
code="audio_tts_no_choices",
|
||||
message="TTS response has no choices.",
|
||||
)
|
||||
message = choices[0].get("message", {})
|
||||
audio_info = message.get("audio")
|
||||
if not audio_info:
|
||||
raise AudioServiceError(
|
||||
status_code=500,
|
||||
code="audio_tts_no_audio_field",
|
||||
message="TTS response missing audio field in message.",
|
||||
)
|
||||
raw = audio_info.get("data")
|
||||
if not raw:
|
||||
raise AudioServiceError(
|
||||
status_code=500,
|
||||
code="audio_tts_empty_audio_data",
|
||||
message="TTS response has empty audio.data.",
|
||||
)
|
||||
return base64.b64decode(raw)
|
||||
except AudioServiceError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
raise AudioServiceError(
|
||||
status_code=500,
|
||||
code="audio_tts_response_parse_error",
|
||||
message=f"Failed to parse TTS response: {exc}",
|
||||
) from exc
|
||||
|
||||
|
||||
def _resolve_api_format(user_format: str) -> str:
|
||||
"""将用户层面的格式名映射为 API 期望的格式。
|
||||
|
||||
MiMo TTS API 接受：``wav``、``mp3``、``pcm``、``pcm16``。
|
||||
"""
|
||||
fmt = user_format.lower()
|
||||
if fmt in ("pcm", "pcm16"):
|
||||
return "pcm16"
|
||||
if fmt in ("wav", "mp3"):
|
||||
return fmt
|
||||
return "wav"
|
||||
@@ -1,231 +0,0 @@
|
||||
"""yuanqi 本地 ASR 适配器 — ``POST /v1/audio/transcriptions`` (multipart)。
|
||||
|
||||
与 MiMo 走 ``/chat/completions`` + ``input_audio`` 不同，元炁本地推理引擎
|
||||
（yuanqi）对外暴露的是 OpenAI Whisper 兼容的 multipart 端点::
|
||||
|
||||
POST {base_url}/v1/audio/transcriptions
|
||||
Content-Type: multipart/form-data
|
||||
Authorization: Bearer {api_key}
|
||||
|
||||
file: <audio bytes>
|
||||
language: zh (可选)
|
||||
|
||||
响应::
|
||||
|
||||
{"text": "转写后的文本"}
|
||||
|
||||
模型实际为 FunASR SenseVoiceSmall，转写文本已在 yuanqi 侧剥离了
|
||||
``<|zh|><|NEUTRAL|><|Speech|><|withitn|>`` 等内部标签，此处无需二次清洗。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import time
|
||||
|
||||
import httpx
|
||||
|
||||
from app.ai_provider.audio.base import (
|
||||
ASRProtocol,
|
||||
ASRRequest,
|
||||
ASRResponse,
|
||||
AudioAdapterBase,
|
||||
AudioServiceError,
|
||||
)
|
||||
from app.core.logging import logger
|
||||
|
||||
# multipart 使用的 MIME 映射（与 yuanqi _SUPPORTED_FORMATS 白名单对齐）
|
||||
_MIME_MAP = {
|
||||
"wav": "audio/wav",
|
||||
"mp3": "audio/mpeg",
|
||||
"m4a": "audio/mp4",
|
||||
"flac": "audio/flac",
|
||||
"ogg": "audio/ogg",
|
||||
"opus": "audio/opus",
|
||||
"aac": "audio/aac",
|
||||
"webm": "audio/webm",
|
||||
}
|
||||
|
||||
# yuanqi ASR 路径前缀（与 src/endpoints/asr.py 的 router 前缀一致）
|
||||
_ASR_PATH = "/v1/audio/transcriptions"
|
||||
|
||||
# ASR 缓存：默认最多 128 条，TTL 1 小时（与 MiMo ASR 适配器一致）
|
||||
_ASR_CACHE_MAX_SIZE = 128
|
||||
_ASR_CACHE_TTL_SECONDS = 3600
|
||||
|
||||
|
||||
class _ASRCacheEntry:
|
||||
__slots__ = ("text", "stored_at")
|
||||
|
||||
def __init__(self, text: str) -> None:
|
||||
self.text = text
|
||||
self.stored_at = time.monotonic()
|
||||
|
||||
|
||||
_asr_cache: dict[str, _ASRCacheEntry] = {}
|
||||
|
||||
|
||||
def _asr_cache_key(model: str, audio_data: bytes, language: str, fmt: str) -> str:
|
||||
"""生成 ASR 缓存键：SHA256(model + audio_bytes + language + format)。"""
|
||||
raw = f"{model}\0{language}\0{fmt}\0"
|
||||
raw_bytes = raw.encode() + audio_data
|
||||
return hashlib.sha256(raw_bytes).hexdigest()
|
||||
|
||||
|
||||
def _asr_cache_get(key: str) -> str | None:
|
||||
entry = _asr_cache.get(key)
|
||||
if entry is None:
|
||||
return None
|
||||
if time.monotonic() - entry.stored_at > _ASR_CACHE_TTL_SECONDS:
|
||||
del _asr_cache[key]
|
||||
return None
|
||||
return entry.text
|
||||
|
||||
|
||||
def _asr_cache_set(key: str, text: str) -> None:
|
||||
if len(_asr_cache) >= _ASR_CACHE_MAX_SIZE:
|
||||
sorted_keys = sorted(_asr_cache.keys(), key=lambda k: _asr_cache[k].stored_at)
|
||||
for old_key in sorted_keys[: _ASR_CACHE_MAX_SIZE // 2]:
|
||||
del _asr_cache[old_key]
|
||||
_asr_cache[key] = _ASRCacheEntry(text)
|
||||
|
||||
|
||||
class YuanqiASRWhisper(AudioAdapterBase, ASRProtocol):
|
||||
"""yuanqi 本地 ASR 适配器，走 OpenAI Whisper 兼容的 multipart 端点。
|
||||
|
||||
将音频字节以 ``multipart/form-data`` 的 ``file`` 字段上传到 yuanqi 的
|
||||
``/v1/audio/transcriptions``，从 JSON 响应中提取转写文本。
|
||||
"""
|
||||
|
||||
_adapter_type = "asr"
|
||||
_default_timeout = 300
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
model: str = "SenseVoiceSmall",
|
||||
request_timeout_seconds: int = 300,
|
||||
) -> None:
|
||||
self._base_url = base_url.rstrip("/")
|
||||
self._api_key = api_key
|
||||
self._model = model
|
||||
self._timeout = request_timeout_seconds
|
||||
self._client: httpx.AsyncClient | None = None
|
||||
|
||||
# ── 客户端 ───────────────────────────────────────────────────────────
|
||||
|
||||
@property
|
||||
def _http(self) -> httpx.AsyncClient:
|
||||
if self._client is None or self._client.is_closed:
|
||||
self._client = httpx.AsyncClient(timeout=httpx.Timeout(self._timeout))
|
||||
return self._client
|
||||
|
||||
async def close(self) -> None:
|
||||
if self._client is not None and not self._client.is_closed:
|
||||
await self._client.aclose()
|
||||
|
||||
# ── ASRProtocol ──────────────────────────────────────────────────────
|
||||
|
||||
async def transcribe(self, request: ASRRequest) -> ASRResponse:
|
||||
"""通过 ``POST /v1/audio/transcriptions`` 转写音频。"""
|
||||
# ── 缓存查找 ──────────────────────────────────────────────────
|
||||
cache_key = _asr_cache_key(
|
||||
self._model, request.audio_data, request.language or "auto", request.format
|
||||
)
|
||||
cached = _asr_cache_get(cache_key)
|
||||
if cached is not None:
|
||||
logger.info("audio_asr.cache_hit", model=self._model, text_length=len(cached))
|
||||
return ASRResponse(text=cached)
|
||||
|
||||
# ── 构造 multipart 请求 ───────────────────────────────────────
|
||||
mime = _MIME_MAP.get(request.format, "audio/wav")
|
||||
filename = f"audio.{request.format}"
|
||||
files = {"file": (filename, request.audio_data, mime)}
|
||||
data: dict[str, str] = {}
|
||||
if request.language:
|
||||
data["language"] = request.language
|
||||
|
||||
# 本地 yuanqi 引擎通常无需鉴权；仅当配置了 api_key 才带 Authorization 头
|
||||
headers: dict[str, str] = {}
|
||||
if self._api_key:
|
||||
headers["Authorization"] = f"Bearer {self._api_key}"
|
||||
|
||||
url = f"{self._base_url}{_ASR_PATH}"
|
||||
|
||||
logger.info(
|
||||
"audio_asr.request",
|
||||
model=self._model,
|
||||
format=request.format,
|
||||
language=request.language,
|
||||
url=url,
|
||||
size_bytes=len(request.audio_data),
|
||||
)
|
||||
|
||||
# ── API 调用 ────────────────────────────────────────────────
|
||||
try:
|
||||
response = await self._http.post(url, files=files, data=data, headers=headers)
|
||||
except httpx.TimeoutException as exc:
|
||||
raise AudioServiceError(
|
||||
status_code=504,
|
||||
code="audio_asr_timeout",
|
||||
message="ASR request timed out calling yuanqi transcriptions.",
|
||||
) from exc
|
||||
except httpx.RequestError as exc:
|
||||
raise AudioServiceError(
|
||||
status_code=503,
|
||||
code="audio_asr_connection_error",
|
||||
message=f"ASR connection error: {exc}",
|
||||
) from exc
|
||||
|
||||
if response.status_code != 200:
|
||||
body = response.text[:500]
|
||||
raise AudioServiceError(
|
||||
status_code=response.status_code,
|
||||
code="audio_asr_api_error",
|
||||
message=f"ASR API returned {response.status_code}: {body}",
|
||||
)
|
||||
|
||||
# 解析响应，提取转写文本
|
||||
text = self._extract_text(response.json())
|
||||
if not text:
|
||||
raise AudioServiceError(
|
||||
status_code=500,
|
||||
code="audio_asr_empty_response",
|
||||
message="ASR API returned empty transcription.",
|
||||
)
|
||||
|
||||
# ── 写入缓存 ────────────────────────────────────────────────
|
||||
_asr_cache_set(cache_key, text)
|
||||
|
||||
logger.info(
|
||||
"audio_asr.success",
|
||||
model=self._model,
|
||||
text_length=len(text),
|
||||
cache_size=len(_asr_cache),
|
||||
)
|
||||
return ASRResponse(text=text)
|
||||
|
||||
# ── 辅助方法 ──────────────────────────────────────────────────────────
|
||||
|
||||
@staticmethod
|
||||
def _extract_text(body: dict) -> str:
|
||||
"""从 Whisper 兼容响应中提取转写文本。"""
|
||||
try:
|
||||
text = body.get("text")
|
||||
if text is None:
|
||||
raise AudioServiceError(
|
||||
status_code=500,
|
||||
code="audio_asr_no_text",
|
||||
message="ASR response has no 'text' field.",
|
||||
)
|
||||
return text or ""
|
||||
except AudioServiceError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
raise AudioServiceError(
|
||||
status_code=500,
|
||||
code="audio_asr_response_parse_error",
|
||||
message=f"Failed to parse ASR response: {exc}",
|
||||
) from exc
|
||||
@@ -1,31 +0,0 @@
|
||||
from app.ai_provider.config.loader import (
|
||||
get_models_config_path,
|
||||
get_pricing_config_path,
|
||||
get_providers_config_path,
|
||||
load_models_config,
|
||||
load_pricing_config,
|
||||
load_providers_config,
|
||||
)
|
||||
from app.ai_provider.config.models import (
|
||||
ModelConfig,
|
||||
ModelEndpointConfig,
|
||||
ParameterMapping,
|
||||
ProviderConfig,
|
||||
derive_capabilities,
|
||||
)
|
||||
from app.ai_provider.core.billing import PricingInfo
|
||||
|
||||
__all__ = [
|
||||
"ModelConfig",
|
||||
"ModelEndpointConfig",
|
||||
"PricingInfo",
|
||||
"ProviderConfig",
|
||||
"ParameterMapping",
|
||||
"derive_capabilities",
|
||||
"load_models_config",
|
||||
"load_providers_config",
|
||||
"load_pricing_config",
|
||||
"get_models_config_path",
|
||||
"get_providers_config_path",
|
||||
"get_pricing_config_path",
|
||||
]
|
||||
@@ -1,80 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import yaml
|
||||
from pydantic import BaseModel
|
||||
|
||||
from app.ai_provider.config.models import ModelConfig, ProviderConfig
|
||||
from app.core.config import settings
|
||||
|
||||
|
||||
class _ModelsFile(BaseModel):
|
||||
version: int = 1
|
||||
models: list[ModelConfig]
|
||||
|
||||
|
||||
class _ProvidersFile(BaseModel):
|
||||
version: int = 1
|
||||
providers: dict[str, ProviderConfig]
|
||||
|
||||
|
||||
class _PricingFile(BaseModel):
|
||||
"""pricing.yaml 顶层结构。"""
|
||||
|
||||
version: int = 1
|
||||
pricing: dict[str, dict] = {}
|
||||
|
||||
|
||||
def get_models_config_path() -> Path:
|
||||
"""返回已解析的 models.yaml 路径（project_root/config/models.yaml）。"""
|
||||
return settings.project_root / "config" / "models.yaml"
|
||||
|
||||
|
||||
def get_providers_config_path() -> Path:
|
||||
"""返回已解析的 providers.yaml 路径（project_root/config/providers.yaml）。"""
|
||||
return settings.project_root / "config" / "providers.yaml"
|
||||
|
||||
|
||||
def load_models_config(path: Path | None = None) -> list[ModelConfig]:
|
||||
"""从 YAML 文件加载模型配置。如果未指定路径，使用配置中的默认路径。"""
|
||||
if path is None:
|
||||
path = get_models_config_path()
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"Models config not found: {path}")
|
||||
raw: dict[str, Any] = yaml.safe_load(path.read_text(encoding="utf-8")) or {}
|
||||
parsed = _ModelsFile.model_validate(raw)
|
||||
return parsed.models
|
||||
|
||||
|
||||
def load_providers_config(path: Path | None = None) -> dict[str, ProviderConfig]:
|
||||
"""从 YAML 文件加载提供商配置。如果未指定路径，使用配置中的默认路径。"""
|
||||
if path is None:
|
||||
path = get_providers_config_path()
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"Providers config not found: {path}")
|
||||
raw: dict[str, Any] = yaml.safe_load(path.read_text(encoding="utf-8")) or {}
|
||||
# 验证前将字典键注入为 provider 的 name 字段
|
||||
providers_raw = raw.get("providers", {})
|
||||
for name, cfg in providers_raw.items():
|
||||
if isinstance(cfg, dict) and "name" not in cfg:
|
||||
cfg["name"] = name
|
||||
parsed = _ProvidersFile.model_validate(raw)
|
||||
return parsed.providers
|
||||
|
||||
|
||||
def get_pricing_config_path() -> Path:
|
||||
"""返回 pricing.yaml 路径（project_root/config/pricing.yaml）。"""
|
||||
return settings.project_root / "config" / "pricing.yaml"
|
||||
|
||||
|
||||
def load_pricing_config(path: Path | None = None) -> dict[str, dict]:
|
||||
"""从 YAML 文件加载定价配置。"""
|
||||
if path is None:
|
||||
path = get_pricing_config_path()
|
||||
if not path.exists():
|
||||
return {}
|
||||
raw: dict[str, Any] = yaml.safe_load(path.read_text(encoding="utf-8")) or {}
|
||||
parsed = _PricingFile.model_validate(raw)
|
||||
return parsed.pricing
|
||||
@@ -1,158 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from enum import StrEnum
|
||||
from typing import Any
|
||||
|
||||
from pydantic import ( # noqa: TC002 — Pydantic 模型
|
||||
BaseModel,
|
||||
Field,
|
||||
field_validator,
|
||||
model_validator,
|
||||
)
|
||||
|
||||
|
||||
class SupportedFormats(StrEnum):
|
||||
OPENAI = "openai"
|
||||
ANTHROPIC = "anthropic"
|
||||
|
||||
|
||||
_BASE_CAPABILITIES = {
|
||||
("openai", "gpt"): {"chat", "tool_use", "structured_output"},
|
||||
("openai", "gpt-4"): {"chat", "tool_use", "structured_output", "vision", "reasoning"},
|
||||
("anthropic", "claude"): {
|
||||
"chat",
|
||||
"tool_use",
|
||||
"structured_output",
|
||||
"vision",
|
||||
"reasoning",
|
||||
"long_context",
|
||||
},
|
||||
("deepseek", "deepseek"): {"chat", "tool_use", "structured_output", "reasoning"},
|
||||
("deepseek", "deepseek-v4"): {
|
||||
"chat",
|
||||
"tool_use",
|
||||
"structured_output",
|
||||
"reasoning",
|
||||
"long_context",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def derive_capabilities(provider: str, model: str) -> set[str]:
|
||||
"""根据 provider + model 前缀自动推导能力。
|
||||
|
||||
调用者可以在调用此函数后，在 ModelConfig 上覆盖个别标记。
|
||||
"""
|
||||
for (prov_prefix, model_prefix), caps in _BASE_CAPABILITIES.items():
|
||||
if provider.startswith(prov_prefix) and model.startswith(model_prefix):
|
||||
return set(caps)
|
||||
return {"chat"}
|
||||
|
||||
|
||||
class ModelEndpointConfig(BaseModel):
|
||||
name: str
|
||||
model: str
|
||||
base_url: str | None = None
|
||||
base_url_env: str | None = None
|
||||
api_key: str | None = None
|
||||
api_key_env: str | None = None
|
||||
weight: int = 100
|
||||
priority: int = 100
|
||||
rpm_limit: int | None = None
|
||||
tpm_limit: int | None = None
|
||||
max_concurrency: int | None = None
|
||||
request_timeout_seconds: int | None = None
|
||||
status: str = "active"
|
||||
extra_params: dict[str, Any] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class ModelConfig(BaseModel):
|
||||
name: str
|
||||
display_name: str | None = None
|
||||
provider: str
|
||||
model: str
|
||||
format: SupportedFormats # "openai" | "anthropic"
|
||||
chat_format: str | None = None
|
||||
tags: list[str] = Field(default_factory=list)
|
||||
supports_tools: bool | None = None
|
||||
supports_vision: bool | None = None
|
||||
supports_reasoning_effort: bool | None = None
|
||||
supports_thinking: bool | None = None
|
||||
supports_streaming: bool | None = True
|
||||
supports_structured_output: bool | None = None
|
||||
context_window: int = 32000
|
||||
response_reserve: int = 4096
|
||||
max_output_tokens: int | None = None
|
||||
max_thinking_tokens: int | None = None
|
||||
reasoning_effort: str | None = None
|
||||
timeout_seconds: int | None = 60
|
||||
dimensions: int | None = None # 向量维度（embedding 模型专用）
|
||||
routing_strategy: str = "weighted_round_robin"
|
||||
endpoint_pool_enabled: bool = False
|
||||
weight: int = 100
|
||||
endpoints: list[ModelEndpointConfig] = Field(default_factory=list)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _validate_embedding_dimensions(self) -> ModelConfig:
|
||||
"""embedding 模型必须声明 dimensions。"""
|
||||
if "embedding" in self.tags and self.dimensions is None:
|
||||
raise ValueError(
|
||||
f"Embedding 模型 '{self.name}' 必须设置 dimensions（如 dimensions: 1024）"
|
||||
)
|
||||
return self
|
||||
|
||||
def resolved_capabilities(self) -> dict[str, bool]:
|
||||
base = derive_capabilities(self.provider, self.model)
|
||||
return {
|
||||
"chat": True,
|
||||
"tool_use": self.supports_tools
|
||||
if self.supports_tools is not None
|
||||
else ("tool_use" in base),
|
||||
"vision": self.supports_vision
|
||||
if self.supports_vision is not None
|
||||
else ("vision" in base),
|
||||
"reasoning": self.supports_reasoning_effort
|
||||
if self.supports_reasoning_effort is not None
|
||||
else ("reasoning" in base),
|
||||
"thinking": self.supports_thinking
|
||||
if self.supports_thinking is not None
|
||||
else ("reasoning" in base),
|
||||
"streaming": self.supports_streaming if self.supports_streaming is not None else True,
|
||||
"structured_output": self.supports_structured_output
|
||||
if self.supports_structured_output is not None
|
||||
else ("structured_output" in base),
|
||||
}
|
||||
|
||||
|
||||
class ParameterMapping(BaseModel):
|
||||
provider_param: str
|
||||
min: float | None = None
|
||||
max: float | None = None
|
||||
default: Any = None
|
||||
transform: str | None = None # "multiply_1000", "int", "str"
|
||||
alias_for: str | None = None # map from another param name
|
||||
handler: str | None = None # named handler for complex param transform
|
||||
|
||||
|
||||
class ProviderConfig(BaseModel):
|
||||
name: str
|
||||
format: SupportedFormats
|
||||
base_url_env: str
|
||||
api_key_env: str
|
||||
base_url: str | None = None
|
||||
timeout_seconds: int = 60
|
||||
max_retries: int = 1
|
||||
parameter_mappings: dict[str, ParameterMapping] = Field(default_factory=dict)
|
||||
parameter_handlers: list[str] = Field(
|
||||
default_factory=list
|
||||
) # ordered global handlers per provider
|
||||
extra_headers: dict[str, str] = Field(default_factory=dict)
|
||||
supports_streaming: bool = True
|
||||
|
||||
@field_validator("format")
|
||||
@classmethod
|
||||
def _validate_format(cls, v: str) -> str:
|
||||
allowed = {e.value for e in SupportedFormats}
|
||||
if v not in allowed:
|
||||
raise ValueError(f"format must be one of {allowed}, got '{v}'")
|
||||
return v
|
||||
@@ -1,17 +0,0 @@
|
||||
from app.ai_provider.core.billing import PricingFactory, PricingInfo, PricingRegistry
|
||||
from app.ai_provider.core.registry import (
|
||||
ModelRegistry,
|
||||
ProviderRegistry,
|
||||
RegistryFactory,
|
||||
RegistryHolder,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"ModelRegistry",
|
||||
"ProviderRegistry",
|
||||
"RegistryFactory",
|
||||
"RegistryHolder",
|
||||
"PricingInfo",
|
||||
"PricingRegistry",
|
||||
"PricingFactory",
|
||||
]
|
||||
@@ -1,253 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import date, datetime
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from app.ai_provider.types.chat import TokenUsage
|
||||
|
||||
|
||||
@dataclass
|
||||
class CostBreakdown:
|
||||
input_cost: float = 0.0
|
||||
output_cost: float = 0.0
|
||||
total: float = 0.0
|
||||
currency: str = "CNY"
|
||||
|
||||
|
||||
class CostCalculator:
|
||||
"""根据价格信息和 token 数量计算成本。
|
||||
|
||||
保持厂商原始单位 —— 不做跨币种换算。
|
||||
|
||||
- 普通 input token 按 ``input_per_1m`` 计费。
|
||||
- 缓存命中的 prompt token 优先按 ``input_cache_read_per_1m`` 计费
|
||||
（若有配置），剩余按普通 input 费率。
|
||||
- 若未配置 cache 费率，所有 prompt token 统一按 input_per_1m 计费。
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def calculate(
|
||||
pricing: PricingInfo | None,
|
||||
prompt_tokens: int,
|
||||
completion_tokens: int,
|
||||
prompt_tokens_cached: int = 0,
|
||||
) -> CostBreakdown:
|
||||
if pricing is None:
|
||||
return CostBreakdown(total=0.0)
|
||||
|
||||
# 缓存 token：优先走缓存费率
|
||||
cached = min(prompt_tokens_cached, prompt_tokens)
|
||||
uncached = prompt_tokens - cached
|
||||
|
||||
input_cost = 0.0
|
||||
if uncached > 0 and pricing.input_per_1m is not None:
|
||||
input_cost += (uncached / 1_000_000) * pricing.input_per_1m
|
||||
if cached > 0 and pricing.input_cache_read_per_1m is not None:
|
||||
input_cost += (cached / 1_000_000) * pricing.input_cache_read_per_1m
|
||||
elif cached > 0 and pricing.input_per_1m is not None:
|
||||
# 无缓存费率时统一按普通 input 计
|
||||
input_cost += (cached / 1_000_000) * pricing.input_per_1m
|
||||
|
||||
output_cost = 0.0
|
||||
if pricing.output_per_1m is not None and completion_tokens > 0:
|
||||
output_cost = (completion_tokens / 1_000_000) * pricing.output_per_1m
|
||||
|
||||
return CostBreakdown(
|
||||
input_cost=round(input_cost, 8),
|
||||
output_cost=round(output_cost, 8),
|
||||
total=round(input_cost + output_cost, 8),
|
||||
currency=pricing.currency,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class UsageRecord:
|
||||
model_name: str
|
||||
provider: str
|
||||
prompt_tokens: int = 0
|
||||
completion_tokens: int = 0
|
||||
total_tokens: int = 0
|
||||
currency: str = "CNY"
|
||||
cost: float = 0.0
|
||||
timestamp: datetime = field(default_factory=datetime.now)
|
||||
endpoint_name: str | None = None
|
||||
status: str = "success"
|
||||
error_message: str | None = None
|
||||
|
||||
|
||||
class BillingTracker:
|
||||
"""记录用量并按模型计算成本。"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._records: list[UsageRecord] = []
|
||||
|
||||
def record(
|
||||
self,
|
||||
model_name: str,
|
||||
provider: str,
|
||||
usage: TokenUsage | None,
|
||||
pricing: PricingInfo | None = None,
|
||||
endpoint_name: str | None = None,
|
||||
status: str = "success",
|
||||
error_message: str | None = None,
|
||||
) -> UsageRecord:
|
||||
prompt_tokens = usage.prompt_tokens if usage else 0
|
||||
completion_tokens = usage.completion_tokens if usage else 0
|
||||
total = usage.total_tokens if usage else 0
|
||||
cached = usage.prompt_tokens_cached if usage else 0
|
||||
|
||||
cost_info = CostCalculator.calculate(pricing, prompt_tokens, completion_tokens, cached)
|
||||
|
||||
record = UsageRecord(
|
||||
model_name=model_name,
|
||||
provider=provider,
|
||||
prompt_tokens=prompt_tokens,
|
||||
completion_tokens=completion_tokens,
|
||||
total_tokens=total,
|
||||
currency=cost_info.currency,
|
||||
cost=cost_info.total,
|
||||
endpoint_name=endpoint_name,
|
||||
status=status,
|
||||
error_message=error_message,
|
||||
)
|
||||
self._records.append(record)
|
||||
return record
|
||||
|
||||
def get_records(
|
||||
self,
|
||||
since: datetime | None = None,
|
||||
model_name: str | None = None,
|
||||
) -> list[UsageRecord]:
|
||||
records = self._records
|
||||
if since:
|
||||
records = [r for r in records if r.timestamp >= since]
|
||||
if model_name:
|
||||
records = [r for r in records if r.model_name == model_name]
|
||||
return records
|
||||
|
||||
def total_cost(self, currency: str | None = None) -> float:
|
||||
if currency:
|
||||
return sum(r.cost for r in self._records if r.currency == currency)
|
||||
return sum(r.cost for r in self._records)
|
||||
|
||||
|
||||
class BudgetController:
|
||||
"""强制执行每日和每月预算限制。"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
daily_budget_cny: float = 0.0,
|
||||
monthly_budget_cny: float = 0.0,
|
||||
) -> None:
|
||||
self.daily_budget_cny = daily_budget_cny
|
||||
self.monthly_budget_cny = monthly_budget_cny
|
||||
self._daily_cost_cny: float = 0.0
|
||||
self._monthly_cost_cny: float = 0.0
|
||||
self._current_date: date = date.today()
|
||||
|
||||
def add_cost_cny(self, cost: float) -> None:
|
||||
self._daily_cost_cny += cost
|
||||
self._monthly_cost_cny += cost
|
||||
|
||||
def allow_request(self) -> bool:
|
||||
return not (
|
||||
(self.daily_budget_cny > 0 and self._daily_cost_cny >= self.daily_budget_cny)
|
||||
or (self.monthly_budget_cny > 0 and self._monthly_cost_cny >= self.monthly_budget_cny)
|
||||
)
|
||||
|
||||
def reset_daily(self) -> None:
|
||||
self._daily_cost_cny = 0.0
|
||||
self._current_date = date.today()
|
||||
|
||||
def reset_monthly(self) -> None:
|
||||
self._monthly_cost_cny = 0.0
|
||||
|
||||
|
||||
class BillingFactory:
|
||||
"""计费组件工厂。"""
|
||||
|
||||
@staticmethod
|
||||
def create_tracker() -> BillingTracker:
|
||||
return BillingTracker()
|
||||
|
||||
@staticmethod
|
||||
def create_budget_controller(
|
||||
daily_budget_cny: float = 0.0,
|
||||
monthly_budget_cny: float = 0.0,
|
||||
) -> BudgetController:
|
||||
return BudgetController(
|
||||
daily_budget_cny=daily_budget_cny,
|
||||
monthly_budget_cny=monthly_budget_cny,
|
||||
)
|
||||
|
||||
|
||||
# ── 定价信息 ──────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class PricingInfo(BaseModel):
|
||||
"""模型定价信息。保持厂商原始单位，不转换不混合。
|
||||
|
||||
各厂商按各自币种/单位记录，展示时原生显示。
|
||||
"""
|
||||
|
||||
currency: str = "CNY"
|
||||
input_per_1m: float | None = None
|
||||
output_per_1m: float | None = None
|
||||
input_cache_read_per_1m: float | None = None
|
||||
|
||||
|
||||
class PricingRegistry:
|
||||
"""模型定价注册表。以模型逻辑名为键查找定价。
|
||||
|
||||
与 ModelRegistry 解耦 —— 改价不影响模型注册，无定价的模型
|
||||
也可以正常使用（不计费）。
|
||||
"""
|
||||
|
||||
def __init__(self, pricing_map: dict[str, PricingInfo] | None = None) -> None:
|
||||
self._pricing: dict[str, PricingInfo] = dict(pricing_map or {})
|
||||
|
||||
def get(self, model_name: str) -> PricingInfo | None:
|
||||
"""获取指定模型的定价信息。无定价时返回 None。"""
|
||||
return self._pricing.get(model_name)
|
||||
|
||||
def list_all(self) -> dict[str, PricingInfo]:
|
||||
"""返回全部定价数据的副本。"""
|
||||
return dict(self._pricing)
|
||||
|
||||
|
||||
class PricingFactory:
|
||||
"""定价注册表工厂。"""
|
||||
|
||||
@staticmethod
|
||||
def from_dict(data: dict[str, dict]) -> PricingRegistry:
|
||||
"""从 YAML 解析后的 dict 构建 PricingRegistry。
|
||||
|
||||
YAML 按厂商嵌套组织（自文档化），工厂将其摊平为按模型逻辑名索引：
|
||||
|
||||
```yaml
|
||||
pricing:
|
||||
deepseek:
|
||||
deepseek-v4-flash:
|
||||
currency: CNY
|
||||
input_per_1m: 1.0
|
||||
anthropic:
|
||||
claude-sonnet:
|
||||
currency: USD
|
||||
input_per_1m: 3.0
|
||||
```
|
||||
"""
|
||||
pricing_map: dict[str, PricingInfo] = {}
|
||||
for _provider_name, models_data in data.items():
|
||||
if not isinstance(models_data, dict):
|
||||
continue
|
||||
for model_name, raw in models_data.items():
|
||||
if raw is None:
|
||||
continue
|
||||
# 跳过所有字段均为 null 的条目（如 mimo-v2.5-asr）
|
||||
values = [v for v in raw.values() if not isinstance(v, bool)]
|
||||
if all(v is None for v in values):
|
||||
continue
|
||||
pricing_map[model_name] = PricingInfo(**raw)
|
||||
return PricingRegistry(pricing_map)
|
||||
@@ -1,164 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import random
|
||||
|
||||
from app.ai_provider.config.models import ModelConfig, ProviderConfig
|
||||
from app.ai_provider.core.billing import PricingFactory, PricingRegistry
|
||||
from app.ai_provider.types.errors import AIProviderError
|
||||
|
||||
|
||||
class ModelRegistry:
|
||||
"""内存模型索引。支持逻辑名和别名查找。"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
models: list[ModelConfig],
|
||||
aliases: dict[str, str] | None = None,
|
||||
) -> None:
|
||||
self._models: dict[str, ModelConfig] = {}
|
||||
self._aliases: dict[str, str] = dict(aliases or {})
|
||||
self._default: str | None = None
|
||||
self._fallback: str | None = None
|
||||
self._capability_index: dict[str, list[ModelConfig]] = {}
|
||||
|
||||
for m in models:
|
||||
self._models[m.name] = m
|
||||
if "default" in m.tags and self._default is None:
|
||||
self._default = m.name
|
||||
if "fallback" in m.tags and self._fallback is None:
|
||||
self._fallback = m.name
|
||||
caps = m.resolved_capabilities()
|
||||
for cap, supported in caps.items():
|
||||
if supported:
|
||||
self._capability_index.setdefault(cap, []).append(m)
|
||||
# 索引 tags（如 tts / asr / voice_clone）——这些不在
|
||||
# resolved_capabilities() 中，但 AudioService 等模块依赖它们查找模型。
|
||||
for tag in m.tags:
|
||||
self._capability_index.setdefault(tag, []).append(m)
|
||||
|
||||
if self._default is None and models:
|
||||
self._default = models[0].name
|
||||
|
||||
def get(self, name: str) -> ModelConfig:
|
||||
model = self._models.get(name)
|
||||
if model is None:
|
||||
raise AIProviderError(
|
||||
message=f"Model '{name}' not found in registry",
|
||||
code="MODEL_NOT_FOUND",
|
||||
status_code=404,
|
||||
)
|
||||
return model
|
||||
|
||||
def resolve(self, name_or_alias: str | None = None) -> ModelConfig:
|
||||
"""解析逻辑名或别名为 ModelConfig。
|
||||
|
||||
无参数时返回默认模型。
|
||||
"""
|
||||
if name_or_alias is None:
|
||||
return self._models[self.get_default_name()]
|
||||
if name_or_alias in self._models:
|
||||
return self._models[name_or_alias]
|
||||
if name_or_alias in self._aliases:
|
||||
logical = self._aliases[name_or_alias]
|
||||
return self._models[logical]
|
||||
raise AIProviderError(
|
||||
message=f"Model or alias '{name_or_alias}' not found",
|
||||
code="MODEL_NOT_FOUND",
|
||||
status_code=404,
|
||||
)
|
||||
|
||||
def resolve_name(self, name_or_alias: str) -> str:
|
||||
"""解析为规范逻辑名（不返回完整配置）。"""
|
||||
if name_or_alias in self._models:
|
||||
return name_or_alias
|
||||
if name_or_alias in self._aliases:
|
||||
return self._aliases[name_or_alias]
|
||||
raise AIProviderError(
|
||||
message=f"Model or alias '{name_or_alias}' not found",
|
||||
code="MODEL_NOT_FOUND",
|
||||
status_code=404,
|
||||
)
|
||||
|
||||
def get_default_name(self) -> str:
|
||||
if self._default is None:
|
||||
raise AIProviderError(
|
||||
message="No default model configured",
|
||||
code="NO_DEFAULT_MODEL",
|
||||
status_code=500,
|
||||
)
|
||||
return self._default
|
||||
|
||||
def get_fallback_name(self) -> str | None:
|
||||
return self._fallback
|
||||
|
||||
def contains(self, name: str) -> bool:
|
||||
return name in self._models
|
||||
|
||||
def list_all(self) -> list[ModelConfig]:
|
||||
return list(self._models.values())
|
||||
|
||||
def list_by_capability(self, capability: str) -> list[ModelConfig]:
|
||||
return list(self._capability_index.get(capability, []))
|
||||
|
||||
def resolve_by_capability(self, capability: str) -> str:
|
||||
candidates = self._capability_index.get(capability)
|
||||
if not candidates:
|
||||
raise AIProviderError(
|
||||
message=f"No active models with capability '{capability}'",
|
||||
code="CAPABILITY_NOT_FOUND",
|
||||
status_code=404,
|
||||
)
|
||||
weights = [max(c.weight, 1) for c in candidates]
|
||||
chosen = random.choices(candidates, weights=weights, k=1)[0]
|
||||
return chosen.name
|
||||
|
||||
|
||||
class ProviderRegistry:
|
||||
"""内存提供商配置索引。"""
|
||||
|
||||
def __init__(self, providers: dict[str, ProviderConfig]) -> None:
|
||||
self._providers = dict(providers)
|
||||
|
||||
def get(self, name: str) -> ProviderConfig:
|
||||
prov = self._providers.get(name)
|
||||
if prov is None:
|
||||
raise AIProviderError(
|
||||
message=f"Provider '{name}' not configured",
|
||||
code="PROVIDER_NOT_FOUND",
|
||||
status_code=404,
|
||||
)
|
||||
return prov
|
||||
|
||||
def list_all(self) -> dict[str, ProviderConfig]:
|
||||
return dict(self._providers)
|
||||
|
||||
|
||||
class RegistryFactory:
|
||||
"""工厂：从配置列表创建注册表。"""
|
||||
|
||||
@staticmethod
|
||||
def create(
|
||||
models: list[ModelConfig],
|
||||
providers: dict[str, ProviderConfig],
|
||||
aliases: dict[str, str] | None = None,
|
||||
pricing_data: dict[str, dict] | None = None,
|
||||
) -> RegistryHolder:
|
||||
return RegistryHolder(
|
||||
model_registry=ModelRegistry(models, aliases=aliases),
|
||||
provider_registry=ProviderRegistry(providers),
|
||||
pricing_registry=PricingFactory.from_dict(pricing_data or {}),
|
||||
)
|
||||
|
||||
|
||||
class RegistryHolder:
|
||||
"""同时持有所有注册表。"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_registry: ModelRegistry,
|
||||
provider_registry: ProviderRegistry,
|
||||
pricing_registry: PricingRegistry | None = None,
|
||||
) -> None:
|
||||
self.model_registry = model_registry
|
||||
self.provider_registry = provider_registry
|
||||
self.pricing_registry = pricing_registry or PricingRegistry()
|
||||
@@ -1,133 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class CircuitBreakerState(Enum):
|
||||
CLOSED = "closed"
|
||||
OPEN = "open"
|
||||
HALF_OPEN = "half_open"
|
||||
|
||||
|
||||
class CircuitBreaker:
|
||||
"""标准断路器，冷却后自动半开。"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
failure_threshold: int = 5,
|
||||
cooldown_seconds: float = 30.0,
|
||||
) -> None:
|
||||
self.name = name
|
||||
self.failure_threshold = failure_threshold
|
||||
self.cooldown_seconds = cooldown_seconds
|
||||
self._state = CircuitBreakerState.CLOSED
|
||||
self._failure_count = 0
|
||||
self._last_failure_time: float = 0.0
|
||||
|
||||
@property
|
||||
def state(self) -> CircuitBreakerState:
|
||||
if (
|
||||
self._state == CircuitBreakerState.OPEN
|
||||
and time.monotonic() - self._last_failure_time >= self.cooldown_seconds
|
||||
):
|
||||
self._state = CircuitBreakerState.HALF_OPEN
|
||||
return self._state
|
||||
|
||||
def record_success(self) -> None:
|
||||
self._failure_count = 0
|
||||
self._state = CircuitBreakerState.CLOSED
|
||||
|
||||
def record_failure(self) -> None:
|
||||
self._failure_count += 1
|
||||
self._last_failure_time = time.monotonic()
|
||||
if self._failure_count >= self.failure_threshold:
|
||||
self._state = CircuitBreakerState.OPEN
|
||||
|
||||
def allow_request(self) -> bool:
|
||||
return self.state != CircuitBreakerState.OPEN
|
||||
|
||||
|
||||
class RetryStrategy:
|
||||
"""指数退避重试。"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
max_retries: int = 3,
|
||||
base_delay: float = 1.0,
|
||||
max_delay: float = 60.0,
|
||||
multiplier: float = 2.0,
|
||||
) -> None:
|
||||
self.max_retries = max_retries
|
||||
self.base_delay = base_delay
|
||||
self.max_delay = max_delay
|
||||
self.multiplier = multiplier
|
||||
|
||||
def should_retry(self, attempt_index: int) -> bool:
|
||||
return attempt_index < self.max_retries
|
||||
|
||||
def get_delay(self, attempt_index: int) -> float:
|
||||
delay = self.base_delay * (self.multiplier**attempt_index)
|
||||
return min(delay, self.max_delay)
|
||||
|
||||
|
||||
class DowngradeAction:
|
||||
"""降级链中的一个步骤。"""
|
||||
|
||||
def __init__(self, type: str, target: str) -> None: # noqa: A002
|
||||
self.type = type # "fallback_endpoint", "fallback_model", "cache", "degrade_quality"
|
||||
self.target = target
|
||||
|
||||
|
||||
class DowngradeChain:
|
||||
"""当模型/端点失败时的降级动作链。"""
|
||||
|
||||
def __init__(self, actions: list[DowngradeAction]) -> None:
|
||||
self._actions = actions
|
||||
self._index = 0
|
||||
|
||||
def next_action(self) -> tuple[str, str] | None:
|
||||
if self._index >= len(self._actions):
|
||||
return None
|
||||
action = self._actions[self._index]
|
||||
return (action.type, action.target)
|
||||
|
||||
def advance(self) -> None:
|
||||
self._index += 1
|
||||
|
||||
def is_exhausted(self) -> bool:
|
||||
return self._index >= len(self._actions)
|
||||
|
||||
def reset(self) -> None:
|
||||
self._index = 0
|
||||
|
||||
|
||||
class ResilienceFactory:
|
||||
"""弹性组件工厂。"""
|
||||
|
||||
@staticmethod
|
||||
def create_circuit_breaker(
|
||||
name: str,
|
||||
failure_threshold: int = 5,
|
||||
cooldown_seconds: float = 30.0,
|
||||
) -> CircuitBreaker:
|
||||
return CircuitBreaker(
|
||||
name=name,
|
||||
failure_threshold=failure_threshold,
|
||||
cooldown_seconds=cooldown_seconds,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def create_retry_strategy(
|
||||
max_retries: int = 3,
|
||||
base_delay: float = 1.0,
|
||||
max_delay: float = 60.0,
|
||||
multiplier: float = 2.0,
|
||||
) -> RetryStrategy:
|
||||
return RetryStrategy(
|
||||
max_retries=max_retries,
|
||||
base_delay=base_delay,
|
||||
max_delay=max_delay,
|
||||
multiplier=multiplier,
|
||||
)
|
||||
@@ -1,149 +0,0 @@
|
||||
"""嵌入服务 — Embedding 适配器外观。
|
||||
|
||||
用法::
|
||||
|
||||
from app.ai_provider.embedding import EmbeddingService
|
||||
|
||||
service = EmbeddingService(manager)
|
||||
result = await service.embed("要嵌入的文本")
|
||||
results = await service.embed(["文本1", "文本2"])
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from app.ai_provider.embedding.base import (
|
||||
EmbeddingProtocol,
|
||||
EmbeddingRequest,
|
||||
EmbeddingResponse,
|
||||
EmbeddingResult,
|
||||
EmbeddingServiceError,
|
||||
)
|
||||
from app.ai_provider.types.errors import AIProviderError
|
||||
|
||||
|
||||
class EmbeddingService:
|
||||
"""外观类，负责解析 Embedding 模型并委托给 ProviderManager。
|
||||
|
||||
接受 ProviderManager 或实现了以下接口的适配器：
|
||||
- ``.model_registry`` → 具有 ``list_by_capability(capability)`` 方法
|
||||
- ``.create_embedding(name)`` → 返回 EmbeddingProtocol 实例
|
||||
"""
|
||||
|
||||
def __init__(self, manager) -> None:
|
||||
self._manager = manager
|
||||
self._adapter: EmbeddingProtocol | None = None
|
||||
|
||||
# ── Registry 兼容适配 ─────────────────────────────────────────
|
||||
|
||||
@property
|
||||
def _registry(self):
|
||||
mgr = self._manager
|
||||
if hasattr(mgr, "model_registry"):
|
||||
return mgr.model_registry
|
||||
return mgr
|
||||
|
||||
@property
|
||||
def _factory(self):
|
||||
mgr = self._manager
|
||||
if hasattr(mgr, "create_embedding"):
|
||||
return mgr
|
||||
return None
|
||||
|
||||
# ── Embedding ──────────────────────────────────────────────────
|
||||
|
||||
@property
|
||||
def adapter(self) -> EmbeddingProtocol:
|
||||
"""解析并缓存 Embedding 适配器。"""
|
||||
if self._adapter is None:
|
||||
self._adapter = self._resolve()
|
||||
return self._adapter
|
||||
|
||||
async def embed(
|
||||
self,
|
||||
input: str | list[str],
|
||||
model_hint: str | None = None,
|
||||
dimensions: int | None = None,
|
||||
) -> EmbeddingResponse:
|
||||
"""快捷方法：构造请求并调用 Embedding 适配器。"""
|
||||
model_name = self._resolve_model_name(model_hint)
|
||||
# 优先从模型配置读 dimensions
|
||||
effective_dimensions = dimensions
|
||||
if effective_dimensions is None:
|
||||
effective_dimensions = self._resolve_model_dimensions(model_hint)
|
||||
request = EmbeddingRequest(
|
||||
input=input if isinstance(input, list) else [input],
|
||||
model=model_name,
|
||||
dimensions=effective_dimensions,
|
||||
)
|
||||
return await self.adapter.embed(request)
|
||||
|
||||
# ── 解析 ───────────────────────────────────────────────────────
|
||||
|
||||
def _resolve(self) -> EmbeddingProtocol:
|
||||
"""查找第一个活跃的 Embedding 模型并通过工厂创建。"""
|
||||
models = self._list_capability_models("embedding")
|
||||
if not models:
|
||||
raise AIProviderError(
|
||||
status_code=503,
|
||||
code="embedding_no_model",
|
||||
message="No embedding model configured in the model registry.",
|
||||
)
|
||||
factory = self._factory
|
||||
if factory is None:
|
||||
raise AIProviderError(
|
||||
status_code=501,
|
||||
code="embedding_no_factory",
|
||||
message="Embedding adapter factory not available (provider manager missing create_embedding).",
|
||||
)
|
||||
return factory.create_embedding(models[0].name)
|
||||
|
||||
def _resolve_model_name(self, model_hint: str | None) -> str:
|
||||
models = self._list_capability_models("embedding")
|
||||
if model_hint:
|
||||
return model_hint
|
||||
if models:
|
||||
return models[0].name
|
||||
return ""
|
||||
|
||||
def _resolve_model_dimensions(self, model_hint: str | None) -> int | None:
|
||||
"""从模型配置读取向量维度。"""
|
||||
models = self._list_capability_models("embedding")
|
||||
if model_hint:
|
||||
for m in models:
|
||||
if m.name == model_hint:
|
||||
return m.dimensions
|
||||
if models:
|
||||
return models[0].dimensions
|
||||
return None
|
||||
|
||||
def _list_capability_models(self, tag: str):
|
||||
reg = self._registry
|
||||
if hasattr(reg, "list_capability_models"):
|
||||
return reg.list_capability_models(tag)
|
||||
if hasattr(reg, "list_by_capability"):
|
||||
return reg.list_by_capability(tag)
|
||||
return []
|
||||
|
||||
# ── 生命周期 ────────────────────────────────────────────────────
|
||||
|
||||
async def close(self) -> None:
|
||||
if self._adapter is not None:
|
||||
close = getattr(self._adapter, "close", None)
|
||||
if close is not None:
|
||||
await close()
|
||||
|
||||
async def __aenter__(self) -> EmbeddingService:
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *exc_info: object) -> None:
|
||||
await self.close()
|
||||
|
||||
|
||||
__all__ = [
|
||||
"EmbeddingProtocol",
|
||||
"EmbeddingRequest",
|
||||
"EmbeddingResponse",
|
||||
"EmbeddingResult",
|
||||
"EmbeddingService",
|
||||
"EmbeddingServiceError",
|
||||
]
|
||||
@@ -1,144 +0,0 @@
|
||||
"""嵌入服务类型 — 请求/响应模型和协议定义。
|
||||
|
||||
Embedding 适配器通过 HTTP 调用嵌入端点（通常是 OpenAI 兼容的 ``POST /v1/embeddings``），
|
||||
经由 ``EmbeddingService`` 或直接从 ``ProviderManager`` 消费。
|
||||
|
||||
用法::
|
||||
|
||||
service = EmbeddingService(manager)
|
||||
result = await service.embed("要嵌入的文本")
|
||||
|
||||
# 或通过工厂直接调用:
|
||||
adapter = manager.create_embedding("text-embedding-3-small")
|
||||
result = await adapter.embed(EmbeddingRequest(input=["文本"]))
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
|
||||
|
||||
from app.ai_provider.types.errors import AIProviderError
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from app.ai_provider.config.models import ModelConfig, ModelEndpointConfig
|
||||
|
||||
|
||||
class EmbeddingServiceError(AIProviderError):
|
||||
"""Embedding API 调用失败的基础错误。"""
|
||||
|
||||
|
||||
# ── 请求 / 响应 ────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@dataclass
|
||||
class EmbeddingRequest:
|
||||
"""将文本转换为向量的请求。
|
||||
|
||||
属性:
|
||||
input: 待嵌入的文本字符串或字符串列表。
|
||||
model: 模型标识符（如 ``"text-embedding-3-small"``）。
|
||||
encoding_format: 返回编码格式（``"float"`` 或 ``"base64"``）。
|
||||
dimensions: 输出向量的维度（仅部分模型支持）。
|
||||
user: 终端用户标识（用于用量追踪）。
|
||||
"""
|
||||
|
||||
input: str | list[str]
|
||||
model: str = ""
|
||||
encoding_format: str = "float"
|
||||
dimensions: int | None = None
|
||||
user: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class EmbeddingResult:
|
||||
"""单条文本的嵌入结果。"""
|
||||
|
||||
embedding: list[float]
|
||||
index: int = 0
|
||||
|
||||
|
||||
@dataclass
|
||||
class EmbeddingResponse:
|
||||
"""嵌入 API 的完整响应。"""
|
||||
|
||||
data: list[EmbeddingResult]
|
||||
model: str = ""
|
||||
usage: dict[str, int] | None = None
|
||||
|
||||
|
||||
# ── Protocol ────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class EmbeddingProtocol(Protocol):
|
||||
"""Embedding 适配器的协议。"""
|
||||
|
||||
async def embed(self, request: EmbeddingRequest) -> EmbeddingResponse:
|
||||
"""将文本转换为向量。
|
||||
|
||||
异常:
|
||||
EmbeddingServiceError: API 调用失败或响应异常时抛出。
|
||||
"""
|
||||
...
|
||||
|
||||
async def close(self) -> None:
|
||||
"""释放适配器持有的 HTTP 连接。"""
|
||||
...
|
||||
|
||||
|
||||
# ── 共享工厂基类 ──────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class EmbeddingAdapterBase:
|
||||
"""Embedding 适配器基类，提供共享的工厂方法。
|
||||
|
||||
子类应在 ``__init__`` 中接受 ``base_url``、``api_key``、``model``、
|
||||
``request_timeout_seconds`` 参数。
|
||||
"""
|
||||
|
||||
_adapter_type: str = "embedding"
|
||||
_default_timeout: int = 60
|
||||
|
||||
@classmethod
|
||||
def from_endpoint(
|
||||
cls,
|
||||
endpoint: ModelEndpointConfig,
|
||||
*,
|
||||
timeout_seconds: int | None = None,
|
||||
) -> Any:
|
||||
"""从已解析的端点配置构建适配器。"""
|
||||
base_url = endpoint.base_url
|
||||
api_key = endpoint.api_key
|
||||
if not base_url:
|
||||
raise EmbeddingServiceError(
|
||||
status_code=500,
|
||||
code="embedding_endpoint_invalid",
|
||||
message="Embedding endpoint has no resolved base_url.",
|
||||
details={"endpoint": endpoint.name},
|
||||
)
|
||||
return cls(
|
||||
base_url=base_url,
|
||||
api_key=api_key,
|
||||
model=endpoint.model,
|
||||
request_timeout_seconds=timeout_seconds
|
||||
or endpoint.request_timeout_seconds
|
||||
or cls._default_timeout,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_model_config(
|
||||
cls,
|
||||
config: ModelConfig,
|
||||
*,
|
||||
timeout_seconds: int | None = None,
|
||||
) -> Any:
|
||||
"""从逻辑模型配置构建适配器（使用第一个活跃端点）。"""
|
||||
for ep in config.endpoints:
|
||||
if ep.status == "active":
|
||||
return cls.from_endpoint(ep, timeout_seconds=timeout_seconds)
|
||||
raise EmbeddingServiceError(
|
||||
status_code=500,
|
||||
code="embedding_no_endpoint",
|
||||
message=f"Embedding model '{config.name}' has no active endpoint.",
|
||||
)
|
||||
@@ -1,169 +0,0 @@
|
||||
"""OpenAI 兼容格式的 Embedding 适配器 — ``POST /v1/embeddings``。
|
||||
|
||||
覆盖 OpenAI、DeepSeek、以及任何兼容 OpenAI 嵌入端点的提供商。
|
||||
|
||||
请求::
|
||||
|
||||
POST {base_url}/v1/embeddings
|
||||
Authorization: Bearer {api_key}
|
||||
Content-Type: application/json
|
||||
|
||||
{
|
||||
"model": "text-embedding-3-small",
|
||||
"input": ["要嵌入的文本"],
|
||||
"encoding_format": "float"
|
||||
}
|
||||
|
||||
响应::
|
||||
|
||||
{
|
||||
"object": "list",
|
||||
"data": [
|
||||
{"object": "embedding", "index": 0, "embedding": [0.xxx, ...]}
|
||||
],
|
||||
"model": "text-embedding-3-small",
|
||||
"usage": {"prompt_tokens": 5, "total_tokens": 5}
|
||||
}
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from app.ai_provider.embedding.base import (
|
||||
EmbeddingAdapterBase,
|
||||
EmbeddingProtocol,
|
||||
EmbeddingRequest,
|
||||
EmbeddingResponse,
|
||||
EmbeddingResult,
|
||||
EmbeddingServiceError,
|
||||
)
|
||||
from app.core.logging import logger
|
||||
|
||||
|
||||
class OpenAIFormatEmbeddingAdapter(EmbeddingAdapterBase, EmbeddingProtocol):
|
||||
"""OpenAI 兼容嵌入适配器，覆盖 95%+ 的嵌入服务提供商。"""
|
||||
|
||||
_adapter_type = "embedding"
|
||||
_default_timeout = 60
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
model: str = "text-embedding-3-small",
|
||||
request_timeout_seconds: int = 60,
|
||||
) -> None:
|
||||
self._base_url = base_url.rstrip("/")
|
||||
self._api_key = api_key
|
||||
self._model = model
|
||||
self._timeout = request_timeout_seconds
|
||||
self._client: httpx.AsyncClient | None = None
|
||||
|
||||
# ── 客户端 ───────────────────────────────────────────────────────────
|
||||
|
||||
@property
|
||||
def _http(self) -> httpx.AsyncClient:
|
||||
if self._client is None or self._client.is_closed:
|
||||
self._client = httpx.AsyncClient(timeout=httpx.Timeout(self._timeout))
|
||||
return self._client
|
||||
|
||||
async def close(self) -> None:
|
||||
if self._client is not None and not self._client.is_closed:
|
||||
await self._client.aclose()
|
||||
|
||||
# ── EmbeddingProtocol ────────────────────────────────────────────────
|
||||
|
||||
async def embed(self, request: EmbeddingRequest) -> EmbeddingResponse:
|
||||
"""通过 ``POST /v1/embeddings`` 将文本转换为向量。"""
|
||||
url = f"{self._base_url}/v1/embeddings"
|
||||
|
||||
payload: dict[str, Any] = {
|
||||
"model": request.model or self._model,
|
||||
"input": request.input,
|
||||
"encoding_format": request.encoding_format,
|
||||
}
|
||||
if request.dimensions is not None:
|
||||
payload["dimensions"] = request.dimensions
|
||||
if request.user is not None:
|
||||
payload["user"] = request.user
|
||||
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
if self._api_key:
|
||||
headers["Authorization"] = f"Bearer {self._api_key}"
|
||||
|
||||
logger.info(
|
||||
"embedding.request",
|
||||
model=payload["model"],
|
||||
input_count=len(request.input) if isinstance(request.input, list) else 1,
|
||||
dimensions=request.dimensions,
|
||||
url=url,
|
||||
)
|
||||
|
||||
try:
|
||||
response = await self._http.post(url, json=payload, headers=headers)
|
||||
except httpx.TimeoutException as exc:
|
||||
raise EmbeddingServiceError(
|
||||
status_code=504,
|
||||
code="embedding_timeout",
|
||||
message="Embedding request timed out.",
|
||||
) from exc
|
||||
except httpx.RequestError as exc:
|
||||
raise EmbeddingServiceError(
|
||||
status_code=503,
|
||||
code="embedding_connection_error",
|
||||
message=f"Embedding connection error: {exc}",
|
||||
) from exc
|
||||
|
||||
if response.status_code != 200:
|
||||
body = response.text[:500]
|
||||
raise EmbeddingServiceError(
|
||||
status_code=response.status_code,
|
||||
code="embedding_api_error",
|
||||
message=f"Embedding API returned {response.status_code}: {body}",
|
||||
)
|
||||
|
||||
result = self._parse_response(response.json())
|
||||
|
||||
logger.info(
|
||||
"embedding.success",
|
||||
model=payload["model"],
|
||||
total_vectors=len(result.data),
|
||||
usage=result.usage,
|
||||
)
|
||||
return result
|
||||
|
||||
# ── 辅助方法 ──────────────────────────────────────────────────────────
|
||||
|
||||
@staticmethod
|
||||
def _parse_response(body: dict) -> EmbeddingResponse:
|
||||
"""从 OpenAI 格式的嵌入响应中提取结果。"""
|
||||
try:
|
||||
raw_data = body.get("data", [])
|
||||
results = []
|
||||
for item in raw_data:
|
||||
results.append(
|
||||
EmbeddingResult(
|
||||
embedding=item.get("embedding", []),
|
||||
index=item.get("index", 0),
|
||||
)
|
||||
)
|
||||
# 按 index 排序确保顺序
|
||||
results.sort(key=lambda r: r.index)
|
||||
|
||||
return EmbeddingResponse(
|
||||
data=results,
|
||||
model=body.get("model", ""),
|
||||
usage=body.get("usage"),
|
||||
)
|
||||
except Exception as exc:
|
||||
raise EmbeddingServiceError(
|
||||
status_code=500,
|
||||
code="embedding_response_parse_error",
|
||||
message=f"Failed to parse embedding response: {exc}",
|
||||
) from exc
|
||||
@@ -1,18 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from app.ai_provider.llm.anthropic_format import (
|
||||
AnthropicFormatAdapter,
|
||||
AnthropicFormatAdapterFactory,
|
||||
)
|
||||
from app.ai_provider.llm.openai_format import OpenAIFormatAdapter, OpenAIFormatAdapterFactory
|
||||
from app.ai_provider.llm.stream import StreamFactory, StreamHandler, ToolCallAccumulator
|
||||
|
||||
__all__ = [
|
||||
"AnthropicFormatAdapter",
|
||||
"AnthropicFormatAdapterFactory",
|
||||
"OpenAIFormatAdapter",
|
||||
"OpenAIFormatAdapterFactory",
|
||||
"StreamFactory",
|
||||
"StreamHandler",
|
||||
"ToolCallAccumulator",
|
||||
]
|
||||
@@ -1,433 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from app.ai_provider.config.models import ModelConfig, ProviderConfig
|
||||
from app.ai_provider.pipeline import ParameterPipeline, PipelineContext, PipelineFactory
|
||||
from app.ai_provider.types.chat import ChatChunk, ChatResult, TokenUsage
|
||||
from app.ai_provider.types.errors import AIProviderError
|
||||
from app.ai_provider.types.messages import (
|
||||
AIMessage,
|
||||
BaseMessage,
|
||||
SystemMessage,
|
||||
ToolMessage,
|
||||
)
|
||||
from app.core.logging import logger
|
||||
|
||||
|
||||
class AnthropicFormatAdapter:
|
||||
"""兼容 Anthropic Messages API 格式的适配器。
|
||||
|
||||
覆盖通过 Anthropic API 访问的 Claude 模型，以及任何镜像
|
||||
Anthropic 消息格式的提供商。
|
||||
"""
|
||||
|
||||
API_VERSION = "2023-06-01"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
provider_config: ProviderConfig,
|
||||
pipeline: ParameterPipeline | None = None,
|
||||
) -> None:
|
||||
self._provider = provider_config
|
||||
self._pipeline = pipeline or PipelineFactory.create(provider_config)
|
||||
self._client: httpx.AsyncClient | None = None
|
||||
|
||||
@property
|
||||
def _http(self) -> httpx.AsyncClient:
|
||||
if self._client is None:
|
||||
self._client = httpx.AsyncClient(
|
||||
base_url=self._provider.base_url or "",
|
||||
timeout=self._provider.timeout_seconds,
|
||||
)
|
||||
return self._client
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 内部方法：消息转换
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _convert_message(self, msg: BaseMessage) -> dict[str, Any]:
|
||||
"""将单条消息转换为 Anthropic 消息格式。
|
||||
|
||||
Anthropic 使用 content blocks 数组而非纯字符串。
|
||||
System 消息**不在此处处理** —— 它们由 :meth:`build_request`
|
||||
提取到顶层的 ``system`` 字段中。
|
||||
"""
|
||||
if isinstance(msg, ToolMessage):
|
||||
return {
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": msg.tool_call_id,
|
||||
"content": msg.content,
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
if isinstance(msg, AIMessage):
|
||||
if msg.raw_content_blocks:
|
||||
return {"role": "assistant", "content": msg.raw_content_blocks}
|
||||
blocks: list[dict[str, Any]] = []
|
||||
if msg.content:
|
||||
blocks.append({"type": "text", "text": msg.content})
|
||||
for tc in msg.tool_calls:
|
||||
blocks.append(
|
||||
{
|
||||
"type": "tool_use",
|
||||
"id": tc.get("id", f"tu_{tc['name']}"),
|
||||
"name": tc["name"],
|
||||
"input": tc.get("args", {}),
|
||||
}
|
||||
)
|
||||
return {"role": "assistant", "content": blocks}
|
||||
|
||||
# HumanMessage
|
||||
return {
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": msg.content}],
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _convert_tools(tools: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
"""将 OpenAI 兼容格式的工具转换为 Anthropic 格式。
|
||||
|
||||
OpenAI 格式::
|
||||
|
||||
{"type": "function", "function": {"name": ..., "parameters": ...}}
|
||||
|
||||
Anthropic 格式::
|
||||
|
||||
{"name": ..., "input_schema": ...}
|
||||
"""
|
||||
result: list[dict[str, Any]] = []
|
||||
for tool in tools:
|
||||
fn = tool.get("function", tool)
|
||||
result.append(
|
||||
{
|
||||
"name": fn.get("name", ""),
|
||||
"description": fn.get("description", ""),
|
||||
"input_schema": fn.get("parameters", {"type": "object"}),
|
||||
}
|
||||
)
|
||||
return result
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 请求构建
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def build_request(
|
||||
self,
|
||||
model: ModelConfig,
|
||||
messages: list[BaseMessage],
|
||||
params: dict[str, Any] | None = None,
|
||||
tools: list[dict[str, Any]] | None = None,
|
||||
tool_choice: str | dict[str, Any] | None = None,
|
||||
stream: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
ctx = PipelineContext(
|
||||
params=params or {},
|
||||
model_config=model,
|
||||
provider_config=self._provider,
|
||||
)
|
||||
pipe_result = self._pipeline.run(ctx)
|
||||
|
||||
# Extract system messages — Anthropic uses a top-level "system" field
|
||||
system_parts: list[str] = []
|
||||
non_system: list[BaseMessage] = []
|
||||
for m in messages:
|
||||
if isinstance(m, SystemMessage):
|
||||
system_parts.append(m.content)
|
||||
else:
|
||||
non_system.append(m)
|
||||
|
||||
body: dict[str, Any] = {
|
||||
"model": model.model,
|
||||
"messages": [self._convert_message(m) for m in non_system],
|
||||
"stream": stream,
|
||||
}
|
||||
|
||||
if system_parts:
|
||||
body["system"] = "\n".join(system_parts)
|
||||
|
||||
body.update(pipe_result.params)
|
||||
|
||||
if tools:
|
||||
body["tools"] = self._convert_tools(tools)
|
||||
if tool_choice:
|
||||
body["tool_choice"] = tool_choice
|
||||
|
||||
# Anthropic requires explicit max_tokens
|
||||
if "max_tokens" not in body:
|
||||
body["max_tokens"] = model.max_output_tokens or 4096
|
||||
|
||||
return body
|
||||
|
||||
def bind_tools(
|
||||
self,
|
||||
request_body: dict[str, Any],
|
||||
tools: list[dict[str, Any]],
|
||||
tool_choice: str | dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
request_body["tools"] = self._convert_tools(tools)
|
||||
if tool_choice:
|
||||
request_body["tool_choice"] = tool_choice
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 认证头
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _headers(self, api_key: str) -> dict[str, str]:
|
||||
return {
|
||||
"x-api-key": api_key,
|
||||
"anthropic-version": self.API_VERSION,
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 响应解析 —— 非流式
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def _parse_content_blocks(raw: dict[str, Any]) -> tuple[str, list[dict[str, Any]], str]:
|
||||
"""解析 Anthropic content blocks 为 (text, tool_calls, reasoning)。"""
|
||||
content_blocks = raw.get("content", [])
|
||||
text_parts: list[str] = []
|
||||
tool_calls: list[dict[str, Any]] = []
|
||||
reasoning = ""
|
||||
|
||||
for block in content_blocks:
|
||||
block_type = block.get("type", "")
|
||||
if block_type == "text":
|
||||
text_parts.append(block.get("text", ""))
|
||||
elif block_type == "tool_use":
|
||||
tool_calls.append(
|
||||
{
|
||||
"name": block.get("name", ""),
|
||||
"args": block.get("input", {}),
|
||||
"id": block.get("id", ""),
|
||||
}
|
||||
)
|
||||
elif block_type == "thinking":
|
||||
reasoning += block.get("thinking", "")
|
||||
# signature 块是 thinking 的加密签名，不是 reasoning 内容，忽略
|
||||
|
||||
return "".join(text_parts), tool_calls, reasoning
|
||||
|
||||
def parse_nonstreaming_response(
|
||||
self,
|
||||
raw: dict[str, Any],
|
||||
model_name: str,
|
||||
) -> ChatResult:
|
||||
content, tool_calls, reasoning = self._parse_content_blocks(raw)
|
||||
|
||||
result = ChatResult(
|
||||
content=content,
|
||||
model_name=model_name,
|
||||
)
|
||||
if tool_calls:
|
||||
result.tool_calls = tool_calls
|
||||
if reasoning:
|
||||
result.reasoning = reasoning
|
||||
|
||||
usage = raw.get("usage")
|
||||
if usage:
|
||||
result.usage = TokenUsage(
|
||||
prompt_tokens=usage.get("input_tokens", 0),
|
||||
prompt_tokens_cached=usage.get("cache_read_input_tokens", 0),
|
||||
completion_tokens=usage.get("output_tokens", 0),
|
||||
total_tokens=usage.get("input_tokens", 0) + usage.get("output_tokens", 0),
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 响应解析 —— 流式（SSE 事件行协议）
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def _parse_stream_event(event_type: str, data: dict[str, Any]) -> ChatChunk:
|
||||
"""解析单条 SSE 流事件为 ChatChunk。"""
|
||||
chunk = ChatChunk()
|
||||
|
||||
if event_type == "message_start":
|
||||
message = data.get("message", {})
|
||||
usage = message.get("usage")
|
||||
if usage:
|
||||
chunk.usage = TokenUsage(
|
||||
prompt_tokens=usage.get("input_tokens", 0),
|
||||
prompt_tokens_cached=usage.get("cache_read_input_tokens", 0),
|
||||
completion_tokens=usage.get("output_tokens", 0),
|
||||
)
|
||||
|
||||
elif event_type == "content_block_start":
|
||||
block = data.get("content_block", {})
|
||||
chunk.content_block_index = data.get("index", 0)
|
||||
chunk.content_block_start = block
|
||||
block_type = block.get("type", "")
|
||||
if block_type == "text":
|
||||
chunk.delta_content = block.get("text", "")
|
||||
elif block_type == "thinking":
|
||||
chunk.delta_reasoning = block.get("thinking", "")
|
||||
elif block_type == "tool_use":
|
||||
chunk.tool_call_chunks.append(
|
||||
{
|
||||
"index": data.get("index", 0),
|
||||
"id": block.get("id"),
|
||||
"name": block.get("name"),
|
||||
"args": "",
|
||||
}
|
||||
)
|
||||
|
||||
elif event_type == "content_block_delta":
|
||||
delta = data.get("delta", {})
|
||||
chunk.content_block_index = data.get("index", 0)
|
||||
chunk.content_block_delta = delta
|
||||
delta_type = delta.get("type", "")
|
||||
if delta_type == "text_delta":
|
||||
chunk.delta_content = delta.get("text", "")
|
||||
elif delta_type == "thinking_delta":
|
||||
chunk.delta_reasoning = delta.get("thinking", "")
|
||||
elif delta_type == "signature_delta":
|
||||
# signature 不展示，但由 Agent Runtime 原样保留供下一轮回放。
|
||||
pass
|
||||
elif delta_type == "input_json_delta":
|
||||
chunk.tool_call_chunks.append(
|
||||
{
|
||||
"index": data.get("index", 0),
|
||||
"id": None,
|
||||
"name": None,
|
||||
"args": delta.get("partial_json", ""),
|
||||
}
|
||||
)
|
||||
|
||||
elif event_type == "message_delta":
|
||||
delta = data.get("delta", {})
|
||||
finish_reason_map = {
|
||||
"end_turn": "stop",
|
||||
"max_tokens": "length",
|
||||
"stop_sequence": "stop",
|
||||
"tool_use": "tool_calls",
|
||||
}
|
||||
stop_reason = delta.get("stop_reason")
|
||||
if stop_reason:
|
||||
chunk.finish_reason = finish_reason_map.get(stop_reason, stop_reason)
|
||||
|
||||
usage = data.get("usage")
|
||||
if usage:
|
||||
chunk.usage = TokenUsage(
|
||||
prompt_tokens=0,
|
||||
completion_tokens=usage.get("output_tokens", 0),
|
||||
)
|
||||
|
||||
# content_block_stop, message_stop, ping — no content to extract
|
||||
|
||||
return chunk
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 公开 API：chat
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def chat(
|
||||
self,
|
||||
model: ModelConfig,
|
||||
messages: list[BaseMessage],
|
||||
api_key: str,
|
||||
base_url: str | None = None,
|
||||
params: dict[str, Any] | None = None,
|
||||
tools: list[dict[str, Any]] | None = None,
|
||||
tool_choice: str | dict[str, Any] | None = None,
|
||||
) -> ChatResult:
|
||||
url = f"{base_url or self._provider.base_url}/v1/messages"
|
||||
body = self.build_request(
|
||||
model,
|
||||
messages,
|
||||
params=params,
|
||||
tools=tools,
|
||||
tool_choice=tool_choice,
|
||||
)
|
||||
|
||||
try:
|
||||
response = await self._http.post(url, json=body, headers=self._headers(api_key))
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
return self.parse_nonstreaming_response(data, model.name)
|
||||
except httpx.HTTPStatusError as exc:
|
||||
raise AIProviderError.from_provider(
|
||||
message=f"Anthropic API error: {exc.response.status_code} {exc.response.text[:200]}",
|
||||
code="PROVIDER_ERROR",
|
||||
provider=self._provider.name,
|
||||
original=exc,
|
||||
) from exc
|
||||
|
||||
async def chat_stream(
|
||||
self,
|
||||
model: ModelConfig,
|
||||
messages: list[BaseMessage],
|
||||
api_key: str,
|
||||
base_url: str | None = None,
|
||||
params: dict[str, Any] | None = None,
|
||||
tools: list[dict[str, Any]] | None = None,
|
||||
tool_choice: str | dict[str, Any] | None = None,
|
||||
) -> AsyncIterator[ChatChunk]:
|
||||
url = f"{base_url or self._provider.base_url}/v1/messages"
|
||||
body = self.build_request(
|
||||
model,
|
||||
messages,
|
||||
params=params,
|
||||
tools=tools,
|
||||
tool_choice=tool_choice,
|
||||
stream=True,
|
||||
)
|
||||
|
||||
async with self._http.stream(
|
||||
"POST",
|
||||
url,
|
||||
json=body,
|
||||
headers=self._headers(api_key),
|
||||
) as response:
|
||||
if response.status_code != 200:
|
||||
error_text = await response.aread()
|
||||
raise AIProviderError.from_provider(
|
||||
message=f"Anthropic stream error: {response.status_code} {error_text[:200]}",
|
||||
code="PROVIDER_ERROR",
|
||||
provider=self._provider.name,
|
||||
)
|
||||
|
||||
current_event = ""
|
||||
async for line in response.aiter_lines():
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
if line.startswith("event: "):
|
||||
current_event = line[7:]
|
||||
elif line.startswith("data: "):
|
||||
data_str = line[6:]
|
||||
if data_str and current_event:
|
||||
try:
|
||||
raw = json.loads(data_str)
|
||||
yield self._parse_stream_event(current_event, raw)
|
||||
except json.JSONDecodeError:
|
||||
logger.warning(
|
||||
"anthropic_stream_parse_error",
|
||||
event=current_event,
|
||||
raw=data_str[:100],
|
||||
)
|
||||
current_event = ""
|
||||
|
||||
async def close(self) -> None:
|
||||
if self._client:
|
||||
await self._client.aclose()
|
||||
self._client = None
|
||||
|
||||
|
||||
class AnthropicFormatAdapterFactory:
|
||||
"""工厂：创建 Anthropic 格式适配器。"""
|
||||
|
||||
@staticmethod
|
||||
def create(provider_config: ProviderConfig) -> AnthropicFormatAdapter:
|
||||
return AnthropicFormatAdapter(provider_config=provider_config)
|
||||
@@ -1,251 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from app.ai_provider.config.models import ModelConfig, ProviderConfig
|
||||
from app.ai_provider.pipeline import ParameterPipeline, PipelineContext, PipelineFactory
|
||||
from app.ai_provider.types.chat import ChatChunk, ChatResult, TokenUsage
|
||||
from app.ai_provider.types.errors import AIProviderError
|
||||
from app.ai_provider.types.messages import AIMessage, BaseMessage, ToolMessage
|
||||
from app.core.logging import logger
|
||||
|
||||
|
||||
class OpenAIFormatAdapter:
|
||||
"""兼容 OpenAI API 格式的适配器（覆盖 deepseek、xiaomi 等）。"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
provider_config: ProviderConfig,
|
||||
pipeline: ParameterPipeline | None = None,
|
||||
) -> None:
|
||||
self._provider = provider_config
|
||||
self._pipeline = pipeline or PipelineFactory.create(provider_config)
|
||||
self._client: httpx.AsyncClient | None = None
|
||||
|
||||
@property
|
||||
def _http(self) -> httpx.AsyncClient:
|
||||
if self._client is None:
|
||||
self._client = httpx.AsyncClient(
|
||||
base_url=self._provider.base_url or "",
|
||||
timeout=self._provider.timeout_seconds,
|
||||
)
|
||||
return self._client
|
||||
|
||||
def _convert_message(self, msg: BaseMessage) -> dict[str, Any]:
|
||||
if isinstance(msg, ToolMessage):
|
||||
return {
|
||||
"role": "tool",
|
||||
"content": msg.content,
|
||||
"tool_call_id": msg.tool_call_id,
|
||||
}
|
||||
if isinstance(msg, AIMessage):
|
||||
d: dict[str, Any] = {"role": "assistant", "content": msg.content}
|
||||
if msg.tool_calls:
|
||||
d["tool_calls"] = [
|
||||
{
|
||||
"id": tc.get("id", f"call_{i}"),
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": tc["name"],
|
||||
"arguments": json.dumps(tc.get("args", {})),
|
||||
},
|
||||
}
|
||||
for i, tc in enumerate(msg.tool_calls)
|
||||
]
|
||||
return d
|
||||
return {"role": msg.role, "content": msg.content}
|
||||
|
||||
def build_request(
|
||||
self,
|
||||
model: ModelConfig,
|
||||
messages: list[BaseMessage],
|
||||
params: dict[str, Any] | None = None,
|
||||
tools: list[dict[str, Any]] | None = None,
|
||||
tool_choice: str | dict[str, Any] | None = None,
|
||||
stream: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
ctx = PipelineContext(
|
||||
params=params or {},
|
||||
model_config=model,
|
||||
provider_config=self._provider,
|
||||
)
|
||||
pipe_result = self._pipeline.run(ctx)
|
||||
|
||||
body: dict[str, Any] = {
|
||||
"model": model.model,
|
||||
"messages": [self._convert_message(m) for m in messages],
|
||||
"stream": stream,
|
||||
}
|
||||
body.update(pipe_result.params)
|
||||
|
||||
if tools:
|
||||
body["tools"] = tools
|
||||
if tool_choice:
|
||||
body["tool_choice"] = tool_choice
|
||||
|
||||
return body
|
||||
|
||||
def bind_tools(
|
||||
self,
|
||||
request_body: dict[str, Any],
|
||||
tools: list[dict[str, Any]],
|
||||
tool_choice: str | dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
request_body["tools"] = tools
|
||||
if tool_choice:
|
||||
request_body["tool_choice"] = tool_choice
|
||||
|
||||
def parse_chunk(self, raw: dict[str, Any]) -> ChatChunk:
|
||||
chunk = ChatChunk()
|
||||
choices = raw.get("choices", [])
|
||||
if choices:
|
||||
delta = choices[0].get("delta", {})
|
||||
chunk.delta_content = delta.get("content") or ""
|
||||
chunk.delta_reasoning = delta.get("reasoning_content") or ""
|
||||
chunk.finish_reason = choices[0].get("finish_reason")
|
||||
|
||||
if delta.get("tool_calls"):
|
||||
for tc in delta["tool_calls"]:
|
||||
fn = tc.get("function", {})
|
||||
chunk.tool_call_chunks.append(
|
||||
{
|
||||
"index": tc.get("index", 0),
|
||||
"id": tc.get("id"),
|
||||
"name": fn.get("name"),
|
||||
"args": fn.get("arguments", ""),
|
||||
}
|
||||
)
|
||||
usage = raw.get("usage")
|
||||
if usage:
|
||||
chunk.usage = TokenUsage(
|
||||
prompt_tokens=usage.get("prompt_tokens", 0),
|
||||
completion_tokens=usage.get("completion_tokens", 0),
|
||||
total_tokens=usage.get("total_tokens", 0),
|
||||
prompt_tokens_cached=usage.get("prompt_tokens_cached", 0),
|
||||
)
|
||||
return chunk
|
||||
|
||||
def parse_nonstreaming_response(
|
||||
self,
|
||||
raw: dict[str, Any],
|
||||
model_name: str,
|
||||
) -> ChatResult:
|
||||
choice = raw.get("choices", [{}])[0]
|
||||
message = choice.get("message", {})
|
||||
content = message.get("content", "")
|
||||
tool_calls_raw = message.get("tool_calls", [])
|
||||
|
||||
result = ChatResult(
|
||||
content=content,
|
||||
reasoning=message.get("reasoning_content", ""),
|
||||
model_name=model_name,
|
||||
)
|
||||
if tool_calls_raw:
|
||||
result.tool_calls = [
|
||||
{
|
||||
"name": tc["function"]["name"],
|
||||
"args": json.loads(tc["function"].get("arguments", "{}")),
|
||||
"id": tc["id"],
|
||||
}
|
||||
for tc in tool_calls_raw
|
||||
]
|
||||
|
||||
usage = raw.get("usage")
|
||||
if usage:
|
||||
result.usage = TokenUsage(
|
||||
prompt_tokens=usage.get("prompt_tokens", 0),
|
||||
completion_tokens=usage.get("completion_tokens", 0),
|
||||
total_tokens=usage.get("total_tokens", 0),
|
||||
prompt_tokens_cached=usage.get("prompt_tokens_cached", 0),
|
||||
)
|
||||
return result
|
||||
|
||||
async def chat(
|
||||
self,
|
||||
model: ModelConfig,
|
||||
messages: list[BaseMessage],
|
||||
api_key: str,
|
||||
base_url: str | None = None,
|
||||
params: dict[str, Any] | None = None,
|
||||
tools: list[dict[str, Any]] | None = None,
|
||||
tool_choice: str | dict[str, Any] | None = None,
|
||||
) -> ChatResult:
|
||||
url = f"{base_url or self._provider.base_url}/chat/completions"
|
||||
body = self.build_request(
|
||||
model, messages, params=params, tools=tools, tool_choice=tool_choice
|
||||
)
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
try:
|
||||
response = await self._http.post(url, json=body, headers=headers)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
return self.parse_nonstreaming_response(data, model.name)
|
||||
except httpx.HTTPStatusError as exc:
|
||||
raise AIProviderError.from_provider(
|
||||
message=f"OpenAI API error: {exc.response.status_code} {exc.response.text[:200]}",
|
||||
code="PROVIDER_ERROR",
|
||||
provider=self._provider.name,
|
||||
original=exc,
|
||||
) from exc
|
||||
|
||||
async def chat_stream(
|
||||
self,
|
||||
model: ModelConfig,
|
||||
messages: list[BaseMessage],
|
||||
api_key: str,
|
||||
base_url: str | None = None,
|
||||
params: dict[str, Any] | None = None,
|
||||
tools: list[dict[str, Any]] | None = None,
|
||||
tool_choice: str | dict[str, Any] | None = None,
|
||||
) -> AsyncIterator[ChatChunk]:
|
||||
url = f"{base_url or self._provider.base_url}/chat/completions"
|
||||
body = self.build_request(
|
||||
model, messages, params=params, tools=tools, tool_choice=tool_choice, stream=True
|
||||
)
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
async with self._http.stream("POST", url, json=body, headers=headers) as response:
|
||||
if response.status_code != 200:
|
||||
error_text = await response.aread()
|
||||
raise AIProviderError.from_provider(
|
||||
message=f"OpenAI stream error: {response.status_code} {error_text[:200]}",
|
||||
code="PROVIDER_ERROR",
|
||||
provider=self._provider.name,
|
||||
)
|
||||
async for line in response.aiter_lines():
|
||||
if line.startswith("data: "):
|
||||
data_str = line[6:].strip()
|
||||
if data_str == "[DONE]":
|
||||
break
|
||||
if data_str:
|
||||
try:
|
||||
raw = json.loads(data_str)
|
||||
yield self.parse_chunk(raw)
|
||||
except json.JSONDecodeError:
|
||||
logger.warning("openai_stream_parse_error", raw=data_str[:100])
|
||||
|
||||
async def close(self) -> None:
|
||||
if self._client:
|
||||
await self._client.aclose()
|
||||
self._client = None
|
||||
|
||||
|
||||
class OpenAIFormatAdapterFactory:
|
||||
"""工厂：创建 OpenAI 格式适配器。"""
|
||||
|
||||
@staticmethod
|
||||
def create(provider_config: ProviderConfig) -> OpenAIFormatAdapter:
|
||||
return OpenAIFormatAdapter(provider_config=provider_config)
|
||||
@@ -1,164 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import random
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Protocol
|
||||
|
||||
|
||||
@dataclass
|
||||
class PoolContext:
|
||||
request_id: str = ""
|
||||
excluded_endpoints: set[str] = field(default_factory=set)
|
||||
|
||||
|
||||
class SchedulingStrategy(Protocol):
|
||||
def select(self, excluded: set[str]) -> dict[str, Any] | None: ...
|
||||
|
||||
|
||||
class WeightedRoundRobinStrategy:
|
||||
"""策略 1：在所有活跃端点间加权轮询。"""
|
||||
|
||||
def __init__(self, endpoints: list[dict[str, Any]]) -> None:
|
||||
self._endpoints = endpoints
|
||||
self._index = 0
|
||||
|
||||
def select(self, excluded: set[str]) -> dict[str, Any] | None:
|
||||
available = [
|
||||
e
|
||||
for e in self._endpoints
|
||||
if e.get("name") not in excluded and e.get("status", "active") == "active"
|
||||
]
|
||||
if not available:
|
||||
return None
|
||||
weights = [e.get("weight", 100) for e in available]
|
||||
self._index = (self._index + 1) % len(available)
|
||||
return random.choices(available, weights=weights, k=1)[0]
|
||||
|
||||
|
||||
class PriorityFailoverStrategy:
|
||||
"""策略 2：选择最高优先级；当排除时顺延。"""
|
||||
|
||||
def __init__(self, endpoints: list[dict[str, Any]]) -> None:
|
||||
self._endpoints = sorted(endpoints, key=lambda e: e.get("priority", 100))
|
||||
|
||||
def select(self, excluded: set[str]) -> dict[str, Any] | None:
|
||||
for ep in self._endpoints:
|
||||
if ep.get("name") not in excluded and ep.get("status", "active") == "active":
|
||||
return ep
|
||||
return None
|
||||
|
||||
|
||||
class LowestLatencyStrategy:
|
||||
"""策略 3：选择最近延迟最低的端点。"""
|
||||
|
||||
def __init__(self, endpoints: list[dict[str, Any]]) -> None:
|
||||
self._endpoints = endpoints
|
||||
|
||||
def select(self, excluded: set[str]) -> dict[str, Any] | None:
|
||||
available = [
|
||||
e
|
||||
for e in self._endpoints
|
||||
if e.get("name") not in excluded and e.get("status", "active") == "active"
|
||||
]
|
||||
if not available:
|
||||
return None
|
||||
return min(available, key=lambda e: e.get("latency_ms", 1000))
|
||||
|
||||
|
||||
class LeastConnectionsStrategy:
|
||||
"""策略 4：选择活跃连接数最少的端点。"""
|
||||
|
||||
def __init__(self, endpoints: list[dict[str, Any]]) -> None:
|
||||
self._endpoints = endpoints
|
||||
|
||||
def select(self, excluded: set[str]) -> dict[str, Any] | None:
|
||||
available = [
|
||||
e
|
||||
for e in self._endpoints
|
||||
if e.get("name") not in excluded and e.get("status", "active") == "active"
|
||||
]
|
||||
if not available:
|
||||
return None
|
||||
return min(available, key=lambda e: e.get("connections", 0))
|
||||
|
||||
|
||||
class CostAwareStrategy:
|
||||
"""策略 5：按每 token 成本选择最便宜的端点。"""
|
||||
|
||||
def __init__(
|
||||
self, endpoints: list[dict[str, Any]], get_cost: Callable[[str], float] | None = None
|
||||
) -> None:
|
||||
self._endpoints = endpoints
|
||||
self._get_cost = get_cost or (lambda _: 0.0)
|
||||
|
||||
def select(self, excluded: set[str]) -> dict[str, Any] | None:
|
||||
available = [
|
||||
e
|
||||
for e in self._endpoints
|
||||
if e.get("name") not in excluded and e.get("status", "active") == "active"
|
||||
]
|
||||
if not available:
|
||||
return None
|
||||
return min(available, key=lambda e: self._get_cost(e.get("name", "")))
|
||||
|
||||
|
||||
_STRATEGIES = {
|
||||
"weighted_round_robin": WeightedRoundRobinStrategy,
|
||||
"priority_failover": PriorityFailoverStrategy,
|
||||
"lowest_latency": LowestLatencyStrategy,
|
||||
"least_connections": LeastConnectionsStrategy,
|
||||
"cost_aware": CostAwareStrategy,
|
||||
}
|
||||
|
||||
|
||||
class EndpointPool:
|
||||
"""模型池，带可插拔的调度策略。"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
logical_model_name: str,
|
||||
endpoints_configs: list[dict[str, Any]],
|
||||
strategy: str = "weighted_round_robin",
|
||||
) -> None:
|
||||
self.logical_model_name = logical_model_name
|
||||
self._endpoints = endpoints_configs
|
||||
strategy_cls = _STRATEGIES.get(strategy, WeightedRoundRobinStrategy)
|
||||
self.strategy = (
|
||||
strategy_cls(endpoints_configs)
|
||||
if strategy == "priority_failover"
|
||||
else strategy_cls(endpoints_configs)
|
||||
)
|
||||
|
||||
def select(self, excluded: set[str]) -> dict[str, Any] | None:
|
||||
"""使用配置的策略选择一个端点。"""
|
||||
return self.strategy.select(excluded)
|
||||
|
||||
def list_endpoints(self) -> list[dict[str, Any]]:
|
||||
return list(self._endpoints)
|
||||
|
||||
|
||||
class PoolFactory:
|
||||
"""工厂：创建端点池。"""
|
||||
|
||||
@staticmethod
|
||||
def create_pool(
|
||||
logical_model_name: str,
|
||||
endpoints: list[dict[str, Any]],
|
||||
strategy_name: str = "weighted_round_robin",
|
||||
) -> EndpointPool:
|
||||
return EndpointPool(
|
||||
logical_model_name=logical_model_name,
|
||||
endpoints_configs=endpoints,
|
||||
strategy=strategy_name,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def create_strategy(
|
||||
name: str,
|
||||
endpoints: list[dict[str, Any]],
|
||||
) -> SchedulingStrategy:
|
||||
strategy_cls = _STRATEGIES.get(name)
|
||||
if strategy_cls is None:
|
||||
raise ValueError(f"Unknown scheduling strategy: {name}")
|
||||
return strategy_cls(endpoints)
|
||||
@@ -1,115 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections import defaultdict
|
||||
from typing import Any
|
||||
|
||||
from app.ai_provider.types.chat import ChatChunk, ChatResult, TokenUsage
|
||||
|
||||
|
||||
class ToolCallAccumulator:
|
||||
"""将流式 tool call 分片累积为完整的 tool call。
|
||||
|
||||
处理流式场景中常见的多分片 tool call 参数，
|
||||
按索引拼接部分 JSON 字符串。
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.content: str = ""
|
||||
self.reasoning: str = ""
|
||||
self._tool_call_slots: dict[int, dict[str, Any]] = defaultdict(
|
||||
lambda: {"index": 0, "id": "", "name": "", "args": ""}
|
||||
)
|
||||
self.usage: TokenUsage | None = None
|
||||
|
||||
def add_chunk(self, chunk: ChatChunk) -> None:
|
||||
"""处理一个 ChatChunk，累积 content/reasoning/tool-calls/usage。"""
|
||||
if chunk.delta_content:
|
||||
self.content += chunk.delta_content
|
||||
if chunk.delta_reasoning:
|
||||
self.reasoning += chunk.delta_reasoning
|
||||
if chunk.usage is not None:
|
||||
self.usage = chunk.usage
|
||||
for tc in chunk.tool_call_chunks:
|
||||
idx = tc.get("index", 0)
|
||||
if idx not in self._tool_call_slots:
|
||||
self._tool_call_slots[idx] = {"index": idx, "id": "", "name": "", "args": ""}
|
||||
slot = self._tool_call_slots[idx]
|
||||
# The first chunk for a given index carries id/name; subsequent
|
||||
# arg-only chunks have them as None and should NOT overwrite.
|
||||
if tc.get("id"):
|
||||
slot["id"] = tc["id"]
|
||||
if tc.get("name"):
|
||||
slot["name"] = tc["name"]
|
||||
slot["args"] += tc.get("args", "")
|
||||
|
||||
def drain(self) -> list[dict[str, Any]]:
|
||||
"""返回所有已累积的 tool call 并重置内部槽位。"""
|
||||
result: list[dict[str, Any]] = []
|
||||
for idx in sorted(self._tool_call_slots.keys()):
|
||||
slot = self._tool_call_slots[idx]
|
||||
result.append(
|
||||
{
|
||||
"index": slot["index"],
|
||||
"id": slot["id"] or f"call_{idx}",
|
||||
"name": slot["name"],
|
||||
"args": slot["args"],
|
||||
}
|
||||
)
|
||||
self._tool_call_slots.clear()
|
||||
return result
|
||||
|
||||
|
||||
class StreamHandler:
|
||||
"""收集流式分片并生成最终的 ChatResult。"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._accumulator = ToolCallAccumulator()
|
||||
self._chunks: list[ChatChunk] = []
|
||||
|
||||
def add_chunk(self, chunk: ChatChunk) -> None:
|
||||
"""收集一个流式分片。"""
|
||||
self._chunks.append(chunk)
|
||||
self._accumulator.add_chunk(chunk)
|
||||
|
||||
@property
|
||||
def chunks(self) -> list[ChatChunk]:
|
||||
"""返回所有已收集分片的副本。"""
|
||||
return list(self._chunks)
|
||||
|
||||
def to_result(self, model_name: str = "") -> ChatResult:
|
||||
"""从所有累积分片构建最终的 ChatResult。"""
|
||||
raw_tool_calls = self._accumulator.drain()
|
||||
parsed_calls: list[dict[str, Any]] = []
|
||||
for tc in raw_tool_calls:
|
||||
args_raw = tc.get("args", "{}")
|
||||
try:
|
||||
parsed_args = json.loads(args_raw)
|
||||
except json.JSONDecodeError:
|
||||
parsed_args = {"_raw": args_raw}
|
||||
parsed_calls.append(
|
||||
{
|
||||
"name": tc["name"],
|
||||
"args": parsed_args,
|
||||
"id": tc["id"],
|
||||
}
|
||||
)
|
||||
return ChatResult(
|
||||
content=self._accumulator.content,
|
||||
reasoning=self._accumulator.reasoning,
|
||||
tool_calls=parsed_calls,
|
||||
usage=self._accumulator.usage,
|
||||
model_name=model_name,
|
||||
)
|
||||
|
||||
|
||||
class StreamFactory:
|
||||
"""流处理器和累积器的工厂。"""
|
||||
|
||||
@staticmethod
|
||||
def create_handler() -> StreamHandler:
|
||||
return StreamHandler()
|
||||
|
||||
@staticmethod
|
||||
def create_accumulator() -> ToolCallAccumulator:
|
||||
return ToolCallAccumulator()
|
||||
@@ -1,463 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from app.ai_provider.config.loader import (
|
||||
load_models_config,
|
||||
load_pricing_config,
|
||||
load_providers_config,
|
||||
)
|
||||
from app.ai_provider.config.models import ModelConfig, ProviderConfig
|
||||
from app.ai_provider.core.billing import (
|
||||
BillingFactory,
|
||||
BillingTracker,
|
||||
BudgetController,
|
||||
PricingInfo,
|
||||
PricingRegistry,
|
||||
)
|
||||
from app.ai_provider.core.registry import ModelRegistry, ProviderRegistry
|
||||
from app.ai_provider.core.resilience import CircuitBreaker
|
||||
from app.ai_provider.llm.anthropic_format import (
|
||||
AnthropicFormatAdapter,
|
||||
AnthropicFormatAdapterFactory,
|
||||
)
|
||||
from app.ai_provider.llm.openai_format import OpenAIFormatAdapter, OpenAIFormatAdapterFactory
|
||||
from app.ai_provider.llm.stream import StreamHandler
|
||||
from app.ai_provider.routing import (
|
||||
CapabilityPermissionController,
|
||||
ModelSelector,
|
||||
MultiModalRouter,
|
||||
)
|
||||
from app.ai_provider.types.chat import ChatResult, TokenUsage
|
||||
from app.ai_provider.types.errors import AIProviderError
|
||||
from app.ai_provider.types.messages import (
|
||||
AIMessage,
|
||||
BaseMessage,
|
||||
HumanMessage,
|
||||
SystemMessage,
|
||||
)
|
||||
from app.core.config import settings
|
||||
|
||||
|
||||
@dataclass
|
||||
class _ResolvedEndpoint:
|
||||
"""端点到适配器和凭证的解析结果。"""
|
||||
|
||||
adapter: OpenAIFormatAdapter | AnthropicFormatAdapter
|
||||
model: ModelConfig
|
||||
provider: ProviderConfig
|
||||
api_key: str
|
||||
base_url: str | None
|
||||
breaker: CircuitBreaker | None = None
|
||||
|
||||
|
||||
class ProviderManager:
|
||||
"""所有 AI 提供商操作的统一入口。
|
||||
|
||||
取代旧的 ModelRouter + get_model_router() 模式。
|
||||
组合了注册表、适配器、池、计费、弹性和路由。
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_registry: ModelRegistry,
|
||||
provider_registry: ProviderRegistry,
|
||||
billing_tracker: BillingTracker | None = None,
|
||||
budget_controller: BudgetController | None = None,
|
||||
pricing_registry: PricingRegistry | None = None,
|
||||
) -> None:
|
||||
self.model_registry = model_registry
|
||||
self.provider_registry = provider_registry
|
||||
self._billing = billing_tracker or BillingFactory.create_tracker()
|
||||
self._budget = budget_controller
|
||||
self._pricing_registry = pricing_registry or PricingRegistry()
|
||||
self._adapter_cache: dict[str, OpenAIFormatAdapter | AnthropicFormatAdapter] = {}
|
||||
self._model_selector = ModelSelector(model_registry)
|
||||
self._multi_modal_router = MultiModalRouter()
|
||||
self._permission_controller = CapabilityPermissionController()
|
||||
|
||||
# ── 模型解析 ──────────────────────────────────────────────
|
||||
|
||||
def resolve_model(self, model_hint: str | None = None) -> str:
|
||||
if model_hint is None:
|
||||
return self.model_registry.get_default_name()
|
||||
return self.model_registry.resolve_name(model_hint)
|
||||
|
||||
def get_model_config(self, name: str) -> ModelConfig:
|
||||
return self.model_registry.get(name)
|
||||
|
||||
def get_provider_config(self, name: str) -> ProviderConfig:
|
||||
return self.provider_registry.get(name)
|
||||
|
||||
# ── 适配器解析 ────────────────────────────────────────────
|
||||
|
||||
def get_adapter(self, model_name: str) -> OpenAIFormatAdapter | AnthropicFormatAdapter:
|
||||
if model_name in self._adapter_cache:
|
||||
return self._adapter_cache[model_name]
|
||||
|
||||
model = self.model_registry.get(model_name)
|
||||
provider = self.provider_registry.get(model.provider)
|
||||
|
||||
if model.format.value == "openai":
|
||||
adapter = OpenAIFormatAdapterFactory.create(provider)
|
||||
else:
|
||||
adapter = AnthropicFormatAdapterFactory.create(provider)
|
||||
|
||||
self._adapter_cache[model_name] = adapter
|
||||
return adapter
|
||||
|
||||
# ── 消息辅助 ───────────────────────────────────────────────
|
||||
|
||||
@staticmethod
|
||||
def to_message_list(
|
||||
system_messages: list[str],
|
||||
history: list[tuple[str, str]],
|
||||
) -> list[BaseMessage]:
|
||||
messages: list[BaseMessage] = []
|
||||
for msg in system_messages:
|
||||
if msg:
|
||||
messages.append(SystemMessage(content=msg))
|
||||
for role, content in history:
|
||||
if role == "system":
|
||||
messages.append(SystemMessage(content=content))
|
||||
elif role == "user":
|
||||
messages.append(HumanMessage(content=content))
|
||||
elif role == "assistant":
|
||||
messages.append(AIMessage(content=content))
|
||||
return messages
|
||||
|
||||
# ── Embedding ──────────────────────────────────────────────────────
|
||||
|
||||
def create_embedding(self, model_name: str) -> Any:
|
||||
"""按模型名创建 Embedding 适配器。
|
||||
|
||||
优先读取 model endpoint 级 base_url/api_key 覆盖；无 endpoint 时
|
||||
回退 provider 级凭证。
|
||||
"""
|
||||
from app.ai_provider.embedding.openai_format import OpenAIFormatEmbeddingAdapter
|
||||
|
||||
model = self.model_registry.get(model_name)
|
||||
provider = self.provider_registry.get(model.provider)
|
||||
|
||||
# 尝试从 model endpoint 读取 base_url / api_key 覆盖
|
||||
base_url: str | None = None
|
||||
api_key: str = ""
|
||||
if model.endpoints and len(model.endpoints) > 0:
|
||||
ep = model.endpoints[0] # 取第一个活跃 endpoint
|
||||
if ep.base_url:
|
||||
base_url = ep.base_url
|
||||
elif ep.base_url_env and hasattr(settings, ep.base_url_env.lower()):
|
||||
base_url = getattr(settings, ep.base_url_env.lower())
|
||||
if ep.api_key:
|
||||
api_key = ep.api_key
|
||||
elif ep.api_key_env and hasattr(settings, ep.api_key_env.lower()):
|
||||
api_key = getattr(settings, ep.api_key_env.lower(), "")
|
||||
|
||||
if not base_url:
|
||||
_, base_url = self._resolve_credentials(provider)
|
||||
if not base_url:
|
||||
base_url = provider.base_url
|
||||
if not api_key:
|
||||
api_key, _ = self._resolve_credentials(provider)
|
||||
|
||||
return OpenAIFormatEmbeddingAdapter(
|
||||
base_url=base_url or "",
|
||||
api_key=api_key or "",
|
||||
model=model.model,
|
||||
request_timeout_seconds=model.timeout_seconds,
|
||||
)
|
||||
|
||||
# ── Reranker ───────────────────────────────────────────────────────
|
||||
|
||||
def create_reranker(self, model_name: str) -> Any:
|
||||
"""按模型名创建 Reranker 适配器。"""
|
||||
from app.ai_provider.reranker.cohere_format import CohereFormatRerankerAdapter
|
||||
|
||||
model = self.model_registry.get(model_name)
|
||||
provider = self.provider_registry.get(model.provider)
|
||||
api_key, base_url = self._resolve_credentials(provider)
|
||||
if not base_url:
|
||||
base_url = provider.base_url
|
||||
return CohereFormatRerankerAdapter(
|
||||
base_url=base_url or "",
|
||||
api_key=api_key,
|
||||
model=model.model,
|
||||
request_timeout_seconds=model.timeout_seconds,
|
||||
)
|
||||
|
||||
# ── 对话 ──────────────────────────────────────────────────────────
|
||||
|
||||
def _resolve_endpoint(self, model_name: str) -> _ResolvedEndpoint:
|
||||
"""解析模型名→(适配器, 凭证, 熔断器)。
|
||||
|
||||
处理池选择、熔断初始化、凭证回退。
|
||||
"""
|
||||
model = self.model_registry.get(model_name)
|
||||
provider = self.provider_registry.get(model.provider)
|
||||
adapter = self.get_adapter(model_name)
|
||||
|
||||
breaker: CircuitBreaker | None = None
|
||||
|
||||
if model.endpoint_pool_enabled and model.endpoints:
|
||||
from app.ai_provider.llm.pool import PoolFactory
|
||||
|
||||
pool = PoolFactory.create_pool(
|
||||
logical_model_name=model_name,
|
||||
endpoints=[ep.model_dump(exclude_none=True) for ep in model.endpoints],
|
||||
strategy_name=model.routing_strategy,
|
||||
)
|
||||
selected = pool.select(set())
|
||||
if selected is not None:
|
||||
api_key = selected.get("api_key") or (
|
||||
settings.__getattribute__(provider.api_key_env.lower())
|
||||
if hasattr(settings, provider.api_key_env.lower())
|
||||
else ""
|
||||
)
|
||||
base_url_env = selected.get("base_url_env") or provider.base_url_env
|
||||
base_url = selected.get("base_url") or (
|
||||
settings.__getattribute__(base_url_env.lower())
|
||||
if hasattr(settings, base_url_env.lower())
|
||||
else None
|
||||
)
|
||||
# 熔断器
|
||||
from app.ai_provider.core.resilience import ResilienceFactory
|
||||
|
||||
breaker = ResilienceFactory.create_circuit_breaker(
|
||||
name=f"{model_name}@{selected.get('name', 'default')}",
|
||||
failure_threshold=3,
|
||||
cooldown_seconds=30,
|
||||
)
|
||||
if not breaker.allow_request():
|
||||
raise AIProviderError(
|
||||
message=f"Circuit breaker open for {selected.get('name', 'unknown')}",
|
||||
code="CIRCUIT_BREAKER_OPEN",
|
||||
status_code=503,
|
||||
)
|
||||
return _ResolvedEndpoint(
|
||||
adapter=adapter,
|
||||
model=model,
|
||||
provider=provider,
|
||||
api_key=api_key,
|
||||
base_url=base_url,
|
||||
breaker=breaker,
|
||||
)
|
||||
# 无活跃 endpoint，回退 provider 凭证（无熔断）
|
||||
api_key, base_url = self._resolve_credentials(provider)
|
||||
return _ResolvedEndpoint(
|
||||
adapter=adapter,
|
||||
model=model,
|
||||
provider=provider,
|
||||
api_key=api_key,
|
||||
base_url=base_url,
|
||||
)
|
||||
|
||||
# 无池，直接走 provider 凭证
|
||||
api_key, base_url = self._resolve_credentials(provider)
|
||||
return _ResolvedEndpoint(
|
||||
adapter=adapter,
|
||||
model=model,
|
||||
provider=provider,
|
||||
api_key=api_key,
|
||||
base_url=base_url,
|
||||
)
|
||||
|
||||
async def chat(
|
||||
self,
|
||||
messages: list[BaseMessage],
|
||||
model_hint: str | None = None,
|
||||
params: dict[str, Any] | None = None,
|
||||
tools: list[dict[str, Any]] | None = None,
|
||||
tool_choice: str | dict[str, Any] | None = None,
|
||||
) -> ChatResult:
|
||||
model_name = self.resolve_model(model_hint)
|
||||
if self._budget and not self._budget.allow_request():
|
||||
raise AIProviderError(
|
||||
message="Budget limit exceeded", code="BUDGET_EXCEEDED", status_code=429
|
||||
)
|
||||
|
||||
ep = self._resolve_endpoint(model_name)
|
||||
|
||||
try:
|
||||
result = await ep.adapter.chat(
|
||||
model=ep.model,
|
||||
messages=messages,
|
||||
api_key=ep.api_key,
|
||||
base_url=ep.base_url,
|
||||
params=params,
|
||||
tools=tools,
|
||||
tool_choice=tool_choice,
|
||||
)
|
||||
if ep.breaker:
|
||||
ep.breaker.record_success()
|
||||
except Exception:
|
||||
if ep.breaker:
|
||||
ep.breaker.record_failure()
|
||||
raise
|
||||
|
||||
self._record_usage(
|
||||
model_name, ep.provider.name, result.usage, self._pricing_registry.get(model_name)
|
||||
)
|
||||
return result
|
||||
|
||||
async def chat_stream(
|
||||
self,
|
||||
messages: list[BaseMessage],
|
||||
model_hint: str | None = None,
|
||||
params: dict[str, Any] | None = None,
|
||||
tools: list[dict[str, Any]] | None = None,
|
||||
tool_choice: str | dict[str, Any] | None = None,
|
||||
):
|
||||
model_name = self.resolve_model(model_hint)
|
||||
if self._budget and not self._budget.allow_request():
|
||||
raise AIProviderError(
|
||||
message="Budget limit exceeded", code="BUDGET_EXCEEDED", status_code=429
|
||||
)
|
||||
|
||||
ep = self._resolve_endpoint(model_name)
|
||||
|
||||
try:
|
||||
stream = ep.adapter.chat_stream(
|
||||
model=ep.model,
|
||||
messages=messages,
|
||||
api_key=ep.api_key,
|
||||
base_url=ep.base_url,
|
||||
params=params,
|
||||
tools=tools,
|
||||
tool_choice=tool_choice,
|
||||
)
|
||||
handler = StreamHandler()
|
||||
async for chunk in stream:
|
||||
handler.add_chunk(chunk)
|
||||
yield chunk
|
||||
|
||||
if ep.breaker:
|
||||
ep.breaker.record_success()
|
||||
except Exception:
|
||||
if ep.breaker:
|
||||
ep.breaker.record_failure()
|
||||
raise
|
||||
|
||||
result = handler.to_result(model_name=model_name)
|
||||
self._record_usage(
|
||||
model_name, ep.provider.name, result.usage, self._pricing_registry.get(model_name)
|
||||
)
|
||||
|
||||
def _resolve_credentials(self, provider: ProviderConfig) -> tuple[str, str | None]:
|
||||
"""从 settings 解析 provider 级别的 API key 和 base URL。"""
|
||||
api_key = (
|
||||
settings.__getattribute__(provider.api_key_env.lower())
|
||||
if hasattr(settings, provider.api_key_env.lower())
|
||||
else ""
|
||||
)
|
||||
base_url = (
|
||||
settings.__getattribute__(provider.base_url_env.lower())
|
||||
if hasattr(settings, provider.base_url_env.lower())
|
||||
else None
|
||||
)
|
||||
return api_key, base_url
|
||||
|
||||
# ── 路由辅助 ───────────────────────────────────────────────
|
||||
|
||||
def select_model_for_context(
|
||||
self,
|
||||
needs_vision: bool = False,
|
||||
needs_thinking: bool = False,
|
||||
) -> str:
|
||||
capabilities: set[str] = set()
|
||||
if needs_vision:
|
||||
capabilities.add("vision")
|
||||
if needs_thinking:
|
||||
capabilities.add("thinking")
|
||||
model = self._model_selector.select({"capabilities": capabilities})
|
||||
return model.name
|
||||
|
||||
def get_adapter_for_tool_binding(
|
||||
self,
|
||||
model_name: str,
|
||||
tools: list[dict[str, Any]],
|
||||
tool_choice: str | dict[str, Any] | None = None,
|
||||
):
|
||||
"""返回 (adapter, request_body_base) 用于工具绑定。
|
||||
|
||||
为 Agent 运行时保留独立的 bind_tools 步骤。
|
||||
"""
|
||||
adapter = self.get_adapter(model_name)
|
||||
body = {"model": self.model_registry.get(model_name).model}
|
||||
adapter.bind_tools(body, tools, tool_choice=tool_choice)
|
||||
return adapter, body
|
||||
|
||||
# ── 用量与计费 ───────────────────────────────────────────────
|
||||
|
||||
def _record_usage(
|
||||
self,
|
||||
model_name: str,
|
||||
provider: str,
|
||||
usage: TokenUsage | None,
|
||||
pricing: PricingInfo | None = None,
|
||||
) -> None:
|
||||
record = self._billing.record(
|
||||
model_name=model_name,
|
||||
provider=provider,
|
||||
usage=usage,
|
||||
pricing=pricing,
|
||||
)
|
||||
if self._budget and record.cost > 0:
|
||||
self._budget.add_cost_cny(
|
||||
record.cost if record.currency == "CNY" else record.cost * 7.3
|
||||
)
|
||||
|
||||
def get_usage_records(self, **kwargs) -> list:
|
||||
return self._billing.get_records(**kwargs)
|
||||
|
||||
def get_budget_status(self) -> dict[str, bool]:
|
||||
if self._budget:
|
||||
return {"within_daily": self._budget.allow_request()}
|
||||
return {"within_daily": True}
|
||||
|
||||
# ── 生命周期 ─────────────────────────────────────────────────────
|
||||
|
||||
async def close(self) -> None:
|
||||
for adapter in self._adapter_cache.values():
|
||||
await adapter.close()
|
||||
|
||||
async def __aenter__(self) -> ProviderManager:
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *exc_info: object) -> None:
|
||||
await self.close()
|
||||
|
||||
|
||||
class ProviderManagerFactory:
|
||||
"""工厂：从配置文件创建 ProviderManager。"""
|
||||
|
||||
@staticmethod
|
||||
def create(
|
||||
models_path: str | None = None,
|
||||
providers_path: str | None = None,
|
||||
aliases: dict[str, str] | None = None,
|
||||
daily_budget_cny: float = 0.0,
|
||||
) -> ProviderManager:
|
||||
from pathlib import Path
|
||||
|
||||
models = load_models_config(Path(models_path) if models_path else None)
|
||||
providers = load_providers_config(Path(providers_path) if providers_path else None)
|
||||
pricing_data = load_pricing_config()
|
||||
from app.ai_provider.core.registry import RegistryFactory
|
||||
|
||||
registries = RegistryFactory.create(
|
||||
models, providers, aliases=aliases, pricing_data=pricing_data
|
||||
)
|
||||
billing = BillingFactory.create_tracker()
|
||||
budget = (
|
||||
BillingFactory.create_budget_controller(daily_budget_cny=daily_budget_cny)
|
||||
if daily_budget_cny > 0
|
||||
else None
|
||||
)
|
||||
return ProviderManager(
|
||||
model_registry=registries.model_registry,
|
||||
provider_registry=registries.provider_registry,
|
||||
billing_tracker=billing,
|
||||
budget_controller=budget,
|
||||
pricing_registry=registries.pricing_registry,
|
||||
)
|
||||
@@ -1,172 +0,0 @@
|
||||
"""配置驱动的 AI 提供商五阶段参数管道。
|
||||
|
||||
根据 ProviderConfig.parameter_mappings 在分发前对参数进行
|
||||
转换、归一化、过滤和注入。第五阶段 HandlerStep 处理跨参数
|
||||
依赖等复杂场景。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
from app.ai_provider.config.models import ModelConfig, ProviderConfig
|
||||
|
||||
|
||||
@dataclass
|
||||
class PipelineContext:
|
||||
"""管道的输入上下文。"""
|
||||
|
||||
params: dict[str, Any] = field(default_factory=dict)
|
||||
model_config: ModelConfig | None = None
|
||||
provider_config: ProviderConfig | None = None
|
||||
messages: list[dict[str, Any]] | None = None
|
||||
tools: list[dict[str, Any]] | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class PipelineResult:
|
||||
"""经过管道处理后的结果。"""
|
||||
|
||||
params: dict[str, Any] = field(default_factory=dict)
|
||||
filtered_params: dict[str, Any] = field(default_factory=dict)
|
||||
messages: list[dict[str, Any]] | None = None
|
||||
tools: list[dict[str, Any]] | None = None
|
||||
|
||||
|
||||
class PipelineStep:
|
||||
"""管道步骤的基类。"""
|
||||
|
||||
def apply(self, ctx: PipelineContext, result: PipelineResult) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class NormalizeStep(PipelineStep):
|
||||
"""步骤 1：归一化参数范围（夹紧到 min/max）。"""
|
||||
|
||||
def apply(self, ctx: PipelineContext, result: PipelineResult) -> None:
|
||||
mappings = ctx.provider_config.parameter_mappings if ctx.provider_config else {}
|
||||
for key, value in ctx.params.items():
|
||||
mapping = mappings.get(key)
|
||||
if mapping:
|
||||
if mapping.min is not None:
|
||||
value = max(mapping.min, value)
|
||||
if mapping.max is not None:
|
||||
value = min(mapping.max, value)
|
||||
result.params[key] = value
|
||||
else:
|
||||
result.params[key] = value
|
||||
|
||||
|
||||
class FilterStep(PipelineStep):
|
||||
"""步骤 2：过滤模型/提供商不支持的参数。"""
|
||||
|
||||
def apply(self, ctx: PipelineContext, result: PipelineResult) -> None:
|
||||
allowed = (
|
||||
set(ctx.provider_config.parameter_mappings.keys())
|
||||
if ctx.provider_config
|
||||
else set(result.params.keys())
|
||||
)
|
||||
filtered = {}
|
||||
for key in list(result.params.keys()):
|
||||
if key not in allowed:
|
||||
filtered[key] = result.params.pop(key)
|
||||
result.filtered_params = filtered
|
||||
|
||||
|
||||
class TransformStep(PipelineStep):
|
||||
"""步骤 3：转换参数名称和值为提供商所需格式。"""
|
||||
|
||||
def apply(self, ctx: PipelineContext, result: PipelineResult) -> None:
|
||||
mappings = ctx.provider_config.parameter_mappings if ctx.provider_config else {}
|
||||
remapped: dict[str, Any] = {}
|
||||
for key, value in result.params.items():
|
||||
mapping = mappings.get(key)
|
||||
if mapping:
|
||||
provider_key = mapping.provider_param
|
||||
if mapping.transform == "int":
|
||||
value = int(value)
|
||||
elif mapping.transform == "str":
|
||||
value = str(value)
|
||||
elif mapping.transform == "multiply_1000":
|
||||
value = int(value * 1000)
|
||||
remapped[provider_key] = value
|
||||
else:
|
||||
remapped[key] = value
|
||||
result.params = remapped
|
||||
|
||||
|
||||
class InjectStep(PipelineStep):
|
||||
"""步骤 4：为缺失参数注入提供商默认值。"""
|
||||
|
||||
def apply(self, ctx: PipelineContext, result: PipelineResult) -> None:
|
||||
mappings = ctx.provider_config.parameter_mappings if ctx.provider_config else {}
|
||||
for mapping in mappings.values():
|
||||
provider_key = mapping.provider_param
|
||||
if provider_key not in result.params and mapping.default is not None:
|
||||
result.params[provider_key] = mapping.default
|
||||
|
||||
|
||||
class HandlerStep(PipelineStep):
|
||||
"""步骤 5：执行已注册的参数处理器。
|
||||
|
||||
处理 YAML 映射规则无法表达的跨参数依赖和复杂变换：
|
||||
- Anthropic thinking 参数要求 max_tokens > budget_tokens
|
||||
- 根据模型能力过滤不支持的参数
|
||||
- 结构性参数重组
|
||||
"""
|
||||
|
||||
def apply(self, ctx: PipelineContext, result: PipelineResult) -> None:
|
||||
if not ctx.provider_config:
|
||||
return
|
||||
from app.ai_provider.pipeline_handlers import ParameterHandlerRegistry
|
||||
|
||||
registry = ParameterHandlerRegistry
|
||||
|
||||
# 阶段 A：逐参数处理器（来自 ParameterMapping.handler）
|
||||
for key, mapping in ctx.provider_config.parameter_mappings.items():
|
||||
if mapping.handler:
|
||||
handler = registry.get(mapping.handler)
|
||||
handler(ctx, result, key) # key is the param key
|
||||
|
||||
# 阶段 B：全局处理器（来自 ProviderConfig.parameter_handlers）
|
||||
for handler_name in ctx.provider_config.parameter_handlers:
|
||||
handler = registry.get(handler_name)
|
||||
handler(ctx, result, None) # global handler: no param_key
|
||||
|
||||
|
||||
class ParameterPipeline:
|
||||
"""配置驱动的五阶段参数管道。"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
steps: list[PipelineStep],
|
||||
provider_config: ProviderConfig | None = None,
|
||||
) -> None:
|
||||
self._steps = steps
|
||||
self._provider_config = provider_config
|
||||
|
||||
def run(self, ctx: PipelineContext) -> PipelineResult:
|
||||
if ctx.provider_config is None and self._provider_config is not None:
|
||||
ctx.provider_config = self._provider_config
|
||||
result = PipelineResult(params=dict(ctx.params))
|
||||
for step in self._steps:
|
||||
step.apply(ctx, result)
|
||||
return result
|
||||
|
||||
|
||||
class PipelineFactory:
|
||||
"""工厂：从提供商配置构建管道。"""
|
||||
|
||||
@staticmethod
|
||||
def create(provider_config: ProviderConfig) -> ParameterPipeline:
|
||||
return ParameterPipeline(
|
||||
provider_config=provider_config,
|
||||
steps=[
|
||||
NormalizeStep(),
|
||||
FilterStep(),
|
||||
TransformStep(),
|
||||
InjectStep(),
|
||||
HandlerStep(),
|
||||
],
|
||||
)
|
||||
@@ -1,153 +0,0 @@
|
||||
"""参数处理器注册中心与内置 handlers。
|
||||
|
||||
处理跨参数依赖、结构性转换等复杂场景——这些无法通过
|
||||
YAML ParameterMapping 的 min/max/transform/default 表达。
|
||||
|
||||
使用方式（在 providers.yaml 中）：
|
||||
```yaml
|
||||
parameter_mappings:
|
||||
thinking: { provider_param: thinking, handler: anthropic_thinking }
|
||||
|
||||
parameter_handlers:
|
||||
- strip_unsupported_capability_params
|
||||
```
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
|
||||
from app.ai_provider.pipeline import PipelineContext, PipelineResult
|
||||
|
||||
# Type alias for a parameter handler function.
|
||||
# param_key is passed for per-param handlers (from ParameterMapping.handler);
|
||||
# for global handlers (from ProviderConfig.parameter_handlers), param_key is None.
|
||||
ParameterHandler = Callable[[PipelineContext, PipelineResult, str | None], None]
|
||||
|
||||
|
||||
class ParameterHandlerRegistry:
|
||||
"""全局参数处理器注册中心。
|
||||
|
||||
使用 register() 装饰器注册 handler，名称字符串即可在 YAML 中引用。
|
||||
"""
|
||||
|
||||
_handlers: dict[str, ParameterHandler] = {}
|
||||
|
||||
@classmethod
|
||||
def register(cls, name: str) -> Callable[[ParameterHandler], ParameterHandler]:
|
||||
"""装饰器：将函数注册为命名 handler。"""
|
||||
|
||||
def decorator(handler: ParameterHandler) -> ParameterHandler:
|
||||
cls._handlers[name] = handler
|
||||
return handler
|
||||
|
||||
return decorator
|
||||
|
||||
@classmethod
|
||||
def get(cls, name: str) -> ParameterHandler:
|
||||
"""按名称获取 handler。未找到时抛出 KeyError。"""
|
||||
if name not in cls._handlers:
|
||||
msg = f"Unknown parameter handler: '{name}'. "
|
||||
msg += f"Available: {', '.join(sorted(cls._handlers))}"
|
||||
raise KeyError(msg)
|
||||
return cls._handlers[name]
|
||||
|
||||
@classmethod
|
||||
def list(cls) -> list[str]:
|
||||
"""返回所有已注册的 handler 名称。"""
|
||||
return list(cls._handlers.keys())
|
||||
|
||||
|
||||
# ── 内建 Handlers ──────────────────────────────────────────────
|
||||
|
||||
|
||||
@ParameterHandlerRegistry.register("anthropic_thinking")
|
||||
def handle_anthropic_thinking(
|
||||
_ctx: PipelineContext,
|
||||
result: PipelineResult,
|
||||
_param_key: str | None = None,
|
||||
) -> None:
|
||||
"""处理 Anthropic thinking 参数的跨参数约束。
|
||||
|
||||
Anthropic API 要求：
|
||||
- max_tokens 必须大于 thinking.budget_tokens
|
||||
- 开启 thinking 时 temperature 必须为 1（不应传此参数）
|
||||
- thinking 参数自身需转换为 {type, budget_tokens} 结构
|
||||
"""
|
||||
thinking = result.params.get("thinking")
|
||||
if thinking is None:
|
||||
return
|
||||
|
||||
# 将标量 thinking=True 转换为启用 extended thinking
|
||||
if thinking is True:
|
||||
result.params["thinking"] = {"type": "enabled"}
|
||||
|
||||
# 提取 budget_tokens
|
||||
if isinstance(thinking, dict):
|
||||
budget = thinking.get("budget_tokens", 0)
|
||||
# 确保 max_tokens > budget_tokens
|
||||
max_tokens = result.params.get("max_tokens", 0)
|
||||
if budget > 0 and max_tokens <= budget:
|
||||
result.params["max_tokens"] = budget + 1
|
||||
|
||||
# 开启 thinking 时移除 temperature 参数（Anthropic 静默忽略，
|
||||
# 但显式移除更安全干净）
|
||||
result.params.pop("temperature", None)
|
||||
|
||||
|
||||
@ParameterHandlerRegistry.register("validate_max_tokens")
|
||||
def validate_max_tokens(
|
||||
ctx: PipelineContext,
|
||||
result: PipelineResult,
|
||||
_param_key: str | None = None,
|
||||
) -> None:
|
||||
"""校验 max_tokens 不超过模型上限。"""
|
||||
model_config = ctx.model_config
|
||||
if model_config is None:
|
||||
return
|
||||
max_tokens = result.params.get("max_tokens")
|
||||
if max_tokens is None:
|
||||
return
|
||||
limit = model_config.max_output_tokens
|
||||
if limit and max_tokens > limit:
|
||||
result.params["max_tokens"] = limit
|
||||
|
||||
|
||||
@ParameterHandlerRegistry.register("validate_temperature")
|
||||
def validate_temperature(
|
||||
_ctx: PipelineContext,
|
||||
_result: PipelineResult,
|
||||
_param_key: str | None = None,
|
||||
) -> None:
|
||||
"""温度参数跨提供商规范校验。
|
||||
|
||||
非 reasoning/thinking 模式下强制温度 1.0 不是必需的规则，
|
||||
OpenAI/DeepSeek 无此要求，此 handler 留作扩展点。
|
||||
"""
|
||||
# 预留：可为特定 provider 添加温度校验逻辑
|
||||
pass
|
||||
|
||||
|
||||
@ParameterHandlerRegistry.register("strip_unsupported_capability_params")
|
||||
def strip_unsupported_capability_params(
|
||||
ctx: PipelineContext,
|
||||
result: PipelineResult,
|
||||
_param_key: str | None = None,
|
||||
) -> None:
|
||||
"""根据模型能力过滤不支持的参数。
|
||||
|
||||
当模型不支持 vision 时，移除 image_url 等视觉参数。
|
||||
当模型不支持 tool_use 时，移除 tools 相关参数。
|
||||
"""
|
||||
model_config = ctx.model_config
|
||||
if model_config is None:
|
||||
return
|
||||
|
||||
caps = model_config.resolved_capabilities()
|
||||
|
||||
if not caps.get("vision", False):
|
||||
result.params.pop("image_url", None)
|
||||
result.params.pop("image", None)
|
||||
|
||||
if not caps.get("thinking", False):
|
||||
result.params.pop("thinking", None)
|
||||
@@ -1,135 +0,0 @@
|
||||
"""重排序服务 — Reranker 适配器外观。
|
||||
|
||||
用法::
|
||||
|
||||
from app.ai_provider.reranker import RerankerService
|
||||
|
||||
service = RerankerService(manager)
|
||||
result = await service.rerank("查询", ["文档1", "文档2"])
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from app.ai_provider.reranker.base import (
|
||||
RerankerProtocol,
|
||||
RerankerServiceError,
|
||||
RerankRequest,
|
||||
RerankResponse,
|
||||
RerankResult,
|
||||
)
|
||||
from app.ai_provider.types.errors import AIProviderError
|
||||
|
||||
|
||||
class RerankerService:
|
||||
"""外观类，负责解析 Reranker 模型并委托给 ProviderManager。
|
||||
|
||||
接受 ProviderManager 或实现了以下接口的适配器：
|
||||
- ``.model_registry`` → 具有 ``list_by_capability(capability)`` 方法
|
||||
- ``.create_reranker(name)`` → 返回 RerankerProtocol 实例
|
||||
"""
|
||||
|
||||
def __init__(self, manager) -> None:
|
||||
self._manager = manager
|
||||
self._adapter: RerankerProtocol | None = None
|
||||
|
||||
# ── Registry 兼容适配 ─────────────────────────────────────────
|
||||
|
||||
@property
|
||||
def _registry(self):
|
||||
mgr = self._manager
|
||||
if hasattr(mgr, "model_registry"):
|
||||
return mgr.model_registry
|
||||
return mgr
|
||||
|
||||
@property
|
||||
def _factory(self):
|
||||
mgr = self._manager
|
||||
if hasattr(mgr, "create_reranker"):
|
||||
return mgr
|
||||
return None
|
||||
|
||||
# ── Reranker ───────────────────────────────────────────────────
|
||||
|
||||
@property
|
||||
def adapter(self) -> RerankerProtocol:
|
||||
"""解析并缓存 Reranker 适配器。"""
|
||||
if self._adapter is None:
|
||||
self._adapter = self._resolve()
|
||||
return self._adapter
|
||||
|
||||
async def rerank(
|
||||
self,
|
||||
query: str,
|
||||
documents: list[str],
|
||||
model_hint: str | None = None,
|
||||
top_n: int | None = None,
|
||||
) -> RerankResponse:
|
||||
"""快捷方法：构造请求并调用 Reranker 适配器。"""
|
||||
model = self._resolve_model_name(model_hint)
|
||||
request = RerankRequest(
|
||||
query=query,
|
||||
documents=documents,
|
||||
model=model,
|
||||
top_n=top_n,
|
||||
)
|
||||
return await self.adapter.rerank(request)
|
||||
|
||||
# ── 解析 ───────────────────────────────────────────────────────
|
||||
|
||||
def _resolve(self) -> RerankerProtocol:
|
||||
"""查找第一个活跃的 Reranker 模型并通过工厂创建。"""
|
||||
models = self._list_capability_models("reranker")
|
||||
if not models:
|
||||
raise AIProviderError(
|
||||
status_code=503,
|
||||
code="reranker_no_model",
|
||||
message="No reranker model configured in the model registry.",
|
||||
)
|
||||
factory = self._factory
|
||||
if factory is None:
|
||||
raise AIProviderError(
|
||||
status_code=501,
|
||||
code="reranker_no_factory",
|
||||
message="Reranker adapter factory not available (provider manager missing create_reranker).",
|
||||
)
|
||||
return factory.create_reranker(models[0].name)
|
||||
|
||||
def _resolve_model_name(self, model_hint: str | None) -> str:
|
||||
models = self._list_capability_models("reranker")
|
||||
if model_hint:
|
||||
return model_hint
|
||||
if models:
|
||||
return models[0].name
|
||||
return ""
|
||||
|
||||
def _list_capability_models(self, tag: str):
|
||||
reg = self._registry
|
||||
if hasattr(reg, "list_capability_models"):
|
||||
return reg.list_capability_models(tag)
|
||||
if hasattr(reg, "list_by_capability"):
|
||||
return reg.list_by_capability(tag)
|
||||
return []
|
||||
|
||||
# ── 生命周期 ────────────────────────────────────────────────────
|
||||
|
||||
async def close(self) -> None:
|
||||
if self._adapter is not None:
|
||||
close = getattr(self._adapter, "close", None)
|
||||
if close is not None:
|
||||
await close()
|
||||
|
||||
async def __aenter__(self) -> RerankerService:
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *exc_info: object) -> None:
|
||||
await self.close()
|
||||
|
||||
|
||||
__all__ = [
|
||||
"RerankerProtocol",
|
||||
"RerankRequest",
|
||||
"RerankResponse",
|
||||
"RerankResult",
|
||||
"RerankerService",
|
||||
"RerankerServiceError",
|
||||
]
|
||||
@@ -1,147 +0,0 @@
|
||||
"""重排序服务类型 — 请求/响应模型和协议定义。
|
||||
|
||||
Reranker 适配器通过 HTTP 调用重排序端点，经由 ``RerankerService``
|
||||
或直接从 ``ProviderManager`` 消费。
|
||||
|
||||
用法::
|
||||
|
||||
service = RerankerService(manager)
|
||||
result = await service.rerank("查询", ["文档1", "文档2"])
|
||||
|
||||
# 或通过工厂直接调用:
|
||||
adapter = manager.create_reranker("rerank-model")
|
||||
result = await adapter.rerank(RerankRequest(query="...", documents=["..."]))
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
|
||||
|
||||
from app.ai_provider.types.errors import AIProviderError
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from app.ai_provider.config.models import ModelConfig, ModelEndpointConfig
|
||||
|
||||
|
||||
class RerankerServiceError(AIProviderError):
|
||||
"""Reranker API 调用失败的基础错误。"""
|
||||
|
||||
|
||||
# ── 请求 / 响应 ────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@dataclass
|
||||
class RerankRequest:
|
||||
"""重排序请求。
|
||||
|
||||
属性:
|
||||
query: 查询文本。
|
||||
documents: 待重排序的文档列表。
|
||||
model: 模型标识符。
|
||||
top_n: 返回的最相关文档数量（nil 表示全部）。
|
||||
max_chunks_per_doc: 每个文档最多切分为多少块。
|
||||
return_documents: 返回结果中是否包含文档文本。
|
||||
"""
|
||||
|
||||
query: str
|
||||
documents: list[str]
|
||||
model: str = ""
|
||||
top_n: int | None = None
|
||||
max_chunks_per_doc: int | None = None
|
||||
return_documents: bool = False
|
||||
|
||||
|
||||
@dataclass
|
||||
class RerankResult:
|
||||
"""单条重排序结果。"""
|
||||
|
||||
index: int
|
||||
relevance_score: float
|
||||
text: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class RerankResponse:
|
||||
"""重排序 API 的完整响应。"""
|
||||
|
||||
results: list[RerankResult]
|
||||
model: str = ""
|
||||
usage: dict[str, int] | None = None
|
||||
|
||||
|
||||
# ── Protocol ────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class RerankerProtocol(Protocol):
|
||||
"""Reranker 适配器的协议。"""
|
||||
|
||||
async def rerank(self, request: RerankRequest) -> RerankResponse:
|
||||
"""对文档进行重排序。
|
||||
|
||||
异常:
|
||||
RerankerServiceError: API 调用失败或响应异常时抛出。
|
||||
"""
|
||||
...
|
||||
|
||||
async def close(self) -> None:
|
||||
"""释放适配器持有的 HTTP 连接。"""
|
||||
...
|
||||
|
||||
|
||||
# ── 共享工厂基类 ──────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class RerankerAdapterBase:
|
||||
"""Reranker 适配器基类，提供共享的工厂方法。
|
||||
|
||||
子类应在 ``__init__`` 中接受 ``base_url``、``api_key``、``model``、
|
||||
``request_timeout_seconds`` 参数。
|
||||
"""
|
||||
|
||||
_adapter_type: str = "reranker"
|
||||
_default_timeout: int = 60
|
||||
|
||||
@classmethod
|
||||
def from_endpoint(
|
||||
cls,
|
||||
endpoint: ModelEndpointConfig,
|
||||
*,
|
||||
timeout_seconds: int | None = None,
|
||||
) -> Any:
|
||||
"""从已解析的端点配置构建适配器。"""
|
||||
base_url = endpoint.base_url
|
||||
api_key = endpoint.api_key
|
||||
if not base_url:
|
||||
raise RerankerServiceError(
|
||||
status_code=500,
|
||||
code="reranker_endpoint_invalid",
|
||||
message="Reranker endpoint has no resolved base_url.",
|
||||
details={"endpoint": endpoint.name},
|
||||
)
|
||||
return cls(
|
||||
base_url=base_url,
|
||||
api_key=api_key,
|
||||
model=endpoint.model,
|
||||
request_timeout_seconds=timeout_seconds
|
||||
or endpoint.request_timeout_seconds
|
||||
or cls._default_timeout,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_model_config(
|
||||
cls,
|
||||
config: ModelConfig,
|
||||
*,
|
||||
timeout_seconds: int | None = None,
|
||||
) -> Any:
|
||||
"""从逻辑模型配置构建适配器（使用第一个活跃端点）。"""
|
||||
for ep in config.endpoints:
|
||||
if ep.status == "active":
|
||||
return cls.from_endpoint(ep, timeout_seconds=timeout_seconds)
|
||||
raise RerankerServiceError(
|
||||
status_code=500,
|
||||
code="reranker_no_endpoint",
|
||||
message=f"Reranker model '{config.name}' has no active endpoint.",
|
||||
)
|
||||
@@ -1,186 +0,0 @@
|
||||
"""Cohere 兼容格式的 Reranker 适配器 — ``POST /v1/rerank``。
|
||||
|
||||
覆盖 Cohere、mixedbread、Jina AI 等采用 Cohere 兼容重排序端点的提供商。
|
||||
|
||||
请求::
|
||||
|
||||
POST {base_url}/v1/rerank
|
||||
Authorization: Bearer {api_key}
|
||||
Content-Type: application/json
|
||||
|
||||
{
|
||||
"model": "rerank-model",
|
||||
"query": "查询文本",
|
||||
"documents": ["文档1", "文档2"],
|
||||
"top_n": 3,
|
||||
"return_documents": false
|
||||
}
|
||||
|
||||
响应::
|
||||
|
||||
{
|
||||
"results": [
|
||||
{"index": 0, "relevance_score": 0.99, "document": {"text": "文档..."}},
|
||||
{"index": 1, "relevance_score": 0.85, "document": {"text": "文档..."}}
|
||||
],
|
||||
"meta": {"billed_units": {"search_units": 2}}
|
||||
}
|
||||
|
||||
另见:
|
||||
|
||||
- Cohere Rerank API: https://docs.cohere.com/reference/rerank
|
||||
- Jina Reranker: https://api.jina.ai/redoc#tag/rerank
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from app.ai_provider.reranker.base import (
|
||||
RerankerAdapterBase,
|
||||
RerankerProtocol,
|
||||
RerankerServiceError,
|
||||
RerankRequest,
|
||||
RerankResponse,
|
||||
RerankResult,
|
||||
)
|
||||
from app.core.logging import logger
|
||||
|
||||
|
||||
class CohereFormatRerankerAdapter(RerankerAdapterBase, RerankerProtocol):
|
||||
"""Cohere 兼容重排序适配器，覆盖多数重排序服务提供商。"""
|
||||
|
||||
_adapter_type = "reranker"
|
||||
_default_timeout = 60
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
model: str = "rerank-model",
|
||||
request_timeout_seconds: int = 60,
|
||||
) -> None:
|
||||
self._base_url = base_url.rstrip("/")
|
||||
self._api_key = api_key
|
||||
self._model = model
|
||||
self._timeout = request_timeout_seconds
|
||||
self._client: httpx.AsyncClient | None = None
|
||||
|
||||
# ── 客户端 ───────────────────────────────────────────────────────────
|
||||
|
||||
@property
|
||||
def _http(self) -> httpx.AsyncClient:
|
||||
if self._client is None or self._client.is_closed:
|
||||
self._client = httpx.AsyncClient(timeout=httpx.Timeout(self._timeout))
|
||||
return self._client
|
||||
|
||||
async def close(self) -> None:
|
||||
if self._client is not None and not self._client.is_closed:
|
||||
await self._client.aclose()
|
||||
|
||||
# ── RerankerProtocol ────────────────────────────────────────────────
|
||||
|
||||
async def rerank(self, request: RerankRequest) -> RerankResponse:
|
||||
"""通过 ``POST /v1/rerank`` 对文档进行重排序。"""
|
||||
url = f"{self._base_url}/v1/rerank"
|
||||
|
||||
payload: dict[str, Any] = {
|
||||
"model": request.model or self._model,
|
||||
"query": request.query,
|
||||
"documents": request.documents,
|
||||
"return_documents": request.return_documents,
|
||||
}
|
||||
if request.top_n is not None:
|
||||
payload["top_n"] = request.top_n
|
||||
if request.max_chunks_per_doc is not None:
|
||||
payload["max_chunks_per_doc"] = request.max_chunks_per_doc
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {self._api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
logger.info(
|
||||
"reranker.request",
|
||||
model=payload["model"],
|
||||
query_length=len(request.query),
|
||||
document_count=len(request.documents),
|
||||
top_n=request.top_n,
|
||||
url=url,
|
||||
)
|
||||
|
||||
try:
|
||||
response = await self._http.post(url, json=payload, headers=headers)
|
||||
except httpx.TimeoutException as exc:
|
||||
raise RerankerServiceError(
|
||||
status_code=504,
|
||||
code="reranker_timeout",
|
||||
message="Reranker request timed out.",
|
||||
) from exc
|
||||
except httpx.RequestError as exc:
|
||||
raise RerankerServiceError(
|
||||
status_code=503,
|
||||
code="reranker_connection_error",
|
||||
message=f"Reranker connection error: {exc}",
|
||||
) from exc
|
||||
|
||||
if response.status_code != 200:
|
||||
body = response.text[:500]
|
||||
raise RerankerServiceError(
|
||||
status_code=response.status_code,
|
||||
code="reranker_api_error",
|
||||
message=f"Reranker API returned {response.status_code}: {body}",
|
||||
)
|
||||
|
||||
result = self._parse_response(response.json())
|
||||
|
||||
logger.info(
|
||||
"reranker.success",
|
||||
model=payload["model"],
|
||||
total_results=len(result.results),
|
||||
)
|
||||
return result
|
||||
|
||||
# ── 辅助方法 ──────────────────────────────────────────────────────────
|
||||
|
||||
@staticmethod
|
||||
def _parse_response(body: dict) -> RerankResponse:
|
||||
"""从 Cohere 格式的重排序响应中提取结果。"""
|
||||
try:
|
||||
raw_results = body.get("results", [])
|
||||
results = []
|
||||
for item in raw_results:
|
||||
results.append(
|
||||
RerankResult(
|
||||
index=item.get("index", 0),
|
||||
relevance_score=float(item.get("relevance_score", 0.0)),
|
||||
text=item.get("document", {}).get("text")
|
||||
if isinstance(item.get("document"), dict)
|
||||
else None,
|
||||
)
|
||||
)
|
||||
# 按相关性分数降序排列
|
||||
results.sort(key=lambda r: r.relevance_score, reverse=True)
|
||||
|
||||
# 提取用量信息（不同提供商格式各异）
|
||||
usage: dict[str, int] | None = None
|
||||
meta = body.get("meta")
|
||||
if isinstance(meta, dict):
|
||||
billed = meta.get("billed_units")
|
||||
if isinstance(billed, dict):
|
||||
usage = {k: int(v) for k, v in billed.items() if isinstance(v, (int, float))}
|
||||
|
||||
return RerankResponse(
|
||||
results=results,
|
||||
model=body.get("model") or body.get("id", ""),
|
||||
usage=usage,
|
||||
)
|
||||
except Exception as exc:
|
||||
raise RerankerServiceError(
|
||||
status_code=500,
|
||||
code="reranker_response_parse_error",
|
||||
message=f"Failed to parse reranker response: {exc}",
|
||||
) from exc
|
||||
@@ -1,14 +0,0 @@
|
||||
from app.ai_provider.routing.multimodal import MultiModalDecision, MultiModalRouter
|
||||
from app.ai_provider.routing.permissions import (
|
||||
CapabilityPermissionController,
|
||||
PermissionRule,
|
||||
)
|
||||
from app.ai_provider.routing.selector import ModelSelector
|
||||
|
||||
__all__ = [
|
||||
"MultiModalDecision",
|
||||
"MultiModalRouter",
|
||||
"ModelSelector",
|
||||
"PermissionRule",
|
||||
"CapabilityPermissionController",
|
||||
]
|
||||
@@ -1,81 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from enum import StrEnum
|
||||
from typing import Any
|
||||
|
||||
|
||||
class MultiModalDecision(StrEnum):
|
||||
"""基于消息内容分析的选路决策。
|
||||
|
||||
每个值代表所选模型必须满足的一项能力需求。
|
||||
PLAIN_TEXT 表示不需要特殊能力。
|
||||
"""
|
||||
|
||||
PLAIN_TEXT = "plain_text"
|
||||
VISION = "vision"
|
||||
THINKING = "thinking"
|
||||
TOOL_USE = "tool_use"
|
||||
STRUCTURED_OUTPUT = "structured_output"
|
||||
|
||||
|
||||
class MultiModalRouter:
|
||||
"""分析消息列表以确定所需的模型能力。
|
||||
|
||||
用法::
|
||||
|
||||
messages = [{"role": "user", "content": "Hello"}]
|
||||
decisions = MultiModalRouter.decide_from_messages(messages)
|
||||
# -> {MultiModalDecision.PLAIN_TEXT}
|
||||
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "What's in this image?"},
|
||||
{"type": "image_url", "image_url": {"url": "..."}},
|
||||
],
|
||||
},
|
||||
]
|
||||
decisions = MultiModalRouter.decide_from_messages(messages)
|
||||
# -> {MultiModalDecision.VISION}
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def decide_from_messages(messages: list[dict[str, Any]]) -> set[MultiModalDecision]:
|
||||
"""分析消息并返回所需能力的集合。
|
||||
|
||||
参数:
|
||||
messages: 消息字典列表，每个字典至少包含 ``role`` 和 ``content`` 键（兼容 OpenAI 格式）。
|
||||
|
||||
返回:
|
||||
:class:`MultiModalDecision` 值的集合。如果未检测到特殊能力，
|
||||
返回 ``{PLAIN_TEXT}``。
|
||||
"""
|
||||
if not messages:
|
||||
return {MultiModalDecision.PLAIN_TEXT}
|
||||
|
||||
decisions: set[MultiModalDecision] = set()
|
||||
|
||||
for msg in messages:
|
||||
role = msg.get("role", "")
|
||||
content = msg.get("content", "")
|
||||
|
||||
# --- VISION: content is a list of blocks with image type ---
|
||||
if isinstance(content, list):
|
||||
for block in content:
|
||||
if isinstance(block, dict):
|
||||
block_type = block.get("type", "")
|
||||
if block_type in ("image_url", "image"):
|
||||
decisions.add(MultiModalDecision.VISION)
|
||||
|
||||
# --- TOOL_USE: assistant sent tool_calls, or message is tool result ---
|
||||
if msg.get("tool_calls"):
|
||||
decisions.add(MultiModalDecision.TOOL_USE)
|
||||
if role == "tool":
|
||||
decisions.add(MultiModalDecision.TOOL_USE)
|
||||
|
||||
# --- THINKING: assistant included reasoning ---
|
||||
if msg.get("reasoning"):
|
||||
decisions.add(MultiModalDecision.THINKING)
|
||||
|
||||
return decisions if decisions else {MultiModalDecision.PLAIN_TEXT}
|
||||
@@ -1,100 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import fnmatch
|
||||
from typing import Any
|
||||
|
||||
|
||||
class PermissionRule:
|
||||
"""模型能力访问控制的单条权限规则。
|
||||
|
||||
规则使用 ``fnmatch`` 风格的通配符模式匹配模型名称。
|
||||
|
||||
参数:
|
||||
model_pattern: ``fnmatch`` 模式，用于匹配模型名（如 ``"deepseek-*"``）。
|
||||
capability: 此规则管控的能力（如 ``"vision"``）。
|
||||
allowed: 能力是否允许（``True`` 允许，``False`` 拒绝）。
|
||||
|
||||
示例::
|
||||
|
||||
# 拒绝 "basic-*" 模型上的 vision 能力
|
||||
PermissionRule(model_pattern="basic-*", capability="vision", allowed=False)
|
||||
|
||||
# 允许所有模型的 thinking 能力
|
||||
PermissionRule(model_pattern="*", capability="thinking", allowed=True)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_pattern: str,
|
||||
capability: str,
|
||||
allowed: bool = True,
|
||||
) -> None:
|
||||
self.model_pattern = model_pattern
|
||||
self.capability = capability
|
||||
self.allowed = allowed
|
||||
|
||||
def match(self, model_name: str) -> bool:
|
||||
"""检查此规则是否适用于给定的模型名。"""
|
||||
return fnmatch.fnmatch(model_name, self.model_pattern)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
verb = "ALLOW" if self.allowed else "DENY"
|
||||
return f"PermissionRule({verb} {self.capability} on {self.model_pattern!r})"
|
||||
|
||||
def __eq__(self, other: Any) -> bool:
|
||||
if not isinstance(other, PermissionRule):
|
||||
return NotImplemented
|
||||
return (
|
||||
self.model_pattern == other.model_pattern
|
||||
and self.capability == other.capability
|
||||
and self.allowed == other.allowed
|
||||
)
|
||||
|
||||
def __hash__(self) -> int:
|
||||
return hash((self.model_pattern, self.capability, self.allowed))
|
||||
|
||||
|
||||
class CapabilityPermissionController:
|
||||
"""控制哪些模型允许哪些能力。
|
||||
|
||||
按顺序评估 :class:`PermissionRule` 列表。
|
||||
**最后匹配的规则** 胜出（最具体的规则应放在较宽泛的规则之后）。
|
||||
如果没有规则匹配某个模型/能力对，该能力**默认允许**。
|
||||
|
||||
用法::
|
||||
|
||||
controller = CapabilityPermissionController([
|
||||
PermissionRule(model_pattern="*", capability="vision", allowed=False),
|
||||
PermissionRule(model_pattern="vision-*", capability="vision", allowed=True),
|
||||
])
|
||||
|
||||
controller.check("fast-model", "vision") # False
|
||||
controller.check("vision-model", "vision") # True
|
||||
"""
|
||||
|
||||
def __init__(self, rules: list[PermissionRule] | None = None) -> None:
|
||||
self._rules = list(rules) if rules else []
|
||||
|
||||
def check(self, model_name: str, capability: str) -> bool:
|
||||
"""检查 *capability* 是否在 *model_name* 上被允许。
|
||||
|
||||
返回 ``True`` 表示允许，``False`` 表示拒绝。
|
||||
"""
|
||||
result: bool | None = None
|
||||
for rule in self._rules:
|
||||
if rule.capability == capability and rule.match(model_name):
|
||||
result = rule.allowed
|
||||
# Default: allow if no rule matched
|
||||
return result if result is not None else True
|
||||
|
||||
def is_allowed(self, model_name: str, capability: str) -> bool:
|
||||
""":meth:`check` 的别名。"""
|
||||
return self.check(model_name, capability)
|
||||
|
||||
def filter_capabilities(
|
||||
self,
|
||||
model_name: str,
|
||||
capabilities: set[str],
|
||||
) -> set[str]:
|
||||
"""返回 *capabilities* 中在 *model_name* 上允许的子集。"""
|
||||
return {c for c in capabilities if self.check(model_name, c)}
|
||||
@@ -1,111 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from app.ai_provider.config.models import ModelConfig
|
||||
from app.ai_provider.core.registry import ModelRegistry
|
||||
from app.ai_provider.routing.multimodal import MultiModalDecision
|
||||
from app.ai_provider.types.errors import AIProviderError
|
||||
|
||||
_SELECTOR_CAPABILITY_MAP: dict[MultiModalDecision, str] = {
|
||||
MultiModalDecision.VISION: "vision",
|
||||
MultiModalDecision.THINKING: "thinking",
|
||||
MultiModalDecision.TOOL_USE: "tool_use",
|
||||
MultiModalDecision.STRUCTURED_OUTPUT: "structured_output",
|
||||
MultiModalDecision.PLAIN_TEXT: "chat",
|
||||
}
|
||||
|
||||
|
||||
class ModelSelector:
|
||||
"""为给定上下文选择最佳模型。
|
||||
|
||||
用法::
|
||||
|
||||
selector = ModelSelector(registry)
|
||||
model = selector.select({"capabilities": {"vision", "thinking"}})
|
||||
model = selector.select({"preferred": "fast-model"})
|
||||
model = selector.select_from_decisions({MultiModalDecision.VISION})
|
||||
"""
|
||||
|
||||
def __init__(self, registry: ModelRegistry) -> None:
|
||||
self._registry = registry
|
||||
|
||||
def select(self, context: dict[str, Any]) -> ModelConfig:
|
||||
"""根据给定的上下文选择模型。
|
||||
|
||||
上下文键:
|
||||
|
||||
* ``capabilities`` (``set[str]``, 可选) — 必需的能力。
|
||||
* ``preferred`` (``str``, 可选) — 偏好的模型名或别名。
|
||||
* ``avoid`` (``set[str]``, 可选) — 要排除的模型名。
|
||||
|
||||
如果给了偏好模型但不能满足所有必需的能力，此方法会回退到
|
||||
能满足要求的模型。
|
||||
"""
|
||||
capabilities: set[str] = set(context.get("capabilities", set()))
|
||||
preferred: str | None = context.get("preferred")
|
||||
avoid: set[str] = set(context.get("avoid", set()))
|
||||
|
||||
# If preferred is given, resolve it and check capabilities
|
||||
if preferred is not None:
|
||||
model = self._registry.resolve(preferred)
|
||||
if (
|
||||
not capabilities or self._model_satisfies_all(model, capabilities)
|
||||
) and model.name not in avoid:
|
||||
return model
|
||||
|
||||
# Find models that satisfy all required capabilities
|
||||
if capabilities:
|
||||
return self._select_by_capabilities(capabilities, avoid)
|
||||
|
||||
# Fall back to default
|
||||
return self._registry.resolve()
|
||||
|
||||
def select_from_decisions(
|
||||
self,
|
||||
decisions: set[MultiModalDecision],
|
||||
) -> ModelConfig:
|
||||
"""基于 :class:`MultiModalDecision` 值选择模型。
|
||||
|
||||
将决策翻译为能力字符串并委托给 :meth:`select`。
|
||||
"""
|
||||
capabilities: set[str] = set()
|
||||
for d in decisions:
|
||||
cap = _SELECTOR_CAPABILITY_MAP.get(d)
|
||||
if cap and cap != "chat":
|
||||
capabilities.add(cap)
|
||||
return self.select({"capabilities": capabilities})
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _select_by_capabilities(
|
||||
self,
|
||||
capabilities: set[str],
|
||||
avoid: set[str],
|
||||
) -> ModelConfig:
|
||||
"""找到支持所有必需能力的模型。"""
|
||||
# Narrow down by each capability (ModelConfig is unhashable, use lists)
|
||||
candidates: list[ModelConfig] = list(self._registry.list_all())
|
||||
|
||||
for cap in capabilities:
|
||||
cap_models = list(self._registry.list_by_capability(cap))
|
||||
cap_names = {m.name for m in cap_models}
|
||||
candidates = [m for m in candidates if m.name in cap_names]
|
||||
|
||||
# Remove avoided models
|
||||
candidates = [m for m in candidates if m.name not in avoid]
|
||||
|
||||
if not candidates:
|
||||
raise AIProviderError(
|
||||
message=f"No active model supports all capabilities: {capabilities}",
|
||||
code="CAPABILITY_NOT_FOUND",
|
||||
status_code=404,
|
||||
)
|
||||
|
||||
# Prefer higher-weight models (higher = better)
|
||||
return max(candidates, key=lambda m: m.weight)
|
||||
|
||||
@staticmethod
|
||||
def _model_satisfies_all(model: ModelConfig, capabilities: set[str]) -> bool:
|
||||
caps = model.resolved_capabilities()
|
||||
return all(caps.get(cap, False) for cap in capabilities)
|
||||
@@ -1,21 +0,0 @@
|
||||
from app.ai_provider.types.chat import ChatChunk, ChatResult, TokenUsage
|
||||
from app.ai_provider.types.errors import AIProviderError
|
||||
from app.ai_provider.types.messages import (
|
||||
AIMessage,
|
||||
BaseMessage,
|
||||
HumanMessage,
|
||||
SystemMessage,
|
||||
ToolMessage,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"AIMessage",
|
||||
"BaseMessage",
|
||||
"HumanMessage",
|
||||
"SystemMessage",
|
||||
"ToolMessage",
|
||||
"ChatChunk",
|
||||
"ChatResult",
|
||||
"TokenUsage",
|
||||
"AIProviderError",
|
||||
]
|
||||
@@ -1,41 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
|
||||
@dataclass
|
||||
class TokenUsage:
|
||||
prompt_tokens: int = 0
|
||||
completion_tokens: int = 0
|
||||
total_tokens: int = 0
|
||||
prompt_tokens_cached: int = 0 # 缓存命中的 prompt token 数（按缓存费率计费）
|
||||
|
||||
def __add__(self, other: TokenUsage) -> TokenUsage:
|
||||
return TokenUsage(
|
||||
prompt_tokens=self.prompt_tokens + other.prompt_tokens,
|
||||
completion_tokens=self.completion_tokens + other.completion_tokens,
|
||||
total_tokens=self.total_tokens + other.total_tokens,
|
||||
prompt_tokens_cached=self.prompt_tokens_cached + other.prompt_tokens_cached,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ChatChunk:
|
||||
delta_content: str = ""
|
||||
delta_reasoning: str = ""
|
||||
tool_call_chunks: list[dict[str, Any]] = field(default_factory=list)
|
||||
usage: TokenUsage | None = None
|
||||
finish_reason: str | None = None
|
||||
content_block_index: int | None = None
|
||||
content_block_start: dict[str, Any] | None = None
|
||||
content_block_delta: dict[str, Any] | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class ChatResult:
|
||||
content: str = ""
|
||||
reasoning: str = ""
|
||||
tool_calls: list[dict[str, Any]] = field(default_factory=list)
|
||||
usage: TokenUsage | None = None
|
||||
model_name: str = ""
|
||||
@@ -1,33 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
class AIProviderError(Exception):
|
||||
def __init__(
|
||||
self,
|
||||
message: str,
|
||||
code: str = "UNKNOWN",
|
||||
status_code: int = 500,
|
||||
provider: str | None = None,
|
||||
original: Exception | None = None,
|
||||
) -> None:
|
||||
self.message = message
|
||||
self.code = code
|
||||
self.status_code = status_code
|
||||
self.provider = provider
|
||||
self.original = original
|
||||
super().__init__(message)
|
||||
|
||||
@classmethod
|
||||
def from_provider(
|
||||
cls,
|
||||
message: str,
|
||||
code: str = "PROVIDER_ERROR",
|
||||
provider: str | None = None,
|
||||
original: Exception | None = None,
|
||||
) -> AIProviderError:
|
||||
return cls(
|
||||
message=message, code=code, status_code=502, provider=provider, original=original
|
||||
)
|
||||
|
||||
def __str__(self) -> str:
|
||||
return self.message
|
||||
@@ -1,81 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Literal
|
||||
|
||||
Role = Literal["system", "user", "assistant", "tool"]
|
||||
|
||||
|
||||
class BaseMessage:
|
||||
role: Role
|
||||
content: str
|
||||
|
||||
def __init__(self, *, content: str, role: Role) -> None:
|
||||
self.content = content
|
||||
self.role = role
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {"role": self.role, "content": self.content}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, d: dict[str, Any]) -> BaseMessage:
|
||||
role = d["role"]
|
||||
content = d.get("content", "")
|
||||
if role == "system":
|
||||
return SystemMessage(content=content)
|
||||
if role == "user":
|
||||
return HumanMessage(content=content)
|
||||
if role == "assistant":
|
||||
return AIMessage(
|
||||
content=content,
|
||||
tool_calls=d.get("tool_calls"),
|
||||
reasoning=d.get("reasoning"),
|
||||
raw_content_blocks=d.get("raw_content_blocks"),
|
||||
)
|
||||
if role == "tool":
|
||||
return ToolMessage(content=content, tool_call_id=d.get("tool_call_id", ""))
|
||||
raise ValueError(f"Unknown role: {role}")
|
||||
|
||||
|
||||
class SystemMessage(BaseMessage):
|
||||
def __init__(self, content: str) -> None:
|
||||
super().__init__(content=content, role="system")
|
||||
|
||||
|
||||
class HumanMessage(BaseMessage):
|
||||
def __init__(self, content: str) -> None:
|
||||
super().__init__(content=content, role="user")
|
||||
|
||||
|
||||
class AIMessage(BaseMessage):
|
||||
def __init__(
|
||||
self,
|
||||
content: str,
|
||||
tool_calls: list[dict[str, Any]] | None = None,
|
||||
reasoning: str | None = None,
|
||||
raw_content_blocks: list[dict[str, Any]] | None = None,
|
||||
) -> None:
|
||||
super().__init__(content=content, role="assistant")
|
||||
self.tool_calls = tool_calls or []
|
||||
self.reasoning = reasoning or ""
|
||||
self.raw_content_blocks = raw_content_blocks or []
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
d = super().to_dict()
|
||||
if self.tool_calls:
|
||||
d["tool_calls"] = self.tool_calls
|
||||
if self.reasoning:
|
||||
d["reasoning"] = self.reasoning
|
||||
if self.raw_content_blocks:
|
||||
d["raw_content_blocks"] = self.raw_content_blocks
|
||||
return d
|
||||
|
||||
|
||||
class ToolMessage(BaseMessage):
|
||||
def __init__(self, content: str, tool_call_id: str) -> None:
|
||||
super().__init__(content=content, role="tool")
|
||||
self.tool_call_id = tool_call_id
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
d = super().to_dict()
|
||||
d["tool_call_id"] = self.tool_call_id
|
||||
return d
|
||||
@@ -4,7 +4,6 @@ from fastapi import APIRouter
|
||||
|
||||
from app.api.v1.endpoints import (
|
||||
admin,
|
||||
audio,
|
||||
auth,
|
||||
avatar,
|
||||
character_import,
|
||||
@@ -57,5 +56,4 @@ api_router.include_router(drive.router)
|
||||
api_router.include_router(events.router)
|
||||
api_router.include_router(goals.router)
|
||||
api_router.include_router(admin.router)
|
||||
api_router.include_router(audio.router)
|
||||
api_router.include_router(avatar.router)
|
||||
@@ -11,7 +11,6 @@ import hmac
|
||||
from fastapi import Depends, HTTPException, Request
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.application.agent.confirmation import ConfirmationRegistry
|
||||
from app.application.tool_runtime.runtime import ToolRuntime
|
||||
from app.core.config import settings
|
||||
from app.db.session import get_session
|
||||
@@ -47,13 +46,6 @@ def get_confirmation_store(request: Request) -> ConfirmationStore:
|
||||
return store
|
||||
|
||||
|
||||
def get_confirmation_registry(request: Request) -> ConfirmationRegistry:
|
||||
registry: ConfirmationRegistry = getattr(request.app.state, "confirmation_registry", None)
|
||||
if registry is None:
|
||||
raise _unavailable("确认注册表未就绪")
|
||||
return registry
|
||||
|
||||
|
||||
def get_event_store(request: Request) -> EventStore:
|
||||
store = getattr(request.app.state, "event_store", None)
|
||||
if store is None:
|
||||
|
||||
@@ -1,275 +0,0 @@
|
||||
"""音频相关 API 端点 — TTS（文本转语音）、ASR（语音转文本）等。
|
||||
|
||||
音频端点使用神识的模型注册表查找可用的音频模型，并调用对应音频适配器合成 /
|
||||
转写。ASR 与 TTS 按 ``tags`` 中的 ``asr`` / ``tts`` 能力区分，ASR 内部再按
|
||||
``provider`` 区分协议：``yuanqi`` 走 multipart ``/v1/audio/transcriptions``，
|
||||
其他走 MiMo 风格的 ``/chat/completions`` + ``input_audio``。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from fastapi import APIRouter, File, Form, HTTPException, Response, UploadFile
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from app.ai_provider.audio import AudioService, AudioServiceError
|
||||
from app.ai_provider.manager import ProviderManagerFactory
|
||||
from app.core.logging import logger
|
||||
|
||||
router = APIRouter(tags=["audio"])
|
||||
|
||||
_MIME_MAP = {
|
||||
"wav": "audio/wav",
|
||||
"mp3": "audio/mpeg",
|
||||
"opus": "audio/opus",
|
||||
"flac": "audio/flac",
|
||||
}
|
||||
|
||||
|
||||
class _AudioModelFactory:
|
||||
"""Adapter: wraps ProviderManager to provide ModelFactory interface for AudioService.
|
||||
|
||||
``create_audio(name)`` 按 ``tags`` 判断 ``asr`` / ``tts`` 能力：
|
||||
- ``tts`` → ``HunjingTTSOpenAI``（``/chat/completions`` + ``audio``）
|
||||
- ``asr`` → 按 ``provider`` 分派：
|
||||
- ``yuanqi`` → ``YuanqiASRWhisper``（multipart ``/v1/audio/transcriptions``）
|
||||
- 其他 → ``HunjingASROpenAI``（``/chat/completions`` + ``input_audio``）
|
||||
"""
|
||||
|
||||
def __init__(self, manager) -> None:
|
||||
self._manager = manager
|
||||
|
||||
@property
|
||||
def registry(self):
|
||||
return self._manager.model_registry
|
||||
|
||||
def create_audio(self, name: str):
|
||||
model = self._manager.get_model_config(name)
|
||||
is_asr = "asr" in model.tags
|
||||
|
||||
if is_asr:
|
||||
return self._create_asr(model)
|
||||
return self._create_tts(model)
|
||||
|
||||
def _provider_or_none(self, model):
|
||||
"""安全获取 provider 配置（provider 未配置时返回 None，走纯端点解析）。"""
|
||||
try:
|
||||
return self._manager.get_provider_config(model.provider)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
# ── ASR ─────────────────────────────────────────────────────────────
|
||||
|
||||
def _create_asr(self, model):
|
||||
"""按 provider 选择 ASR 协议适配器。"""
|
||||
# yuanqi 本地引擎走 multipart /v1/audio/transcriptions
|
||||
if model.provider == "yuanqi":
|
||||
from app.ai_provider.audio.yuanqi_asr import YuanqiASRWhisper
|
||||
|
||||
if model.endpoints:
|
||||
return YuanqiASRWhisper.from_model_config(
|
||||
model, provider=self._provider_or_none(model)
|
||||
)
|
||||
|
||||
provider = self._manager.get_provider_config(model.provider)
|
||||
api_key, base_url = self._resolve_credentials(provider)
|
||||
if not base_url:
|
||||
raise AudioServiceError(
|
||||
status_code=500,
|
||||
code="audio_asr_no_base_url",
|
||||
message=f"Provider '{model.provider}' has no base_url configured.",
|
||||
)
|
||||
return YuanqiASRWhisper(
|
||||
base_url=base_url,
|
||||
api_key=api_key or "",
|
||||
model=model.model,
|
||||
request_timeout_seconds=model.timeout_seconds or 300,
|
||||
)
|
||||
|
||||
# MiMo 等走 /chat/completions + input_audio
|
||||
from app.ai_provider.audio.asr import HunjingASROpenAI
|
||||
|
||||
if model.endpoints:
|
||||
return HunjingASROpenAI.from_model_config(
|
||||
model, provider=self._provider_or_none(model)
|
||||
)
|
||||
|
||||
provider = self._manager.get_provider_config(model.provider)
|
||||
api_key, base_url = self._resolve_credentials(provider)
|
||||
if not base_url:
|
||||
raise AudioServiceError(
|
||||
status_code=500,
|
||||
code="audio_asr_no_base_url",
|
||||
message=f"Provider '{model.provider}' has no base_url configured.",
|
||||
)
|
||||
return HunjingASROpenAI(
|
||||
base_url=base_url,
|
||||
api_key=api_key or "",
|
||||
model=model.model,
|
||||
request_timeout_seconds=model.timeout_seconds or 300,
|
||||
)
|
||||
|
||||
# ── TTS ─────────────────────────────────────────────────────────────
|
||||
|
||||
def _create_tts(self, model):
|
||||
from app.ai_provider.audio.tts import HunjingTTSOpenAI
|
||||
|
||||
# 有 endpoints 时交给 from_model_config（按端点的 status/model/超时）
|
||||
if model.endpoints:
|
||||
return HunjingTTSOpenAI.from_model_config(
|
||||
model, provider=self._provider_or_none(model)
|
||||
)
|
||||
|
||||
# 无 endpoint 时回退到 provider 级别凭证
|
||||
provider = self._manager.get_provider_config(model.provider)
|
||||
api_key, base_url = self._resolve_credentials(provider)
|
||||
if not base_url:
|
||||
raise AudioServiceError(
|
||||
status_code=500,
|
||||
code="audio_tts_no_base_url",
|
||||
message=f"Provider '{model.provider}' has no base_url configured.",
|
||||
)
|
||||
return HunjingTTSOpenAI(
|
||||
base_url=base_url,
|
||||
api_key=api_key or "",
|
||||
model=model.model,
|
||||
request_timeout_seconds=model.timeout_seconds or 120,
|
||||
)
|
||||
|
||||
def _resolve_credentials(self, provider):
|
||||
"""从 settings 解析 provider 级别的 API key 和 base URL。"""
|
||||
from app.core.config import settings
|
||||
|
||||
api_key = (
|
||||
settings.__getattribute__(provider.api_key_env.lower())
|
||||
if hasattr(settings, provider.api_key_env.lower())
|
||||
else ""
|
||||
)
|
||||
base_url = (
|
||||
settings.__getattribute__(provider.base_url_env.lower())
|
||||
if hasattr(settings, provider.base_url_env.lower())
|
||||
else None
|
||||
)
|
||||
return api_key, base_url
|
||||
|
||||
|
||||
class TTSRequest(BaseModel):
|
||||
"""TTS 合成请求。"""
|
||||
|
||||
text: str = Field(..., min_length=1, max_length=4096, description="要合成语音的文本。")
|
||||
voice: str = Field("default", description="语音角色标识（取决于模型，如 default 或声纹 ID）。")
|
||||
format: str = Field("wav", description="音频格式：wav / mp3 / opus / flac。")
|
||||
|
||||
|
||||
@router.post("/audio/tts")
|
||||
async def synthesize_speech(body: TTSRequest) -> Response:
|
||||
"""合成文本为语音，返回音频二进制流。
|
||||
|
||||
使用神识注册表中第一个活跃的 TTS 模型（capability: tts）。
|
||||
未配置 TTS 模型时返回 503。
|
||||
"""
|
||||
if not body.text.strip():
|
||||
raise HTTPException(status_code=400, detail="text 不能为空。")
|
||||
|
||||
mime = _MIME_MAP.get(body.format, "audio/wav")
|
||||
|
||||
try:
|
||||
manager = ProviderManagerFactory.create()
|
||||
service = AudioService(_AudioModelFactory(manager))
|
||||
async with service:
|
||||
result = await service.synthesize(
|
||||
text=body.text,
|
||||
voice=body.voice,
|
||||
response_format=body.format,
|
||||
)
|
||||
except AudioServiceError as exc:
|
||||
logger.error(
|
||||
"audio.tts.api_error",
|
||||
voice=body.voice,
|
||||
format=body.format,
|
||||
error_code=exc.code,
|
||||
error_message=exc.message,
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code=exc.status_code or 502,
|
||||
detail=exc.message,
|
||||
) from exc
|
||||
|
||||
logger.info(
|
||||
"audio.tts.success",
|
||||
voice=body.voice,
|
||||
format=body.format,
|
||||
audio_bytes=len(result.audio_data),
|
||||
)
|
||||
|
||||
return Response(
|
||||
content=result.audio_data,
|
||||
media_type=mime,
|
||||
headers={
|
||||
"Content-Disposition": f'inline; filename="speech.{body.format}"',
|
||||
"X-Audio-Format": body.format,
|
||||
"Cache-Control": "public, max-age=3600",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
# ── ASR（语音转文本）───────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TranscriptionResponse(BaseModel):
|
||||
"""OpenAI Whisper 兼容的转写响应。"""
|
||||
|
||||
text: str = Field(description="转写后的文本")
|
||||
|
||||
|
||||
@router.post("/audio/transcriptions", response_model=TranscriptionResponse)
|
||||
async def transcribe_audio(
|
||||
file: UploadFile = File(..., description="音频文件（wav/mp3/m4a/flac/ogg/opus/webm 等）"),
|
||||
language: str | None = Form(default=None, description="语言提示（ISO-639-1：zh/en/ja/ko 等），省略则自动检测"),
|
||||
model: str | None = Form(default=None, description="可选，指定使用的 ASR 模型逻辑名（如 funasr-sensevoice / mimo-v2.5-asr），省略则用注册表中第一个"),
|
||||
) -> TranscriptionResponse:
|
||||
"""转写音频为文本，返回 OpenAI Whisper 兼容的 JSON。
|
||||
|
||||
不指定 ``model`` 时用神识注册表中第一个活跃的 ASR 模型（capability: asr），
|
||||
按其 ``provider`` 自动选择协议适配器：``yuanqi`` 走本地 multipart 端点，
|
||||
其他走 MiMo 风格的 chat completions。未配置 ASR 模型时返回 503。
|
||||
"""
|
||||
audio_bytes = await file.read()
|
||||
if not audio_bytes:
|
||||
raise HTTPException(status_code=400, detail="音频文件为空。")
|
||||
|
||||
fmt = (file.filename or "audio.wav").rsplit(".", 1)[-1].lower() if "." in (file.filename or "") else "wav"
|
||||
|
||||
try:
|
||||
manager = ProviderManagerFactory.create()
|
||||
service = AudioService(_AudioModelFactory(manager))
|
||||
async with service:
|
||||
result = await service.transcribe(
|
||||
audio_data=audio_bytes,
|
||||
format=fmt,
|
||||
language=language,
|
||||
model_name=model,
|
||||
)
|
||||
except AudioServiceError as exc:
|
||||
logger.error(
|
||||
"audio.asr.api_error",
|
||||
filename=file.filename,
|
||||
format=fmt,
|
||||
language=language,
|
||||
error_code=exc.code,
|
||||
error_message=exc.message,
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code=exc.status_code or 502,
|
||||
detail=exc.message,
|
||||
) from exc
|
||||
|
||||
logger.info(
|
||||
"audio.asr.success",
|
||||
filename=file.filename,
|
||||
format=fmt,
|
||||
language=language,
|
||||
text_length=len(result.text),
|
||||
text=result.text,
|
||||
)
|
||||
|
||||
return TranscriptionResponse(text=result.text)
|
||||
@@ -1,32 +1,23 @@
|
||||
"""Run 与 SSE API。"""
|
||||
"""Run 与 SSE API（Run 的启动契约见 kernel_runs.py）。"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from collections.abc import AsyncGenerator
|
||||
from typing import Any
|
||||
|
||||
from fastapi import APIRouter, Header, Query, Request
|
||||
from fastapi import APIRouter, Header, Query
|
||||
from sse_starlette.sse import EventSourceResponse
|
||||
|
||||
from app.application.agent.confirmation import ConfirmationRegistry
|
||||
from app.application.agent.run_registry import AgentRunRegistry
|
||||
from app.core.config import settings
|
||||
from app.core.errors import AppError
|
||||
from app.core.time import to_iso, utc_now
|
||||
from app.db.session import AsyncSessionLocal
|
||||
from app.repositories import runtime as runtime_repo
|
||||
from app.schemas.common import Page
|
||||
from app.schemas.runtime import (
|
||||
AgentControlRequest,
|
||||
AgentControlResponse,
|
||||
ConfirmationResolutionResponse,
|
||||
EventResponse,
|
||||
RunResponse,
|
||||
)
|
||||
from app.schemas.runtime import EventResponse, RunResponse
|
||||
from app.services.event_bus import bus
|
||||
from app.services.runtime import (
|
||||
cancel_run,
|
||||
confirmation_exists_in_run,
|
||||
get_run_response,
|
||||
list_run_event_responses,
|
||||
list_runs_response,
|
||||
@@ -66,107 +57,11 @@ async def list_run_events(
|
||||
return await list_run_event_responses(run_id, limit=limit, offset=offset)
|
||||
|
||||
|
||||
def _get_agent_run_registry(request: Request) -> AgentRunRegistry:
|
||||
registry: AgentRunRegistry | None = getattr(request.app.state, "agent_run_registry", None)
|
||||
if registry is None:
|
||||
raise AppError("service_unavailable", "Agent 运行控制未启用。", status_code=503)
|
||||
return registry
|
||||
|
||||
|
||||
@router.post("/{run_id}/cancel", response_model=RunResponse)
|
||||
async def cancel_run_api(run_id: str, request: Request) -> RunResponse:
|
||||
registry = getattr(request.app.state, "agent_run_registry", None)
|
||||
if registry is not None:
|
||||
await registry.abort(run_id)
|
||||
async def cancel_run_api(run_id: str) -> RunResponse:
|
||||
return await cancel_run(run_id)
|
||||
|
||||
|
||||
@router.post("/{run_id}/steer", response_model=AgentControlResponse)
|
||||
async def steer_run(
|
||||
run_id: str,
|
||||
payload: AgentControlRequest,
|
||||
request: Request,
|
||||
) -> AgentControlResponse:
|
||||
if not _get_agent_run_registry(request).steer(run_id, payload.content):
|
||||
raise AppError("conflict", "Run 当前不在本进程执行，无法 steer。", status_code=409)
|
||||
return AgentControlResponse()
|
||||
|
||||
|
||||
@router.post("/{run_id}/follow-up", response_model=AgentControlResponse)
|
||||
async def follow_up_run(
|
||||
run_id: str,
|
||||
payload: AgentControlRequest,
|
||||
request: Request,
|
||||
) -> AgentControlResponse:
|
||||
if not _get_agent_run_registry(request).follow_up(run_id, payload.content):
|
||||
raise AppError("conflict", "Run 当前不在本进程执行，无法追加 follow-up。", status_code=409)
|
||||
return AgentControlResponse()
|
||||
|
||||
|
||||
def _get_confirmation_registry(request: Request) -> ConfirmationRegistry:
|
||||
"""从 app.state 取 ConfirmationRegistry；未挂载时 503。"""
|
||||
registry: ConfirmationRegistry | None = getattr(
|
||||
request.app.state, "confirmation_registry", None
|
||||
)
|
||||
if registry is None:
|
||||
raise AppError(
|
||||
"service_unavailable",
|
||||
"确认机制未启用（进程内 Future 注册表未挂载）",
|
||||
status_code=503,
|
||||
)
|
||||
return registry
|
||||
|
||||
|
||||
async def _resolve_confirmation(
|
||||
request: Request,
|
||||
run_id: str,
|
||||
confirmation_id: str,
|
||||
approved: bool,
|
||||
) -> ConfirmationResolutionResponse:
|
||||
"""approve/deny 共用逻辑：通过 ConfirmationRegistry 更新确认状态。"""
|
||||
registry = _get_confirmation_registry(request)
|
||||
if await registry.resolve(run_id, approved, confirmation_id=confirmation_id):
|
||||
return ConfirmationResolutionResponse(approved=approved)
|
||||
|
||||
# 记录已处理/不存在——查 events 表区分 410 与 404
|
||||
event_store = getattr(request.app.state, "event_store", None)
|
||||
if event_store is not None and await confirmation_exists_in_run(
|
||||
event_store, run_id, confirmation_id
|
||||
):
|
||||
raise AppError(
|
||||
"confirmation_expired",
|
||||
"确认已失效（服务可能已重启或已处理），请重新发起对话",
|
||||
status_code=410,
|
||||
)
|
||||
raise AppError("confirmation_not_found", "确认项不存在", status_code=404)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/{run_id}/confirmations/{confirmation_id}/approve",
|
||||
response_model=ConfirmationResolutionResponse,
|
||||
)
|
||||
async def approve_confirmation(
|
||||
run_id: str,
|
||||
confirmation_id: str,
|
||||
request: Request,
|
||||
) -> ConfirmationResolutionResponse:
|
||||
"""批准高风险工具执行：触发 Agent 循环恢复（confirmed=True 重试）。"""
|
||||
return await _resolve_confirmation(request, run_id, confirmation_id, approved=True)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/{run_id}/confirmations/{confirmation_id}/deny",
|
||||
response_model=ConfirmationResolutionResponse,
|
||||
)
|
||||
async def deny_confirmation(
|
||||
run_id: str,
|
||||
confirmation_id: str,
|
||||
request: Request,
|
||||
) -> ConfirmationResolutionResponse:
|
||||
"""拒绝高风险工具执行：Agent 循环回填拒绝 ToolMessage。"""
|
||||
return await _resolve_confirmation(request, run_id, confirmation_id, approved=False)
|
||||
|
||||
|
||||
@router.get("/{run_id}/stream")
|
||||
async def stream_run(
|
||||
run_id: str,
|
||||
|
||||
@@ -1,16 +1,20 @@
|
||||
"""安全 API（api-spec §16）。
|
||||
|
||||
audit-events 按人物/枢过滤查询审计时间线；confirmations 的 approve/deny 负责
|
||||
高风险工具确认：通过 ConfirmationRegistry 更新工具确认状态唤醒 Agent 循环。
|
||||
audit-events 按人物/枢过滤查询审计时间线；confirmations 的 approve/deny 落库裁决
|
||||
高风险工具确认（pending → approved/denied），客户端随后以 confirmed=True 重新调用。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query, Response, status
|
||||
from datetime import UTC, datetime
|
||||
|
||||
from app.api.v1.deps import get_confirmation_registry, get_optional_tool_runtime
|
||||
from app.application.agent.confirmation import ConfirmationRegistry
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query, Response, status
|
||||
from sqlalchemy import select, update
|
||||
|
||||
from app.api.v1.deps import get_optional_tool_runtime
|
||||
from app.application.tool_runtime.runtime import ToolRuntime
|
||||
from app.db.session import AsyncSessionLocal
|
||||
from app.models.safety import PendingConfirmationRecord
|
||||
from app.schemas.admin import AuditEntry
|
||||
from app.schemas.common import Page
|
||||
from app.schemas.safety import (
|
||||
@@ -102,47 +106,22 @@ def _cancelled(confirmation_id: str) -> dict[str, str]:
|
||||
|
||||
|
||||
@router.post("/confirmations/{confirmation_id}/approve")
|
||||
async def approve(
|
||||
confirmation_id: str,
|
||||
registry: ConfirmationRegistry = Depends(get_confirmation_registry),
|
||||
) -> dict[str, bool | str]:
|
||||
"""批准工具执行。通过 ConfirmationRegistry 唤醒等待的 Agent 循环。"""
|
||||
from sqlalchemy import select
|
||||
|
||||
from app.db.session import AsyncSessionLocal
|
||||
from app.models.safety import PendingConfirmationRecord
|
||||
|
||||
async with AsyncSessionLocal() as session:
|
||||
result = await session.execute(
|
||||
select(PendingConfirmationRecord).where(PendingConfirmationRecord.id == confirmation_id)
|
||||
)
|
||||
record = result.scalar_one_or_none()
|
||||
if record is None or record.run_id is None:
|
||||
raise _not_found(confirmation_id)
|
||||
|
||||
# 更新 resolution 并唤醒进程内 Event
|
||||
ok = await registry.resolve(
|
||||
record.run_id,
|
||||
approved=True,
|
||||
confirmation_id=confirmation_id,
|
||||
)
|
||||
if not ok:
|
||||
raise _not_found(confirmation_id)
|
||||
|
||||
async def approve(confirmation_id: str) -> dict[str, bool | str]:
|
||||
"""批准工具执行：pending 记录落库为 approved，客户端以 confirmed=True 重新调用。"""
|
||||
record = await _resolve_confirmation_record(confirmation_id, approved=True)
|
||||
return {"status": "resolved", "approved": True, "run_id": record.run_id}
|
||||
|
||||
|
||||
@router.post("/confirmations/{confirmation_id}/deny")
|
||||
async def deny(
|
||||
confirmation_id: str,
|
||||
registry: ConfirmationRegistry = Depends(get_confirmation_registry),
|
||||
) -> dict[str, str]:
|
||||
"""拒绝工具执行。通过 ConfirmationRegistry 通知 Agent 循环。"""
|
||||
from sqlalchemy import select
|
||||
async def deny(confirmation_id: str) -> dict[str, str]:
|
||||
"""拒绝工具执行：pending 记录落库为 denied。"""
|
||||
await _resolve_confirmation_record(confirmation_id, approved=False)
|
||||
return {"status": "cancelled", "confirmation_id": confirmation_id}
|
||||
|
||||
from app.db.session import AsyncSessionLocal
|
||||
from app.models.safety import PendingConfirmationRecord
|
||||
|
||||
async def _resolve_confirmation_record(
|
||||
confirmation_id: str, *, approved: bool
|
||||
) -> PendingConfirmationRecord:
|
||||
async with AsyncSessionLocal() as session:
|
||||
result = await session.execute(
|
||||
select(PendingConfirmationRecord).where(PendingConfirmationRecord.id == confirmation_id)
|
||||
@@ -151,12 +130,20 @@ async def deny(
|
||||
if record is None or record.run_id is None:
|
||||
raise _not_found(confirmation_id)
|
||||
|
||||
ok = await registry.resolve(
|
||||
record.run_id,
|
||||
approved=False,
|
||||
confirmation_id=confirmation_id,
|
||||
)
|
||||
if not ok:
|
||||
raise _not_found(confirmation_id)
|
||||
|
||||
return {"status": "cancelled", "confirmation_id": confirmation_id}
|
||||
async with AsyncSessionLocal() as session:
|
||||
stmt = (
|
||||
update(PendingConfirmationRecord)
|
||||
.where(
|
||||
PendingConfirmationRecord.id == confirmation_id,
|
||||
PendingConfirmationRecord.status == "pending",
|
||||
)
|
||||
.values(
|
||||
status="approved" if approved else "denied",
|
||||
resolved_at=datetime.now(UTC),
|
||||
)
|
||||
)
|
||||
updated = await session.execute(stmt)
|
||||
await session.commit()
|
||||
if updated.rowcount == 0:
|
||||
raise _not_found(confirmation_id)
|
||||
return record
|
||||
@@ -1,60 +1,25 @@
|
||||
"""设置 API（模型网关状态 + 生成偏好）。
|
||||
"""设置 API（生成偏好）。
|
||||
|
||||
对应前端 settingsApi。模型网关为只读总览，生成偏好可读可写；
|
||||
真实密钥不进入响应。业务逻辑在 services/preferences（与系统配置 core/config/settings 区分）。
|
||||
生成偏好可读可写；模型网关状态由神识承接（用户面见神识阁）。
|
||||
业务逻辑在 services/preferences（与系统配置 core/config/settings 区分）。
|
||||
"""
|
||||
|
||||
from fastapi import APIRouter, Depends
|
||||
|
||||
from app.api.v1.deps import get_current_user
|
||||
from app.core.errors import AppError
|
||||
from app.schemas.auth import CurrentUser
|
||||
from app.schemas.preferences import (
|
||||
GenerationPreferencesResponse,
|
||||
ModelEndpointStatus,
|
||||
ModelGatewayStatus,
|
||||
UpdateEndpointOverrideRequest,
|
||||
UpdateGenerationPreferencesRequest,
|
||||
)
|
||||
from app.services.ai_provider_admin import update_endpoint_override
|
||||
from app.services.preferences import (
|
||||
get_generation_preferences,
|
||||
get_provider_gateway,
|
||||
update_generation_preferences,
|
||||
)
|
||||
|
||||
router = APIRouter(prefix="/settings", tags=["settings"])
|
||||
|
||||
|
||||
@router.get("/models", response_model=ModelGatewayStatus)
|
||||
async def read_provider_gateway() -> ModelGatewayStatus:
|
||||
return get_provider_gateway()
|
||||
|
||||
|
||||
@router.patch(
|
||||
"/models/{model}/endpoints/{endpoint}",
|
||||
response_model=ModelEndpointStatus,
|
||||
)
|
||||
async def update_model_endpoint(
|
||||
model: str,
|
||||
endpoint: str,
|
||||
payload: UpdateEndpointOverrideRequest,
|
||||
) -> ModelEndpointStatus:
|
||||
"""更新 endpoint 覆盖（base_url / api_key / status / weight / priority，DB 优先 env）。
|
||||
|
||||
写后热重载 registry（reset_model_router）。响应脱敏，不回 base_url / api_key 原值。
|
||||
"""
|
||||
await update_endpoint_override(model, endpoint, payload)
|
||||
gateway = get_provider_gateway()
|
||||
for m in gateway.models:
|
||||
if m.name != model:
|
||||
continue
|
||||
for ep in m.endpoints or []:
|
||||
if ep.name == endpoint:
|
||||
return ep
|
||||
raise AppError("not_found", "endpoint 不存在。")
|
||||
|
||||
|
||||
@router.get("/preferences", response_model=GenerationPreferencesResponse)
|
||||
async def read_generation_preferences(
|
||||
current_user: CurrentUser = Depends(get_current_user),
|
||||
|
||||
@@ -1,82 +0,0 @@
|
||||
"""非流式 Agent 应用层：工具调用闭环编排。
|
||||
|
||||
入口：
|
||||
- ``run_agent_turn`` 执行一轮带工具调用的对话循环
|
||||
- ``AgentTurnResult`` 循环产出
|
||||
- ``to_ai_provider_tools`` 把 ToolSpec 转成 OpenAI function schema dict
|
||||
- ``default_tools`` 返回全部内置工具（依赖司命 + 确认流控风险）
|
||||
"""
|
||||
|
||||
from app.application.agent.agent import (
|
||||
Agent,
|
||||
AgentConfig,
|
||||
AgentState,
|
||||
AgentStatus,
|
||||
PendingMessageQueue,
|
||||
QueueMode,
|
||||
)
|
||||
from app.application.agent.confirmation import ConfirmationRegistry
|
||||
from app.application.agent.extensions import (
|
||||
AfterToolCallDecision,
|
||||
AgentContext,
|
||||
AgentEvent,
|
||||
AgentRunContext,
|
||||
BeforeToolCallDecision,
|
||||
CancellationToken,
|
||||
ExtensionRegistry,
|
||||
ToolCallContext,
|
||||
TurnSnapshot,
|
||||
)
|
||||
from app.application.agent.prompt_composer import compose_system_messages
|
||||
from app.application.agent.python_skills import (
|
||||
PythonSkillActivation,
|
||||
PythonSkillCatalog,
|
||||
PythonSkillDefinition,
|
||||
activate_python_skills,
|
||||
discover_python_skills,
|
||||
)
|
||||
from app.application.agent.run_registry import AgentRunRegistry
|
||||
from app.application.agent.runtime import AgentTurnResult, run_agent_turn
|
||||
from app.application.agent.skill_loader import load_skill_manifests
|
||||
from app.application.agent.tasks import (
|
||||
next_step_for_tool,
|
||||
step_event_payload,
|
||||
task_from_deliberative_plan,
|
||||
task_plan_payload,
|
||||
)
|
||||
from app.application.agent.tools import default_tools, to_ai_provider_tools
|
||||
|
||||
__all__ = [
|
||||
"AfterToolCallDecision",
|
||||
"Agent",
|
||||
"AgentConfig",
|
||||
"AgentContext",
|
||||
"AgentEvent",
|
||||
"AgentRunContext",
|
||||
"AgentRunRegistry",
|
||||
"AgentState",
|
||||
"AgentStatus",
|
||||
"AgentTurnResult",
|
||||
"BeforeToolCallDecision",
|
||||
"CancellationToken",
|
||||
"ConfirmationRegistry",
|
||||
"ExtensionRegistry",
|
||||
"PendingMessageQueue",
|
||||
"PythonSkillActivation",
|
||||
"PythonSkillCatalog",
|
||||
"PythonSkillDefinition",
|
||||
"QueueMode",
|
||||
"ToolCallContext",
|
||||
"TurnSnapshot",
|
||||
"activate_python_skills",
|
||||
"compose_system_messages",
|
||||
"default_tools",
|
||||
"discover_python_skills",
|
||||
"load_skill_manifests",
|
||||
"next_step_for_tool",
|
||||
"run_agent_turn",
|
||||
"step_event_payload",
|
||||
"task_from_deliberative_plan",
|
||||
"task_plan_payload",
|
||||
"to_ai_provider_tools",
|
||||
]
|
||||
@@ -1,303 +0,0 @@
|
||||
"""Pi agent-core 风格的有状态 Agent 门面。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Awaitable, Callable
|
||||
from contextlib import suppress
|
||||
from dataclasses import dataclass, field, replace
|
||||
from enum import StrEnum
|
||||
from typing import Any
|
||||
|
||||
from app.ai_provider.types.chat import ChatChunk, TokenUsage
|
||||
from app.ai_provider.types.messages import BaseMessage, HumanMessage
|
||||
from app.application.agent.confirmation import ConfirmationRegistry
|
||||
from app.application.agent.extensions import (
|
||||
AgentEvent,
|
||||
AgentRunContext,
|
||||
CancellationToken,
|
||||
ExtensionRegistry,
|
||||
)
|
||||
from app.application.agent.runtime import AgentTurnResult, run_agent_turn
|
||||
from app.modules.tool_runtime.definition import ToolRuntimeLike
|
||||
|
||||
|
||||
class AgentStatus(StrEnum):
|
||||
idle = "idle"
|
||||
running = "running"
|
||||
completed = "completed"
|
||||
failed = "failed"
|
||||
cancelled = "cancelled"
|
||||
|
||||
|
||||
class QueueMode(StrEnum):
|
||||
all = "all"
|
||||
one_at_a_time = "one-at-a-time"
|
||||
|
||||
|
||||
@dataclass
|
||||
class AgentState:
|
||||
"""有状态 Agent 的进程内快照。"""
|
||||
|
||||
messages: list[BaseMessage] = field(default_factory=list)
|
||||
status: AgentStatus = AgentStatus.idle
|
||||
streaming_message: str = ""
|
||||
pending_tool_calls: set[str] = field(default_factory=set)
|
||||
error_message: str | None = None
|
||||
usage: TokenUsage = field(default_factory=TokenUsage)
|
||||
|
||||
|
||||
class PendingMessageQueue:
|
||||
"""steering/follow-up 共用的小型 FIFO。"""
|
||||
|
||||
def __init__(self, mode: QueueMode = QueueMode.one_at_a_time) -> None:
|
||||
self.mode = mode
|
||||
self._messages: list[BaseMessage] = []
|
||||
|
||||
def put(self, message: BaseMessage) -> None:
|
||||
self._messages.append(message)
|
||||
|
||||
def drain(self) -> list[BaseMessage]:
|
||||
if not self._messages:
|
||||
return []
|
||||
if self.mode is QueueMode.all:
|
||||
messages, self._messages = self._messages, []
|
||||
return messages
|
||||
return [self._messages.pop(0)]
|
||||
|
||||
def clear(self) -> None:
|
||||
self._messages.clear()
|
||||
|
||||
def __bool__(self) -> bool:
|
||||
return bool(self._messages)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class AgentConfig:
|
||||
"""Stateful Agent 的稳定运行配置。"""
|
||||
|
||||
tools: tuple[dict[str, Any], ...]
|
||||
tool_runtime: ToolRuntimeLike | None
|
||||
model_hint: str | None = None
|
||||
thinking_enabled: bool | None = None
|
||||
reasoning_effort: str | None = None
|
||||
max_iterations: int = 8
|
||||
confirmation_registry: ConfirmationRegistry | None = None
|
||||
confirmation_timeout: float = 300.0
|
||||
on_turn: Callable[[], Awaitable[None]] | None = None
|
||||
on_chunk: Callable[[ChatChunk], Awaitable[None]] | None = None
|
||||
on_tool_event: Callable[[str, dict[str, Any]], Awaitable[None]] | None = None
|
||||
on_confirm_timeout: Callable[[str], Awaitable[None]] | None = None
|
||||
|
||||
|
||||
class Agent:
|
||||
"""管理状态、消息队列、取消与扩展 Hook；低层循环仍由 run_agent_turn 负责。"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: AgentConfig,
|
||||
*,
|
||||
messages: list[BaseMessage] | None = None,
|
||||
run_context: AgentRunContext | None = None,
|
||||
extensions: ExtensionRegistry | None = None,
|
||||
runner: Callable[..., Awaitable[AgentTurnResult]] = run_agent_turn,
|
||||
) -> None:
|
||||
self.config = config
|
||||
self.state = AgentState(messages=list(messages or []))
|
||||
self.run_context = run_context or AgentRunContext()
|
||||
self.extensions = extensions or ExtensionRegistry()
|
||||
self._runner = runner
|
||||
self.steering = PendingMessageQueue()
|
||||
self.follow_ups = PendingMessageQueue()
|
||||
self._active_task: asyncio.Task[AgentTurnResult] | None = None
|
||||
self._accepting_messages = False
|
||||
self._idle = asyncio.Event()
|
||||
self._idle.set()
|
||||
|
||||
def steer(self, message: str | BaseMessage) -> bool:
|
||||
if not self._accepting_messages:
|
||||
return False
|
||||
self.steering.put(self._as_message(message))
|
||||
return True
|
||||
|
||||
def follow_up(self, message: str | BaseMessage) -> bool:
|
||||
if not self._accepting_messages:
|
||||
return False
|
||||
self.follow_ups.put(self._as_message(message))
|
||||
return True
|
||||
|
||||
async def prompt(self, message: str | BaseMessage) -> AgentTurnResult:
|
||||
if self.state.status is AgentStatus.running:
|
||||
raise RuntimeError("Agent 正在运行，请使用 steer 或 follow_up")
|
||||
self.state.messages.append(self._as_message(message))
|
||||
return await self._run()
|
||||
|
||||
async def continue_run(self) -> AgentTurnResult:
|
||||
if self.state.status is AgentStatus.running:
|
||||
raise RuntimeError("Agent 正在运行")
|
||||
return await self._run()
|
||||
|
||||
async def abort(self) -> bool:
|
||||
task = self._active_task
|
||||
if task is None or task.done():
|
||||
return False
|
||||
self.run_context.cancellation.cancel()
|
||||
task.cancel()
|
||||
return True
|
||||
|
||||
async def wait_for_idle(self) -> None:
|
||||
await self._idle.wait()
|
||||
|
||||
async def reset(self) -> None:
|
||||
await self.abort()
|
||||
await self.wait_for_idle()
|
||||
self.state = AgentState()
|
||||
self.steering.clear()
|
||||
self.follow_ups.clear()
|
||||
|
||||
async def _run(self) -> AgentTurnResult:
|
||||
if self.run_context.cancellation.cancelled:
|
||||
self.run_context = replace(self.run_context, cancellation=CancellationToken())
|
||||
self.state.status = AgentStatus.running
|
||||
self.state.streaming_message = ""
|
||||
self.state.pending_tool_calls.clear()
|
||||
self.state.error_message = None
|
||||
self._accepting_messages = True
|
||||
self._idle.clear()
|
||||
try:
|
||||
await self.extensions.publish(AgentEvent("agent.started"), self.run_context)
|
||||
task = asyncio.create_task(self._run_loop())
|
||||
self._active_task = task
|
||||
result = await task
|
||||
self._accepting_messages = False
|
||||
self.state.status = AgentStatus.completed
|
||||
self.state.streaming_message = result.content
|
||||
self.state.usage = result.usage
|
||||
await self.extensions.publish(
|
||||
AgentEvent("agent.ended", {"iterations": result.iterations}), self.run_context
|
||||
)
|
||||
return result
|
||||
except asyncio.CancelledError:
|
||||
self._accepting_messages = False
|
||||
self.state.status = AgentStatus.cancelled
|
||||
await self.extensions.publish(AgentEvent("agent.cancelled"), self.run_context)
|
||||
raise
|
||||
except Exception as exc:
|
||||
self._accepting_messages = False
|
||||
self.state.status = AgentStatus.failed
|
||||
self.state.error_message = str(exc)
|
||||
with suppress(Exception):
|
||||
await self.extensions.publish(
|
||||
AgentEvent("agent.failed", {"error": str(exc)}), self.run_context
|
||||
)
|
||||
raise
|
||||
finally:
|
||||
self._accepting_messages = False
|
||||
with suppress(Exception):
|
||||
await self.extensions.publish(
|
||||
AgentEvent("agent.settled", {"status": self.state.status.value}),
|
||||
self.run_context,
|
||||
)
|
||||
self._active_task = None
|
||||
self._idle.set()
|
||||
|
||||
async def _run_loop(self) -> AgentTurnResult:
|
||||
remaining = self.config.max_iterations
|
||||
total_iterations = 0
|
||||
total_usage = TokenUsage()
|
||||
total_tool_history: list[dict[str, Any]] = []
|
||||
final = AgentTurnResult(content="", reasoning="")
|
||||
while remaining > 0:
|
||||
working_messages = list(self.state.messages)
|
||||
final = await self._runner(
|
||||
messages=working_messages,
|
||||
tools=list(self.config.tools),
|
||||
tool_runtime=self.config.tool_runtime,
|
||||
model_hint=self.config.model_hint,
|
||||
thinking_enabled=self.config.thinking_enabled,
|
||||
reasoning_effort=self.config.reasoning_effort,
|
||||
max_iterations=remaining,
|
||||
run_id=self.run_context.run_id,
|
||||
character_id=self.run_context.character_id,
|
||||
trace_id=self.run_context.trace_id,
|
||||
on_turn=self.config.on_turn,
|
||||
on_chunk=self._on_chunk,
|
||||
on_tool_event=self._on_tool_event,
|
||||
confirmation_registry=self.config.confirmation_registry,
|
||||
confirmation_timeout=self.config.confirmation_timeout,
|
||||
on_confirm_timeout=self.config.on_confirm_timeout,
|
||||
extensions=self.extensions,
|
||||
run_context=self.run_context,
|
||||
get_steering_messages=self.steering.drain,
|
||||
)
|
||||
self.state.messages = working_messages
|
||||
total_iterations += final.iterations
|
||||
total_usage += final.usage
|
||||
total_tool_history.extend(final.tool_history)
|
||||
remaining -= final.iterations
|
||||
if final.truncated:
|
||||
return AgentTurnResult(
|
||||
content=final.content,
|
||||
reasoning=final.reasoning,
|
||||
tool_history=total_tool_history,
|
||||
iterations=total_iterations,
|
||||
truncated=True,
|
||||
usage=total_usage,
|
||||
)
|
||||
if remaining <= 0:
|
||||
return AgentTurnResult(
|
||||
content=final.content,
|
||||
reasoning=final.reasoning,
|
||||
tool_history=total_tool_history,
|
||||
iterations=total_iterations,
|
||||
truncated=bool(self.follow_ups),
|
||||
usage=total_usage,
|
||||
)
|
||||
follow_ups = self.follow_ups.drain()
|
||||
if not follow_ups:
|
||||
return AgentTurnResult(
|
||||
content=final.content,
|
||||
reasoning=final.reasoning,
|
||||
tool_history=total_tool_history,
|
||||
iterations=total_iterations,
|
||||
truncated=final.truncated,
|
||||
usage=total_usage,
|
||||
)
|
||||
self.state.messages.extend(follow_ups)
|
||||
return final
|
||||
|
||||
async def _on_chunk(self, chunk: ChatChunk) -> None:
|
||||
if chunk.delta_content:
|
||||
self.state.streaming_message += chunk.delta_content
|
||||
if self.config.on_chunk is not None:
|
||||
await self.config.on_chunk(chunk)
|
||||
|
||||
async def _on_tool_event(self, event_type: str, payload: dict[str, Any]) -> None:
|
||||
call_id = str(payload.get("tool_call_id", ""))
|
||||
if event_type == "tool.started" and call_id:
|
||||
self.state.pending_tool_calls.add(call_id)
|
||||
elif event_type in ("tool.completed", "tool.failed") and call_id:
|
||||
self.state.pending_tool_calls.discard(call_id)
|
||||
# 先转发 SSE（用户可见流优先），再把同一事件发布进扩展事件流，
|
||||
# 让 AgentEventListener 也能观测工具级生命周期。
|
||||
# 扩展 listener 异常被吞掉，避免观测副菜阻断工具事件主流程。
|
||||
if self.config.on_tool_event is not None:
|
||||
await self.config.on_tool_event(event_type, payload)
|
||||
with suppress(Exception):
|
||||
await self.extensions.publish(
|
||||
AgentEvent(event_type, payload), self.run_context
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _as_message(message: str | BaseMessage) -> BaseMessage:
|
||||
return HumanMessage(content=message) if isinstance(message, str) else message
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Agent",
|
||||
"AgentConfig",
|
||||
"AgentState",
|
||||
"AgentStatus",
|
||||
"PendingMessageQueue",
|
||||
"QueueMode",
|
||||
]
|
||||
@@ -1,148 +0,0 @@
|
||||
"""混合确认注册表 — 同进程 Event 唤醒 + 重启后 DB polling 兜底。
|
||||
|
||||
Agent 循环在工具需要确认时：
|
||||
1. 写 PendingConfirmationRecord 到 DB（持久化挂起状态）
|
||||
2. 调 ``ConfirmationRegistry.wait(run_id)`` — 同进程用 Event 0 延迟唤醒，
|
||||
重启后 polling DB 兜底
|
||||
|
||||
approve/deny 端点统一调 ``registry.resolve(run_id, approved)``：
|
||||
- 更新 PendingConfirmationRecord.status + resolved_at
|
||||
- 同一进程直接 set Event
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from datetime import UTC, datetime
|
||||
|
||||
from sqlalchemy import or_, select, update
|
||||
|
||||
from app.models.safety import PendingConfirmationRecord
|
||||
|
||||
|
||||
class ConfirmationRegistry:
|
||||
"""混合确认注册表。
|
||||
|
||||
- 同进程：``asyncio.Event`` 直接唤醒，0 延迟
|
||||
- 重启后：Event 不存在 → polling ``PendingConfirmationRecord.status`` 兜底
|
||||
- approve/deny 端点统一调用 ``resolve(run_id, approved)``
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._events: dict[str, asyncio.Event] = {}
|
||||
|
||||
async def wait(
|
||||
self,
|
||||
run_id: str,
|
||||
*,
|
||||
poll_interval: float = 0.5,
|
||||
timeout: float = 60,
|
||||
confirmation_id: str | None = None,
|
||||
) -> bool:
|
||||
"""等待确认。
|
||||
|
||||
同一进程的事件直接唤醒；跨进程或重启后 polling 兜底。
|
||||
|
||||
``confirmation_id`` 可选，传入时会同时注册为该 Event 的别名，
|
||||
使得外部以 ``confirmation_id`` 调 ``resolve`` 时也能立即唤醒同进程的 Event
|
||||
（无需等 DB polling 兜底）。
|
||||
|
||||
返回：
|
||||
True=批准, False=拒绝 ; 超时抛出 TimeoutError
|
||||
"""
|
||||
event = asyncio.Event()
|
||||
self._events[run_id] = event
|
||||
if confirmation_id:
|
||||
self._events[confirmation_id] = event
|
||||
|
||||
try:
|
||||
deadline = time.monotonic() + timeout
|
||||
while time.monotonic() < deadline:
|
||||
# 1) polling DB 获取真实状态（Event 只用于唤醒，不携带批准/拒绝值）
|
||||
from app.db.session import AsyncSessionLocal
|
||||
|
||||
async with AsyncSessionLocal() as session:
|
||||
if confirmation_id is None:
|
||||
stmt = select(PendingConfirmationRecord).where(
|
||||
or_(
|
||||
PendingConfirmationRecord.id == run_id,
|
||||
PendingConfirmationRecord.run_id == run_id,
|
||||
)
|
||||
)
|
||||
else:
|
||||
stmt = select(PendingConfirmationRecord).where(
|
||||
PendingConfirmationRecord.id == confirmation_id,
|
||||
PendingConfirmationRecord.run_id == run_id,
|
||||
)
|
||||
result = await session.execute(stmt)
|
||||
record = result.scalar_one_or_none()
|
||||
if record is not None and record.status != "pending":
|
||||
return record.status == "approved"
|
||||
|
||||
# 2) 等一个间隔，同时监听 Event（Event set 只表示「有变化」，实际值靠 DB）
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
asyncio.shield(event.wait()),
|
||||
timeout=poll_interval,
|
||||
)
|
||||
# Event 被唤醒 → 回到循环顶部重新查 DB
|
||||
continue
|
||||
except TimeoutError:
|
||||
continue
|
||||
|
||||
# 超时
|
||||
raise TimeoutError(f"确认超时（等待 {timeout:.0f}s 未裁决）")
|
||||
finally:
|
||||
self._events.pop(run_id, None)
|
||||
if confirmation_id:
|
||||
self._events.pop(confirmation_id, None)
|
||||
|
||||
async def resolve(
|
||||
self,
|
||||
run_id: str,
|
||||
approved: bool,
|
||||
*,
|
||||
confirmation_id: str | None = None,
|
||||
) -> bool:
|
||||
"""精确裁决一个待确认工具调用，并唤醒对应等待者。"""
|
||||
from app.db.session import AsyncSessionLocal
|
||||
|
||||
target_id = confirmation_id or run_id
|
||||
async with AsyncSessionLocal() as session:
|
||||
status = "approved" if approved else "denied"
|
||||
stmt = update(PendingConfirmationRecord).where(
|
||||
PendingConfirmationRecord.status == "pending"
|
||||
)
|
||||
if confirmation_id is None:
|
||||
stmt = stmt.where(
|
||||
or_(
|
||||
PendingConfirmationRecord.id == run_id,
|
||||
PendingConfirmationRecord.run_id == run_id,
|
||||
)
|
||||
)
|
||||
else:
|
||||
stmt = stmt.where(
|
||||
PendingConfirmationRecord.id == confirmation_id,
|
||||
PendingConfirmationRecord.run_id == run_id,
|
||||
)
|
||||
result = await session.execute(
|
||||
stmt.values(status=status, resolved_at=datetime.now(UTC))
|
||||
)
|
||||
await session.commit()
|
||||
|
||||
if result.rowcount == 0:
|
||||
return False
|
||||
|
||||
event = self._events.get(target_id) or self._events.get(run_id)
|
||||
if event is not None:
|
||||
event.set()
|
||||
return True
|
||||
|
||||
def has(self, run_id: str) -> bool:
|
||||
"""是否仍在等待（同进程检查）。"""
|
||||
event = self._events.get(run_id)
|
||||
return event is not None and not event.is_set()
|
||||
|
||||
|
||||
__all__ = ["ConfirmationRegistry"]
|
||||
@@ -1,229 +0,0 @@
|
||||
"""Agent 扩展契约与注册表。
|
||||
|
||||
扩展只接触本次运行的快照与 ToolResult，不持有数据库会话、密钥或安全治理对象。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass, field, replace
|
||||
from typing import Any
|
||||
|
||||
from app.ai_provider.types.chat import TokenUsage
|
||||
from app.ai_provider.types.messages import BaseMessage
|
||||
from app.modules.tool_runtime.definition import ToolResult
|
||||
|
||||
|
||||
class CancellationToken:
|
||||
"""协作式取消信号；阻塞 await 仍由活动 asyncio.Task.cancel() 负责中断。"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._cancelled = False
|
||||
|
||||
@property
|
||||
def cancelled(self) -> bool:
|
||||
return self._cancelled
|
||||
|
||||
def cancel(self) -> None:
|
||||
self._cancelled = True
|
||||
|
||||
def raise_if_cancelled(self) -> None:
|
||||
if self._cancelled:
|
||||
raise asyncio.CancelledError
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class AgentRunContext:
|
||||
"""扩展可见的最小运行上下文。"""
|
||||
|
||||
run_id: str | None = None
|
||||
session_id: str | None = None
|
||||
character_id: str | None = None
|
||||
trace_id: str | None = None
|
||||
metadata: dict[str, Any] = field(default_factory=dict)
|
||||
cancellation: CancellationToken = field(default_factory=CancellationToken)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class AgentContext:
|
||||
"""每轮发送给模型前可转换的上下文快照。"""
|
||||
|
||||
messages: tuple[BaseMessage, ...]
|
||||
tools: tuple[dict[str, Any], ...]
|
||||
model_hint: str | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class AgentEvent:
|
||||
"""Agent Core 内部生命周期事件。"""
|
||||
|
||||
type: str
|
||||
payload: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ToolCallContext:
|
||||
"""工具调用 Hook 的稳定输入。"""
|
||||
|
||||
tool_call_id: str
|
||||
name: str
|
||||
input: dict[str, Any]
|
||||
messages: tuple[BaseMessage, ...]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BeforeToolCallDecision:
|
||||
"""执行前决策；扩展只能收紧或改参，不能绕过 ToolRuntime。"""
|
||||
|
||||
allowed: bool = True
|
||||
reason: str | None = None
|
||||
input: dict[str, Any] | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class AfterToolCallDecision:
|
||||
"""执行后结果覆盖与提前终止提示。"""
|
||||
|
||||
result: ToolResult | None = None
|
||||
terminate: bool = False
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TurnSnapshot:
|
||||
"""一轮模型调用及其工具执行结果摘要。"""
|
||||
|
||||
iteration: int
|
||||
content: str
|
||||
reasoning: str
|
||||
tool_calls: tuple[dict[str, Any], ...]
|
||||
tool_results: tuple[dict[str, Any], ...]
|
||||
usage: TokenUsage = field(default_factory=TokenUsage)
|
||||
|
||||
|
||||
ContextTransformer = Callable[[AgentContext, AgentRunContext], Awaitable[AgentContext]]
|
||||
BeforeToolCallHook = Callable[
|
||||
[ToolCallContext, AgentRunContext], Awaitable[BeforeToolCallDecision | None]
|
||||
]
|
||||
AfterToolCallHook = Callable[
|
||||
[ToolCallContext, ToolResult, AgentRunContext], Awaitable[AfterToolCallDecision | None]
|
||||
]
|
||||
TurnPolicy = Callable[[TurnSnapshot, AgentRunContext], Awaitable[bool]]
|
||||
AgentEventListener = Callable[[AgentEvent, AgentRunContext], Awaitable[None]]
|
||||
|
||||
|
||||
class ExtensionRegistry:
|
||||
"""受约束的 Agent 扩展注册表；按注册顺序执行 Hook。"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._context_transformers: list[ContextTransformer] = []
|
||||
self._before_tool_hooks: list[BeforeToolCallHook] = []
|
||||
self._after_tool_hooks: list[AfterToolCallHook] = []
|
||||
self._turn_policies: list[TurnPolicy] = []
|
||||
self._event_listeners: list[AgentEventListener] = []
|
||||
|
||||
def add_context_transformer(self, transformer: ContextTransformer) -> None:
|
||||
self._context_transformers.append(transformer)
|
||||
|
||||
def add_before_tool_hook(self, hook: BeforeToolCallHook) -> None:
|
||||
self._before_tool_hooks.append(hook)
|
||||
|
||||
def add_after_tool_hook(self, hook: AfterToolCallHook) -> None:
|
||||
self._after_tool_hooks.append(hook)
|
||||
|
||||
def add_turn_policy(self, policy: TurnPolicy) -> None:
|
||||
self._turn_policies.append(policy)
|
||||
|
||||
def add_event_listener(self, listener: AgentEventListener) -> Callable[[], None]:
|
||||
self._event_listeners.append(listener)
|
||||
|
||||
def unsubscribe() -> None:
|
||||
if listener in self._event_listeners:
|
||||
self._event_listeners.remove(listener)
|
||||
|
||||
return unsubscribe
|
||||
|
||||
def extend(self, registry: ExtensionRegistry) -> None:
|
||||
"""按原注册顺序合并另一组扩展。"""
|
||||
self._context_transformers.extend(registry._context_transformers)
|
||||
self._before_tool_hooks.extend(registry._before_tool_hooks)
|
||||
self._after_tool_hooks.extend(registry._after_tool_hooks)
|
||||
self._turn_policies.extend(registry._turn_policies)
|
||||
self._event_listeners.extend(registry._event_listeners)
|
||||
|
||||
async def transform_context(
|
||||
self, context: AgentContext, run: AgentRunContext
|
||||
) -> AgentContext:
|
||||
for transformer in self._context_transformers:
|
||||
run.cancellation.raise_if_cancelled()
|
||||
context = await transformer(context, run)
|
||||
return context
|
||||
|
||||
async def before_tool_call(
|
||||
self, call: ToolCallContext, run: AgentRunContext
|
||||
) -> BeforeToolCallDecision:
|
||||
current = call
|
||||
for hook in self._before_tool_hooks:
|
||||
run.cancellation.raise_if_cancelled()
|
||||
decision = await hook(current, run)
|
||||
if decision is None:
|
||||
continue
|
||||
if not decision.allowed:
|
||||
return decision
|
||||
if decision.input is not None:
|
||||
current = replace(current, input=decision.input)
|
||||
return BeforeToolCallDecision(input=current.input)
|
||||
|
||||
async def after_tool_call(
|
||||
self,
|
||||
call: ToolCallContext,
|
||||
result: ToolResult,
|
||||
run: AgentRunContext,
|
||||
) -> AfterToolCallDecision:
|
||||
current = result
|
||||
terminate = False
|
||||
for hook in self._after_tool_hooks:
|
||||
run.cancellation.raise_if_cancelled()
|
||||
decision = await hook(call, current, run)
|
||||
if decision is None:
|
||||
continue
|
||||
if decision.result is not None:
|
||||
current = decision.result
|
||||
terminate = terminate or decision.terminate
|
||||
return AfterToolCallDecision(result=current, terminate=terminate)
|
||||
|
||||
async def should_stop(self, turn: TurnSnapshot, run: AgentRunContext) -> bool:
|
||||
for policy in self._turn_policies:
|
||||
run.cancellation.raise_if_cancelled()
|
||||
if await policy(turn, run):
|
||||
return True
|
||||
return False
|
||||
|
||||
async def publish(self, event: AgentEvent, run: AgentRunContext) -> None:
|
||||
for listener in tuple(self._event_listeners):
|
||||
await listener(event, run)
|
||||
|
||||
@property
|
||||
def enabled(self) -> bool:
|
||||
return any(
|
||||
(
|
||||
self._context_transformers,
|
||||
self._before_tool_hooks,
|
||||
self._after_tool_hooks,
|
||||
self._turn_policies,
|
||||
self._event_listeners,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"AfterToolCallDecision",
|
||||
"AgentContext",
|
||||
"AgentEvent",
|
||||
"AgentRunContext",
|
||||
"BeforeToolCallDecision",
|
||||
"CancellationToken",
|
||||
"ExtensionRegistry",
|
||||
"ToolCallContext",
|
||||
"TurnSnapshot",
|
||||
]
|
||||
@@ -1,21 +0,0 @@
|
||||
"""通用 Agent 提示组合。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from app.modules.agent import AgentDefinition
|
||||
from app.modules.skills import SkillManifest
|
||||
|
||||
|
||||
def compose_system_messages(
|
||||
base_messages: list[str],
|
||||
agent: AgentDefinition | None,
|
||||
skills: list[SkillManifest],
|
||||
) -> list[str]:
|
||||
"""按稳定基础提示、Agent 身份、技能说明的顺序组合提示。"""
|
||||
messages = list(base_messages)
|
||||
if agent is not None and agent.system_prompt:
|
||||
messages.append(f"【Agent：{agent.name}】\n{agent.system_prompt}")
|
||||
for skill in skills:
|
||||
if skill.prompt:
|
||||
messages.append(f"【{skill.name}·能力注入】\n{skill.prompt}")
|
||||
return messages
|
||||
@@ -1,132 +0,0 @@
|
||||
"""受信任 Python Skill 的发现、匹配与 Run 级激活。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
import inspect
|
||||
import pkgutil
|
||||
from collections.abc import Awaitable, Callable, Iterable
|
||||
from dataclasses import dataclass
|
||||
from types import ModuleType
|
||||
from typing import Any
|
||||
|
||||
from app.application.agent.extensions import ExtensionRegistry
|
||||
|
||||
SkillCleanup = Callable[[], Awaitable[None] | None]
|
||||
SkillSetup = Callable[[ExtensionRegistry], SkillCleanup | None]
|
||||
SkillMatcher = Callable[[str, dict[str, Any]], bool]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PythonSkillDefinition:
|
||||
"""部署期 Python Skill 契约。"""
|
||||
|
||||
name: str
|
||||
description: str
|
||||
triggers: tuple[str, ...]
|
||||
instructions: str
|
||||
setup: SkillSetup
|
||||
version: str = "1"
|
||||
enabled: bool = True
|
||||
match: SkillMatcher | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class PythonSkillActivation:
|
||||
"""单个 Run 的 Skill 激活结果。"""
|
||||
|
||||
skills: tuple[PythonSkillDefinition, ...]
|
||||
cleanups: list[SkillCleanup]
|
||||
|
||||
async def close(self) -> None:
|
||||
"""按激活逆序释放 Skill 资源。"""
|
||||
while self.cleanups:
|
||||
result = self.cleanups.pop()()
|
||||
if inspect.isawaitable(result):
|
||||
await result
|
||||
|
||||
|
||||
class PythonSkillCatalog:
|
||||
"""只保存 Skill 描述与 setup，不持有 Run 状态。"""
|
||||
|
||||
def __init__(self, skills: Iterable[PythonSkillDefinition] = ()) -> None:
|
||||
self._skills: dict[str, PythonSkillDefinition] = {}
|
||||
for skill in skills:
|
||||
if skill.name in self._skills:
|
||||
raise ValueError(f"Python Skill 名称重复：{skill.name}")
|
||||
self._skills[skill.name] = skill
|
||||
|
||||
@property
|
||||
def skills(self) -> tuple[PythonSkillDefinition, ...]:
|
||||
return tuple(self._skills.values())
|
||||
|
||||
def match(
|
||||
self,
|
||||
context_text: str,
|
||||
*,
|
||||
skill_refs: tuple[str, ...] = (),
|
||||
metadata: dict[str, Any] | None = None,
|
||||
) -> tuple[PythonSkillDefinition, ...]:
|
||||
"""返回显式引用、关键词或自定义匹配器命中的 enabled Skills。"""
|
||||
refs = set(skill_refs)
|
||||
text = context_text.casefold()
|
||||
match_metadata = metadata or {}
|
||||
return tuple(
|
||||
skill
|
||||
for skill in self._skills.values()
|
||||
if skill.enabled
|
||||
and (
|
||||
skill.name in refs
|
||||
or any(trigger.casefold() in text for trigger in skill.triggers if trigger)
|
||||
or (skill.match is not None and skill.match(context_text, match_metadata))
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
async def activate_python_skills(
|
||||
registry: ExtensionRegistry,
|
||||
skills: Iterable[PythonSkillDefinition],
|
||||
) -> PythonSkillActivation:
|
||||
"""把命中的 Skills 注册到当前 Run，并在部分失败时清理已激活项。"""
|
||||
activation = PythonSkillActivation(tuple(skills), [])
|
||||
try:
|
||||
for skill in activation.skills:
|
||||
cleanup = skill.setup(registry)
|
||||
if cleanup is not None:
|
||||
activation.cleanups.append(cleanup)
|
||||
except Exception:
|
||||
await activation.close()
|
||||
raise
|
||||
return activation
|
||||
|
||||
|
||||
def discover_python_skills(packages: Iterable[str] = ()) -> PythonSkillCatalog:
|
||||
"""从受信任 packages 发现定义，但不执行 Skill setup。"""
|
||||
modules: list[ModuleType] = []
|
||||
for package_name in dict.fromkeys(packages):
|
||||
package = importlib.import_module(package_name)
|
||||
package_path = getattr(package, "__path__", None)
|
||||
if package_path is None:
|
||||
raise TypeError(f"Python Skill 包必须是 package：{package_name}")
|
||||
modules.extend(
|
||||
importlib.import_module(module.name)
|
||||
for module in pkgutil.iter_modules(package_path, f"{package.__name__}.")
|
||||
if not module.name.rsplit(".", 1)[-1].startswith("_")
|
||||
)
|
||||
|
||||
definitions: list[PythonSkillDefinition] = []
|
||||
for module in sorted(modules, key=lambda item: item.__name__):
|
||||
skill = getattr(module, "SKILL", None)
|
||||
if not isinstance(skill, PythonSkillDefinition):
|
||||
raise TypeError(f"Python Skill 缺少 SKILL 定义：{module.__name__}")
|
||||
definitions.append(skill)
|
||||
return PythonSkillCatalog(definitions)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"PythonSkillActivation",
|
||||
"PythonSkillCatalog",
|
||||
"PythonSkillDefinition",
|
||||
"activate_python_skills",
|
||||
"discover_python_skills",
|
||||
]
|
||||
@@ -1,49 +0,0 @@
|
||||
"""活动 Agent Run 注册表：为 steer、follow-up、abort 提供进程内控制面。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
|
||||
from app.application.agent.agent import Agent
|
||||
|
||||
|
||||
class AgentRunRegistry:
|
||||
"""跟踪当前进程中的活动 Agent；持久状态仍以 Run/EventStore 为准。"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._agents: dict[str, Agent] = {}
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
async def register(self, run_id: str, agent: Agent) -> None:
|
||||
async with self._lock:
|
||||
self._agents[run_id] = agent
|
||||
|
||||
async def unregister(self, run_id: str, agent: Agent) -> None:
|
||||
async with self._lock:
|
||||
if self._agents.get(run_id) is agent:
|
||||
self._agents.pop(run_id, None)
|
||||
|
||||
def get(self, run_id: str) -> Agent | None:
|
||||
return self._agents.get(run_id)
|
||||
|
||||
def steer(self, run_id: str, content: str) -> bool:
|
||||
agent = self.get(run_id)
|
||||
if agent is None:
|
||||
return False
|
||||
return agent.steer(content)
|
||||
|
||||
def follow_up(self, run_id: str, content: str) -> bool:
|
||||
agent = self.get(run_id)
|
||||
if agent is None:
|
||||
return False
|
||||
return agent.follow_up(content)
|
||||
|
||||
async def abort(self, run_id: str) -> bool:
|
||||
agent = self.get(run_id)
|
||||
if agent is None or not await agent.abort():
|
||||
return False
|
||||
await agent.wait_for_idle()
|
||||
return True
|
||||
|
||||
|
||||
__all__ = ["AgentRunRegistry"]
|
||||
@@ -1,840 +0,0 @@
|
||||
"""流式 Agent 循环：stream → 累积 tool_calls → execute → 回填 ToolMessage → 再 stream。
|
||||
|
||||
阶段2b 步骤A：把 Agent 工具调用闭环从非流式 ``router.chat`` 改成流式 ``router.stream``，
|
||||
并在主循环中暴露两个可选 async 回调（``on_chunk`` / ``on_tool_event``）供上层（步骤B
|
||||
send_message）接入 SSE 增量冲刷与工具事件。本模块不感知 SSE / 前端，回调签名仅约定：
|
||||
|
||||
- ``on_chunk(chunk: ChatChunk)``：每个流式 chunk 透传（含 delta_content / delta_reasoning
|
||||
/ tool_call_chunks / usage 等），上层可据此写 ``message.delta``。
|
||||
- ``on_tool_event(event_type: str, payload: dict)``：工具生命周期事件：
|
||||
- ``"tool.started"``：执行前
|
||||
- ``"tool.completed"``：执行后，含 status（success / error / denied）
|
||||
- ``"tool.requires_confirmation"``（阶段3 新增）：高风险工具抛 RequiresConfirmation 时
|
||||
发出，带 confirmation_id / risk_level / args，供前端展示确认卡片、approve/deny 端点
|
||||
据此恢复循环。
|
||||
|
||||
阶段3 步骤1：工具确认挂起/恢复机制。
|
||||
- 主循环遇到 ``RequiresConfirmation`` 时，通过 ``confirmation_registry.register(cid)`` 取
|
||||
一个 ``Future[bool]`` 并 ``await``；外部 approve/deny 端点（步骤2）调
|
||||
``registry.resolve(cid, approved)`` 触发恢复。
|
||||
- 批准：以 ``confirmed=True`` 重试 execute；拒绝：回填"用户拒绝"ToolMessage，循环继续。
|
||||
- 超时：``confirmation_timeout``（默认 300s）未裁决按拒绝处理；该超时受外层
|
||||
``run_timeout``（send_message 的 ``asyncio.timeout``）约束——确认必须在 run 超时内完成。
|
||||
- ``confirmation_registry`` 为 None 时降级：回填"需确认但未启用"消息，循环不卡（向下兼容）。
|
||||
|
||||
设计要点：
|
||||
- 流式 ``router.stream(...)`` 返回 ``AsyncIterator[ChatChunk]``；流结束（async for 正常
|
||||
退出）即本轮模型输出完毕，``ToolCallAccumulator.drain()`` 给出完整 tool_calls——不依赖
|
||||
``finish_reason`` 判断。
|
||||
- 工具执行经 ``ToolRuntime.execute``——司命评估与事件落库由其内部完成，本层只负责调用与
|
||||
错误兜底。
|
||||
- 工具失败**不中断循环**：异常被捕获、错误信息以 ToolMessage content 回填，让模型自行决策。
|
||||
RequiresConfirmation 单独处理（不走错误兜底分支）。
|
||||
- usage 累计每轮 prompt/completion/total_tokens。
|
||||
- ``on_chunk`` / ``on_tool_event`` 默认 None；不传时本循环与纯后端逻辑等价。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
from app.ai_provider import ProviderManagerFactory, TokenUsage
|
||||
from app.ai_provider.llm import ToolCallAccumulator
|
||||
from app.ai_provider.types.chat import ChatChunk
|
||||
from app.ai_provider.types.messages import AIMessage, BaseMessage, ToolMessage
|
||||
from app.application.agent.confirmation import ConfirmationRegistry
|
||||
from app.application.agent.extensions import (
|
||||
AgentContext,
|
||||
AgentEvent,
|
||||
AgentRunContext,
|
||||
CancellationToken,
|
||||
ExtensionRegistry,
|
||||
ToolCallContext,
|
||||
TurnSnapshot,
|
||||
)
|
||||
from app.core.logging import logger
|
||||
from app.models.safety import PendingConfirmationRecord
|
||||
from app.modules.safety.confirmation import new_confirmation_id
|
||||
from app.modules.safety.exceptions import RequiresConfirmation
|
||||
from app.modules.tool_runtime.definition import ToolResult, ToolRuntimeLike
|
||||
|
||||
# 默认最大循环次数：覆盖常见多步工具调用，超限则截断返回，避免无限循环。
|
||||
DEFAULT_MAX_ITERATIONS = 8
|
||||
|
||||
# 默认确认超时：高风险工具挂起等待用户裁决的最长时间（秒）；超时按拒绝处理。
|
||||
DEFAULT_CONFIRMATION_TIMEOUT = 300.0
|
||||
|
||||
# 回调类型别名：给上层（步骤B send_message）类型检查用
|
||||
OnChunkCallback = Callable[[ChatChunk], Awaitable[None]]
|
||||
OnToolEventCallback = Callable[[str, dict[str, Any]], Awaitable[None]]
|
||||
|
||||
|
||||
@dataclass
|
||||
class AgentTurnResult:
|
||||
"""单轮 Agent 循环的产出。"""
|
||||
|
||||
content: str # 最终回复正文（最后一次模型输出）
|
||||
reasoning: str # 最终回复的思考内容（如有）
|
||||
tool_history: list[dict[str, Any]] = field(
|
||||
default_factory=list
|
||||
) # [{name, args, result, status, tool_call_id?}]
|
||||
iterations: int = 0 # 实际循环次数
|
||||
truncated: bool = False # 是否因达上限截断
|
||||
usage: TokenUsage = field(default_factory=TokenUsage) # 累计 token 用量
|
||||
|
||||
|
||||
def _sum_usage(acc: TokenUsage, delta: TokenUsage) -> TokenUsage:
|
||||
return acc + delta
|
||||
|
||||
|
||||
def _parse_tool_args(args_str: str) -> dict[str, Any]:
|
||||
"""把累积到的 args JSON 字符串解析成 dict；解析失败时回退 ``{"_raw": ...}``。"""
|
||||
if not args_str:
|
||||
return {}
|
||||
try:
|
||||
parsed = json.loads(args_str)
|
||||
except json.JSONDecodeError:
|
||||
return {"_raw": args_str}
|
||||
if isinstance(parsed, dict):
|
||||
return parsed
|
||||
# 非对象 JSON（数组/字符串/数字）——包成 _raw 让模型自己理解
|
||||
return {"_raw": args_str}
|
||||
|
||||
|
||||
def _ai_message_from_stream(
|
||||
content: str,
|
||||
reasoning: str,
|
||||
tool_calls: list[dict[str, Any]],
|
||||
raw_content_blocks: list[dict[str, Any]] | None = None,
|
||||
) -> AIMessage:
|
||||
"""把流式累积的 content/reasoning/tool_calls 还原成 AIMessage 以便追加到 messages 历史。
|
||||
|
||||
``tool_calls`` 元素结构为 ``{"index", "id"?, "name"?, "args": str}``（由
|
||||
``ToolCallAccumulator.drain`` 产出，``args`` 为完整 JSON 字符串）。AIMessage 期望
|
||||
``{"name": ..., "args": dict, "id": ...}``，故需反序列化 args 并补全缺失 id。
|
||||
"""
|
||||
lc_tool_calls: list[dict[str, Any]] = []
|
||||
for entry in tool_calls:
|
||||
name = entry.get("name")
|
||||
if not name:
|
||||
# 没有 name 的 tool_call 无法执行，跳过（保留在累积器视图但不进历史）
|
||||
continue
|
||||
tc_id = entry.get("id") or f"call_{name}"
|
||||
lc_tool_calls.append(
|
||||
{
|
||||
"name": name,
|
||||
"args": _parse_tool_args(entry.get("args", "")),
|
||||
"id": tc_id,
|
||||
}
|
||||
)
|
||||
return AIMessage(
|
||||
content=content,
|
||||
tool_calls=lc_tool_calls,
|
||||
reasoning=reasoning,
|
||||
raw_content_blocks=raw_content_blocks,
|
||||
)
|
||||
|
||||
|
||||
def _truncate(text: str, limit: int = 4000) -> str:
|
||||
"""截断超长文本（沙箱 stdout/stderr 可能很大），避免喂回模型时爆 token。
|
||||
|
||||
措辞明确这是「回显截断」而非数据损坏，并告知取后续内容的方法：否则模型读到大文件
|
||||
被 read_file 截断的结果，会反复换方法重读，陷入循环直到 ``max_iterations`` 截断。
|
||||
阈值 4000 兼顾 token 成本与单次可读量；大文件配合 read_file 的 offset/limit 分段读取。
|
||||
"""
|
||||
if len(text) <= limit:
|
||||
return text
|
||||
return (
|
||||
f"{text[:limit]}\n…[输出过长，仅回显前 {limit} 字符；原始内容完整。"
|
||||
f"如需后续内容：read_file 用 offset/limit 分段，或 shell 用 sed -n / tail / head 限定输出]"
|
||||
)
|
||||
|
||||
|
||||
async def _publish_message_lifecycle(
|
||||
registry: ExtensionRegistry,
|
||||
run: AgentRunContext,
|
||||
message: BaseMessage,
|
||||
) -> None:
|
||||
"""发布一条完整消息的 start/end 生命周期。"""
|
||||
if not registry.enabled:
|
||||
return
|
||||
payload = {"role": message.role, "message": message.to_dict()}
|
||||
await registry.publish(AgentEvent("message.started", payload), run)
|
||||
await registry.publish(AgentEvent("message.ended", payload), run)
|
||||
|
||||
|
||||
def _tool_result_text(result: ToolResult) -> str:
|
||||
"""把 ToolResult 转成回填给模型的紧凑文本。
|
||||
|
||||
显式标注 ``status=成功/失败``：否则成功路径（如 write_file 落盘后 stdout/stderr 全空）
|
||||
只剩孤零零的 ``exit_code=0``，模型缺乏「已成功」的明确语义，倾向再调 read_file 验证，
|
||||
叠加回显截断极易陷入反复重试。
|
||||
"""
|
||||
status = "成功" if result.ok else "失败"
|
||||
parts = [f"status={status} exit_code={result.exit_code}"]
|
||||
if result.stdout:
|
||||
parts.append(f"stdout:\n{_truncate(result.stdout)}")
|
||||
if result.stderr:
|
||||
parts.append(f"stderr:\n{_truncate(result.stderr)}")
|
||||
if result.timeout:
|
||||
parts.append("timeout=true")
|
||||
if result.truncated:
|
||||
parts.append("truncated=true（仅截断回显给模型的输出，不影响已落盘等实际副作用）")
|
||||
return "\n".join(parts)
|
||||
|
||||
|
||||
# 控制流哨兵异常：只在本模块内传递，不泄漏到 run_agent_turn 调用方。
|
||||
class _ToolDenied(Exception):
|
||||
"""用户拒绝 / 确认超时。``message`` 为回填文本，``reason`` 为内部归因（denied/timeout）。"""
|
||||
|
||||
def __init__(self, message: str, *, reason: str = "denied") -> None:
|
||||
super().__init__(message)
|
||||
self.message = message
|
||||
self.reason = reason
|
||||
|
||||
|
||||
class _ToolSkipped(Exception):
|
||||
"""confirmation_registry 未启用时的降级：回填"需确认但未启用"消息。"""
|
||||
|
||||
def __init__(self, message: str) -> None:
|
||||
super().__init__(message)
|
||||
self.message = message
|
||||
|
||||
|
||||
async def _execute_tool_with_confirmation(
|
||||
*,
|
||||
tool_runtime: ToolRuntimeLike | None,
|
||||
name: str,
|
||||
tool_input: dict[str, Any],
|
||||
args_str: str,
|
||||
tool_call_id: str,
|
||||
run_id: str | None,
|
||||
character_id: str | None,
|
||||
trace_id: str | None,
|
||||
base_confirmed: bool,
|
||||
confirmation_registry: ConfirmationRegistry | None,
|
||||
confirmation_timeout: float,
|
||||
on_tool_event: OnToolEventCallback | None,
|
||||
cancellation: CancellationToken | None = None,
|
||||
on_timeout: Callable[[str], Awaitable[None]] | None = None,
|
||||
) -> ToolResult:
|
||||
"""执行单次工具调用，处理 RequiresConfirmation 挂起/恢复。
|
||||
|
||||
- 正常返回 ToolResult；任意终态（拒绝/超时/降级）以本模块哨兵异常向调用方信号。
|
||||
- 批准路径以 ``confirmed=True`` 重试 execute（绕过司命 ask），失败仍正常返回 ToolResult
|
||||
（由调用方按 ok=False 回填 error 文本）。
|
||||
- ``on_timeout(confirmation_id)``：确认超时时回调，供调用方标记超时。
|
||||
"""
|
||||
# base_confirmed=True 表示调用方已保证确认（如 API 已先确认），直接执行
|
||||
use_confirmed = base_confirmed
|
||||
if tool_runtime is None:
|
||||
raise _ToolSkipped("工具运行时未启用，已跳过工具调用。")
|
||||
|
||||
while True:
|
||||
try:
|
||||
return await tool_runtime.execute(
|
||||
name,
|
||||
tool_input,
|
||||
run_id=run_id,
|
||||
character_id=character_id,
|
||||
trace_id=trace_id,
|
||||
confirmed=use_confirmed,
|
||||
context={"cancellation": cancellation} if cancellation is not None else None,
|
||||
)
|
||||
except RequiresConfirmation as exc:
|
||||
if use_confirmed:
|
||||
raise _ToolSkipped(f"工具 {name} 在已确认状态下仍要求确认，已跳过。") from exc
|
||||
|
||||
# ---- registry 未启用：降级 ----
|
||||
if confirmation_registry is None:
|
||||
raise _ToolSkipped(
|
||||
f"工具 {name} 需要用户确认（风险等级 {exc.risk_level or 'L2'}），"
|
||||
"但当前未启用确认流程，已跳过。"
|
||||
) from exc
|
||||
|
||||
# ---- registry 已启用：登记 + 持久化待确认 + 发事件 + await ----
|
||||
confirmation_id = new_confirmation_id()
|
||||
|
||||
# 先持久化待确认记录（多 worker 共享确认状态），再发事件。
|
||||
# 顺序保证：resolve/deny 回调（on_tool_event 触发的 create_task）看到记录时，
|
||||
# DB 里已有该行，不会空转超时。
|
||||
if run_id:
|
||||
from app.db.session import AsyncSessionLocal
|
||||
|
||||
async with AsyncSessionLocal() as session:
|
||||
record = PendingConfirmationRecord(
|
||||
id=confirmation_id,
|
||||
tool=name,
|
||||
input_json=json.dumps(
|
||||
{
|
||||
"name": name,
|
||||
"args": tool_input,
|
||||
"tool_call_id": tool_call_id,
|
||||
"args_str": args_str,
|
||||
}
|
||||
),
|
||||
risk_level=exc.risk_level or "L2",
|
||||
reason=exc.message,
|
||||
run_id=run_id,
|
||||
character_id=character_id,
|
||||
status="pending",
|
||||
)
|
||||
session.add(record)
|
||||
await session.commit()
|
||||
|
||||
if on_tool_event is not None:
|
||||
await on_tool_event(
|
||||
"tool.requires_confirmation",
|
||||
{
|
||||
"confirmation_id": confirmation_id,
|
||||
"run_id": run_id,
|
||||
"tool_name": name,
|
||||
"tool_call_id": tool_call_id,
|
||||
"risk_level": exc.risk_level,
|
||||
"reason": exc.message,
|
||||
"args": args_str,
|
||||
},
|
||||
)
|
||||
|
||||
# 使用混合确认注册表（同进程 Event + 重启后 polling）。无 run_id 的
|
||||
# 独立调用仍可等待到超时，避免把“未裁决”误报成“确认机制未启用”。
|
||||
confirmation_run_id = run_id or confirmation_id
|
||||
try:
|
||||
approved = await confirmation_registry.wait(
|
||||
confirmation_run_id,
|
||||
confirmation_id=confirmation_id,
|
||||
timeout=confirmation_timeout,
|
||||
)
|
||||
except TimeoutError:
|
||||
if on_timeout is not None:
|
||||
await on_timeout(confirmation_id)
|
||||
raise _ToolDenied(
|
||||
f"工具 {name} 的确认超时（等待 {confirmation_timeout:.0f}s 未裁决），已按拒绝处理。",
|
||||
reason="timeout",
|
||||
) from None
|
||||
|
||||
if not approved:
|
||||
raise _ToolDenied(f"用户拒绝了工具 {name} 的执行。") from exc
|
||||
|
||||
# 批准：以 confirmed=True 重试 execute（循环回到顶部 try）
|
||||
logger.info(
|
||||
"agent_tool_confirmed",
|
||||
tool=name,
|
||||
confirmation_id=confirmation_id,
|
||||
run_id=run_id,
|
||||
)
|
||||
use_confirmed = True
|
||||
if on_tool_event is not None:
|
||||
await on_tool_event(
|
||||
"tool.started",
|
||||
{
|
||||
"tool_name": name,
|
||||
"tool_call_id": tool_call_id,
|
||||
"args": args_str,
|
||||
"confirmed": True,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
async def run_agent_turn(
|
||||
*,
|
||||
messages: list[BaseMessage],
|
||||
tools: list[dict[str, Any]],
|
||||
tool_runtime: ToolRuntimeLike | None,
|
||||
model_hint: str | None = None,
|
||||
thinking_enabled: bool | None = None,
|
||||
reasoning_effort: str | None = None,
|
||||
max_iterations: int = DEFAULT_MAX_ITERATIONS,
|
||||
# 可选上下文：透传给 ToolRuntime.execute 做事件归属与工作目录隔离
|
||||
run_id: str | None = None,
|
||||
character_id: str | None = None,
|
||||
trace_id: str | None = None,
|
||||
# 可选：高风险工具是否已由上层确认（透传给 ToolRuntime.execute）
|
||||
confirmed: bool = False,
|
||||
# 可选：每轮开始时回调（供外层重置 per-turn 状态如 reasoning_finalized）
|
||||
on_turn: Callable[[], Awaitable[None]] | None = None,
|
||||
# 可选回调：把流式 chunk 透传给上层做 SSE 增量冲刷（默认 None 表示不回调）
|
||||
on_chunk: OnChunkCallback | None = None,
|
||||
# 可选回调：工具生命周期事件（"tool.started" / "tool.completed" /
|
||||
# "tool.requires_confirmation"）
|
||||
on_tool_event: OnToolEventCallback | None = None,
|
||||
# 可选：确认 Future 注册表（阶段3 步骤1）。缺省 None 时遇 RequiresConfirmation 降级。
|
||||
confirmation_registry: ConfirmationRegistry | None = None,
|
||||
# 可选：单次确认挂起超时（秒），默认 300s；超时按拒绝处理。
|
||||
confirmation_timeout: float = DEFAULT_CONFIRMATION_TIMEOUT,
|
||||
# 可选：确认超时回调（confirmation_id → None），供调用方标记 DB 超时状态。
|
||||
on_confirm_timeout: Callable[[str], Awaitable[None]] | None = None,
|
||||
# 受约束扩展 Hook 与本次运行上下文。
|
||||
extensions: ExtensionRegistry | None = None,
|
||||
run_context: AgentRunContext | None = None,
|
||||
# Stateful Agent 的 steering 队列；每轮工具批次完成后消费。
|
||||
get_steering_messages: Callable[[], list[BaseMessage]] | None = None,
|
||||
) -> AgentTurnResult:
|
||||
"""执行一轮带工具调用的 Agent 循环（流式版）。
|
||||
|
||||
循环逻辑见模块 docstring；返回值始终是 ``AgentTurnResult``，工具失败不抛异常。
|
||||
``max_iterations`` 为硬上限：达到上限且模型仍返回 tool_calls 时返回 truncated=True。
|
||||
``on_chunk`` / ``on_tool_event`` 缺省 None 时与纯后端逻辑等价。
|
||||
``on_confirm_timeout(confirmation_id)``：确认超时时回调，用于标记 DB 记录。
|
||||
|
||||
确认挂起/恢复（阶段3 步骤1）：``confirmation_registry`` 缺省 None 时，遇
|
||||
``RequiresConfirmation`` 走降级路径（回填"需确认但未启用"消息，循环继续）；否则按
|
||||
register → 发 ``tool.requires_confirmation`` 事件 → ``await Future``（带超时）→
|
||||
批准重试 / 拒绝回填 的流程处理。``confirmation_timeout`` 须小于上层 run_timeout，否则
|
||||
run 先超时。
|
||||
"""
|
||||
if max_iterations < 1:
|
||||
raise ValueError("max_iterations 必须 >= 1")
|
||||
|
||||
manager = ProviderManagerFactory.create()
|
||||
extension_registry = extensions or ExtensionRegistry()
|
||||
extension_context = run_context or AgentRunContext(
|
||||
run_id=run_id,
|
||||
character_id=character_id,
|
||||
trace_id=trace_id,
|
||||
)
|
||||
usage = TokenUsage()
|
||||
tool_history: list[dict[str, Any]] = []
|
||||
iterations = 0
|
||||
# 最后一轮累积值（用于 truncated 分支返回 content/reasoning）
|
||||
last_content = ""
|
||||
last_reasoning = ""
|
||||
|
||||
while iterations < max_iterations:
|
||||
iterations += 1
|
||||
if on_turn is not None:
|
||||
await on_turn()
|
||||
if extension_registry.enabled:
|
||||
await extension_registry.publish(
|
||||
AgentEvent("agent.turn.started", {"iteration": iterations}),
|
||||
extension_context,
|
||||
)
|
||||
|
||||
# ---- 本轮流式消费 ----
|
||||
acc = ToolCallAccumulator()
|
||||
content_buf = ""
|
||||
reasoning_buf = ""
|
||||
finish_reason: str | None = None
|
||||
raw_blocks: dict[int, dict[str, Any]] = {}
|
||||
params: dict[str, object] = {}
|
||||
if thinking_enabled is not None:
|
||||
params["thinking_enabled"] = thinking_enabled
|
||||
if reasoning_effort is not None:
|
||||
params["reasoning_effort"] = reasoning_effort
|
||||
extension_context.cancellation.raise_if_cancelled()
|
||||
request_context = AgentContext(tuple(messages), tuple(tools), model_hint)
|
||||
if extension_registry.enabled:
|
||||
request_context = await extension_registry.transform_context(
|
||||
request_context,
|
||||
extension_context,
|
||||
)
|
||||
if extension_registry.enabled:
|
||||
await extension_registry.publish(
|
||||
AgentEvent("message.started", {"role": "assistant", "iteration": iterations}),
|
||||
extension_context,
|
||||
)
|
||||
stream = manager.chat_stream(
|
||||
list(request_context.messages),
|
||||
model_hint=request_context.model_hint,
|
||||
params=params,
|
||||
tools=list(request_context.tools) or None,
|
||||
)
|
||||
async for chunk in stream:
|
||||
if on_chunk is not None:
|
||||
await on_chunk(chunk)
|
||||
if extension_registry.enabled and (chunk.delta_content or chunk.delta_reasoning):
|
||||
await extension_registry.publish(
|
||||
AgentEvent(
|
||||
"message.updated",
|
||||
{
|
||||
"role": "assistant",
|
||||
"iteration": iterations,
|
||||
"delta": chunk.delta_content,
|
||||
"reasoning_delta": chunk.delta_reasoning,
|
||||
},
|
||||
),
|
||||
extension_context,
|
||||
)
|
||||
if chunk.delta_content:
|
||||
content_buf += chunk.delta_content
|
||||
if chunk.delta_reasoning:
|
||||
reasoning_buf += chunk.delta_reasoning
|
||||
if chunk.finish_reason is not None:
|
||||
finish_reason = chunk.finish_reason
|
||||
if chunk.content_block_start is not None and chunk.content_block_index is not None:
|
||||
raw_blocks[chunk.content_block_index] = dict(chunk.content_block_start)
|
||||
if chunk.content_block_delta is not None and chunk.content_block_index is not None:
|
||||
block = raw_blocks.get(chunk.content_block_index)
|
||||
if block is not None:
|
||||
delta = chunk.content_block_delta
|
||||
delta_type = delta.get("type")
|
||||
if delta_type == "text_delta":
|
||||
block["text"] = f"{block.get('text', '')}{delta.get('text', '')}"
|
||||
elif delta_type == "thinking_delta":
|
||||
block["thinking"] = (
|
||||
f"{block.get('thinking', '')}{delta.get('thinking', '')}"
|
||||
)
|
||||
elif delta_type == "signature_delta":
|
||||
block["signature"] = (
|
||||
f"{block.get('signature', '')}{delta.get('signature', '')}"
|
||||
)
|
||||
elif delta_type == "input_json_delta":
|
||||
block["_partial_json"] = (
|
||||
f"{block.get('_partial_json', '')}{delta.get('partial_json', '')}"
|
||||
)
|
||||
if chunk.usage is not None:
|
||||
usage = _sum_usage(usage, chunk.usage)
|
||||
acc.add_chunk(chunk)
|
||||
|
||||
# 流结束即本轮模型输出完毕
|
||||
tool_calls = acc.drain()
|
||||
last_content = content_buf
|
||||
last_reasoning = reasoning_buf
|
||||
|
||||
raw_content_blocks: list[dict[str, Any]] = []
|
||||
for index in sorted(raw_blocks):
|
||||
block = raw_blocks[index]
|
||||
partial_json = block.pop("_partial_json", None)
|
||||
if partial_json is not None:
|
||||
block["input"] = _parse_tool_args(partial_json)
|
||||
raw_content_blocks.append(block)
|
||||
|
||||
# 把模型输出（含 thinking signature / tool_calls）原样追加到历史。
|
||||
ai_message = _ai_message_from_stream(
|
||||
content_buf,
|
||||
reasoning_buf,
|
||||
tool_calls,
|
||||
raw_content_blocks or None,
|
||||
)
|
||||
messages.append(ai_message)
|
||||
if extension_registry.enabled:
|
||||
await extension_registry.publish(
|
||||
AgentEvent(
|
||||
"message.ended",
|
||||
{
|
||||
"role": "assistant",
|
||||
"iteration": iterations,
|
||||
"message": ai_message.to_dict(),
|
||||
"finish_reason": finish_reason,
|
||||
},
|
||||
),
|
||||
extension_context,
|
||||
)
|
||||
|
||||
# 无工具调用：先消费 steering，再决定是否结束。
|
||||
if not any(tc.get("name") for tc in tool_calls):
|
||||
steering = get_steering_messages() if get_steering_messages is not None else []
|
||||
turn = TurnSnapshot(
|
||||
iteration=iterations,
|
||||
content=content_buf,
|
||||
reasoning=reasoning_buf,
|
||||
tool_calls=(),
|
||||
tool_results=tuple(tool_history),
|
||||
usage=usage,
|
||||
)
|
||||
if extension_registry.enabled:
|
||||
await extension_registry.publish(
|
||||
AgentEvent("agent.turn.ended", {"iteration": iterations}),
|
||||
extension_context,
|
||||
)
|
||||
# TurnPolicy 在此分支不参与终止决策：模型无工具调用时已由 steering /
|
||||
# finish_reason / 默认 return 收尾，强行接入会覆盖 pause_turn / refusal
|
||||
# 既有处理；其真实终止作用在下方工具分支（should_stop or terminate_after_tools）。
|
||||
await extension_registry.should_stop(turn, extension_context)
|
||||
if steering:
|
||||
messages.extend(steering)
|
||||
continue
|
||||
if finish_reason == "pause_turn":
|
||||
continue
|
||||
if finish_reason == "refusal":
|
||||
raise RuntimeError("模型拒绝了当前请求。")
|
||||
return AgentTurnResult(
|
||||
content=content_buf,
|
||||
reasoning=reasoning_buf,
|
||||
tool_history=tool_history,
|
||||
iterations=iterations,
|
||||
truncated=finish_reason == "length",
|
||||
usage=usage,
|
||||
)
|
||||
|
||||
# 逐个执行 tool_call 并把结果以 ToolMessage 回填。工具可能有副作用或需要确认，
|
||||
# 在 ToolSpec 未声明并行安全前保持确定性串行执行。
|
||||
turn_history_start = len(tool_history)
|
||||
terminate_after_tools = False
|
||||
for call_index, entry in enumerate(tool_calls):
|
||||
name = entry.get("name")
|
||||
if not name:
|
||||
continue
|
||||
args_str = entry.get("args", "")
|
||||
tool_input = _parse_tool_args(args_str)
|
||||
tool_call_id = entry.get("id") or f"call_{iterations}_{call_index}_{name}"
|
||||
hook_call = ToolCallContext(
|
||||
tool_call_id=tool_call_id,
|
||||
name=name,
|
||||
input=tool_input,
|
||||
messages=tuple(messages),
|
||||
)
|
||||
if extension_registry.enabled:
|
||||
decision = await extension_registry.before_tool_call(
|
||||
hook_call,
|
||||
extension_context,
|
||||
)
|
||||
if not decision.allowed:
|
||||
denied_text = decision.reason or f"扩展阻止了工具 {name} 的执行。"
|
||||
tool_history.append(
|
||||
{
|
||||
"name": name,
|
||||
"args": tool_input,
|
||||
"result": denied_text,
|
||||
"status": "denied",
|
||||
"tool_call_id": tool_call_id,
|
||||
}
|
||||
)
|
||||
tool_message = ToolMessage(content=denied_text, tool_call_id=tool_call_id)
|
||||
messages.append(tool_message)
|
||||
await _publish_message_lifecycle(
|
||||
extension_registry, extension_context, tool_message
|
||||
)
|
||||
if on_tool_event is not None:
|
||||
await on_tool_event(
|
||||
"tool.completed",
|
||||
{
|
||||
"tool_name": name,
|
||||
"tool_call_id": tool_call_id,
|
||||
"status": "denied",
|
||||
"result": denied_text[:500],
|
||||
},
|
||||
)
|
||||
continue
|
||||
if decision.input is not None:
|
||||
tool_input = decision.input
|
||||
args_str = json.dumps(tool_input, ensure_ascii=False)
|
||||
hook_call = ToolCallContext(
|
||||
tool_call_id=tool_call_id,
|
||||
name=name,
|
||||
input=tool_input,
|
||||
messages=tuple(messages),
|
||||
)
|
||||
|
||||
# 执行前事件
|
||||
if on_tool_event is not None:
|
||||
await on_tool_event(
|
||||
"tool.started",
|
||||
{
|
||||
"tool_name": name,
|
||||
"tool_call_id": tool_call_id,
|
||||
"args": args_str,
|
||||
},
|
||||
)
|
||||
|
||||
try:
|
||||
tool_result = await _execute_tool_with_confirmation(
|
||||
tool_runtime=tool_runtime,
|
||||
name=name,
|
||||
tool_input=tool_input,
|
||||
args_str=args_str,
|
||||
tool_call_id=tool_call_id,
|
||||
run_id=run_id,
|
||||
character_id=character_id,
|
||||
trace_id=trace_id,
|
||||
base_confirmed=confirmed,
|
||||
confirmation_registry=confirmation_registry,
|
||||
confirmation_timeout=confirmation_timeout,
|
||||
on_tool_event=on_tool_event,
|
||||
cancellation=extension_context.cancellation,
|
||||
on_timeout=on_confirm_timeout,
|
||||
)
|
||||
except _ToolDenied as denied:
|
||||
# 用户拒绝或确认超时：回填拒绝消息，循环继续让模型自决策
|
||||
logger.info(
|
||||
"agent_tool_denied",
|
||||
tool=name,
|
||||
reason=denied.reason,
|
||||
run_id=run_id,
|
||||
)
|
||||
denied_text = denied.message
|
||||
tool_history.append(
|
||||
{
|
||||
"name": name,
|
||||
"args": tool_input,
|
||||
"result": denied_text,
|
||||
"status": "denied",
|
||||
"tool_call_id": tool_call_id,
|
||||
}
|
||||
)
|
||||
tool_message = ToolMessage(content=denied_text, tool_call_id=tool_call_id)
|
||||
messages.append(tool_message)
|
||||
await _publish_message_lifecycle(
|
||||
extension_registry, extension_context, tool_message
|
||||
)
|
||||
if on_tool_event is not None:
|
||||
await on_tool_event(
|
||||
"tool.completed",
|
||||
{
|
||||
"tool_name": name,
|
||||
"tool_call_id": tool_call_id,
|
||||
"status": "denied",
|
||||
"result": denied_text[:500],
|
||||
},
|
||||
)
|
||||
continue
|
||||
except _ToolSkipped as skipped:
|
||||
# registry 未启用：回填"需确认但未启用"，循环继续
|
||||
skipped_text = skipped.message
|
||||
tool_history.append(
|
||||
{
|
||||
"name": name,
|
||||
"args": tool_input,
|
||||
"result": skipped_text,
|
||||
"status": "skipped",
|
||||
"tool_call_id": tool_call_id,
|
||||
}
|
||||
)
|
||||
tool_message = ToolMessage(content=skipped_text, tool_call_id=tool_call_id)
|
||||
messages.append(tool_message)
|
||||
await _publish_message_lifecycle(
|
||||
extension_registry, extension_context, tool_message
|
||||
)
|
||||
if on_tool_event is not None:
|
||||
await on_tool_event(
|
||||
"tool.completed",
|
||||
{
|
||||
"tool_name": name,
|
||||
"tool_call_id": tool_call_id,
|
||||
"status": "skipped",
|
||||
"result": skipped_text[:500],
|
||||
},
|
||||
)
|
||||
continue
|
||||
except Exception as exc: # noqa: BLE001 - 任意工具错误都回填、不中断循环
|
||||
# 司命 deny / 沙箱错误 / 未知工具等：错误信息回填，循环继续
|
||||
logger.warning(
|
||||
"agent_tool_failed",
|
||||
tool=name,
|
||||
error=str(exc),
|
||||
run_id=run_id,
|
||||
)
|
||||
error_text = f"工具执行失败：{type(exc).__name__}: {_truncate(str(exc))}"
|
||||
tool_history.append(
|
||||
{
|
||||
"name": name,
|
||||
"args": tool_input,
|
||||
"result": error_text,
|
||||
"status": "error",
|
||||
"tool_call_id": tool_call_id,
|
||||
}
|
||||
)
|
||||
tool_message = ToolMessage(content=error_text, tool_call_id=tool_call_id)
|
||||
messages.append(tool_message)
|
||||
await _publish_message_lifecycle(
|
||||
extension_registry, extension_context, tool_message
|
||||
)
|
||||
if on_tool_event is not None:
|
||||
await on_tool_event(
|
||||
"tool.completed",
|
||||
{
|
||||
"tool_name": name,
|
||||
"tool_call_id": tool_call_id,
|
||||
"status": "error",
|
||||
"result": error_text[:500],
|
||||
},
|
||||
)
|
||||
continue
|
||||
|
||||
# 成功路径：扩展可变换结果或提示终止；安全裁定已在 ToolRuntime 内完成。
|
||||
if extension_registry.enabled:
|
||||
after = await extension_registry.after_tool_call(
|
||||
hook_call,
|
||||
tool_result,
|
||||
extension_context,
|
||||
)
|
||||
tool_result = after.result or tool_result
|
||||
terminate_after_tools = terminate_after_tools or after.terminate
|
||||
outcome = _tool_result_text(tool_result)
|
||||
status = "success" if tool_result.ok else "error"
|
||||
if not tool_result.ok:
|
||||
# 沙箱执行返回非 ok（exit_code != 0 等）——不抛异常，回填让模型自决策
|
||||
logger.info(
|
||||
"agent_tool_returned_error",
|
||||
tool=name,
|
||||
exit_code=tool_result.exit_code,
|
||||
stderr=tool_result.stderr[:500] if tool_result.stderr else "",
|
||||
run_id=run_id,
|
||||
)
|
||||
tool_history.append(
|
||||
{
|
||||
"name": name,
|
||||
"args": tool_input,
|
||||
"result": outcome,
|
||||
"status": status,
|
||||
"tool_call_id": tool_call_id,
|
||||
}
|
||||
)
|
||||
tool_message = ToolMessage(content=outcome, tool_call_id=tool_call_id)
|
||||
messages.append(tool_message)
|
||||
await _publish_message_lifecycle(extension_registry, extension_context, tool_message)
|
||||
if on_tool_event is not None:
|
||||
await on_tool_event(
|
||||
"tool.completed",
|
||||
{
|
||||
"tool_name": name,
|
||||
"tool_call_id": tool_call_id,
|
||||
"status": status,
|
||||
"result": outcome[:500],
|
||||
},
|
||||
)
|
||||
|
||||
steering = get_steering_messages() if get_steering_messages is not None else []
|
||||
if steering:
|
||||
messages.extend(steering)
|
||||
turn = TurnSnapshot(
|
||||
iteration=iterations,
|
||||
content=content_buf,
|
||||
reasoning=reasoning_buf,
|
||||
tool_calls=tuple(tool_calls),
|
||||
tool_results=tuple(tool_history[turn_history_start:]),
|
||||
usage=usage,
|
||||
)
|
||||
should_stop = terminate_after_tools
|
||||
if extension_registry.enabled:
|
||||
await extension_registry.publish(
|
||||
AgentEvent("agent.turn.ended", {"iteration": iterations}),
|
||||
extension_context,
|
||||
)
|
||||
should_stop = should_stop or await extension_registry.should_stop(
|
||||
turn,
|
||||
extension_context,
|
||||
)
|
||||
if should_stop and not steering:
|
||||
return AgentTurnResult(
|
||||
content=content_buf,
|
||||
reasoning=reasoning_buf,
|
||||
tool_history=tool_history,
|
||||
iterations=iterations,
|
||||
truncated=False,
|
||||
usage=usage,
|
||||
)
|
||||
|
||||
# 达上限仍带 tool_calls：截断返回，content 用最后一次模型输出或说明已截断
|
||||
content = last_content or f"Agent 循环达到最大次数 {max_iterations}，已截断。"
|
||||
return AgentTurnResult(
|
||||
content=content,
|
||||
reasoning=last_reasoning,
|
||||
tool_history=tool_history,
|
||||
iterations=iterations,
|
||||
truncated=True,
|
||||
usage=usage,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"DEFAULT_CONFIRMATION_TIMEOUT",
|
||||
"DEFAULT_MAX_ITERATIONS",
|
||||
"AgentTurnResult",
|
||||
"OnChunkCallback",
|
||||
"OnToolEventCallback",
|
||||
"run_agent_turn",
|
||||
]
|
||||
@@ -1,22 +0,0 @@
|
||||
"""从既有 Skill 存储加载运行时 manifest。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from app.modules.skills import SkillManifest, parse_skill_manifest
|
||||
from app.repositories.json_utils import loads_list
|
||||
from app.services.skills import get_enabled_skills_by_names, match_active_skills
|
||||
|
||||
|
||||
async def load_skill_manifests(
|
||||
context_text: str,
|
||||
*,
|
||||
skill_refs: tuple[str, ...] = (),
|
||||
) -> list[SkillManifest]:
|
||||
"""加载关键词命中或显式引用的 enabled Skill。"""
|
||||
skills = await match_active_skills(context_text)
|
||||
if skill_refs:
|
||||
by_name = {skill.name: skill for skill in skills}
|
||||
for skill in await get_enabled_skills_by_names(skill_refs):
|
||||
by_name.setdefault(skill.name, skill)
|
||||
skills = list(by_name.values())
|
||||
return [parse_skill_manifest(skill, loads_list(skill.tool_bindings_json)) for skill in skills]
|
||||
@@ -1,95 +0,0 @@
|
||||
"""将洛书规划结果映射为通用 Agent 任务，并写入 Run 事件。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from app.modules.agent import AgentTask, Plan, PlanStep
|
||||
from app.modules.planner import DeliberativePlan
|
||||
|
||||
|
||||
def task_from_deliberative_plan(
|
||||
plan: DeliberativePlan,
|
||||
*,
|
||||
run_id: str,
|
||||
agent_name: str,
|
||||
objective: str,
|
||||
) -> AgentTask:
|
||||
"""把一次规划结果转换为仅在当前 Run 存活的 AgentTask。"""
|
||||
return AgentTask(
|
||||
id=f"task_{run_id}",
|
||||
objective=objective,
|
||||
agent_name=agent_name,
|
||||
run_id=run_id,
|
||||
status="running",
|
||||
plan=Plan(
|
||||
steps=tuple(
|
||||
PlanStep(
|
||||
id=str(step.id),
|
||||
title=step.title,
|
||||
description=step.description,
|
||||
tool=step.tool,
|
||||
risk_level=step.risk_level,
|
||||
)
|
||||
for step in plan.steps
|
||||
),
|
||||
assumptions=tuple(plan.assumptions),
|
||||
risks=tuple(plan.risks),
|
||||
status="running",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def next_step_for_tool(task: AgentTask, tool_name: str, completed_step_ids: set[str]) -> PlanStep | None:
|
||||
"""返回尚未处理且明确绑定到工具的首个计划步骤。"""
|
||||
if task.plan is None:
|
||||
return None
|
||||
return next(
|
||||
(
|
||||
step
|
||||
for step in task.plan.steps
|
||||
if step.id not in completed_step_ids and step.tool == tool_name
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
|
||||
def step_event_payload(
|
||||
task: AgentTask,
|
||||
step: PlanStep,
|
||||
*,
|
||||
status: str,
|
||||
result_summary: str | None = None,
|
||||
) -> dict[str, object]:
|
||||
"""生成计划步骤生命周期事件 payload。"""
|
||||
return {
|
||||
"task_id": task.id,
|
||||
"step_id": step.id,
|
||||
"title": step.title,
|
||||
"tool": step.tool,
|
||||
"risk_level": step.risk_level,
|
||||
"status": status,
|
||||
**({"result_summary": result_summary} if result_summary else {}),
|
||||
}
|
||||
|
||||
|
||||
def task_plan_payload(task: AgentTask) -> dict[str, object]:
|
||||
"""生成安全且可直接展示的计划事件 payload。"""
|
||||
assert task.plan is not None
|
||||
return {
|
||||
"task_id": task.id,
|
||||
"agent_name": task.agent_name,
|
||||
"objective": task.objective,
|
||||
"status": task.plan.status,
|
||||
"assumptions": list(task.plan.assumptions),
|
||||
"risks": list(task.plan.risks),
|
||||
"steps": [
|
||||
{
|
||||
"id": step.id,
|
||||
"title": step.title,
|
||||
"description": step.description,
|
||||
"tool": step.tool,
|
||||
"risk_level": step.risk_level,
|
||||
"status": step.status,
|
||||
}
|
||||
for step in task.plan.steps
|
||||
],
|
||||
}
|
||||
@@ -1,137 +0,0 @@
|
||||
"""ToolSpec → OpenAI 函数 schema 的转换与默认工具集选取。
|
||||
|
||||
直接产出 OpenAI function schema dict（``{"type": "function", "function": {...}}``），
|
||||
与 ChatResult.tool_calls 里的 name/args 回填闭环天然对齐。
|
||||
|
||||
ToolSpec 的 name/description 直接映射；参数 schema 因 ToolSpec 未内嵌 JSON Schema，
|
||||
这里按工具 name 查表给出最小参数 schema（builtins 内工具参数固定，可选参数也显式列出）。
|
||||
未知工具默认空参数 schema（``{"type": "object", "properties": {}}``），模型按需传参。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from app.modules.tool_runtime.definition import ToolSpec
|
||||
|
||||
# 内置工具的参数 schema（name → JSON Schema）。新增工具时扩展此表即可。
|
||||
# 保持最小集合：只声明 known 工具的参数；调用方传入未登记工具时回退空对象 schema。
|
||||
_PARAM_SCHEMAS: dict[str, dict[str, Any]] = {
|
||||
"datetime": {"type": "object", "properties": {}, "required": []},
|
||||
"read_file": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"path": {"type": "string", "description": "workspace 内的文件相对路径。"},
|
||||
"offset": {
|
||||
"type": "integer",
|
||||
"description": "起始行号（1-based），默认 1。大文件分段读取后续内容时使用。",
|
||||
},
|
||||
"limit": {
|
||||
"type": "integer",
|
||||
"description": "读取的行数上限；不传则读到末尾。大文件分段读取时使用。",
|
||||
},
|
||||
},
|
||||
"required": ["path"],
|
||||
},
|
||||
"list_dir": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"path": {
|
||||
"type": "string",
|
||||
"description": "workspace 内的目录相对路径，缺省为 '.'。",
|
||||
},
|
||||
},
|
||||
"required": [],
|
||||
},
|
||||
"write_file": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"path": {"type": "string", "description": "workspace 内文件相对路径。"},
|
||||
"content": {"type": "string", "description": "要写入的文本内容。"},
|
||||
},
|
||||
"required": ["path"],
|
||||
},
|
||||
"shell": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"command": {
|
||||
"type": "string",
|
||||
"description": "要在沙箱中执行的 shell 命令。工作目录为沙箱隔离的 workspace 目录（并非 /home/user 或 /tmp 等 Linux 路径，这些路径不存在，不要 cd 到它们）。read_file/write_file 操作的也是这个 workspace 目录，shell 里 cd 到别的目录将看不到这些文件。",
|
||||
},
|
||||
"timeout_ms": {"type": "integer", "description": "可选超时（毫秒），默认 30000。"},
|
||||
"stdin": {"type": "string", "description": "可选标准输入。"},
|
||||
},
|
||||
"required": ["command"],
|
||||
},
|
||||
"fetch_url": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"url": {"type": "string", "description": "要抓取的网页 URL。"},
|
||||
"max_chars": {
|
||||
"type": "integer",
|
||||
"description": "正文最大字符数，默认 8000。",
|
||||
},
|
||||
},
|
||||
"required": ["url"],
|
||||
},
|
||||
}
|
||||
|
||||
_EMPTY_SCHEMA: dict[str, Any] = {"type": "object", "properties": {}, "required": []}
|
||||
|
||||
|
||||
def to_ai_provider_tools(tool_specs: list[ToolSpec]) -> list[dict[str, Any]]:
|
||||
"""把 ToolSpec 列表转成 OpenAI function schema dict 列表。
|
||||
|
||||
先查 ``_PARAM_SCHEMAS``（内置工具），未命中则用 ``spec.input_schema``（MCP 工具），
|
||||
仍为空则回退无参占位 ``_EMPTY_SCHEMA``。
|
||||
"""
|
||||
tools: list[dict[str, Any]] = []
|
||||
for spec in tool_specs:
|
||||
params = _PARAM_SCHEMAS.get(spec.name)
|
||||
if not params and spec.name.startswith("skill:"):
|
||||
params = _PARAM_SCHEMAS.get(spec.name.rsplit("/", 1)[-1])
|
||||
if not params:
|
||||
params = spec.input_schema or _EMPTY_SCHEMA
|
||||
tools.append(
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": spec.name,
|
||||
"description": spec.description,
|
||||
"parameters": params,
|
||||
},
|
||||
}
|
||||
)
|
||||
return tools
|
||||
|
||||
|
||||
# 核心工具恒绑（沙箱类，始终可用）。
|
||||
_CORE_TOOL_NAMES = ["datetime", "read_file", "list_dir", "write_file", "shell"]
|
||||
|
||||
|
||||
def default_tools(
|
||||
tool_specs_by_name: dict[str, ToolSpec] | None = None,
|
||||
*,
|
||||
allowed_tools: tuple[str, ...] = (),
|
||||
) -> list[dict[str, Any]]:
|
||||
"""返回工具对应的 OpenAI function schema 列表。
|
||||
|
||||
- 未传 ``tool_specs_by_name`` 时，从 ``builtin_specs()`` 加载仅返回核心工具
|
||||
（datetime/read_file/list_dir/write_file/shell），共 5 个。
|
||||
- 传入自定义工具池时（如 ``tool_runtime._resolve_specs()``），返回池中所有工具，
|
||||
含 MCP 注册的工具，不限于核心工具。
|
||||
"""
|
||||
if tool_specs_by_name is None:
|
||||
from app.application.tool_runtime.builtins import builtin_specs
|
||||
|
||||
tool_specs_by_name = builtin_specs()
|
||||
chosen = [
|
||||
tool_specs_by_name[name] for name in _CORE_TOOL_NAMES if name in tool_specs_by_name
|
||||
]
|
||||
else:
|
||||
chosen = list(tool_specs_by_name.values())
|
||||
|
||||
if allowed_tools:
|
||||
allowed = set(allowed_tools)
|
||||
chosen = [spec for spec in chosen if spec.name in allowed]
|
||||
return to_ai_provider_tools(chosen)
|
||||
@@ -1,198 +0,0 @@
|
||||
"""资料导入候选抽取用例（异步后台任务，跨 services + ai_provider）。
|
||||
|
||||
触发端点在校验并置 job 为 ``extracting_*`` 后即返回；实际模型调用在后台任务中
|
||||
进行（独立 session），完成后落候选并置 job 为 ``reviewing`` / ``parsed``。
|
||||
仅适用单进程；多 worker 部署需替换为跨进程任务队列。
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.ai_provider import ProviderManager, ProviderManagerFactory
|
||||
from app.core.logging import logger
|
||||
from app.db.session import AsyncSessionLocal
|
||||
from app.models.candidates import MemoryCandidate, ProfileCandidate
|
||||
from app.models.imports import CharacterImportJob, ImportItem
|
||||
from app.models.personality import PersonalityProposal
|
||||
from app.modules.imports.extract import (
|
||||
parse_memory_candidates,
|
||||
parse_personality_proposals,
|
||||
parse_profile_candidates,
|
||||
)
|
||||
from app.modules.imports.prompts import (
|
||||
build_memory_extraction_messages,
|
||||
build_personality_extraction_messages,
|
||||
build_profile_extraction_messages,
|
||||
)
|
||||
from app.repositories import candidates as candidates_repo
|
||||
from app.repositories import imports as imports_repo
|
||||
from app.repositories.json_utils import dumps
|
||||
|
||||
# 后台抽取任务强引用：asyncio.create_task 仅保留弱引用，需强引用防 GC。
|
||||
_running_extraction_tasks: set[asyncio.Task[None]] = set()
|
||||
|
||||
_provider_manager: ProviderManager | None = None
|
||||
|
||||
|
||||
def _get_provider_manager() -> ProviderManager:
|
||||
global _provider_manager
|
||||
if _provider_manager is None:
|
||||
_provider_manager = ProviderManagerFactory.create()
|
||||
return _provider_manager
|
||||
|
||||
|
||||
_PROFILE_MODEL_HINT = "personality-analysis"
|
||||
_MEMORY_MODEL_HINT = "memory-extraction"
|
||||
_PERSONALITY_MODEL_HINT = "personality-analysis"
|
||||
|
||||
|
||||
async def drain_pending_extractions() -> None:
|
||||
"""等待所有进行中的后台抽取任务完成（测试断言 / 优雅停机用）。"""
|
||||
if _running_extraction_tasks:
|
||||
await asyncio.gather(*_running_extraction_tasks, return_exceptions=True)
|
||||
|
||||
|
||||
def start_profile_extraction(job_id: str) -> None:
|
||||
task = asyncio.create_task(_run_profile_extraction(job_id))
|
||||
_running_extraction_tasks.add(task)
|
||||
task.add_done_callback(_running_extraction_tasks.discard)
|
||||
|
||||
|
||||
def start_memory_extraction(job_id: str) -> None:
|
||||
task = asyncio.create_task(_run_memory_extraction(job_id))
|
||||
_running_extraction_tasks.add(task)
|
||||
task.add_done_callback(_running_extraction_tasks.discard)
|
||||
|
||||
|
||||
def start_personality_extraction(job_id: str) -> None:
|
||||
task = asyncio.create_task(_run_personality_extraction(job_id))
|
||||
_running_extraction_tasks.add(task)
|
||||
task.add_done_callback(_running_extraction_tasks.discard)
|
||||
|
||||
|
||||
async def _load_job_and_items(
|
||||
session: AsyncSession, job_id: str
|
||||
) -> tuple[CharacterImportJob | None, list[ImportItem]]:
|
||||
job = await imports_repo.get_job(session, job_id)
|
||||
if job is None:
|
||||
return None, []
|
||||
result = await session.execute(
|
||||
select(ImportItem)
|
||||
.where(ImportItem.import_job_id == job_id)
|
||||
.order_by(ImportItem.sequence.asc())
|
||||
)
|
||||
return job, list(result.scalars().all())
|
||||
|
||||
|
||||
async def _run_profile_extraction(job_id: str) -> None:
|
||||
async with AsyncSessionLocal() as session:
|
||||
try:
|
||||
job, items = await _load_job_and_items(session, job_id)
|
||||
if job is None:
|
||||
return
|
||||
system, user = build_profile_extraction_messages(_pairs(items))
|
||||
raw = await _call_model(system, user, _PROFILE_MODEL_HINT, job.character_id)
|
||||
candidates = [
|
||||
ProfileCandidate(
|
||||
character_id=job.character_id,
|
||||
import_job_id=job.id,
|
||||
field=draft.field_name,
|
||||
value=draft.value,
|
||||
confidence=draft.confidence,
|
||||
reason=draft.reason,
|
||||
status="pending",
|
||||
evidence_ids_json=dumps(_evidence_ids(items, draft.evidence)),
|
||||
)
|
||||
for draft in parse_profile_candidates(raw)
|
||||
]
|
||||
await candidates_repo.add_profile_candidates(session, candidates)
|
||||
job.status = "reviewing" if candidates else "parsed"
|
||||
await session.commit()
|
||||
except Exception as exc:
|
||||
await _fail_job(session, job_id, exc)
|
||||
|
||||
|
||||
async def _run_memory_extraction(job_id: str) -> None:
|
||||
async with AsyncSessionLocal() as session:
|
||||
try:
|
||||
job, items = await _load_job_and_items(session, job_id)
|
||||
if job is None:
|
||||
return
|
||||
system, user = build_memory_extraction_messages(_pairs(items))
|
||||
raw = await _call_model(system, user, _MEMORY_MODEL_HINT, job.character_id)
|
||||
candidates = [
|
||||
MemoryCandidate(
|
||||
character_id=job.character_id,
|
||||
import_job_id=job.id,
|
||||
type=draft.type,
|
||||
title=draft.title,
|
||||
content=draft.content,
|
||||
confidence=0.6,
|
||||
importance=draft.importance,
|
||||
sensitivity=draft.sensitivity,
|
||||
reason=draft.reason,
|
||||
status="pending",
|
||||
evidence_ids_json=dumps(_evidence_ids(items, draft.evidence)),
|
||||
)
|
||||
for draft in parse_memory_candidates(raw)
|
||||
]
|
||||
await candidates_repo.add_memory_candidates(session, candidates)
|
||||
job.status = "reviewing" if candidates else "parsed"
|
||||
await session.commit()
|
||||
except Exception as exc:
|
||||
await _fail_job(session, job_id, exc)
|
||||
|
||||
|
||||
async def _run_personality_extraction(job_id: str) -> None:
|
||||
async with AsyncSessionLocal() as session:
|
||||
try:
|
||||
job, items = await _load_job_and_items(session, job_id)
|
||||
if job is None:
|
||||
return
|
||||
system, user = build_personality_extraction_messages(_pairs(items), "")
|
||||
raw = await _call_model(system, user, _PERSONALITY_MODEL_HINT, job.character_id)
|
||||
proposals = [
|
||||
PersonalityProposal(
|
||||
character_id=job.character_id,
|
||||
fields_json=dumps({draft.field_name: draft.change}),
|
||||
reason=draft.reason or "资料导入抽取的人格完善建议",
|
||||
status="pending",
|
||||
)
|
||||
for draft in parse_personality_proposals(raw)
|
||||
]
|
||||
session.add_all(proposals)
|
||||
await session.flush()
|
||||
job.status = "reviewing" if proposals else "parsed"
|
||||
await session.commit()
|
||||
except Exception as exc:
|
||||
await _fail_job(session, job_id, exc)
|
||||
|
||||
|
||||
def _pairs(items: list[ImportItem]) -> list[tuple[str, str]]:
|
||||
return [(item.role, item.content) for item in items]
|
||||
|
||||
|
||||
def _evidence_ids(items: list[ImportItem], indices: list[int]) -> list[str]:
|
||||
return [items[i].id for i in indices if 0 <= i < len(items)]
|
||||
|
||||
|
||||
async def _call_model(system: str, user: str, model_hint: str, _character_id: str) -> str:
|
||||
manager = _get_provider_manager()
|
||||
messages = manager.to_message_list([system], [("user", user)])
|
||||
response = await manager.chat(messages, model_hint=model_hint)
|
||||
return response.content
|
||||
|
||||
|
||||
async def _fail_job(session: AsyncSession, job_id: str, exc: Exception) -> None:
|
||||
logger.error(f"导入候选抽取失败 job={job_id}: {exc}")
|
||||
try:
|
||||
job = await imports_repo.get_job(session, job_id)
|
||||
if job is not None:
|
||||
job.status = "failed"
|
||||
job.error_code = "extraction_failed"
|
||||
job.error_message = str(exc)[:500]
|
||||
await session.commit()
|
||||
except Exception: # noqa: BLE001 — 失败处理本身不能再抛
|
||||
pass
|
||||
@@ -1,83 +0,0 @@
|
||||
"""声明式 Skill 激活：匹配、提示注入、Run 级工具别名和事件。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from app.application.agent.prompt_composer import compose_system_messages
|
||||
from app.application.agent.skill_loader import load_skill_manifests
|
||||
from app.modules.runtime.events import Event
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from app.application.tool_runtime.view import RunScopedToolView
|
||||
from app.modules.runtime.events import EventStore
|
||||
|
||||
|
||||
class SkillExecutor:
|
||||
"""把数据库 Skill 激活到当前 Run 的工具视图。"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
tool_view: RunScopedToolView | None = None,
|
||||
event_store: EventStore | None = None,
|
||||
) -> None:
|
||||
self._tool_view = tool_view
|
||||
self._event_store = event_store
|
||||
|
||||
async def match_and_activate(
|
||||
self,
|
||||
context_text: str,
|
||||
system_messages: list[str],
|
||||
*,
|
||||
run_id: str | None = None,
|
||||
character_id: str | None = None,
|
||||
skill_refs: tuple[str, ...] = (),
|
||||
) -> list[str]:
|
||||
"""匹配技能并把提示与别名写入当前 Run。"""
|
||||
if not context_text:
|
||||
return system_messages
|
||||
|
||||
manifests = await load_skill_manifests(context_text, skill_refs=skill_refs)
|
||||
system_messages.extend(compose_system_messages([], None, manifests))
|
||||
activated_names: list[str] = []
|
||||
for manifest in manifests:
|
||||
if self._tool_view is not None:
|
||||
for tool_name in manifest.allowed_tools:
|
||||
self._tool_view.bind_alias(
|
||||
f"skill:{manifest.name}/{tool_name}",
|
||||
tool_name,
|
||||
description_prefix=f"[{manifest.name}] ",
|
||||
)
|
||||
activated_names.append(manifest.name)
|
||||
|
||||
if activated_names and self._event_store:
|
||||
await self._emit_activation_event(
|
||||
activated_names,
|
||||
run_id=run_id,
|
||||
character_id=character_id,
|
||||
)
|
||||
return system_messages
|
||||
|
||||
async def _emit_activation_event(
|
||||
self,
|
||||
skill_names: list[str],
|
||||
*,
|
||||
run_id: str | None = None,
|
||||
character_id: str | None = None,
|
||||
) -> None:
|
||||
event_store = self._event_store
|
||||
if event_store is None:
|
||||
return
|
||||
await event_store.append(
|
||||
Event(
|
||||
type="skill.activated",
|
||||
source="skill_executor",
|
||||
visibility="public",
|
||||
run_id=run_id,
|
||||
character_id=character_id,
|
||||
payload={"skills": skill_names},
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["SkillExecutor"]
|
||||
@@ -5,9 +5,6 @@
|
||||
|
||||
- ``.env`` 永远放在项目根目录（backend 父级），``Settings`` 初始化时以**绝对路径**
|
||||
通过 ``_env_file`` 传入，不使用相对路径，不依赖 ``python-dotenv``。
|
||||
- ``ai_provider/`` 只**消费**本目录的配置，不持有自己的
|
||||
``settings.py`` / ``config.py``；领域数据结构（``ModelConfig`` 等）保留在
|
||||
``ai_provider/config/models.py``。
|
||||
"""
|
||||
|
||||
from app.core.config.sandbox_config import SandboxSettings
|
||||
|
||||
@@ -1,46 +0,0 @@
|
||||
"""模型网关（神识）env 配置：各 provider 的 API Key / Base URL 与 endpoint 池策略。
|
||||
|
||||
注意区分**配置字段**与**领域数据结构**：
|
||||
- 本文件只承载「从 env 读的开关 / 参数 / 密钥」。
|
||||
- 领域数据结构（``ModelConfig`` / ``ModelEndpointConfig``，从 ``config/models.yaml`` 加载）
|
||||
保留在 ``ai_provider/config/models.py``，不外迁到 ``core/config``。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pydantic_settings import BaseSettings
|
||||
|
||||
|
||||
class ModelGatewayConfig(BaseSettings):
|
||||
"""模型网关 env 配置。"""
|
||||
|
||||
# ==================== 模型配置文件 ====================
|
||||
model_config_path: str = "../config/models.yaml" # 模型配置文件路径（相对 backend）
|
||||
|
||||
# ==================== provider API Key / Base URL ====================
|
||||
openai_api_key: str = "" # OpenAI API Key
|
||||
openai_base_url: str = "https://api.openai.com/v1" # OpenAI API 地址
|
||||
deepseek_api_key: str = "" # DeepSeek API Key
|
||||
deepseek_base_url: str = "https://api.deepseek.com" # DeepSeek API 地址（OpenAI 兼容）
|
||||
xiaomi_mimo_api_key: str = "" # 小米 MiMo API Key
|
||||
xiaomi_mimo_base_url: str = "https://api.xiaomimimo.com/v1" # 小米 MiMo API 地址
|
||||
anthropic_api_key: str = "" # Anthropic API Key
|
||||
anthropic_base_url: str = "https://api.anthropic.com" # Anthropic API 地址
|
||||
volcengine_ark_api_key: str = "" # 火山方舟 / 豆包 API Key
|
||||
volcengine_ark_base_url: str = "https://ark.cn-beijing.volces.com/api/v3" # 火山方舟 API 地址
|
||||
custom_llm_api_key: str = "" # 自定义 LLM API Key
|
||||
custom_llm_base_url: str = "" # 自定义 LLM API 地址
|
||||
yuanqi_api_key: str = "" # 元炁本地推理引擎 API Key
|
||||
yuanqi_base_url: str = "http://localhost:11434" # 元炁本地推理引擎 API 地址
|
||||
|
||||
# ==================== endpoint 池配置 ====================
|
||||
# 物理 endpoint 池路由、failover 与重试策略（见 ai_provider/pool.py 与 resilience.py）
|
||||
model_endpoint_pool_enabled: bool = True # 是否启用 endpoint 池路由（False 时走逻辑模型直连）
|
||||
model_endpoint_pool_dry_run_selection: bool = False # dry-run：只预选并记录，不实际路由
|
||||
model_endpoint_pool_expose_debug_headers: bool = False # 是否在响应头暴露 endpoint 路由调试信息
|
||||
model_endpoint_pool_agent_failover_enabled: bool = (
|
||||
False # Agent 调用是否允许跨 endpoint failover
|
||||
)
|
||||
model_endpoint_pool_per_endpoint_max_retries: int = 1 # 单 endpoint 失败时的本地重试次数
|
||||
model_endpoint_pool_retry_base_delay_seconds: float = 0.25 # endpoint 本地重试基础退避（秒）
|
||||
model_endpoint_pool_retry_max_delay_seconds: float = 2.0 # endpoint 本地重试最大退避（秒）
|
||||
@@ -1,7 +1,7 @@
|
||||
"""主 Settings：多继承所有子配置，统一以绝对路径 ``_env_file`` 加载 ``.env``。
|
||||
|
||||
子配置见同目录下的 ``app_config`` / ``auth_config`` / ``database_config`` / ``log_config`` /
|
||||
``runtime_config`` / ``model_config`` / ``sandbox_config``。本文件聚合字段，并承载
|
||||
``runtime_config`` / ``sandbox_config``。本文件聚合字段，并承载
|
||||
跨子配置的路径解析与计算属性。
|
||||
|
||||
配置优先级（自高到低）：命令行参数 > 系统环境变量 > ``.env`` > 类属性默认值。
|
||||
@@ -18,7 +18,6 @@ from app.core.config.app_config import AppConfig
|
||||
from app.core.config.auth_config import AuthConfig
|
||||
from app.core.config.database_config import DatabaseConfig
|
||||
from app.core.config.log_config import LogConfig
|
||||
from app.core.config.model_config import ModelGatewayConfig
|
||||
from app.core.config.runtime_config import RuntimeConfig
|
||||
from app.core.config.sandbox_config import SandboxSettings
|
||||
from app.core.config.time_config import TimeConfig
|
||||
@@ -35,7 +34,6 @@ class Settings(
|
||||
DatabaseConfig,
|
||||
LogConfig,
|
||||
RuntimeConfig,
|
||||
ModelGatewayConfig,
|
||||
SandboxSettings,
|
||||
TimeConfig,
|
||||
BaseSettings,
|
||||
@@ -117,21 +115,6 @@ class Settings(
|
||||
return example
|
||||
return configured
|
||||
|
||||
@cached_property
|
||||
def resolved_model_config_path(self) -> Path:
|
||||
"""解析模型配置文件路径。
|
||||
|
||||
如果配置文件不存在，尝试使用示例配置文件。
|
||||
"""
|
||||
configured = self.resolve_from_backend(self.model_config_path)
|
||||
if configured.exists():
|
||||
return configured
|
||||
# 如果配置文件不存在，尝试使用示例配置
|
||||
example = self.project_root / "config" / "models.example.yaml"
|
||||
if example.exists():
|
||||
return example
|
||||
return configured
|
||||
|
||||
@cached_property
|
||||
def resolved_avatar_dir(self) -> Path:
|
||||
"""解析角色头像存储目录为绝对路径。"""
|
||||
|
||||
@@ -15,7 +15,6 @@ from fastapi.middleware.cors import CORSMiddleware
|
||||
from sqlalchemy.exc import OperationalError
|
||||
|
||||
from app.api.router import api_router
|
||||
from app.application.agent.confirmation import ConfirmationRegistry
|
||||
from app.application.tool_runtime.runtime import ToolRuntime
|
||||
from app.core.config import settings
|
||||
from app.core.errors import register_exception_handlers, trace_id_var
|
||||
@@ -71,14 +70,6 @@ async def lifespan(_app: FastAPI) -> AsyncGenerator[None, None]:
|
||||
except Exception: # noqa: BLE001
|
||||
log_with_context("SQLite WAL 模式设置失败，已跳过", level="WARNING")
|
||||
|
||||
# 加载模型网关 endpoint 覆盖缓存（DB 优先于 yaml/env）；迁移未完或读失败则降级跳过
|
||||
try:
|
||||
from app.services.ai_provider_admin import reload_provider_manager
|
||||
|
||||
await reload_provider_manager()
|
||||
except Exception: # noqa: BLE001
|
||||
log_with_context("模型网关覆盖缓存加载失败，已跳过", level="WARNING")
|
||||
|
||||
# 运行时注入：事件存储始终就绪（DbEventStore 持久化到 events 表，重启不丢）；
|
||||
# 工具运行时依赖沙箱配置，未配置（缺 SANDBOX_AUTH_SECRET 等）则降级 None（工具 API 503）。
|
||||
# 事件流与历史查询不依赖沙箱，始终可用。
|
||||
@@ -122,16 +113,6 @@ async def lifespan(_app: FastAPI) -> AsyncGenerator[None, None]:
|
||||
|
||||
_app.state.confirmation_store = DbConfirmationStore(AsyncSessionLocal)
|
||||
log_with_context("confirmation_store_set_to_db")
|
||||
# 进程内检查点确认注册表：混合唤醒（同进程 Event + 重启后 DB polling）。
|
||||
_app.state.confirmation_registry = ConfirmationRegistry()
|
||||
from app.application.agent import AgentRunRegistry, discover_python_skills
|
||||
|
||||
_app.state.agent_run_registry = AgentRunRegistry()
|
||||
_app.state.python_skill_catalog = discover_python_skills(("app.skills",))
|
||||
log_with_context(
|
||||
"Python Skills 目录就绪",
|
||||
skills=[skill.name for skill in _app.state.python_skill_catalog.skills],
|
||||
)
|
||||
sandbox_http: httpx.AsyncClient | None = None
|
||||
tools_http: httpx.AsyncClient | None = None
|
||||
if settings.enable_tool_runtime:
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
包含角色档案处理、状态转换规则、模型策略提取等不依赖外部系统的领域逻辑。
|
||||
设计见 ``docs/domain/identity-continuity.md``。
|
||||
|
||||
模块导入限制：禁止引入 ``repositories`` / ``db`` / ``integrations`` / ``ai_provider`` / ``httpx`` / ``sqlalchemy``。
|
||||
模块导入限制：禁止引入 ``repositories`` / ``db`` / ``integrations`` / ``httpx`` / ``sqlalchemy``。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
"""辰正时态感知纯计算核心。
|
||||
|
||||
本包为零 IO 纯函数与数据结构，禁止 import repositories/db/integrations/
|
||||
ai_provider/httpx/sqlalchemy/services。需要 IO 时由上层 services/application
|
||||
httpx/sqlalchemy/services。需要 IO 时由上层 services/application
|
||||
调入本包函数。
|
||||
"""
|
||||
@@ -1,83 +1,11 @@
|
||||
"""设置 API schema（模型网关与生成偏好）。
|
||||
"""设置 API schema（生成偏好）。
|
||||
|
||||
前端只见逻辑模型 / provider 密钥脱敏状态 / 生成参数偏好；
|
||||
真实 model 名、endpoint URL、API Key 由后端持有，不进入响应（environment-config §3）。
|
||||
前端只见生成参数偏好；模型网关状态由神识承接（用户面见神识阁）。
|
||||
"""
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
|
||||
class ModelEndpointStatus(BaseModel):
|
||||
"""endpoint 池状态。base_url / api_key 原值不回传（脱敏），仅回覆盖标记 + env 变量名。"""
|
||||
|
||||
name: str = Field(description="endpoint 名称（如 primary / fallback）。")
|
||||
provider: str = Field(description="所属 provider。")
|
||||
status: str = Field(description="endpoint 状态：active / inactive / unhealthy / cooldown。")
|
||||
weight: int | None = Field(default=None, description="权重（weighted_random 路由用）。")
|
||||
priority: int | None = Field(
|
||||
default=None, description="优先级（priority_failover 用，数值越小越优先）。"
|
||||
)
|
||||
is_overridden: bool = Field(default=False, description="是否存在 DB 覆盖。")
|
||||
has_url_override: bool = Field(
|
||||
default=False, description="是否用 DB 覆盖了 base_url（优先于 .env）。"
|
||||
)
|
||||
has_key_override: bool = Field(
|
||||
default=False, description="是否用 DB 覆盖了 api_key（优先于 .env）。"
|
||||
)
|
||||
base_url_env: str | None = Field(
|
||||
default=None, description="base_url 的 .env 变量名（只读展示）。"
|
||||
)
|
||||
api_key_env: str | None = Field(
|
||||
default=None, description="api_key 的 .env 变量名（只读展示）。"
|
||||
)
|
||||
|
||||
|
||||
class ProviderKeyStatus(BaseModel):
|
||||
"""供应商密钥脱敏状态：configured 表示 .env 是否已配置该 provider 的 API Key。"""
|
||||
|
||||
provider: str = Field(description="provider 标识。")
|
||||
configured: bool = Field(description="是否已配置 API Key（仅 true/false，绝不暴露值）。")
|
||||
|
||||
|
||||
class LogicalModel(BaseModel):
|
||||
"""逻辑模型（契约字典 §16）。业务层 model_hint 引用 name。"""
|
||||
|
||||
name: str = Field(description="逻辑模型名，业务层 model_hint 引用。")
|
||||
display_name: str = Field(description="展示名。")
|
||||
provider: str = Field(description="默认 provider。")
|
||||
capabilities: list[str] = Field(
|
||||
description="能力清单：chat / tool_use / structured_output / reasoning / long_context / vision。"
|
||||
)
|
||||
routing_strategy: str = Field(description="路由策略。")
|
||||
supports_streaming: bool = Field(description="是否支持流式。")
|
||||
supports_thinking: bool = Field(description="是否支持思考。")
|
||||
supports_structured_output: bool = Field(description="是否支持结构化输出。")
|
||||
context_window: int = Field(description="上下文窗口大小。")
|
||||
max_output_tokens: int | None = Field(default=None, description="最大输出 token 数。")
|
||||
endpoints: list[ModelEndpointStatus] | None = Field(
|
||||
default=None, description="endpoint 池脱敏状态。"
|
||||
)
|
||||
is_default: bool | None = Field(default=None, description="是否为默认对话模型槽位。")
|
||||
|
||||
|
||||
class ModelGatewayStatus(BaseModel):
|
||||
model_config = ConfigDict(
|
||||
json_schema_extra={
|
||||
"examples": [
|
||||
{
|
||||
"default_model_hint": "default-chat",
|
||||
"models": [],
|
||||
"providers": [{"provider": "deepseek", "configured": True}],
|
||||
}
|
||||
]
|
||||
}
|
||||
)
|
||||
|
||||
models: list[LogicalModel] = Field(description="逻辑模型清单。")
|
||||
providers: list[ProviderKeyStatus] = Field(description="provider 密钥脱敏状态。")
|
||||
default_model_hint: str = Field(description="默认逻辑模型名。")
|
||||
|
||||
|
||||
class GenerationPreferencesResponse(BaseModel):
|
||||
"""生成参数偏好（业务层，不含密钥）。"""
|
||||
|
||||
@@ -114,21 +42,3 @@ class UpdateGenerationPreferencesRequest(BaseModel):
|
||||
default=None, ge=1, le=200000, description="最大输出 token 数。"
|
||||
)
|
||||
streaming: bool | None = Field(default=None, description="是否启用流式输出。")
|
||||
|
||||
|
||||
class UpdateEndpointOverrideRequest(BaseModel):
|
||||
"""更新 endpoint 覆盖（PATCH /settings/models/{model}/endpoints/{endpoint}）。
|
||||
|
||||
字段全可选；传 None 表示清除该字段覆盖、回退 yaml/env。base_url/api_key 经 DB
|
||||
覆盖优先于 .env。响应脱敏，不回 base_url/api_key 原值。
|
||||
"""
|
||||
|
||||
base_url: str | None = Field(
|
||||
default=None, description="覆盖 base_url（DB 优先 env）；None 清除。"
|
||||
)
|
||||
api_key: str | None = Field(
|
||||
default=None, description="覆盖 api_key（DB 优先 env）；None 清除。"
|
||||
)
|
||||
status: str | None = Field(default=None, description="覆盖 endpoint 状态。")
|
||||
weight: int | None = Field(default=None, ge=0, description="覆盖权重。")
|
||||
priority: int | None = Field(default=None, ge=0, description="覆盖优先级。")
|
||||
@@ -1,6 +1,6 @@
|
||||
"""运行时 API schema。"""
|
||||
|
||||
from typing import Any, Literal
|
||||
from typing import Any
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
@@ -85,31 +85,3 @@ class RunResponse(BaseModel):
|
||||
error_message: str | None = Field(description="失败错误信息；成功时为空。")
|
||||
created_at: str = Field(description="创建时间，UTC ISO 8601 字符串。")
|
||||
updated_at: str = Field(description="更新时间，UTC ISO 8601 字符串。")
|
||||
|
||||
|
||||
class AgentControlRequest(BaseModel):
|
||||
"""运行中 steer / follow-up 消息。"""
|
||||
|
||||
content: str = Field(min_length=1, max_length=20_000)
|
||||
|
||||
|
||||
class AgentControlResponse(BaseModel):
|
||||
accepted: bool = True
|
||||
|
||||
|
||||
class ConfirmationResolutionResponse(BaseModel):
|
||||
"""工具确认 resolve 结果（approve/deny 端点成功响应）。"""
|
||||
|
||||
model_config = ConfigDict(
|
||||
json_schema_extra={
|
||||
"examples": [
|
||||
{"status": "resolved", "approved": True},
|
||||
{"status": "resolved", "approved": False},
|
||||
]
|
||||
}
|
||||
)
|
||||
|
||||
status: Literal["resolved"] = Field(
|
||||
default="resolved", description="恒为 resolved；失败路径由错误响应覆盖。"
|
||||
)
|
||||
approved: bool = Field(description="裁决结果：True=批准，False=拒绝。")
|
||||
@@ -1,108 +0,0 @@
|
||||
"""AI provider 运行时治理：endpoint 覆盖（DB 优先于 .env/yaml）+ 热重载。
|
||||
|
||||
override 写 DB → 读 DB 全量刷 spec 同步缓存 → ProviderManager 下次构建读新缓存。
|
||||
api_key 明文存 DB（本地单用户边界），响应层脱敏不回值（见 schemas）。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from app.ai_provider.config.loader import load_models_config
|
||||
from app.core.errors import AppError
|
||||
from app.db.session import AsyncSessionLocal
|
||||
from app.models.governance import ModelEndpointOverride
|
||||
from app.repositories import model_endpoints as overrides_repo
|
||||
from app.schemas.preferences import UpdateEndpointOverrideRequest
|
||||
from app.services.audit import record_audit_event
|
||||
|
||||
# —— endpoint 运行时覆盖（DB 优先于 .env/yaml）——
|
||||
# reload_provider_manager 时（异步 service）读 DB 全量刷此同步缓存，
|
||||
# load_models_config 仍可独立使用。None 字段表示不覆盖、回退 yaml/env。
|
||||
_endpoint_overrides_cache: dict[tuple[str, str], dict[str, Any]] = {}
|
||||
|
||||
|
||||
def set_endpoint_overrides_cache(overrides: list[dict[str, Any]]) -> None:
|
||||
"""用 DB 全量覆盖列表刷新模块缓存（同步）。"""
|
||||
global _endpoint_overrides_cache
|
||||
_endpoint_overrides_cache = {
|
||||
(o["logical_model_name"], o["endpoint_name"]): o for o in overrides
|
||||
}
|
||||
|
||||
|
||||
def get_endpoint_overrides_cache() -> dict[tuple[str, str], dict[str, Any]]:
|
||||
"""读取当前覆盖缓存快照。"""
|
||||
return dict(_endpoint_overrides_cache)
|
||||
|
||||
|
||||
async def list_endpoint_overrides() -> list[dict[str, Any]]:
|
||||
"""读 DB 全量覆盖，序列化为 dict 列表（供刷缓存 + 响应）。"""
|
||||
async with AsyncSessionLocal() as session:
|
||||
rows = await overrides_repo.list_overrides(session)
|
||||
return [_override_to_dict(r) for r in rows]
|
||||
|
||||
|
||||
async def reload_provider_manager() -> None:
|
||||
"""读 DB 全量 → 刷 endpoint 覆盖同步缓存。"""
|
||||
overrides = await list_endpoint_overrides()
|
||||
set_endpoint_overrides_cache(overrides)
|
||||
|
||||
|
||||
async def update_endpoint_override(
|
||||
logical_model_name: str,
|
||||
endpoint_name: str,
|
||||
payload: UpdateEndpointOverrideRequest,
|
||||
) -> dict[str, Any]:
|
||||
"""upsert 覆盖（payload 字段全量写入，None 清除字段）+ reload + 审计。
|
||||
|
||||
校验 (model, endpoint) 存在于 yaml（避免覆盖幽灵 endpoint）。
|
||||
"""
|
||||
if not _endpoint_exists(logical_model_name, endpoint_name):
|
||||
raise AppError("not_found", "endpoint 不存在于模型配置中。")
|
||||
async with AsyncSessionLocal() as session:
|
||||
override = ModelEndpointOverride(
|
||||
logical_model_name=logical_model_name,
|
||||
endpoint_name=endpoint_name,
|
||||
base_url=payload.base_url,
|
||||
api_key=payload.api_key,
|
||||
status=payload.status,
|
||||
weight=payload.weight,
|
||||
priority=payload.priority,
|
||||
)
|
||||
saved = await overrides_repo.upsert_override(session, override)
|
||||
await record_audit_event(
|
||||
session,
|
||||
action="model_endpoint.override_updated",
|
||||
actor="user",
|
||||
input_summary=f"{logical_model_name}/{endpoint_name}",
|
||||
resource_type="model_endpoint_override",
|
||||
resource_id=f"{logical_model_name}/{endpoint_name}",
|
||||
)
|
||||
await session.commit()
|
||||
await session.refresh(saved)
|
||||
result = _override_to_dict(saved)
|
||||
await reload_provider_manager()
|
||||
return result
|
||||
|
||||
|
||||
def _endpoint_exists(logical_model_name: str, endpoint_name: str) -> bool:
|
||||
"""校验 (model, endpoint) 在 yaml 配置中存在（避免覆盖幽灵 endpoint）。"""
|
||||
models = load_models_config()
|
||||
for model in models:
|
||||
if model.name != logical_model_name:
|
||||
continue
|
||||
return any(ep.name == endpoint_name for ep in model.endpoints)
|
||||
return False
|
||||
|
||||
|
||||
def _override_to_dict(row: ModelEndpointOverride) -> dict[str, Any]:
|
||||
return {
|
||||
"logical_model_name": row.logical_model_name,
|
||||
"endpoint_name": row.endpoint_name,
|
||||
"base_url": row.base_url,
|
||||
"api_key": row.api_key,
|
||||
"status": row.status,
|
||||
"weight": row.weight,
|
||||
"priority": row.priority,
|
||||
"updated_at": row.updated_at,
|
||||
}
|
||||
@@ -1,11 +1,8 @@
|
||||
"""设置服务：模型网关状态聚合与生成偏好。
|
||||
"""设置服务：生成偏好。
|
||||
|
||||
模型网关状态从 config/models.yaml 派生，provider 密钥按 settings 是否配置脱敏为
|
||||
configured 布尔（绝不暴露值）；生成偏好为本地单用户单行，无则返回默认值。
|
||||
生成偏好为本地单用户单行，无则返回默认值；模型网关状态由神识承接（用户面见神识阁）。
|
||||
"""
|
||||
|
||||
from app.ai_provider.config.loader import load_models_config
|
||||
from app.ai_provider.config.models import ModelConfig
|
||||
from app.core.config import settings
|
||||
from app.core.time import utc_now
|
||||
from app.db.session import AsyncSessionLocal
|
||||
@@ -13,96 +10,10 @@ from app.models.governance import GenerationPreferences
|
||||
from app.repositories import preferences as preferences_repo
|
||||
from app.schemas.preferences import (
|
||||
GenerationPreferencesResponse,
|
||||
LogicalModel,
|
||||
ModelEndpointStatus,
|
||||
ModelGatewayStatus,
|
||||
ProviderKeyStatus,
|
||||
UpdateGenerationPreferencesRequest,
|
||||
)
|
||||
from app.services.ai_provider_admin import get_endpoint_overrides_cache
|
||||
from app.services.serializers import generation_preferences_to_response
|
||||
|
||||
# provider → 持有其 API Key 的环境变量名；未登记的 provider 回退 <PROVIDER>_API_KEY。
|
||||
_PROVIDER_KEY_ENV: dict[str, str] = {
|
||||
"openai": "OPENAI_API_KEY",
|
||||
"deepseek": "DEEPSEEK_API_KEY",
|
||||
"anthropic": "ANTHROPIC_API_KEY",
|
||||
"google": "GOOGLE_API_KEY",
|
||||
"minimax": "MINIMAX_API_KEY",
|
||||
"zhipu": "ZHIPU_API_KEY",
|
||||
"tencentcloud": "TENCENTCLOUD_API_KEY",
|
||||
"volcengine": "VOLCENGINE_ARK_API_KEY",
|
||||
"bailian": "BAILIAN_API_KEY",
|
||||
}
|
||||
|
||||
|
||||
def provider_configured(provider: str | None) -> bool:
|
||||
"""provider 的 API Key 是否已在 settings 配置（仅返回布尔，不暴露值）。"""
|
||||
if not provider:
|
||||
return False
|
||||
env_name = _PROVIDER_KEY_ENV.get(provider, f"{provider.upper()}_API_KEY")
|
||||
return bool(getattr(settings, env_name.lower(), None))
|
||||
|
||||
|
||||
def _endpoint_status(
|
||||
endpoint, model_name: str, provider: str, overrides_cache
|
||||
) -> ModelEndpointStatus:
|
||||
ov = overrides_cache.get((model_name, endpoint.name))
|
||||
return ModelEndpointStatus(
|
||||
name=endpoint.name,
|
||||
provider=provider,
|
||||
status=(ov.get("status") or endpoint.status) if ov else endpoint.status,
|
||||
weight=endpoint.weight,
|
||||
priority=endpoint.priority,
|
||||
is_overridden=ov is not None,
|
||||
has_url_override=bool(ov and ov.get("base_url")),
|
||||
has_key_override=bool(ov and ov.get("api_key")),
|
||||
base_url_env=endpoint.base_url_env,
|
||||
api_key_env=endpoint.api_key_env,
|
||||
)
|
||||
|
||||
|
||||
def _logical_model(cfg: ModelConfig, overrides_cache) -> LogicalModel:
|
||||
provider = cfg.provider or "unknown"
|
||||
endpoints = [
|
||||
_endpoint_status(ep, cfg.name, provider, overrides_cache) for ep in cfg.endpoints
|
||||
] or None
|
||||
return LogicalModel(
|
||||
name=cfg.name,
|
||||
display_name=cfg.display_name or cfg.name,
|
||||
provider=provider,
|
||||
capabilities=[k for k, v in cfg.resolved_capabilities().items() if v],
|
||||
routing_strategy=cfg.routing_strategy,
|
||||
supports_streaming=bool(cfg.supports_streaming),
|
||||
supports_thinking=bool(cfg.supports_thinking),
|
||||
supports_structured_output=bool(cfg.supports_structured_output),
|
||||
context_window=cfg.context_window,
|
||||
max_output_tokens=cfg.max_output_tokens,
|
||||
endpoints=endpoints,
|
||||
is_default=True if cfg.name == settings.default_model_hint else None,
|
||||
)
|
||||
|
||||
|
||||
def get_provider_gateway() -> ModelGatewayStatus:
|
||||
"""从 config/models.yaml 派生模型网关总览（含 DB 覆盖标记，脱敏不回原值）。"""
|
||||
models_cfg = load_models_config()
|
||||
overrides_cache = get_endpoint_overrides_cache()
|
||||
providers_seen: dict[str, None] = {}
|
||||
logical_models: list[LogicalModel] = []
|
||||
for cfg in models_cfg:
|
||||
logical_models.append(_logical_model(cfg, overrides_cache))
|
||||
provider = cfg.provider or "unknown"
|
||||
providers_seen.setdefault(provider, None)
|
||||
providers = [
|
||||
ProviderKeyStatus(provider=name, configured=provider_configured(name))
|
||||
for name in providers_seen
|
||||
]
|
||||
return ModelGatewayStatus(
|
||||
models=logical_models,
|
||||
providers=providers,
|
||||
default_model_hint=settings.default_model_hint,
|
||||
)
|
||||
|
||||
|
||||
def _default_preferences() -> GenerationPreferences:
|
||||
return GenerationPreferences(
|
||||
|
||||
@@ -10,7 +10,6 @@ from app.core.security import redact_value
|
||||
from app.core.time import utc_now
|
||||
from app.db.session import AsyncSessionLocal
|
||||
from app.models.runtime import Event, Run
|
||||
from app.modules.runtime.events import EventStore
|
||||
from app.repositories import runtime as runtime_repo
|
||||
from app.repositories.json_utils import dumps
|
||||
from app.schemas.common import Page
|
||||
@@ -21,7 +20,6 @@ from app.services.serializers import event_to_response, run_to_response
|
||||
|
||||
TERMINAL_RUN_STATUS = {"completed", "failed", "cancelled", "interrupted"}
|
||||
TERMINAL_EVENT_TYPES = {"run.completed", "run.failed", "run.cancelled", "run.interrupted"}
|
||||
REQUIRES_CONFIRMATION_EVENT_TYPE = "tool.requires_confirmation"
|
||||
|
||||
|
||||
async def create_run(
|
||||
@@ -215,20 +213,3 @@ async def list_run_event_responses(
|
||||
session, run_id, limit=limit, offset=offset
|
||||
)
|
||||
return [event_to_response(event) for event in events]
|
||||
|
||||
|
||||
async def confirmation_exists_in_run(
|
||||
event_store: EventStore, run_id: str, confirmation_id: str
|
||||
) -> bool:
|
||||
"""失效感知：run 下是否曾经发过 ``tool.requires_confirmation`` 且 payload 含该 confirmation_id。
|
||||
|
||||
实现走 EventStore.list_by_run + Python 侧过滤，避免 PostgreSQL 专用 JSON 操作符，
|
||||
兼容本地 SQLite。用于 approve/deny 端点在 Future 已失效时区分 410 与 404。
|
||||
"""
|
||||
events = await event_store.list_by_run(run_id)
|
||||
for event in events:
|
||||
if event.type != REQUIRES_CONFIRMATION_EVENT_TYPE:
|
||||
continue
|
||||
if event.payload.get("confirmation_id") == confirmation_id:
|
||||
return True
|
||||
return False
|
||||
@@ -1,5 +0,0 @@
|
||||
"""随应用部署的受信任 Python Skills。
|
||||
|
||||
新增 Skill 时创建同级模块并导出 ``setup(registry)``。目录仅接收经过代码审查的模块；
|
||||
数据库 Skill 与用户上传内容不得写入这里或在 FastAPI 进程内动态执行。
|
||||
"""
|
||||
@@ -1,39 +0,0 @@
|
||||
"""工具参数守卫：拒绝参数过大的工具调用，防止模型生成超大输入。"""
|
||||
|
||||
import json
|
||||
|
||||
from app.application.agent.extensions import (
|
||||
AgentRunContext,
|
||||
BeforeToolCallDecision,
|
||||
ExtensionRegistry,
|
||||
ToolCallContext,
|
||||
)
|
||||
from app.application.agent.python_skills import PythonSkillDefinition
|
||||
|
||||
# 单次工具调用参数序列化后的字符上限；超出则拒绝并要求模型精简。
|
||||
_MAX_ARG_CHARS = 8000
|
||||
|
||||
|
||||
async def _reject_oversized_args(
|
||||
call: ToolCallContext, _run: AgentRunContext
|
||||
) -> BeforeToolCallDecision | None:
|
||||
size = len(json.dumps(call.input, ensure_ascii=False))
|
||||
if size <= _MAX_ARG_CHARS:
|
||||
return None
|
||||
return BeforeToolCallDecision(
|
||||
allowed=False,
|
||||
reason=f"工具 {call.name} 的参数过大（{size} 字符超过 {_MAX_ARG_CHARS}），请精简后重试。",
|
||||
)
|
||||
|
||||
|
||||
def _setup(registry: ExtensionRegistry) -> None:
|
||||
registry.add_before_tool_hook(_reject_oversized_args)
|
||||
|
||||
|
||||
SKILL = PythonSkillDefinition(
|
||||
name="arg_guard",
|
||||
description="拒绝参数过大的工具调用，避免超大输入挤占上下文或触发限制。",
|
||||
triggers=("大参数", "arg guard"),
|
||||
instructions="发起工具调用时保持参数精简；超大参数会被拒绝。",
|
||||
setup=_setup,
|
||||
)
|
||||
@@ -1,34 +0,0 @@
|
||||
"""Token 预算守卫：累计消耗超阈值时优雅收尾。"""
|
||||
|
||||
from app.application.agent.extensions import (
|
||||
AgentRunContext,
|
||||
ExtensionRegistry,
|
||||
TurnSnapshot,
|
||||
)
|
||||
from app.application.agent.python_skills import PythonSkillDefinition
|
||||
|
||||
# 累计 total_tokens 软上限。达到后 TurnPolicy 要求停止，让 Agent 用当前内容自然
|
||||
# 收尾——区别于 max_iterations 的硬截断（truncated=True，回复可能只说了一半）。
|
||||
_BUDGET_TOTAL_TOKENS = 20000
|
||||
|
||||
|
||||
async def _stop_when_budget_exhausted(
|
||||
turn: TurnSnapshot, _run: AgentRunContext
|
||||
) -> bool:
|
||||
return turn.usage.total_tokens >= _BUDGET_TOTAL_TOKENS
|
||||
|
||||
|
||||
def _setup(registry: ExtensionRegistry) -> None:
|
||||
registry.add_turn_policy(_stop_when_budget_exhausted)
|
||||
|
||||
|
||||
SKILL = PythonSkillDefinition(
|
||||
name="budget_guard",
|
||||
description="累计 token 消耗超预算时优雅停止 Agent 循环，避免成本失控。",
|
||||
triggers=("预算", "budget", "长任务"),
|
||||
instructions=(
|
||||
"本轮受 token 预算约束。若已接近上限，请用当前已有信息给出最终回复，"
|
||||
"不要发起新的工具调用或长篇推理。"
|
||||
),
|
||||
setup=_setup,
|
||||
)
|
||||
@@ -1,47 +0,0 @@
|
||||
"""上下文预算守卫：历史过长时折叠中间消息，保护 prompt 预算。"""
|
||||
|
||||
from app.ai_provider.types.messages import BaseMessage
|
||||
from app.application.agent.extensions import (
|
||||
AgentContext,
|
||||
AgentRunContext,
|
||||
ExtensionRegistry,
|
||||
)
|
||||
from app.application.agent.python_skills import PythonSkillDefinition
|
||||
|
||||
# messages 总字符数软上限；超出则保留首条与末尾 _TAIL_KEEP 条，中间折叠。
|
||||
_MAX_CONTEXT_CHARS = 12000
|
||||
_TAIL_KEEP = 8
|
||||
|
||||
|
||||
async def _trim_history(context: AgentContext, _run: AgentRunContext) -> AgentContext:
|
||||
messages = context.messages
|
||||
total = sum(len(message.content) for message in messages)
|
||||
if total <= _MAX_CONTEXT_CHARS or len(messages) <= _TAIL_KEEP + 1:
|
||||
return context
|
||||
first = messages[0]
|
||||
tail = messages[-_TAIL_KEEP:]
|
||||
dropped = len(messages) - 1 - _TAIL_KEEP
|
||||
# 在首条追加折叠提示（新建对象避免污染原 message）；不在历史中间插入新消息，
|
||||
# 以免破坏 provider 对 system 位置或多轮 tool 消息配对的预期。
|
||||
folded = BaseMessage(
|
||||
content=f"{first.content}\n\n[已省略 {dropped} 条较早历史]",
|
||||
role=first.role,
|
||||
)
|
||||
return AgentContext(
|
||||
messages=(folded, *tail),
|
||||
tools=context.tools,
|
||||
model_hint=context.model_hint,
|
||||
)
|
||||
|
||||
|
||||
def _setup(registry: ExtensionRegistry) -> None:
|
||||
registry.add_context_transformer(_trim_history)
|
||||
|
||||
|
||||
SKILL = PythonSkillDefinition(
|
||||
name="context_budget",
|
||||
description="对话历史过长时折叠中间消息，控制 prompt 预算。",
|
||||
triggers=("长对话", "context", "历史"),
|
||||
instructions="历史较长时优先依赖近期消息，不必回溯很早的内容。",
|
||||
setup=_setup,
|
||||
)
|
||||
@@ -1,44 +0,0 @@
|
||||
"""工具输出防护示例 Skill。"""
|
||||
|
||||
from dataclasses import replace
|
||||
|
||||
from app.application.agent.extensions import (
|
||||
AfterToolCallDecision,
|
||||
AgentRunContext,
|
||||
ExtensionRegistry,
|
||||
ToolCallContext,
|
||||
)
|
||||
from app.application.agent.python_skills import PythonSkillDefinition
|
||||
from app.modules.tool_runtime.definition import ToolResult
|
||||
|
||||
_MAX_OUTPUT_CHARS = 4000
|
||||
|
||||
|
||||
async def _limit_tool_output(
|
||||
_call: ToolCallContext,
|
||||
result: ToolResult,
|
||||
_run: AgentRunContext,
|
||||
) -> AfterToolCallDecision | None:
|
||||
if len(result.stdout) <= _MAX_OUTPUT_CHARS and len(result.stderr) <= _MAX_OUTPUT_CHARS:
|
||||
return None
|
||||
return AfterToolCallDecision(
|
||||
result=replace(
|
||||
result,
|
||||
stdout=result.stdout[:_MAX_OUTPUT_CHARS],
|
||||
stderr=result.stderr[:_MAX_OUTPUT_CHARS],
|
||||
truncated=True,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _setup(registry: ExtensionRegistry) -> None:
|
||||
registry.add_after_tool_hook(_limit_tool_output)
|
||||
|
||||
|
||||
SKILL = PythonSkillDefinition(
|
||||
name="response_guard",
|
||||
description="限制过长工具输出，避免无关内容挤占模型上下文。",
|
||||
triggers=("长输出", "大量输出", "response guard"),
|
||||
instructions="工具输出可能很长；优先分段读取，只保留完成当前任务所需的内容。",
|
||||
setup=_setup,
|
||||
)
|
||||
@@ -1,53 +0,0 @@
|
||||
"""运行指标采集：监听 Agent 生命周期事件，结束时记录摘要日志。"""
|
||||
|
||||
from loguru import logger
|
||||
|
||||
from app.application.agent.extensions import (
|
||||
AgentEvent,
|
||||
AgentRunContext,
|
||||
ExtensionRegistry,
|
||||
)
|
||||
from app.application.agent.python_skills import PythonSkillDefinition
|
||||
|
||||
|
||||
def _setup(registry: ExtensionRegistry) -> None:
|
||||
turns = 0
|
||||
tool_calls = 0
|
||||
tool_errors = 0
|
||||
tool_denied = 0
|
||||
|
||||
async def listen(event: AgentEvent, run: AgentRunContext) -> None:
|
||||
nonlocal turns, tool_calls, tool_errors, tool_denied
|
||||
# 工具事件经 Agent._on_tool_event 桥接进扩展事件流，故此处可直接观测
|
||||
# tool.started / tool.completed（status: success / error / denied）。
|
||||
if event.type == "agent.turn.started":
|
||||
turns += 1
|
||||
elif event.type == "tool.started":
|
||||
tool_calls += 1
|
||||
elif event.type == "tool.completed":
|
||||
status = event.payload.get("status")
|
||||
if status == "error":
|
||||
tool_errors += 1
|
||||
elif status == "denied":
|
||||
tool_denied += 1
|
||||
elif event.type == "agent.settled":
|
||||
logger.info(
|
||||
"agent_run_metrics",
|
||||
run_id=run.run_id,
|
||||
turns=turns,
|
||||
tool_calls=tool_calls,
|
||||
tool_errors=tool_errors,
|
||||
tool_denied=tool_denied,
|
||||
status=event.payload.get("status"),
|
||||
)
|
||||
|
||||
registry.add_event_listener(listen)
|
||||
|
||||
|
||||
SKILL = PythonSkillDefinition(
|
||||
name="run_metrics",
|
||||
description="采集 Agent 轮次指标并在运行结束时记录摘要日志。",
|
||||
triggers=("指标", "metrics", "观测"),
|
||||
instructions="",
|
||||
setup=_setup,
|
||||
)
|
||||
Whitespace-only changes.
@@ -1,41 +0,0 @@
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from app.ai_provider.config.loader import load_models_config, load_providers_config
|
||||
|
||||
|
||||
class TestLoader:
|
||||
def test_load_models_config(self, tmp_path: Path):
|
||||
yaml_content = """
|
||||
version: 1
|
||||
models:
|
||||
- name: test-model
|
||||
provider: deepseek
|
||||
model: deepseek-v4-flash
|
||||
format: openai
|
||||
"""
|
||||
config_file = tmp_path / "models.yaml"
|
||||
config_file.write_text(yaml_content)
|
||||
models = load_models_config(config_file)
|
||||
assert len(models) == 1
|
||||
assert models[0].name == "test-model"
|
||||
|
||||
def test_load_providers_config(self, tmp_path: Path):
|
||||
yaml_content = """
|
||||
version: 1
|
||||
providers:
|
||||
deepseek:
|
||||
format: openai
|
||||
base_url_env: DEEPSEEK_BASE_URL
|
||||
api_key_env: DEEPSEEK_API_KEY
|
||||
"""
|
||||
config_file = tmp_path / "providers.yaml"
|
||||
config_file.write_text(yaml_content)
|
||||
providers = load_providers_config(config_file)
|
||||
assert "deepseek" in providers
|
||||
assert providers["deepseek"].format == "openai"
|
||||
|
||||
def test_missing_file(self):
|
||||
with pytest.raises(FileNotFoundError):
|
||||
load_models_config(Path("/nonexistent/models.yaml"))
|
||||
@@ -1,69 +0,0 @@
|
||||
import pytest
|
||||
|
||||
from app.ai_provider.config.models import (
|
||||
ModelConfig,
|
||||
ProviderConfig,
|
||||
derive_capabilities,
|
||||
)
|
||||
|
||||
|
||||
class TestModelConfig:
|
||||
def test_minimal_config(self):
|
||||
cfg = ModelConfig(
|
||||
name="deepseek-v4-flash",
|
||||
provider="deepseek",
|
||||
model="deepseek-v4-flash",
|
||||
format="openai",
|
||||
)
|
||||
assert cfg.name == "deepseek-v4-flash"
|
||||
assert cfg.format == "openai"
|
||||
assert not cfg.supports_vision # defaulted
|
||||
|
||||
def test_capability_override(self):
|
||||
cfg = ModelConfig(
|
||||
name="claude-sonnet",
|
||||
provider="anthropic",
|
||||
model="claude-sonnet-4-6",
|
||||
format="anthropic",
|
||||
supports_vision=True,
|
||||
supports_tools=True,
|
||||
)
|
||||
assert cfg.supports_vision
|
||||
assert cfg.supports_tools
|
||||
|
||||
def test_capability_auto_derive(self):
|
||||
caps = derive_capabilities(provider="anthropic", model="claude-fable-5")
|
||||
assert "chat" in caps
|
||||
assert "tool_use" in caps
|
||||
assert "vision" in caps
|
||||
assert "reasoning" in caps
|
||||
|
||||
|
||||
class TestProviderConfig:
|
||||
def test_minimal(self):
|
||||
cfg = ProviderConfig(
|
||||
name="deepseek",
|
||||
format="openai",
|
||||
base_url_env="DEEPSEEK_BASE_URL",
|
||||
api_key_env="DEEPSEEK_API_KEY",
|
||||
)
|
||||
assert cfg.name == "deepseek"
|
||||
|
||||
def test_with_parameter_mapping(self):
|
||||
cfg = ProviderConfig(
|
||||
name="deepseek",
|
||||
format="openai",
|
||||
base_url_env="DEEPSEEK_BASE_URL",
|
||||
api_key_env="DEEPSEEK_API_KEY",
|
||||
parameter_mappings={
|
||||
"max_tokens": {"provider_param": "max_tokens"},
|
||||
"temperature": {"provider_param": "temperature", "min": 0, "max": 2},
|
||||
},
|
||||
)
|
||||
mapping = cfg.parameter_mappings["temperature"]
|
||||
assert mapping.min == 0
|
||||
assert mapping.max == 2
|
||||
|
||||
def test_format_validation(self):
|
||||
with pytest.raises(ValueError, match="format"):
|
||||
ProviderConfig(name="bad", format="unknown", base_url_env="X", api_key_env="Y")
|
||||
Whitespace-only changes.
@@ -1,67 +0,0 @@
|
||||
import pytest
|
||||
|
||||
from app.ai_provider.core.billing import (
|
||||
BudgetController,
|
||||
CostCalculator,
|
||||
PricingInfo,
|
||||
UsageRecord,
|
||||
)
|
||||
|
||||
|
||||
class TestCostCalculator:
|
||||
def test_openai_pricing(self):
|
||||
pricing = PricingInfo(currency="USD", input_per_1m=5.0, output_per_1m=15.0)
|
||||
cost = CostCalculator.calculate(pricing, 1000, 500)
|
||||
assert cost.input_cost == 0.005 # 1000 / 1M * 5.0
|
||||
assert cost.output_cost == 0.0075 # 500 / 1M * 15.0
|
||||
assert cost.total == 0.0125
|
||||
assert cost.currency == "USD"
|
||||
|
||||
def test_no_output_pricing(self):
|
||||
pricing = PricingInfo(currency="CNY", input_per_1m=1.0, output_per_1m=None)
|
||||
cost = CostCalculator.calculate(pricing, 1000, 0)
|
||||
assert cost.input_cost == 0.001
|
||||
assert cost.output_cost == 0.0
|
||||
|
||||
def test_no_pricing(self):
|
||||
cost = CostCalculator.calculate(None, 1000, 500)
|
||||
assert cost.total == 0.0 # gracefully handles missing pricing
|
||||
|
||||
|
||||
class TestUsageRecord:
|
||||
def test_create_record(self):
|
||||
record = UsageRecord(
|
||||
model_name="deepseek-v4-flash",
|
||||
provider="deepseek",
|
||||
prompt_tokens=100,
|
||||
completion_tokens=50,
|
||||
total_tokens=150,
|
||||
currency="CNY",
|
||||
cost=0.0002,
|
||||
)
|
||||
assert record.total_tokens == 150
|
||||
|
||||
|
||||
class TestBudgetController:
|
||||
@pytest.fixture
|
||||
def controller(self):
|
||||
return BudgetController(
|
||||
daily_budget_cny=10.0,
|
||||
monthly_budget_cny=200.0,
|
||||
)
|
||||
|
||||
def test_within_budget(self, controller):
|
||||
controller.add_cost_cny(1.0)
|
||||
assert controller.allow_request()
|
||||
|
||||
def test_exceed_daily_budget(self, controller):
|
||||
controller.add_cost_cny(11.0)
|
||||
assert not controller.allow_request()
|
||||
|
||||
def test_reset_daily(self):
|
||||
controller = BudgetController(daily_budget_cny=10.0, monthly_budget_cny=200.0)
|
||||
controller.add_cost_cny(11.0)
|
||||
assert not controller.allow_request()
|
||||
# Simulate next day
|
||||
controller.reset_daily()
|
||||
assert controller.allow_request()
|
||||
@@ -1,225 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from app.ai_provider.config.models import ModelConfig, ProviderConfig
|
||||
from app.ai_provider.core.registry import (
|
||||
ModelRegistry,
|
||||
ProviderRegistry,
|
||||
RegistryFactory,
|
||||
RegistryHolder,
|
||||
)
|
||||
from app.ai_provider.types.errors import AIProviderError
|
||||
|
||||
|
||||
class TestModelRegistry:
|
||||
@pytest.fixture
|
||||
def sample_models(self):
|
||||
return [
|
||||
ModelConfig(
|
||||
name="fast-model",
|
||||
provider="deepseek",
|
||||
model="deepseek-v4-flash",
|
||||
format="openai",
|
||||
tags=["default"],
|
||||
weight=100,
|
||||
),
|
||||
ModelConfig(
|
||||
name="pro-model",
|
||||
provider="deepseek",
|
||||
model="deepseek-v4-pro",
|
||||
format="openai",
|
||||
tags=["fallback"],
|
||||
weight=50,
|
||||
),
|
||||
ModelConfig(
|
||||
name="vision-model",
|
||||
provider="anthropic",
|
||||
model="claude-sonnet-4-6",
|
||||
format="anthropic",
|
||||
supports_vision=True,
|
||||
),
|
||||
]
|
||||
|
||||
def test_get_model(self, sample_models):
|
||||
reg = ModelRegistry(sample_models)
|
||||
model = reg.get("fast-model")
|
||||
assert model.name == "fast-model"
|
||||
|
||||
def test_get_nonexistent(self, sample_models):
|
||||
reg = ModelRegistry(sample_models)
|
||||
with pytest.raises(AIProviderError):
|
||||
reg.get("nonexistent")
|
||||
|
||||
def test_default_model(self, sample_models):
|
||||
reg = ModelRegistry(sample_models)
|
||||
assert reg.get_default_name() == "fast-model"
|
||||
|
||||
def test_fallback_model(self, sample_models):
|
||||
reg = ModelRegistry(sample_models)
|
||||
assert reg.get_fallback_name() == "pro-model"
|
||||
|
||||
def test_resolve_by_alias(self, sample_models):
|
||||
reg = ModelRegistry(sample_models, aliases={"deepseek:deepseek-v4-flash": "fast-model"})
|
||||
model = reg.resolve("deepseek:deepseek-v4-flash")
|
||||
assert model.name == "fast-model"
|
||||
|
||||
def test_resolve_by_logical_name(self, sample_models):
|
||||
reg = ModelRegistry(sample_models)
|
||||
model = reg.resolve("fast-model")
|
||||
assert model.name == "fast-model"
|
||||
|
||||
def test_list_by_capability(self, sample_models):
|
||||
reg = ModelRegistry(sample_models)
|
||||
models = reg.list_by_capability("vision")
|
||||
assert len(models) == 1
|
||||
assert models[0].name == "vision-model"
|
||||
|
||||
def test_resolve_by_capability(self, sample_models):
|
||||
reg = ModelRegistry(sample_models)
|
||||
name = reg.resolve_by_capability("vision")
|
||||
assert name == "vision-model"
|
||||
|
||||
def test_weighted_random_resolve(self, sample_models):
|
||||
reg = ModelRegistry(sample_models)
|
||||
name = reg.resolve()
|
||||
assert name.name == "fast-model"
|
||||
|
||||
def test_contains_returns_true_for_registered(self, sample_models):
|
||||
reg = ModelRegistry(sample_models)
|
||||
assert reg.contains("fast-model") is True
|
||||
|
||||
def test_contains_returns_false_for_missing(self, sample_models):
|
||||
reg = ModelRegistry(sample_models)
|
||||
assert reg.contains("phantom-model") is False
|
||||
|
||||
def test_list_all_returns_all_models(self, sample_models):
|
||||
reg = ModelRegistry(sample_models)
|
||||
assert len(reg.list_all()) == 3
|
||||
|
||||
|
||||
class TestProviderRegistry:
|
||||
@pytest.fixture
|
||||
def sample_providers(self):
|
||||
return {
|
||||
"deepseek": ProviderConfig(
|
||||
name="deepseek",
|
||||
format="openai",
|
||||
base_url_env="DEEPSEEK_BASE_URL",
|
||||
api_key_env="DEEPSEEK_API_KEY",
|
||||
),
|
||||
"anthropic": ProviderConfig(
|
||||
name="anthropic",
|
||||
format="anthropic",
|
||||
base_url_env="ANTHROPIC_BASE_URL",
|
||||
api_key_env="ANTHROPIC_API_KEY",
|
||||
),
|
||||
}
|
||||
|
||||
def test_get_provider(self, sample_providers):
|
||||
reg = ProviderRegistry(sample_providers)
|
||||
prov = reg.get("deepseek")
|
||||
assert prov.format == "openai"
|
||||
|
||||
def test_get_nonexistent(self, sample_providers):
|
||||
reg = ProviderRegistry(sample_providers)
|
||||
with pytest.raises(AIProviderError):
|
||||
reg.get("nonexistent")
|
||||
|
||||
def test_list_all(self, sample_providers):
|
||||
reg = ProviderRegistry(sample_providers)
|
||||
provs = reg.list_all()
|
||||
assert set(provs.keys()) == {"deepseek", "anthropic"}
|
||||
|
||||
|
||||
class TestRegistryFactory:
|
||||
def test_from_configs(self, tmp_path):
|
||||
models_yaml = tmp_path / "models.yaml"
|
||||
models_yaml.write_text(
|
||||
"""
|
||||
version: 1
|
||||
models:
|
||||
- name: test-model
|
||||
provider: deepseek
|
||||
model: deepseek-v4-flash
|
||||
format: openai
|
||||
"""
|
||||
)
|
||||
providers_yaml = tmp_path / "providers.yaml"
|
||||
providers_yaml.write_text(
|
||||
"""
|
||||
version: 1
|
||||
providers:
|
||||
deepseek:
|
||||
format: openai
|
||||
base_url_env: DEEPSEEK_BASE_URL
|
||||
api_key_env: DEEPSEEK_API_KEY
|
||||
"""
|
||||
)
|
||||
from app.ai_provider.config.loader import load_models_config, load_providers_config
|
||||
|
||||
models = load_models_config(models_yaml)
|
||||
providers = load_providers_config(providers_yaml)
|
||||
registry = RegistryFactory.create(models, providers)
|
||||
assert registry.model_registry.get("test-model") is not None
|
||||
assert registry.provider_registry.get("deepseek") is not None
|
||||
|
||||
def test_create_with_aliases(self, tmp_path):
|
||||
models_yaml = tmp_path / "models.yaml"
|
||||
models_yaml.write_text(
|
||||
"""
|
||||
version: 1
|
||||
models:
|
||||
- name: fast-model
|
||||
provider: deepseek
|
||||
model: deepseek-v4-flash
|
||||
format: openai
|
||||
tags: ["default"]
|
||||
"""
|
||||
)
|
||||
providers_yaml = tmp_path / "providers.yaml"
|
||||
providers_yaml.write_text(
|
||||
"""
|
||||
version: 1
|
||||
providers:
|
||||
deepseek:
|
||||
format: openai
|
||||
base_url_env: DEEPSEEK_BASE_URL
|
||||
api_key_env: DEEPSEEK_API_KEY
|
||||
"""
|
||||
)
|
||||
from app.ai_provider.config.loader import load_models_config, load_providers_config
|
||||
|
||||
models = load_models_config(models_yaml)
|
||||
providers = load_providers_config(providers_yaml)
|
||||
registry = RegistryFactory.create(
|
||||
models, providers, aliases={"deepseek:deepseek-v4-flash": "fast-model"}
|
||||
)
|
||||
assert registry.model_registry.resolve("deepseek:deepseek-v4-flash").name == "fast-model"
|
||||
|
||||
|
||||
class TestRegistryHolder:
|
||||
def test_holds_both_registries(self):
|
||||
model_reg = ModelRegistry(
|
||||
[
|
||||
ModelConfig(
|
||||
name="m1",
|
||||
provider="deepseek",
|
||||
model="deepseek-v4-flash",
|
||||
format="openai",
|
||||
)
|
||||
]
|
||||
)
|
||||
prov_reg = ProviderRegistry(
|
||||
{
|
||||
"deepseek": ProviderConfig(
|
||||
name="deepseek",
|
||||
format="openai",
|
||||
base_url_env="DEEPSEEK_BASE_URL",
|
||||
api_key_env="DEEPSEEK_API_KEY",
|
||||
)
|
||||
}
|
||||
)
|
||||
holder = RegistryHolder(model_reg, prov_reg)
|
||||
assert holder.model_registry.get("m1") is not None
|
||||
assert holder.provider_registry.get("deepseek") is not None
|
||||
@@ -1,109 +0,0 @@
|
||||
# backend/tests/ai_provider/core/test_resilience.py
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
|
||||
from app.ai_provider.core.resilience import (
|
||||
CircuitBreaker,
|
||||
CircuitBreakerState,
|
||||
DowngradeAction,
|
||||
DowngradeChain,
|
||||
ResilienceFactory,
|
||||
RetryStrategy,
|
||||
)
|
||||
|
||||
|
||||
class TestCircuitBreaker:
|
||||
@pytest.fixture
|
||||
def breaker(self):
|
||||
return CircuitBreaker(
|
||||
name="test-breaker",
|
||||
failure_threshold=3,
|
||||
cooldown_seconds=1.0,
|
||||
)
|
||||
|
||||
def test_initial_state(self, breaker):
|
||||
assert breaker.state == CircuitBreakerState.CLOSED
|
||||
|
||||
def test_trip_after_threshold(self, breaker):
|
||||
breaker.record_failure()
|
||||
assert breaker.state == CircuitBreakerState.CLOSED
|
||||
breaker.record_failure()
|
||||
assert breaker.state == CircuitBreakerState.CLOSED
|
||||
breaker.record_failure()
|
||||
assert breaker.state == CircuitBreakerState.OPEN
|
||||
|
||||
def test_reset_on_success(self, breaker):
|
||||
breaker.record_failure()
|
||||
breaker.record_failure()
|
||||
breaker.record_success()
|
||||
assert breaker.state == CircuitBreakerState.CLOSED
|
||||
|
||||
def test_half_open_after_cooldown(self):
|
||||
breaker = CircuitBreaker(name="test", failure_threshold=1, cooldown_seconds=0.01)
|
||||
breaker.record_failure()
|
||||
assert breaker.state == CircuitBreakerState.OPEN
|
||||
asyncio.run(asyncio.sleep(0.02))
|
||||
assert breaker.state == CircuitBreakerState.HALF_OPEN
|
||||
|
||||
def test_trip_from_half_open(self):
|
||||
breaker = CircuitBreaker(name="test", failure_threshold=1, cooldown_seconds=0.01)
|
||||
breaker.record_failure()
|
||||
asyncio.run(asyncio.sleep(0.02))
|
||||
assert breaker.state == CircuitBreakerState.HALF_OPEN
|
||||
breaker.record_failure()
|
||||
assert breaker.state == CircuitBreakerState.OPEN
|
||||
|
||||
|
||||
class TestRetryStrategy:
|
||||
@pytest.fixture
|
||||
def strategy(self):
|
||||
return RetryStrategy(max_retries=3, base_delay=0.01, max_delay=0.1)
|
||||
|
||||
def test_should_retry(self, strategy):
|
||||
assert strategy.should_retry(0)
|
||||
assert strategy.should_retry(1)
|
||||
assert strategy.should_retry(2)
|
||||
assert not strategy.should_retry(3)
|
||||
|
||||
def test_delay_exponential(self, strategy):
|
||||
d0 = strategy.get_delay(0)
|
||||
d1 = strategy.get_delay(1)
|
||||
d2 = strategy.get_delay(2)
|
||||
assert d1 > d0
|
||||
assert d2 > d1
|
||||
|
||||
|
||||
class TestDowngradeChain:
|
||||
def test_fallback_chain(self):
|
||||
chain = DowngradeChain(
|
||||
[
|
||||
DowngradeAction(type="fallback_endpoint", target="backup-1"),
|
||||
DowngradeAction(type="fallback_endpoint", target="backup-2"),
|
||||
DowngradeAction(type="fallback_model", target="cheap-model"),
|
||||
]
|
||||
)
|
||||
assert chain.next_action() == ("fallback_endpoint", "backup-1")
|
||||
chain.advance()
|
||||
assert chain.next_action() == ("fallback_endpoint", "backup-2")
|
||||
chain.advance()
|
||||
assert chain.next_action() == ("fallback_model", "cheap-model")
|
||||
chain.advance()
|
||||
assert chain.is_exhausted()
|
||||
|
||||
|
||||
class TestResilienceFactory:
|
||||
def test_create_breaker(self):
|
||||
breaker = ResilienceFactory.create_circuit_breaker(
|
||||
name="test",
|
||||
failure_threshold=5,
|
||||
cooldown_seconds=30,
|
||||
)
|
||||
assert breaker.failure_threshold == 5
|
||||
|
||||
def test_create_retry(self):
|
||||
strategy = ResilienceFactory.create_retry_strategy(
|
||||
max_retries=3,
|
||||
base_delay=1.0,
|
||||
)
|
||||
assert strategy.max_retries == 3
|
||||
@@ -1,543 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from app.ai_provider.config.models import ModelConfig, ParameterMapping, ProviderConfig
|
||||
from app.ai_provider.llm.anthropic_format import AnthropicFormatAdapter
|
||||
from app.ai_provider.types.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage
|
||||
|
||||
|
||||
class TestAnthropicMessageConversion:
|
||||
"""Conversion of internal message types to Anthropic content-block format."""
|
||||
|
||||
@pytest.fixture
|
||||
def adapter(self):
|
||||
provider = ProviderConfig(
|
||||
name="anthropic",
|
||||
format="anthropic",
|
||||
base_url_env="ANTHROPIC_BASE_URL",
|
||||
api_key_env="ANTHROPIC_API_KEY",
|
||||
parameter_mappings={
|
||||
"temperature": ParameterMapping(provider_param="temperature", min=0, max=1),
|
||||
"max_tokens": ParameterMapping(provider_param="max_tokens"),
|
||||
"top_p": ParameterMapping(provider_param="top_p", min=0, max=1),
|
||||
},
|
||||
)
|
||||
return AnthropicFormatAdapter(provider_config=provider)
|
||||
|
||||
def test_convert_human_message(self, adapter):
|
||||
"""HumanMessage becomes user role with a text content block."""
|
||||
msg = HumanMessage(content="Hello!")
|
||||
result = adapter._convert_message(msg)
|
||||
assert result["role"] == "user"
|
||||
assert result["content"] == [{"type": "text", "text": "Hello!"}]
|
||||
|
||||
def test_convert_ai_message(self, adapter):
|
||||
"""AIMessage becomes assistant role with a text content block."""
|
||||
msg = AIMessage(content="Sure, let me help.")
|
||||
result = adapter._convert_message(msg)
|
||||
assert result["role"] == "assistant"
|
||||
assert result["content"] == [{"type": "text", "text": "Sure, let me help."}]
|
||||
|
||||
def test_convert_ai_message_replays_raw_content_blocks(self, adapter):
|
||||
blocks = [
|
||||
{"type": "thinking", "thinking": "summary", "signature": "sig_1"},
|
||||
{"type": "tool_use", "id": "call_1", "name": "read_file", "input": {}},
|
||||
]
|
||||
msg = AIMessage(content="", raw_content_blocks=blocks)
|
||||
|
||||
assert adapter._convert_message(msg) == {"role": "assistant", "content": blocks}
|
||||
|
||||
def test_convert_ai_message_with_tool_calls(self, adapter):
|
||||
"""AIMessage with tool_calls produces tool_use content blocks."""
|
||||
msg = AIMessage(
|
||||
content="",
|
||||
tool_calls=[
|
||||
{"name": "read_file", "args": {"path": "x.txt"}, "id": "call_1"},
|
||||
],
|
||||
)
|
||||
result = adapter._convert_message(msg)
|
||||
assert result["role"] == "assistant"
|
||||
assert result["content"] == [
|
||||
{"type": "tool_use", "id": "call_1", "name": "read_file", "input": {"path": "x.txt"}},
|
||||
]
|
||||
|
||||
def test_convert_ai_message_with_text_and_tool_calls(self, adapter):
|
||||
"""AIMessage with both text and tool_calls produces multiple content blocks."""
|
||||
msg = AIMessage(
|
||||
content="I will check that file.",
|
||||
tool_calls=[
|
||||
{"name": "read_file", "args": {"path": "x.txt"}, "id": "call_1"},
|
||||
{"name": "list_dir", "args": {"path": "/tmp"}, "id": "call_2"},
|
||||
],
|
||||
)
|
||||
result = adapter._convert_message(msg)
|
||||
assert result["role"] == "assistant"
|
||||
assert len(result["content"]) == 3
|
||||
assert result["content"][0] == {"type": "text", "text": "I will check that file."}
|
||||
assert result["content"][1] == {
|
||||
"type": "tool_use",
|
||||
"id": "call_1",
|
||||
"name": "read_file",
|
||||
"input": {"path": "x.txt"},
|
||||
}
|
||||
assert result["content"][2] == {
|
||||
"type": "tool_use",
|
||||
"id": "call_2",
|
||||
"name": "list_dir",
|
||||
"input": {"path": "/tmp"},
|
||||
}
|
||||
|
||||
def test_convert_tool_message(self, adapter):
|
||||
"""ToolMessage becomes user role with a tool_result content block."""
|
||||
msg = ToolMessage(content="File contents", tool_call_id="call_1")
|
||||
result = adapter._convert_message(msg)
|
||||
assert result["role"] == "user"
|
||||
assert result["content"] == [
|
||||
{"type": "tool_result", "tool_use_id": "call_1", "content": "File contents"},
|
||||
]
|
||||
|
||||
|
||||
class TestAnthropicRequestBuilding:
|
||||
"""Building the full request body for the Anthropic API."""
|
||||
|
||||
@pytest.fixture
|
||||
def adapter(self):
|
||||
provider = ProviderConfig(
|
||||
name="anthropic",
|
||||
format="anthropic",
|
||||
base_url_env="ANTHROPIC_BASE_URL",
|
||||
api_key_env="ANTHROPIC_API_KEY",
|
||||
parameter_mappings={
|
||||
"temperature": ParameterMapping(provider_param="temperature", min=0, max=1),
|
||||
"max_tokens": ParameterMapping(provider_param="max_tokens"),
|
||||
},
|
||||
)
|
||||
return AnthropicFormatAdapter(provider_config=provider)
|
||||
|
||||
@pytest.fixture
|
||||
def model(self):
|
||||
return ModelConfig(
|
||||
name="claude-sonnet-4",
|
||||
provider="anthropic",
|
||||
model="claude-sonnet-4-20250514",
|
||||
format="anthropic",
|
||||
max_output_tokens=8192,
|
||||
)
|
||||
|
||||
def test_build_request_body(self, adapter, model):
|
||||
"""Request body includes model, messages, and pipeline-processed params."""
|
||||
messages = [HumanMessage(content="Hi")]
|
||||
body = adapter.build_request(
|
||||
model=model,
|
||||
messages=messages,
|
||||
params={"temperature": 0.7, "max_tokens": 4096},
|
||||
)
|
||||
assert body["model"] == "claude-sonnet-4-20250514"
|
||||
assert body["temperature"] == 0.7
|
||||
assert body["max_tokens"] == 4096
|
||||
assert len(body["messages"]) == 1
|
||||
assert body["messages"][0]["role"] == "user"
|
||||
|
||||
def test_build_request_with_system_message(self, adapter, model):
|
||||
"""System message is placed in top-level 'system' field, not in messages."""
|
||||
messages = [
|
||||
SystemMessage(content="You are a helpful assistant."),
|
||||
HumanMessage(content="Hi"),
|
||||
]
|
||||
body = adapter.build_request(model=model, messages=messages)
|
||||
assert body["system"] == "You are a helpful assistant."
|
||||
# system message is removed from the messages array
|
||||
assert len(body["messages"]) == 1
|
||||
assert body["messages"][0]["role"] == "user"
|
||||
|
||||
def test_build_request_with_multiple_system_messages(self, adapter, model):
|
||||
"""Multiple system messages are concatenated with newlines."""
|
||||
messages = [
|
||||
SystemMessage(content="Be concise."),
|
||||
SystemMessage(content="Use Chinese."),
|
||||
HumanMessage(content="Hi"),
|
||||
]
|
||||
body = adapter.build_request(model=model, messages=messages)
|
||||
assert body["system"] == "Be concise.\nUse Chinese."
|
||||
assert len(body["messages"]) == 1
|
||||
|
||||
def test_build_request_sets_default_max_tokens(self, adapter, model):
|
||||
"""max_tokens defaults to model.max_output_tokens when not in params."""
|
||||
messages = [HumanMessage(content="Hi")]
|
||||
body = adapter.build_request(model=model, messages=messages, params={})
|
||||
assert body["max_tokens"] == 8192
|
||||
|
||||
def test_build_request_with_tools(self, adapter, model):
|
||||
"""Tools are converted to Anthropic format in the request."""
|
||||
messages = [HumanMessage(content="Read file x.txt")]
|
||||
openai_tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "read_file",
|
||||
"description": "Read a file's contents",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"path": {"type": "string"}},
|
||||
"required": ["path"],
|
||||
},
|
||||
},
|
||||
},
|
||||
]
|
||||
body = adapter.build_request(
|
||||
model=model,
|
||||
messages=messages,
|
||||
tools=openai_tools,
|
||||
tool_choice="auto",
|
||||
)
|
||||
assert "tools" in body
|
||||
assert len(body["tools"]) == 1
|
||||
anthropic_tool = body["tools"][0]
|
||||
assert anthropic_tool["name"] == "read_file"
|
||||
assert anthropic_tool["description"] == "Read a file's contents"
|
||||
assert anthropic_tool["input_schema"] == {
|
||||
"type": "object",
|
||||
"properties": {"path": {"type": "string"}},
|
||||
"required": ["path"],
|
||||
}
|
||||
assert body["tool_choice"] == "auto"
|
||||
|
||||
def test_build_request_streaming(self, adapter, model):
|
||||
"""Streaming flag is passed through in the request."""
|
||||
messages = [HumanMessage(content="Hi")]
|
||||
body = adapter.build_request(model=model, messages=messages, stream=True)
|
||||
assert body["stream"] is True
|
||||
|
||||
|
||||
class TestAnthropicToolConversion:
|
||||
"""Converting tools between OpenAI and Anthropic formats."""
|
||||
|
||||
def test_convert_tools_standard(self):
|
||||
"""OpenAI format tools are converted to Anthropic format."""
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "read_file",
|
||||
"description": "Read a file",
|
||||
"parameters": {"type": "object"},
|
||||
},
|
||||
},
|
||||
]
|
||||
result = AnthropicFormatAdapter._convert_tools(tools)
|
||||
assert len(result) == 1
|
||||
assert result[0] == {
|
||||
"name": "read_file",
|
||||
"description": "Read a file",
|
||||
"input_schema": {"type": "object"},
|
||||
}
|
||||
|
||||
def test_convert_tools_already_anthropic(self):
|
||||
"""Tools already in Anthropic format pass through correctly."""
|
||||
tools = [
|
||||
{
|
||||
"name": "read_file",
|
||||
"description": "Read a file",
|
||||
"input_schema": {"type": "object"},
|
||||
},
|
||||
]
|
||||
result = AnthropicFormatAdapter._convert_tools(tools)
|
||||
assert len(result) == 1
|
||||
assert result[0] == tools[0]
|
||||
|
||||
|
||||
class TestAnthropicResponseParsing:
|
||||
"""Parsing Anthropic API responses (non-streaming)."""
|
||||
|
||||
@pytest.fixture
|
||||
def adapter(self):
|
||||
return AnthropicFormatAdapter(
|
||||
provider_config=ProviderConfig(
|
||||
name="anthropic",
|
||||
format="anthropic",
|
||||
base_url_env="ANTHROPIC_BASE_URL",
|
||||
api_key_env="ANTHROPIC_API_KEY",
|
||||
)
|
||||
)
|
||||
|
||||
def test_parse_nonstreaming_text(self, adapter):
|
||||
"""Simple text response is parsed correctly."""
|
||||
raw = {
|
||||
"id": "msg_123",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "text", "text": "Hello, world!"}],
|
||||
"model": "claude-sonnet-4-20250514",
|
||||
"stop_reason": "end_turn",
|
||||
"usage": {"input_tokens": 10, "output_tokens": 5},
|
||||
}
|
||||
result = adapter.parse_nonstreaming_response(raw, "claude-sonnet-4")
|
||||
assert result.content == "Hello, world!"
|
||||
assert result.model_name == "claude-sonnet-4"
|
||||
assert result.usage is not None
|
||||
assert result.usage.prompt_tokens == 10
|
||||
assert result.usage.completion_tokens == 5
|
||||
assert result.usage.total_tokens == 15
|
||||
|
||||
def test_parse_nonstreaming_with_tool_calls(self, adapter):
|
||||
"""Response with tool_use blocks is parsed with tool_calls."""
|
||||
raw = {
|
||||
"id": "msg_456",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "text", "text": "I will read the file."},
|
||||
{
|
||||
"type": "tool_use",
|
||||
"id": "tu_123",
|
||||
"name": "read_file",
|
||||
"input": {"path": "x.txt"},
|
||||
},
|
||||
],
|
||||
"model": "claude-sonnet-4-20250514",
|
||||
"stop_reason": "tool_use",
|
||||
"usage": {"input_tokens": 20, "output_tokens": 15},
|
||||
}
|
||||
result = adapter.parse_nonstreaming_response(raw, "claude-sonnet-4")
|
||||
assert result.content == "I will read the file."
|
||||
assert len(result.tool_calls) == 1
|
||||
assert result.tool_calls[0] == {
|
||||
"name": "read_file",
|
||||
"args": {"path": "x.txt"},
|
||||
"id": "tu_123",
|
||||
}
|
||||
|
||||
def test_parse_nonstreaming_with_reasoning(self, adapter):
|
||||
"""Response with thinking blocks populates reasoning field."""
|
||||
raw = {
|
||||
"id": "msg_789",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "thinking", "thinking": "Let me think step by step..."},
|
||||
{"type": "text", "text": "Here is the answer."},
|
||||
{"type": "signature", "signature": "sig_12345"},
|
||||
],
|
||||
"model": "claude-sonnet-4-20250514",
|
||||
"stop_reason": "end_turn",
|
||||
"usage": {"input_tokens": 5, "output_tokens": 10},
|
||||
}
|
||||
result = adapter.parse_nonstreaming_response(raw, "claude-sonnet-4")
|
||||
assert result.content == "Here is the answer."
|
||||
# signature 块是加密签名，不进入 reasoning
|
||||
assert result.reasoning == "Let me think step by step..."
|
||||
|
||||
def test_parse_nonstreaming_no_usage(self, adapter):
|
||||
"""Response without usage data still parses correctly."""
|
||||
raw = {
|
||||
"id": "msg_000",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "text", "text": "No usage info."}],
|
||||
"model": "claude-sonnet-4-20250514",
|
||||
"stop_reason": "end_turn",
|
||||
}
|
||||
result = adapter.parse_nonstreaming_response(raw, "claude-sonnet-4")
|
||||
assert result.content == "No usage info."
|
||||
assert result.usage is None
|
||||
|
||||
def test_parse_content_blocks_empty(self):
|
||||
"""Empty content blocks produce empty text and no tool calls."""
|
||||
raw = {"content": []}
|
||||
text, tools, reasoning = AnthropicFormatAdapter._parse_content_blocks(raw)
|
||||
assert text == ""
|
||||
assert tools == []
|
||||
assert reasoning == ""
|
||||
|
||||
def test_bind_tools(self):
|
||||
"""bind_tools adds converted tools to the request body."""
|
||||
adapter = AnthropicFormatAdapter(
|
||||
provider_config=ProviderConfig(
|
||||
name="anthropic",
|
||||
format="anthropic",
|
||||
base_url_env="ANTHROPIC_BASE_URL",
|
||||
api_key_env="ANTHROPIC_API_KEY",
|
||||
)
|
||||
)
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "read_file",
|
||||
"description": "Read a file",
|
||||
"parameters": {"type": "object"},
|
||||
},
|
||||
},
|
||||
]
|
||||
body = {}
|
||||
adapter.bind_tools(body, tools, tool_choice="auto")
|
||||
assert body["tools"] == [
|
||||
{
|
||||
"name": "read_file",
|
||||
"description": "Read a file",
|
||||
"input_schema": {"type": "object"},
|
||||
},
|
||||
]
|
||||
assert body["tool_choice"] == "auto"
|
||||
|
||||
|
||||
class TestAnthropicStreamParsing:
|
||||
"""Parsing Anthropic SSE stream events."""
|
||||
|
||||
def test_parse_message_start(self):
|
||||
"""message_start event sets initial token usage."""
|
||||
data = {
|
||||
"type": "message_start",
|
||||
"message": {
|
||||
"id": "msg_123",
|
||||
"usage": {"input_tokens": 15, "output_tokens": 0},
|
||||
},
|
||||
}
|
||||
chunk = AnthropicFormatAdapter._parse_stream_event("message_start", data)
|
||||
assert chunk.usage is not None
|
||||
assert chunk.usage.prompt_tokens == 15
|
||||
assert chunk.usage.completion_tokens == 0
|
||||
assert chunk.delta_content == ""
|
||||
|
||||
def test_parse_content_block_start_text(self):
|
||||
"""content_block_start for text sets delta_content."""
|
||||
data = {
|
||||
"type": "content_block_start",
|
||||
"index": 0,
|
||||
"content_block": {"type": "text", "text": "Hello"},
|
||||
}
|
||||
chunk = AnthropicFormatAdapter._parse_stream_event("content_block_start", data)
|
||||
assert chunk.delta_content == "Hello"
|
||||
|
||||
def test_parse_content_block_start_thinking(self):
|
||||
"""content_block_start for thinking sets delta_reasoning."""
|
||||
data = {
|
||||
"type": "content_block_start",
|
||||
"index": 0,
|
||||
"content_block": {"type": "thinking", "thinking": "Let me think..."},
|
||||
}
|
||||
chunk = AnthropicFormatAdapter._parse_stream_event("content_block_start", data)
|
||||
assert chunk.delta_reasoning == "Let me think..."
|
||||
|
||||
def test_parse_content_block_start_tool_use(self):
|
||||
"""content_block_start for tool_use creates a tool call chunk stub."""
|
||||
data = {
|
||||
"type": "content_block_start",
|
||||
"index": 1,
|
||||
"content_block": {
|
||||
"type": "tool_use",
|
||||
"id": "tu_123",
|
||||
"name": "read_file",
|
||||
"input": {},
|
||||
},
|
||||
}
|
||||
chunk = AnthropicFormatAdapter._parse_stream_event("content_block_start", data)
|
||||
assert len(chunk.tool_call_chunks) == 1
|
||||
assert chunk.tool_call_chunks[0]["index"] == 1
|
||||
assert chunk.tool_call_chunks[0]["id"] == "tu_123"
|
||||
assert chunk.tool_call_chunks[0]["name"] == "read_file"
|
||||
assert chunk.tool_call_chunks[0]["args"] == ""
|
||||
|
||||
def test_parse_content_block_delta_text(self):
|
||||
"""content_block_delta for text_delta appends content."""
|
||||
data = {
|
||||
"type": "content_block_delta",
|
||||
"index": 0,
|
||||
"delta": {"type": "text_delta", "text": " world"},
|
||||
}
|
||||
chunk = AnthropicFormatAdapter._parse_stream_event("content_block_delta", data)
|
||||
assert chunk.delta_content == " world"
|
||||
|
||||
def test_parse_content_block_delta_thinking(self):
|
||||
"""content_block_delta for thinking_delta appends reasoning."""
|
||||
data = {
|
||||
"type": "content_block_delta",
|
||||
"index": 0,
|
||||
"delta": {"type": "thinking_delta", "thinking": " more..."},
|
||||
}
|
||||
chunk = AnthropicFormatAdapter._parse_stream_event("content_block_delta", data)
|
||||
assert chunk.delta_reasoning == " more..."
|
||||
|
||||
def test_parse_content_block_delta_signature(self):
|
||||
"""content_block_delta for signature_delta is ignored (not reasoning)."""
|
||||
data = {
|
||||
"type": "content_block_delta",
|
||||
"index": 0,
|
||||
"delta": {"type": "signature_delta", "signature": "sig_abc"},
|
||||
}
|
||||
chunk = AnthropicFormatAdapter._parse_stream_event("content_block_delta", data)
|
||||
# signature 不展示为 reasoning，但保留原始 delta 供 Agent 下一轮回放。
|
||||
assert chunk.delta_reasoning == ""
|
||||
assert chunk.content_block_index == 0
|
||||
assert chunk.content_block_delta == data["delta"]
|
||||
|
||||
def test_parse_content_block_delta_input_json(self):
|
||||
"""content_block_delta for input_json_delta creates a tool call arg chunk."""
|
||||
data = {
|
||||
"type": "content_block_delta",
|
||||
"index": 0,
|
||||
"delta": {"type": "input_json_delta", "partial_json": '{"path": "'},
|
||||
}
|
||||
chunk = AnthropicFormatAdapter._parse_stream_event("content_block_delta", data)
|
||||
assert len(chunk.tool_call_chunks) == 1
|
||||
assert chunk.tool_call_chunks[0]["args"] == '{"path": "'
|
||||
|
||||
def test_parse_message_delta_stop(self):
|
||||
"""message_delta with end_turn sets finish_reason to 'stop'."""
|
||||
data = {
|
||||
"type": "message_delta",
|
||||
"delta": {"stop_reason": "end_turn", "stop_sequence": None},
|
||||
"usage": {"output_tokens": 42},
|
||||
}
|
||||
chunk = AnthropicFormatAdapter._parse_stream_event("message_delta", data)
|
||||
assert chunk.finish_reason == "stop"
|
||||
assert chunk.usage is not None
|
||||
assert chunk.usage.completion_tokens == 42
|
||||
|
||||
def test_parse_message_delta_tool_use(self):
|
||||
"""message_delta with tool_use sets finish_reason to 'tool_calls'."""
|
||||
data = {
|
||||
"type": "message_delta",
|
||||
"delta": {"stop_reason": "tool_use"},
|
||||
"usage": {"output_tokens": 10},
|
||||
}
|
||||
chunk = AnthropicFormatAdapter._parse_stream_event("message_delta", data)
|
||||
assert chunk.finish_reason == "tool_calls"
|
||||
|
||||
def test_parse_message_delta_max_tokens(self):
|
||||
"""message_delta with max_tokens stop reason."""
|
||||
data = {
|
||||
"type": "message_delta",
|
||||
"delta": {"stop_reason": "max_tokens"},
|
||||
"usage": {"output_tokens": 5},
|
||||
}
|
||||
chunk = AnthropicFormatAdapter._parse_stream_event("message_delta", data)
|
||||
assert chunk.finish_reason == "length"
|
||||
|
||||
def test_parse_content_block_stop_noop(self):
|
||||
"""content_block_stop produces an empty chunk."""
|
||||
data = {"type": "content_block_stop", "index": 0}
|
||||
chunk = AnthropicFormatAdapter._parse_stream_event("content_block_stop", data)
|
||||
assert chunk.delta_content == ""
|
||||
assert chunk.delta_reasoning == ""
|
||||
assert chunk.usage is None
|
||||
|
||||
|
||||
class TestAnthropicAuthHeaders:
|
||||
"""Anthropic-specific authentication headers."""
|
||||
|
||||
def test_headers_contain_api_key_and_version(self):
|
||||
"""Headers include x-api-key, anthropic-version, and content type."""
|
||||
adapter = AnthropicFormatAdapter(
|
||||
provider_config=ProviderConfig(
|
||||
name="anthropic",
|
||||
format="anthropic",
|
||||
base_url_env="ANTHROPIC_BASE_URL",
|
||||
api_key_env="ANTHROPIC_API_KEY",
|
||||
)
|
||||
)
|
||||
headers = adapter._headers("sk-ant-12345")
|
||||
assert headers["x-api-key"] == "sk-ant-12345"
|
||||
assert headers["anthropic-version"] == "2023-06-01"
|
||||
assert headers["Content-Type"] == "application/json"
|
||||
@@ -1,114 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from app.ai_provider.config.models import ModelConfig, ParameterMapping, ProviderConfig
|
||||
from app.ai_provider.llm.openai_format import OpenAIFormatAdapter
|
||||
from app.ai_provider.types.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage
|
||||
|
||||
|
||||
class TestOpenAIMessageConversion:
|
||||
@pytest.fixture
|
||||
def adapter(self):
|
||||
provider = ProviderConfig(
|
||||
name="deepseek",
|
||||
format="openai",
|
||||
base_url_env="DEEPSEEK_BASE_URL",
|
||||
api_key_env="DEEPSEEK_API_KEY",
|
||||
parameter_mappings={
|
||||
"temperature": ParameterMapping(provider_param="temperature", min=0, max=2),
|
||||
"max_tokens": ParameterMapping(provider_param="max_tokens"),
|
||||
},
|
||||
)
|
||||
return OpenAIFormatAdapter(provider_config=provider)
|
||||
|
||||
def test_convert_system_message(self, adapter):
|
||||
msg = SystemMessage(content="You are a helpful assistant.")
|
||||
result = adapter._convert_message(msg)
|
||||
assert result["role"] == "system"
|
||||
assert result["content"] == "You are a helpful assistant."
|
||||
|
||||
def test_convert_human_message(self, adapter):
|
||||
msg = HumanMessage(content="Hello!")
|
||||
result = adapter._convert_message(msg)
|
||||
assert result["role"] == "user"
|
||||
|
||||
def test_convert_ai_message(self, adapter):
|
||||
msg = AIMessage(
|
||||
content="Sure!",
|
||||
tool_calls=[{"name": "read_file", "args": {"path": "x.txt"}, "id": "call_1"}],
|
||||
)
|
||||
result = adapter._convert_message(msg)
|
||||
assert result["role"] == "assistant"
|
||||
assert "tool_calls" in result
|
||||
|
||||
def test_convert_tool_message(self, adapter):
|
||||
msg = ToolMessage(content="File contents", tool_call_id="call_1")
|
||||
result = adapter._convert_message(msg)
|
||||
assert result["role"] == "tool"
|
||||
assert result["tool_call_id"] == "call_1"
|
||||
|
||||
def test_build_request_body(self, adapter):
|
||||
model = ModelConfig(
|
||||
name="test", provider="deepseek", model="deepseek-v4-flash", format="openai"
|
||||
)
|
||||
messages = [HumanMessage(content="Hi")]
|
||||
body = adapter.build_request(
|
||||
model=model,
|
||||
messages=messages,
|
||||
params={"temperature": 0.7, "max_tokens": 4096},
|
||||
)
|
||||
assert body["model"] == "deepseek-v4-flash"
|
||||
assert body["temperature"] == 0.7
|
||||
assert body["max_tokens"] == 4096
|
||||
assert len(body["messages"]) == 1
|
||||
|
||||
def test_parse_chunk(self, adapter):
|
||||
raw = {"choices": [{"delta": {"content": "Hello"}, "index": 0}], "usage": None}
|
||||
chunk = adapter.parse_chunk(raw)
|
||||
assert chunk.delta_content == "Hello"
|
||||
|
||||
def test_parse_chunk_with_usage(self, adapter):
|
||||
raw = {
|
||||
"choices": [{"delta": {}, "index": 0, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
|
||||
}
|
||||
chunk = adapter.parse_chunk(raw)
|
||||
assert chunk.usage is not None
|
||||
assert chunk.usage.prompt_tokens == 10
|
||||
|
||||
def test_parse_tool_call_chunk(self, adapter):
|
||||
raw = {
|
||||
"choices": [
|
||||
{
|
||||
"delta": {
|
||||
"tool_calls": [
|
||||
{
|
||||
"index": 0,
|
||||
"id": "call_1",
|
||||
"function": {
|
||||
"name": "read_file",
|
||||
"arguments": '{"path": "test.txt"}',
|
||||
},
|
||||
}
|
||||
],
|
||||
},
|
||||
"index": 0,
|
||||
}
|
||||
],
|
||||
}
|
||||
chunk = adapter.parse_chunk(raw)
|
||||
assert len(chunk.tool_call_chunks) == 1
|
||||
assert chunk.tool_call_chunks[0]["name"] == "read_file"
|
||||
|
||||
def test_bind_tools(self, adapter):
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {"name": "read_file", "parameters": {"type": "object"}},
|
||||
}
|
||||
]
|
||||
body = {}
|
||||
adapter.bind_tools(body, tools, tool_choice="auto")
|
||||
assert body["tools"] == tools
|
||||
assert body["tool_choice"] == "auto"
|
||||
@@ -1,100 +0,0 @@
|
||||
from app.ai_provider.llm.pool import (
|
||||
CostAwareStrategy,
|
||||
EndpointPool,
|
||||
LeastConnectionsStrategy,
|
||||
LowestLatencyStrategy,
|
||||
PoolFactory,
|
||||
PriorityFailoverStrategy,
|
||||
WeightedRoundRobinStrategy,
|
||||
)
|
||||
|
||||
|
||||
class TestWeightedRoundRobin:
|
||||
def test_distribution(self):
|
||||
endpoints = [{"name": "a", "weight": 100}, {"name": "b", "weight": 50}]
|
||||
strategy = WeightedRoundRobinStrategy(endpoints)
|
||||
picks = [strategy.select([])["name"] for _ in range(15)]
|
||||
assert picks.count("a") > picks.count("b")
|
||||
|
||||
def test_skip_excluded(self):
|
||||
endpoints = [{"name": "a", "weight": 100}, {"name": "b", "weight": 100}]
|
||||
strategy = WeightedRoundRobinStrategy(endpoints)
|
||||
pick = strategy.select(["a"])
|
||||
assert pick["name"] == "b"
|
||||
|
||||
|
||||
class TestPriorityFailover:
|
||||
def test_primary_selected(self):
|
||||
endpoints = [{"name": "primary", "priority": 10}, {"name": "backup", "priority": 20}]
|
||||
strategy = PriorityFailoverStrategy(endpoints)
|
||||
pick = strategy.select([])
|
||||
assert pick["name"] == "primary"
|
||||
|
||||
def test_failover_on_excluded(self):
|
||||
endpoints = [{"name": "primary", "priority": 10}, {"name": "backup", "priority": 20}]
|
||||
strategy = PriorityFailoverStrategy(endpoints)
|
||||
pick = strategy.select(["primary"])
|
||||
assert pick["name"] == "backup"
|
||||
|
||||
def test_exhausted(self):
|
||||
endpoints = [{"name": "a", "priority": 10}]
|
||||
strategy = PriorityFailoverStrategy(endpoints)
|
||||
pick = strategy.select(["a"])
|
||||
assert pick is None
|
||||
|
||||
|
||||
class TestLowestLatency:
|
||||
def test_pick_lowest_latency(self):
|
||||
endpoints = [{"name": "fast", "latency_ms": 100}, {"name": "slow", "latency_ms": 500}]
|
||||
strategy = LowestLatencyStrategy(endpoints)
|
||||
pick = strategy.select([])
|
||||
assert pick["name"] == "fast"
|
||||
|
||||
|
||||
class TestCostAware:
|
||||
def test_pick_cheapest(self):
|
||||
endpoints = [{"name": "cheap"}, {"name": "expensive"}]
|
||||
|
||||
def _cost(name: str) -> float:
|
||||
return 0.1 if name == "cheap" else 1.0
|
||||
|
||||
strategy = CostAwareStrategy(endpoints, get_cost=_cost)
|
||||
pick = strategy.select([])
|
||||
assert pick["name"] == "cheap"
|
||||
|
||||
|
||||
class TestLeastConnections:
|
||||
def test_pick_least_connections(self):
|
||||
endpoints = [{"name": "busy", "connections": 5}, {"name": "idle", "connections": 0}]
|
||||
strategy = LeastConnectionsStrategy(endpoints)
|
||||
pick = strategy.select([])
|
||||
assert pick["name"] == "idle"
|
||||
|
||||
|
||||
class TestEndpointPool:
|
||||
def test_select_with_weighted_round_robin(self):
|
||||
pool = EndpointPool(
|
||||
logical_model_name="test",
|
||||
endpoints_configs=[
|
||||
{"name": "a", "weight": 100, "status": "active"},
|
||||
{"name": "b", "weight": 50, "status": "active"},
|
||||
],
|
||||
strategy="weighted_round_robin",
|
||||
)
|
||||
selected = pool.select([])
|
||||
assert selected["name"] in ("a", "b")
|
||||
|
||||
def test_select_all_excluded_returns_none(self):
|
||||
pool = EndpointPool(
|
||||
logical_model_name="test",
|
||||
endpoints_configs=[{"name": "a", "weight": 100, "status": "active"}],
|
||||
)
|
||||
assert pool.select(["a"]) is None
|
||||
|
||||
def test_strategy_factory(self):
|
||||
pool = PoolFactory.create_pool(
|
||||
logical_model_name="test",
|
||||
endpoints=[{"name": "a", "weight": 100, "status": "active"}],
|
||||
strategy_name="priority_failover",
|
||||
)
|
||||
assert isinstance(pool.strategy, PriorityFailoverStrategy)
|
||||
@@ -1,168 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from app.ai_provider.llm.stream import StreamFactory, StreamHandler, ToolCallAccumulator
|
||||
from app.ai_provider.types.chat import ChatChunk, TokenUsage
|
||||
|
||||
|
||||
class TestToolCallAccumulator:
|
||||
def test_accumulate_simple(self):
|
||||
acc = ToolCallAccumulator()
|
||||
acc.add_chunk(ChatChunk(delta_content="Hello"))
|
||||
assert acc.content == "Hello"
|
||||
|
||||
def test_accumulate_reasoning(self):
|
||||
acc = ToolCallAccumulator()
|
||||
acc.add_chunk(ChatChunk(delta_reasoning="I think"))
|
||||
acc.add_chunk(ChatChunk(delta_reasoning=" therefore"))
|
||||
assert acc.reasoning == "I think therefore"
|
||||
|
||||
def test_accumulate_tool_call(self):
|
||||
acc = ToolCallAccumulator()
|
||||
acc.add_chunk(
|
||||
ChatChunk(
|
||||
tool_call_chunks=[
|
||||
{
|
||||
"index": 0,
|
||||
"id": "call_1",
|
||||
"name": "read_file",
|
||||
"args": '{"path": "test.txt"}',
|
||||
},
|
||||
]
|
||||
)
|
||||
)
|
||||
result = acc.drain()
|
||||
assert len(result) == 1
|
||||
assert result[0]["name"] == "read_file"
|
||||
|
||||
def test_accumulate_tool_call_args(self):
|
||||
acc = ToolCallAccumulator()
|
||||
acc.add_chunk(ChatChunk(tool_call_chunks=[{"index": 0, "args": '{"path": "'}]))
|
||||
acc.add_chunk(ChatChunk(tool_call_chunks=[{"index": 0, "args": 'test.txt"}'}]))
|
||||
result = acc.drain()
|
||||
assert result[0]["args"] == '{"path": "test.txt"}'
|
||||
|
||||
def test_accumulate_usage(self):
|
||||
acc = ToolCallAccumulator()
|
||||
acc.add_chunk(
|
||||
ChatChunk(usage=TokenUsage(prompt_tokens=10, completion_tokens=5, total_tokens=15))
|
||||
)
|
||||
assert acc.usage is not None
|
||||
assert acc.usage.prompt_tokens == 10
|
||||
|
||||
def test_empty_drain(self):
|
||||
acc = ToolCallAccumulator()
|
||||
result = acc.drain()
|
||||
assert result == []
|
||||
|
||||
def test_accumulate_multiple_tool_calls(self):
|
||||
acc = ToolCallAccumulator()
|
||||
acc.add_chunk(
|
||||
ChatChunk(
|
||||
tool_call_chunks=[
|
||||
{"index": 0, "id": "call_1", "name": "read_file", "args": '{"path": "'},
|
||||
{"index": 1, "id": "call_2", "name": "list_dir", "args": '{"path": "'},
|
||||
]
|
||||
)
|
||||
)
|
||||
acc.add_chunk(
|
||||
ChatChunk(
|
||||
tool_call_chunks=[
|
||||
{"index": 0, "args": 'x.txt"}'},
|
||||
{"index": 1, "args": '/tmp"}'},
|
||||
]
|
||||
)
|
||||
)
|
||||
result = acc.drain()
|
||||
assert len(result) == 2
|
||||
assert result[0]["name"] == "read_file"
|
||||
assert result[0]["args"] == '{"path": "x.txt"}'
|
||||
assert result[1]["name"] == "list_dir"
|
||||
assert result[1]["args"] == '{"path": "/tmp"}'
|
||||
|
||||
|
||||
class TestStreamHandler:
|
||||
def test_empty_result(self):
|
||||
handler = StreamHandler()
|
||||
result = handler.to_result(model_name="test-model")
|
||||
assert result.content == ""
|
||||
assert result.reasoning == ""
|
||||
assert result.tool_calls == []
|
||||
assert result.usage is None
|
||||
assert result.model_name == "test-model"
|
||||
|
||||
def test_stream_to_result(self):
|
||||
handler = StreamHandler()
|
||||
handler.add_chunk(ChatChunk(delta_content="Hello world"))
|
||||
handler.add_chunk(ChatChunk(delta_reasoning="Thinking..."))
|
||||
handler.add_chunk(
|
||||
ChatChunk(usage=TokenUsage(prompt_tokens=5, completion_tokens=3, total_tokens=8))
|
||||
)
|
||||
result = handler.to_result(model_name="test-model")
|
||||
assert result.content == "Hello world"
|
||||
assert result.reasoning == "Thinking..."
|
||||
assert result.usage is not None
|
||||
assert result.usage.total_tokens == 8
|
||||
assert result.model_name == "test-model"
|
||||
|
||||
def test_stream_with_tool_calls(self):
|
||||
handler = StreamHandler()
|
||||
handler.add_chunk(
|
||||
ChatChunk(
|
||||
tool_call_chunks=[
|
||||
{"index": 0, "id": "call_1", "name": "read_file", "args": '{"path": "'},
|
||||
]
|
||||
)
|
||||
)
|
||||
handler.add_chunk(
|
||||
ChatChunk(
|
||||
tool_call_chunks=[
|
||||
{"index": 0, "args": 'test.txt"}'},
|
||||
]
|
||||
)
|
||||
)
|
||||
result = handler.to_result(model_name="test-model")
|
||||
assert len(result.tool_calls) == 1
|
||||
assert result.tool_calls[0]["name"] == "read_file"
|
||||
assert result.tool_calls[0]["args"] == {"path": "test.txt"}
|
||||
|
||||
def test_chunks_property(self):
|
||||
handler = StreamHandler()
|
||||
handler.add_chunk(ChatChunk(delta_content="Hello"))
|
||||
handler.add_chunk(ChatChunk(delta_content=" world"))
|
||||
assert len(handler.chunks) == 2
|
||||
|
||||
def test_stream_content_and_tool_calls_together(self):
|
||||
handler = StreamHandler()
|
||||
handler.add_chunk(ChatChunk(delta_content="I will read the file."))
|
||||
handler.add_chunk(
|
||||
ChatChunk(
|
||||
tool_call_chunks=[
|
||||
{"index": 0, "id": "call_1", "name": "read_file", "args": "{}"},
|
||||
]
|
||||
)
|
||||
)
|
||||
result = handler.to_result()
|
||||
assert result.content == "I will read the file."
|
||||
assert len(result.tool_calls) == 1
|
||||
|
||||
def test_non_json_args_fallback(self):
|
||||
handler = StreamHandler()
|
||||
handler.add_chunk(
|
||||
ChatChunk(
|
||||
tool_call_chunks=[
|
||||
{"index": 0, "id": "call_1", "name": "read_file", "args": "not valid json"},
|
||||
]
|
||||
)
|
||||
)
|
||||
result = handler.to_result()
|
||||
assert result.tool_calls[0]["args"] == {"_raw": "not valid json"}
|
||||
|
||||
|
||||
class TestStreamFactory:
|
||||
def test_create_handler(self):
|
||||
handler = StreamFactory.create_handler()
|
||||
assert isinstance(handler, StreamHandler)
|
||||
|
||||
def test_create_accumulator(self):
|
||||
acc = StreamFactory.create_accumulator()
|
||||
assert isinstance(acc, ToolCallAccumulator)
|
||||
@@ -1,139 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from app.ai_provider.routing.multimodal import MultiModalDecision, MultiModalRouter
|
||||
|
||||
|
||||
class TestMultiModalDecision:
|
||||
def test_enum_values(self):
|
||||
assert MultiModalDecision.PLAIN_TEXT == "plain_text"
|
||||
assert MultiModalDecision.VISION == "vision"
|
||||
assert MultiModalDecision.THINKING == "thinking"
|
||||
assert MultiModalDecision.TOOL_USE == "tool_use"
|
||||
assert MultiModalDecision.STRUCTURED_OUTPUT == "structured_output"
|
||||
|
||||
def test_enum_is_str_enum(self):
|
||||
assert issubclass(MultiModalDecision, str)
|
||||
|
||||
|
||||
class TestMultiModalRouter:
|
||||
def test_plain_text_message(self):
|
||||
messages = [{"role": "user", "content": "Hello, world!"}]
|
||||
decisions = MultiModalRouter.decide_from_messages(messages)
|
||||
assert decisions == {MultiModalDecision.PLAIN_TEXT}
|
||||
|
||||
def test_plain_text_multiple_messages(self):
|
||||
messages = [
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": "What is Python?"},
|
||||
{"role": "assistant", "content": "Python is a programming language."},
|
||||
]
|
||||
decisions = MultiModalRouter.decide_from_messages(messages)
|
||||
assert decisions == {MultiModalDecision.PLAIN_TEXT}
|
||||
|
||||
def test_vision_with_openai_image_url_block(self):
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "What's in this image?"},
|
||||
{"type": "image_url", "image_url": {"url": "https://example.com/img.png"}},
|
||||
],
|
||||
}
|
||||
]
|
||||
decisions = MultiModalRouter.decide_from_messages(messages)
|
||||
assert MultiModalDecision.VISION in decisions
|
||||
|
||||
def test_vision_with_image_type_block(self):
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Describe this"},
|
||||
{
|
||||
"type": "image",
|
||||
"source": {"type": "base64", "media_type": "image/png", "data": "..."},
|
||||
},
|
||||
],
|
||||
}
|
||||
]
|
||||
decisions = MultiModalRouter.decide_from_messages(messages)
|
||||
assert MultiModalDecision.VISION in decisions
|
||||
|
||||
def test_plain_text_with_single_content_block(self):
|
||||
"""A single text-only content block should not trigger VISION."""
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": "Hello!"}],
|
||||
}
|
||||
]
|
||||
decisions = MultiModalRouter.decide_from_messages(messages)
|
||||
assert MultiModalDecision.VISION not in decisions
|
||||
|
||||
def test_tool_use_from_assistant_tool_calls(self):
|
||||
messages = [
|
||||
{"role": "user", "content": "What's the weather?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"tool_calls": [
|
||||
{"id": "call_1", "function": {"name": "get_weather", "arguments": "{}"}}
|
||||
],
|
||||
},
|
||||
]
|
||||
decisions = MultiModalRouter.decide_from_messages(messages)
|
||||
assert MultiModalDecision.TOOL_USE in decisions
|
||||
|
||||
def test_tool_use_from_tool_role(self):
|
||||
messages = [
|
||||
{"role": "user", "content": "What's the weather?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"tool_calls": [{"id": "call_1", "name": "get_weather"}],
|
||||
},
|
||||
{"role": "tool", "content": "25°C", "tool_call_id": "call_1"},
|
||||
]
|
||||
decisions = MultiModalRouter.decide_from_messages(messages)
|
||||
assert MultiModalDecision.TOOL_USE in decisions
|
||||
|
||||
def test_thinking_from_reasoning_field(self):
|
||||
messages = [
|
||||
{"role": "user", "content": "Think step by step"},
|
||||
{"role": "assistant", "content": "Final answer", "reasoning": "Let me think..."},
|
||||
]
|
||||
decisions = MultiModalRouter.decide_from_messages(messages)
|
||||
assert MultiModalDecision.THINKING in decisions
|
||||
|
||||
def test_empty_messages(self):
|
||||
decisions = MultiModalRouter.decide_from_messages([])
|
||||
assert decisions == {MultiModalDecision.PLAIN_TEXT}
|
||||
|
||||
def test_multiple_capabilities_combined(self):
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Analyze this image step by step"},
|
||||
{"type": "image_url", "image_url": {"url": "https://example.com/img.png"}},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "Let me think about this image...",
|
||||
"reasoning": "First, I see an image. Let me analyze it carefully.",
|
||||
"tool_calls": [
|
||||
{"id": "call_1", "function": {"name": "analyze_image", "arguments": "{}"}}
|
||||
],
|
||||
},
|
||||
{"role": "tool", "content": "Analysis complete", "tool_call_id": "call_1"},
|
||||
]
|
||||
decisions = MultiModalRouter.decide_from_messages(messages)
|
||||
assert MultiModalDecision.VISION in decisions
|
||||
assert MultiModalDecision.THINKING in decisions
|
||||
assert MultiModalDecision.TOOL_USE in decisions
|
||||
|
||||
def test_decide_with_empty_content(self):
|
||||
messages = [{"role": "user", "content": ""}]
|
||||
decisions = MultiModalRouter.decide_from_messages(messages)
|
||||
assert decisions == {MultiModalDecision.PLAIN_TEXT}
|
||||
@@ -1,141 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from app.ai_provider.routing.permissions import CapabilityPermissionController, PermissionRule
|
||||
|
||||
|
||||
class TestPermissionRule:
|
||||
def test_exact_match(self):
|
||||
rule = PermissionRule(model_pattern="fast-model", capability="vision", allowed=True)
|
||||
assert rule.match("fast-model") is True
|
||||
assert rule.match("pro-model") is False
|
||||
|
||||
def test_wildcard_match(self):
|
||||
rule = PermissionRule(model_pattern="deepseek-*", capability="thinking", allowed=False)
|
||||
assert rule.match("deepseek-v4-flash") is True
|
||||
assert rule.match("deepseek-v4-pro") is True
|
||||
assert rule.match("claude-sonnet-4-6") is False
|
||||
|
||||
def test_universal_match(self):
|
||||
rule = PermissionRule(model_pattern="*", capability="vision", allowed=True)
|
||||
assert rule.match("any-model") is True
|
||||
assert rule.match("another-one") is True
|
||||
|
||||
def test_suffix_wildcard(self):
|
||||
rule = PermissionRule(model_pattern="*-pro", capability="thinking", allowed=True)
|
||||
assert rule.match("deepseek-v4-pro") is True
|
||||
assert rule.match("fast-model") is False
|
||||
|
||||
def test_match_method_called_explicitly(self):
|
||||
rule = PermissionRule(model_pattern="vision-*", capability="vision")
|
||||
assert rule.match("vision-model") is True
|
||||
|
||||
def test_allowed_defaults_to_true(self):
|
||||
rule = PermissionRule(model_pattern="*", capability="vision")
|
||||
assert rule.allowed is True
|
||||
|
||||
def test_allowed_set_to_false(self):
|
||||
rule = PermissionRule(model_pattern="*", capability="vision", allowed=False)
|
||||
assert rule.allowed is False
|
||||
|
||||
def test_repr(self):
|
||||
rule = PermissionRule(model_pattern="*", capability="vision", allowed=True)
|
||||
r = repr(rule)
|
||||
assert "PermissionRule" in r
|
||||
assert "*" in r
|
||||
assert "vision" in r
|
||||
|
||||
def test_eq_same_rules(self):
|
||||
rule1 = PermissionRule(model_pattern="foo-*", capability="vision", allowed=True)
|
||||
rule2 = PermissionRule(model_pattern="foo-*", capability="vision", allowed=True)
|
||||
assert rule1 == rule2
|
||||
|
||||
def test_eq_different_rules(self):
|
||||
rule1 = PermissionRule(model_pattern="foo-*", capability="vision", allowed=True)
|
||||
rule2 = PermissionRule(model_pattern="bar-*", capability="vision", allowed=True)
|
||||
assert rule1 != rule2
|
||||
|
||||
|
||||
class TestCapabilityPermissionController:
|
||||
def test_allow_all_by_default(self):
|
||||
controller = CapabilityPermissionController([])
|
||||
assert controller.check("fast-model", "vision") is True
|
||||
|
||||
def test_deny_specific_capability_on_model(self):
|
||||
controller = CapabilityPermissionController(
|
||||
[
|
||||
PermissionRule(model_pattern="fast-*", capability="vision", allowed=False),
|
||||
]
|
||||
)
|
||||
assert controller.check("fast-model", "vision") is False
|
||||
assert controller.check("pro-model", "vision") is True
|
||||
|
||||
def test_deny_then_allow_more_specific(self):
|
||||
"""A more specific deny should override a broader allow."""
|
||||
controller = CapabilityPermissionController(
|
||||
[
|
||||
PermissionRule(model_pattern="*", capability="thinking", allowed=True),
|
||||
PermissionRule(model_pattern="fast-*", capability="thinking", allowed=False),
|
||||
]
|
||||
)
|
||||
assert controller.check("fast-model", "thinking") is False
|
||||
assert controller.check("pro-model", "thinking") is True
|
||||
|
||||
def test_filter_capabilities(self):
|
||||
controller = CapabilityPermissionController(
|
||||
[
|
||||
PermissionRule(model_pattern="*", capability="vision", allowed=False),
|
||||
PermissionRule(model_pattern="vision-*", capability="vision", allowed=True),
|
||||
]
|
||||
)
|
||||
filtered = controller.filter_capabilities("fast-model", {"chat", "vision", "thinking"})
|
||||
assert filtered == {"chat", "thinking"}
|
||||
|
||||
filtered_vision = controller.filter_capabilities(
|
||||
"vision-model", {"chat", "vision", "thinking"}
|
||||
)
|
||||
assert filtered_vision == {"chat", "vision", "thinking"}
|
||||
|
||||
def test_filter_with_no_rules(self):
|
||||
controller = CapabilityPermissionController([])
|
||||
filtered = controller.filter_capabilities("any-model", {"chat", "vision"})
|
||||
assert filtered == {"chat", "vision"}
|
||||
|
||||
def test_check_nonexistent_model(self):
|
||||
"""Non-existent models should still be checked against rules."""
|
||||
controller = CapabilityPermissionController(
|
||||
[
|
||||
PermissionRule(model_pattern="*", capability="vision", allowed=False),
|
||||
]
|
||||
)
|
||||
assert controller.check("unknown-model", "vision") is False
|
||||
|
||||
def test_multiple_rules_matching(self):
|
||||
"""When multiple rules match, the last one wins (most specific wins by order)."""
|
||||
controller = CapabilityPermissionController(
|
||||
[
|
||||
PermissionRule(model_pattern="*", capability="tool_use", allowed=True),
|
||||
PermissionRule(model_pattern="basic-*", capability="tool_use", allowed=False),
|
||||
PermissionRule(model_pattern="basic-pro", capability="tool_use", allowed=True),
|
||||
]
|
||||
)
|
||||
assert controller.check("basic-lite", "tool_use") is False
|
||||
assert controller.check("basic-pro", "tool_use") is True
|
||||
assert controller.check("pro-model", "tool_use") is True
|
||||
|
||||
def test_check_capability_not_in_rules(self):
|
||||
controller = CapabilityPermissionController(
|
||||
[
|
||||
PermissionRule(model_pattern="*", capability="vision", allowed=False),
|
||||
]
|
||||
)
|
||||
# 'chat' is not in any rule, so it's allowed by default
|
||||
assert controller.check("fast-model", "chat") is True
|
||||
|
||||
def test_is_allowed_model(self):
|
||||
controller = CapabilityPermissionController(
|
||||
[
|
||||
PermissionRule(model_pattern="*", capability="vision", allowed=False),
|
||||
]
|
||||
)
|
||||
assert controller.is_allowed("fast-model", "vision") is False
|
||||
assert controller.is_allowed("fast-model", "chat") is True
|
||||
@@ -1,135 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from app.ai_provider.config.models import ModelConfig
|
||||
from app.ai_provider.core.registry import ModelRegistry
|
||||
from app.ai_provider.routing.multimodal import MultiModalDecision
|
||||
from app.ai_provider.routing.selector import ModelSelector
|
||||
from app.ai_provider.types.errors import AIProviderError
|
||||
|
||||
|
||||
class TestModelSelector:
|
||||
@pytest.fixture
|
||||
def registry(self):
|
||||
models = [
|
||||
ModelConfig(
|
||||
name="fast-model",
|
||||
provider="deepseek",
|
||||
model="deepseek-v4-flash",
|
||||
format="openai",
|
||||
tags=["default"],
|
||||
weight=100,
|
||||
),
|
||||
ModelConfig(
|
||||
name="pro-model",
|
||||
provider="deepseek",
|
||||
model="deepseek-v4-pro",
|
||||
format="openai",
|
||||
tags=["fallback"],
|
||||
weight=50,
|
||||
),
|
||||
ModelConfig(
|
||||
name="vision-model",
|
||||
provider="anthropic",
|
||||
model="claude-sonnet-4-6",
|
||||
format="anthropic",
|
||||
supports_vision=True,
|
||||
supports_thinking=True,
|
||||
supports_tools=True,
|
||||
),
|
||||
ModelConfig(
|
||||
name="thinking-model",
|
||||
provider="deepseek",
|
||||
model="deepseek-v4-thinking",
|
||||
format="openai",
|
||||
supports_reasoning_effort=True,
|
||||
supports_thinking=True,
|
||||
),
|
||||
]
|
||||
return ModelRegistry(models)
|
||||
|
||||
def test_select_default_model(self, registry):
|
||||
selector = ModelSelector(registry)
|
||||
model = selector.select({})
|
||||
assert model.name == "fast-model"
|
||||
|
||||
def test_select_with_preferred_model(self, registry):
|
||||
selector = ModelSelector(registry)
|
||||
model = selector.select({"preferred": "pro-model"})
|
||||
assert model.name == "pro-model"
|
||||
|
||||
def test_select_preferred_with_capability(self, registry):
|
||||
selector = ModelSelector(registry)
|
||||
model = selector.select({"preferred": "vision-model", "capabilities": {"vision"}})
|
||||
assert model.name == "vision-model"
|
||||
|
||||
def test_select_preferred_missing_required_capability(self, registry):
|
||||
"""Preferred model doesn't have required capability — falls back."""
|
||||
selector = ModelSelector(registry)
|
||||
model = selector.select({"preferred": "fast-model", "capabilities": {"vision"}})
|
||||
assert model.name != "fast-model"
|
||||
assert "vision" in model.resolved_capabilities()
|
||||
assert model.resolved_capabilities()["vision"] is True
|
||||
|
||||
def test_select_for_vision(self, registry):
|
||||
selector = ModelSelector(registry)
|
||||
model = selector.select({"capabilities": {"vision"}})
|
||||
assert model.name == "vision-model"
|
||||
|
||||
def test_select_for_thinking(self, registry):
|
||||
selector = ModelSelector(registry)
|
||||
model = selector.select({"capabilities": {"thinking"}})
|
||||
assert model.resolved_capabilities()["thinking"] is True
|
||||
|
||||
def test_select_for_tool_use(self, registry):
|
||||
selector = ModelSelector(registry)
|
||||
model = selector.select({"capabilities": {"tool_use"}})
|
||||
assert model.resolved_capabilities()["tool_use"] is True
|
||||
|
||||
def test_select_for_multiple_capabilities(self, registry):
|
||||
selector = ModelSelector(registry)
|
||||
model = selector.select({"capabilities": {"vision", "thinking"}})
|
||||
assert model.resolved_capabilities()["vision"] is True
|
||||
assert model.resolved_capabilities()["thinking"] is True
|
||||
|
||||
def test_select_from_multi_modal_decision(self, registry):
|
||||
selector = ModelSelector(registry)
|
||||
model = selector.select_from_decisions(
|
||||
{MultiModalDecision.VISION, MultiModalDecision.THINKING}
|
||||
)
|
||||
assert model.resolved_capabilities()["vision"] is True
|
||||
assert model.resolved_capabilities()["thinking"] is True
|
||||
|
||||
def test_select_from_plain_text_decision(self, registry):
|
||||
selector = ModelSelector(registry)
|
||||
model = selector.select_from_decisions({MultiModalDecision.PLAIN_TEXT})
|
||||
assert model.name == "fast-model"
|
||||
|
||||
def test_select_with_invalid_preferred(self, registry):
|
||||
selector = ModelSelector(registry)
|
||||
with pytest.raises(AIProviderError):
|
||||
selector.select({"preferred": "nonexistent-model"})
|
||||
|
||||
def test_select_with_capability_no_match(self, registry):
|
||||
"""No model supports an unknown capability."""
|
||||
selector = ModelSelector(registry)
|
||||
with pytest.raises(AIProviderError):
|
||||
selector.select({"capabilities": {"telepathy"}})
|
||||
|
||||
def test_select_allows_preferred_with_superset_capabilities(self, registry):
|
||||
"""Preferred model has MORE capabilities than required — should still match."""
|
||||
selector = ModelSelector(registry)
|
||||
model = selector.select({"preferred": "vision-model", "capabilities": {"tool_use"}})
|
||||
assert model.name == "vision-model"
|
||||
|
||||
def test_select_from_decisions_with_empty_set(self, registry):
|
||||
selector = ModelSelector(registry)
|
||||
model = selector.select_from_decisions(set())
|
||||
assert model.name == "fast-model"
|
||||
|
||||
def test_select_for_avoid_list(self, registry):
|
||||
"""Avoid certain models — prefer a non-avoided model."""
|
||||
selector = ModelSelector(registry)
|
||||
model = selector.select({"capabilities": {"chat"}, "avoid": {"fast-model"}})
|
||||
assert model.name != "fast-model"
|
||||
@@ -1,61 +0,0 @@
|
||||
import pytest
|
||||
|
||||
from app.ai_provider.config.models import ModelConfig, ProviderConfig
|
||||
from app.ai_provider.core.registry import ModelRegistry, ProviderRegistry
|
||||
from app.ai_provider.manager import ProviderManager
|
||||
from app.ai_provider.types.messages import HumanMessage, SystemMessage
|
||||
|
||||
|
||||
class TestProviderManager:
|
||||
@pytest.fixture
|
||||
def manager(self):
|
||||
models = [
|
||||
ModelConfig(
|
||||
name="test",
|
||||
provider="deepseek",
|
||||
model="deepseek-v4-flash",
|
||||
format="openai",
|
||||
tags=["default"],
|
||||
),
|
||||
]
|
||||
providers = {
|
||||
"deepseek": ProviderConfig(
|
||||
name="deepseek",
|
||||
format="openai",
|
||||
base_url_env="DEEPSEEK_BASE_URL",
|
||||
api_key_env="DEEPSEEK_API_KEY",
|
||||
),
|
||||
}
|
||||
model_reg = ModelRegistry(models)
|
||||
provider_reg = ProviderRegistry(providers)
|
||||
return ProviderManager(
|
||||
model_registry=model_reg,
|
||||
provider_registry=provider_reg,
|
||||
)
|
||||
|
||||
def test_resolve_model(self, manager):
|
||||
name = manager.resolve_model("test")
|
||||
assert name == "test"
|
||||
|
||||
def test_resolve_model_default(self, manager):
|
||||
name = manager.resolve_model()
|
||||
assert name == "test"
|
||||
|
||||
def test_get_adapter_for_model(self, manager):
|
||||
adapter = manager.get_adapter("test")
|
||||
from app.ai_provider.llm.openai_format import OpenAIFormatAdapter
|
||||
|
||||
assert isinstance(adapter, OpenAIFormatAdapter)
|
||||
|
||||
def test_get_provider_config(self, manager):
|
||||
prov = manager.get_provider_config("deepseek")
|
||||
assert prov.name == "deepseek"
|
||||
|
||||
def test_to_message_list(self, manager):
|
||||
msgs = manager.to_message_list(
|
||||
system_messages=["You are helpful."],
|
||||
history=[("user", "Hello")],
|
||||
)
|
||||
assert len(msgs) == 2
|
||||
assert isinstance(msgs[0], SystemMessage)
|
||||
assert isinstance(msgs[1], HumanMessage)
|
||||
@@ -1,91 +0,0 @@
|
||||
"""Tests for the parameter pipeline (Task 4)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from app.ai_provider.config.models import ModelConfig, ParameterMapping, ProviderConfig
|
||||
from app.ai_provider.pipeline import PipelineContext, PipelineFactory
|
||||
|
||||
|
||||
class TestParameterPipeline:
|
||||
"""Test the config-driven 4-stage parameter pipeline."""
|
||||
|
||||
@pytest.fixture
|
||||
def pipeline(self):
|
||||
provider_cfg = ProviderConfig(
|
||||
name="deepseek",
|
||||
format="openai",
|
||||
base_url_env="DEEPSEEK_BASE_URL",
|
||||
api_key_env="DEEPSEEK_API_KEY",
|
||||
parameter_mappings={
|
||||
"temperature": ParameterMapping(provider_param="temperature", min=0, max=2),
|
||||
"max_tokens": ParameterMapping(provider_param="max_tokens", min=1),
|
||||
"top_p": ParameterMapping(provider_param="top_p", min=0, max=1),
|
||||
},
|
||||
)
|
||||
return PipelineFactory.create(provider_cfg)
|
||||
|
||||
def test_normalize_temperature(self, pipeline):
|
||||
ctx = PipelineContext(params={"temperature": 0.7})
|
||||
result = pipeline.run(ctx)
|
||||
assert result.params["temperature"] == 0.7
|
||||
|
||||
def test_clamp_temperature(self, pipeline):
|
||||
ctx = PipelineContext(params={"temperature": 5.0})
|
||||
result = pipeline.run(ctx)
|
||||
assert result.params["temperature"] == 2.0 # clamped to max
|
||||
|
||||
def test_filter_unknown_param(self, pipeline):
|
||||
ctx = PipelineContext(params={"temperature": 0.5, "unknown_param": "value"})
|
||||
result = pipeline.run(ctx)
|
||||
assert "unknown_param" not in result.params
|
||||
assert "temperature" in result.params
|
||||
|
||||
def test_inject_provider_defaults(self):
|
||||
provider_cfg = ProviderConfig(
|
||||
name="anthropic",
|
||||
format="anthropic",
|
||||
base_url_env="ANTHROPIC_BASE_URL",
|
||||
api_key_env="ANTHROPIC_API_KEY",
|
||||
parameter_mappings={
|
||||
"max_tokens": ParameterMapping(provider_param="max_tokens", default=4096),
|
||||
},
|
||||
)
|
||||
pipeline = PipelineFactory.create(provider_cfg)
|
||||
ctx = PipelineContext(params={})
|
||||
result = pipeline.run(ctx)
|
||||
assert result.params["max_tokens"] == 4096
|
||||
|
||||
def test_transform_param(self):
|
||||
provider_cfg = ProviderConfig(
|
||||
name="openai",
|
||||
format="openai",
|
||||
base_url_env="OPENAI_BASE_URL",
|
||||
api_key_env="OPENAI_API_KEY",
|
||||
parameter_mappings={
|
||||
"max_tokens": ParameterMapping(provider_param="max_tokens", transform="int"),
|
||||
},
|
||||
)
|
||||
pipeline = PipelineFactory.create(provider_cfg)
|
||||
ctx = PipelineContext(params={"max_tokens": 4096.0})
|
||||
result = pipeline.run(ctx)
|
||||
assert isinstance(result.params["max_tokens"], int)
|
||||
|
||||
def test_capability_filter_vision(self):
|
||||
model = ModelConfig(
|
||||
name="test",
|
||||
provider="openai",
|
||||
model="gpt-4o",
|
||||
format="openai",
|
||||
supports_vision=False,
|
||||
)
|
||||
provider_cfg = ProviderConfig(
|
||||
name="openai", format="openai", base_url_env="X", api_key_env="Y"
|
||||
)
|
||||
pipeline = PipelineFactory.create(provider_cfg)
|
||||
ctx = PipelineContext(
|
||||
params={"image_url": "http://example.com/img.png"}, model_config=model
|
||||
)
|
||||
result = pipeline.run(ctx)
|
||||
assert "image_url" in result.filtered_params
|
||||
@@ -1,305 +0,0 @@
|
||||
"""Tests for the parameter handler mechanism (Task 4 extension)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from app.ai_provider.config.models import ModelConfig, ParameterMapping, ProviderConfig
|
||||
from app.ai_provider.pipeline import PipelineContext, PipelineFactory
|
||||
from app.ai_provider.pipeline_handlers import ParameterHandlerRegistry
|
||||
|
||||
|
||||
class TestParameterHandlerRegistry:
|
||||
"""Test the handler registry itself."""
|
||||
|
||||
def test_register_and_get(self):
|
||||
@ParameterHandlerRegistry.register("test_handler_1")
|
||||
def my_handler(_ctx, _result, _param_key=None):
|
||||
pass
|
||||
|
||||
handler = ParameterHandlerRegistry.get("test_handler_1")
|
||||
assert handler is my_handler
|
||||
|
||||
def test_get_unknown_handler_raises(self):
|
||||
with pytest.raises(KeyError, match="Unknown parameter handler"):
|
||||
ParameterHandlerRegistry.get("nonexistent_handler")
|
||||
|
||||
def test_list_handlers(self):
|
||||
handlers = ParameterHandlerRegistry.list()
|
||||
assert "anthropic_thinking" in handlers
|
||||
assert "validate_max_tokens" in handlers
|
||||
assert "strip_unsupported_capability_params" in handlers
|
||||
|
||||
|
||||
class TestAnthropicThinkingHandler:
|
||||
"""Test the anthropic_thinking handler via the pipeline."""
|
||||
|
||||
def test_thinking_true_converts_to_dict(self):
|
||||
"""标量 thinking=True → {type: enabled}"""
|
||||
provider = ProviderConfig(
|
||||
name="anthropic",
|
||||
format="anthropic",
|
||||
base_url_env="ANTHROPIC_BASE_URL",
|
||||
api_key_env="ANTHROPIC_API_KEY",
|
||||
parameter_mappings={
|
||||
"thinking": ParameterMapping(
|
||||
provider_param="thinking",
|
||||
handler="anthropic_thinking",
|
||||
),
|
||||
"max_tokens": ParameterMapping(provider_param="max_tokens"),
|
||||
},
|
||||
)
|
||||
pipeline = PipelineFactory.create(provider)
|
||||
ctx = PipelineContext(params={"thinking": True, "max_tokens": 4096})
|
||||
result = pipeline.run(ctx)
|
||||
assert result.params["thinking"] == {"type": "enabled"}
|
||||
assert result.params["max_tokens"] == 4096
|
||||
|
||||
def test_thinking_budget_enforces_max_tokens(self):
|
||||
"""thinking.budget_tokens 确保 max_tokens > budget_tokens"""
|
||||
provider = ProviderConfig(
|
||||
name="anthropic",
|
||||
format="anthropic",
|
||||
base_url_env="X",
|
||||
api_key_env="Y",
|
||||
parameter_mappings={
|
||||
"thinking": ParameterMapping(
|
||||
provider_param="thinking",
|
||||
handler="anthropic_thinking",
|
||||
),
|
||||
"max_tokens": ParameterMapping(provider_param="max_tokens"),
|
||||
},
|
||||
)
|
||||
pipeline = PipelineFactory.create(provider)
|
||||
# budget=8192, max_tokens=4096 → 不满足，应被抬到 8193
|
||||
ctx = PipelineContext(params={"thinking": {"budget_tokens": 8192}, "max_tokens": 4096})
|
||||
result = pipeline.run(ctx)
|
||||
assert result.params["max_tokens"] >= 8193
|
||||
|
||||
def test_thinking_removes_temperature(self):
|
||||
"""开启 thinking 时 temperature 被移除"""
|
||||
provider = ProviderConfig(
|
||||
name="anthropic",
|
||||
format="anthropic",
|
||||
base_url_env="X",
|
||||
api_key_env="Y",
|
||||
parameter_mappings={
|
||||
"thinking": ParameterMapping(
|
||||
provider_param="thinking",
|
||||
handler="anthropic_thinking",
|
||||
),
|
||||
"temperature": ParameterMapping(
|
||||
provider_param="temperature",
|
||||
min=0,
|
||||
max=1,
|
||||
),
|
||||
"max_tokens": ParameterMapping(provider_param="max_tokens"),
|
||||
},
|
||||
)
|
||||
pipeline = PipelineFactory.create(provider)
|
||||
ctx = PipelineContext(params={"thinking": True, "temperature": 0.7, "max_tokens": 4096})
|
||||
result = pipeline.run(ctx)
|
||||
assert "temperature" not in result.params
|
||||
|
||||
|
||||
class TestValidateMaxTokensHandler:
|
||||
"""Test the validate_max_tokens handler."""
|
||||
|
||||
def test_clamps_to_model_limit(self):
|
||||
"""max_tokens 应被裁剪到模型上限"""
|
||||
model = ModelConfig(
|
||||
name="test-model",
|
||||
provider="anthropic",
|
||||
model="claude-sonnet-4-6",
|
||||
format="anthropic",
|
||||
max_output_tokens=16000,
|
||||
)
|
||||
provider = ProviderConfig(
|
||||
name="anthropic",
|
||||
format="anthropic",
|
||||
base_url_env="X",
|
||||
api_key_env="Y",
|
||||
parameter_mappings={
|
||||
"max_tokens": ParameterMapping(provider_param="max_tokens"),
|
||||
},
|
||||
parameter_handlers=["validate_max_tokens"],
|
||||
)
|
||||
pipeline = PipelineFactory.create(provider)
|
||||
ctx = PipelineContext(
|
||||
params={"max_tokens": 999999},
|
||||
model_config=model,
|
||||
provider_config=provider,
|
||||
)
|
||||
result = pipeline.run(ctx)
|
||||
assert result.params["max_tokens"] == 16000
|
||||
|
||||
def test_within_limit_unchanged(self):
|
||||
"""max_tokens 在合理范围内时不变"""
|
||||
model = ModelConfig(
|
||||
name="test-model",
|
||||
provider="anthropic",
|
||||
model="claude-sonnet-4-6",
|
||||
format="anthropic",
|
||||
max_output_tokens=16000,
|
||||
)
|
||||
provider = ProviderConfig(
|
||||
name="anthropic",
|
||||
format="anthropic",
|
||||
base_url_env="X",
|
||||
api_key_env="Y",
|
||||
parameter_mappings={
|
||||
"max_tokens": ParameterMapping(provider_param="max_tokens"),
|
||||
},
|
||||
parameter_handlers=["validate_max_tokens"],
|
||||
)
|
||||
pipeline = PipelineFactory.create(provider)
|
||||
ctx = PipelineContext(
|
||||
params={"max_tokens": 8000},
|
||||
model_config=model,
|
||||
provider_config=provider,
|
||||
)
|
||||
result = pipeline.run(ctx)
|
||||
assert result.params["max_tokens"] == 8000
|
||||
|
||||
|
||||
class TestStripUnsupportedCapabilityParams:
|
||||
"""Test the strip_unsupported_capability_params handler."""
|
||||
|
||||
def test_strips_image_params_when_no_vision(self):
|
||||
"""模型不支持 vision 时移除 image_url"""
|
||||
model = ModelConfig(
|
||||
name="test-model",
|
||||
provider="openai",
|
||||
model="gpt-4o",
|
||||
format="openai",
|
||||
supports_vision=False,
|
||||
)
|
||||
provider = ProviderConfig(
|
||||
name="openai",
|
||||
format="openai",
|
||||
base_url_env="X",
|
||||
api_key_env="Y",
|
||||
parameter_mappings={
|
||||
"max_tokens": ParameterMapping(provider_param="max_tokens"),
|
||||
"image_url": ParameterMapping(provider_param="image_url"),
|
||||
},
|
||||
parameter_handlers=["strip_unsupported_capability_params"],
|
||||
)
|
||||
pipeline = PipelineFactory.create(provider)
|
||||
ctx = PipelineContext(
|
||||
params={"max_tokens": 100, "image_url": "http://example.com/img.png"},
|
||||
model_config=model,
|
||||
provider_config=provider,
|
||||
)
|
||||
result = pipeline.run(ctx)
|
||||
assert "image_url" not in result.params
|
||||
assert "max_tokens" in result.params
|
||||
|
||||
def test_keeps_image_params_when_vision_supported(self):
|
||||
"""模型支持 vision 时保留 image_url"""
|
||||
model = ModelConfig(
|
||||
name="test-model",
|
||||
provider="openai",
|
||||
model="gpt-4o",
|
||||
format="openai",
|
||||
supports_vision=True,
|
||||
)
|
||||
provider = ProviderConfig(
|
||||
name="openai",
|
||||
format="openai",
|
||||
base_url_env="X",
|
||||
api_key_env="Y",
|
||||
parameter_mappings={
|
||||
"max_tokens": ParameterMapping(provider_param="max_tokens"),
|
||||
"image_url": ParameterMapping(provider_param="image_url"),
|
||||
},
|
||||
parameter_handlers=["strip_unsupported_capability_params"],
|
||||
)
|
||||
pipeline = PipelineFactory.create(provider)
|
||||
ctx = PipelineContext(
|
||||
params={"max_tokens": 100, "image_url": "http://example.com/img.png"},
|
||||
model_config=model,
|
||||
provider_config=provider,
|
||||
)
|
||||
result = pipeline.run(ctx)
|
||||
assert "image_url" in result.params
|
||||
|
||||
def test_strips_thinking_when_not_supported(self):
|
||||
"""模型不支持 thinking 时移除 thinking 参数"""
|
||||
model = ModelConfig(
|
||||
name="test-model",
|
||||
provider="openai",
|
||||
model="gpt-4o",
|
||||
format="openai",
|
||||
supports_thinking=False,
|
||||
)
|
||||
provider = ProviderConfig(
|
||||
name="openai",
|
||||
format="openai",
|
||||
base_url_env="X",
|
||||
api_key_env="Y",
|
||||
parameter_mappings={
|
||||
"thinking": ParameterMapping(provider_param="thinking"),
|
||||
},
|
||||
parameter_handlers=["strip_unsupported_capability_params"],
|
||||
)
|
||||
pipeline = PipelineFactory.create(provider)
|
||||
ctx = PipelineContext(
|
||||
params={"thinking": True},
|
||||
model_config=model,
|
||||
provider_config=provider,
|
||||
)
|
||||
result = pipeline.run(ctx)
|
||||
assert "thinking" not in result.params
|
||||
|
||||
|
||||
class TestHandlerIntegration:
|
||||
"""Test full pipeline integration with handlers."""
|
||||
|
||||
def test_anthropic_pipeline_end_to_end(self):
|
||||
"""模拟完整的 Anthropic pipeline 流程"""
|
||||
provider = ProviderConfig(
|
||||
name="anthropic",
|
||||
format="anthropic",
|
||||
base_url_env="ANTHROPIC_BASE_URL",
|
||||
api_key_env="ANTHROPIC_API_KEY",
|
||||
parameter_mappings={
|
||||
"max_tokens": ParameterMapping(provider_param="max_tokens", default=4096),
|
||||
"thinking": ParameterMapping(
|
||||
provider_param="thinking", handler="anthropic_thinking"
|
||||
),
|
||||
"temperature": ParameterMapping(
|
||||
provider_param="temperature", min=0, max=1, default=1.0
|
||||
),
|
||||
},
|
||||
parameter_handlers=[
|
||||
"strip_unsupported_capability_params",
|
||||
"validate_max_tokens",
|
||||
],
|
||||
)
|
||||
model = ModelConfig(
|
||||
name="claude-sonnet",
|
||||
provider="anthropic",
|
||||
model="claude-sonnet-4-6",
|
||||
format="anthropic",
|
||||
supports_vision=True,
|
||||
supports_thinking=True,
|
||||
max_output_tokens=16000,
|
||||
)
|
||||
pipeline = PipelineFactory.create(provider)
|
||||
ctx = PipelineContext(
|
||||
params={
|
||||
"thinking": {"budget_tokens": 8000},
|
||||
"temperature": 0.7,
|
||||
},
|
||||
model_config=model,
|
||||
provider_config=provider,
|
||||
)
|
||||
result = pipeline.run(ctx)
|
||||
|
||||
# thinking → provider_param name
|
||||
assert "thinking" in result.params
|
||||
# max_tokens injected via default, then lifted by handler
|
||||
assert result.params["max_tokens"] >= 8001
|
||||
# temperature removed by handler
|
||||
assert "temperature" not in result.params
|
||||
@@ -1,46 +0,0 @@
|
||||
# backend/tests/ai_provider/types/test_messages.py
|
||||
from app.ai_provider.types.messages import (
|
||||
AIMessage,
|
||||
HumanMessage,
|
||||
SystemMessage,
|
||||
ToolMessage,
|
||||
)
|
||||
|
||||
|
||||
class TestMessageCreation:
|
||||
def test_system_message(self):
|
||||
msg = SystemMessage(content="You are a helpful assistant.")
|
||||
assert msg.content == "You are a helpful assistant."
|
||||
assert msg.role == "system"
|
||||
|
||||
def test_human_message(self):
|
||||
msg = HumanMessage(content="Hello!")
|
||||
assert msg.content == "Hello!"
|
||||
assert msg.role == "user"
|
||||
|
||||
def test_ai_message_with_tool_calls(self):
|
||||
msg = AIMessage(
|
||||
content="I'll check that.",
|
||||
tool_calls=[{"name": "read_file", "args": {"path": "test.txt"}, "id": "call_1"}],
|
||||
)
|
||||
assert msg.tool_calls is not None
|
||||
assert len(msg.tool_calls) == 1
|
||||
assert msg.tool_calls[0]["name"] == "read_file"
|
||||
|
||||
def test_tool_message(self):
|
||||
msg = ToolMessage(content="File contents here", tool_call_id="call_1")
|
||||
assert msg.content == "File contents here"
|
||||
assert msg.tool_call_id == "call_1"
|
||||
|
||||
def test_ai_message_with_reasoning(self):
|
||||
msg = AIMessage(content="Answer", reasoning="I think therefore...")
|
||||
assert msg.reasoning == "I think therefore..."
|
||||
|
||||
def test_message_to_dict_roundtrip(self):
|
||||
original = HumanMessage(content="Hello")
|
||||
d = original.to_dict()
|
||||
assert d["role"] == "user"
|
||||
assert d["content"] == "Hello"
|
||||
restored = HumanMessage.from_dict(d)
|
||||
assert restored.content == original.content
|
||||
assert restored.role == original.role
|
||||
@@ -1,21 +0,0 @@
|
||||
# backend/tests/ai_provider/types/test_errors.py
|
||||
from app.ai_provider.types.errors import AIProviderError
|
||||
|
||||
|
||||
class TestErrors:
|
||||
def test_error_with_code(self):
|
||||
err = AIProviderError(message="Model not found", code="MODEL_NOT_FOUND", status_code=404)
|
||||
assert err.code == "MODEL_NOT_FOUND"
|
||||
assert err.status_code == 404
|
||||
assert str(err) == "Model not found"
|
||||
|
||||
def test_error_from_provider(self):
|
||||
cause = ValueError("API returned 429")
|
||||
err = AIProviderError.from_provider(
|
||||
message="Rate limited",
|
||||
code="RATE_LIMITED",
|
||||
provider="deepseek",
|
||||
original=cause,
|
||||
)
|
||||
assert err.provider == "deepseek"
|
||||
assert err.original is cause
|
||||
@@ -1,78 +0,0 @@
|
||||
"""模型网关 endpoint 覆盖 API 测试（api-spec §4.4，DB 优先于 .env/yaml）。"""
|
||||
|
||||
from collections.abc import AsyncGenerator
|
||||
|
||||
import pytest
|
||||
from httpx import ASGITransport, AsyncClient
|
||||
|
||||
from app.main import app
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def client() -> AsyncGenerator[AsyncClient, None]:
|
||||
transport = ASGITransport(app=app)
|
||||
async with AsyncClient(transport=transport, base_url="http://test") as http_client:
|
||||
yield http_client
|
||||
|
||||
|
||||
def _first_endpoint(gateway: dict) -> tuple[str, str]:
|
||||
model = gateway["models"][0]
|
||||
return model["name"], model["endpoints"][0]["name"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_models_has_endpoint_override_fields(client: AsyncClient) -> None:
|
||||
resp = await client.get("/api/v1/settings/models")
|
||||
assert resp.status_code == 200
|
||||
ep = resp.json()["models"][0]["endpoints"][0]
|
||||
assert "priority" in ep
|
||||
assert "is_overridden" in ep
|
||||
assert "has_url_override" in ep
|
||||
assert "has_key_override" in ep
|
||||
assert ep["is_overridden"] is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_patch_endpoint_status_override(client: AsyncClient) -> None:
|
||||
gateway = (await client.get("/api/v1/settings/models")).json()
|
||||
model, endpoint = _first_endpoint(gateway)
|
||||
|
||||
resp = await client.patch(
|
||||
f"/api/v1/settings/models/{model}/endpoints/{endpoint}",
|
||||
json={"status": "inactive"},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
assert body["is_overridden"] is True
|
||||
assert body["status"] == "inactive"
|
||||
|
||||
# GET 反映覆盖
|
||||
gateway2 = (await client.get("/api/v1/settings/models")).json()
|
||||
eps = [m for m in gateway2["models"] if m["name"] == model][0]["endpoints"]
|
||||
ep = [e for e in eps if e["name"] == endpoint][0]
|
||||
assert ep["is_overridden"] is True
|
||||
assert ep["status"] == "inactive"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_patch_endpoint_key_override_redacted(client: AsyncClient) -> None:
|
||||
gateway = (await client.get("/api/v1/settings/models")).json()
|
||||
model, endpoint = _first_endpoint(gateway)
|
||||
|
||||
resp = await client.patch(
|
||||
f"/api/v1/settings/models/{model}/endpoints/{endpoint}",
|
||||
json={"api_key": "sk-test-secret-value"},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
assert resp.json()["has_key_override"] is True
|
||||
# 脱敏：响应体不含 api_key 原值
|
||||
assert "sk-test-secret-value" not in resp.text
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_patch_endpoint_not_found(client: AsyncClient) -> None:
|
||||
resp = await client.patch(
|
||||
"/api/v1/settings/models/no_such_model/endpoints/primary",
|
||||
json={"status": "inactive"},
|
||||
)
|
||||
assert resp.status_code == 404
|
||||
@@ -1,203 +0,0 @@
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from app.application.agent import (
|
||||
compose_system_messages,
|
||||
next_step_for_tool,
|
||||
step_event_payload,
|
||||
task_from_deliberative_plan,
|
||||
task_plan_payload,
|
||||
)
|
||||
from app.application.agent.skill_loader import load_skill_manifests
|
||||
from app.modules.agent import (
|
||||
AgentDefinition,
|
||||
AgentTask,
|
||||
ExecutionPolicy,
|
||||
Plan,
|
||||
PlanStep,
|
||||
ToolPolicy,
|
||||
)
|
||||
from app.modules.planner import DeliberativePlan
|
||||
from app.modules.planner import PlanStep as DeliberativePlanStep
|
||||
from app.modules.skills import parse_skill_manifest
|
||||
|
||||
|
||||
def test_agent_definition_keeps_minimal_execution_and_tool_policy() -> None:
|
||||
agent = AgentDefinition(
|
||||
name="工师",
|
||||
system_prompt="完成已授权的工程任务。",
|
||||
skill_refs=("代码审查",),
|
||||
tool_policy=ToolPolicy(("read_file", "shell")),
|
||||
execution_policy=ExecutionPolicy(max_iterations=3),
|
||||
)
|
||||
|
||||
assert agent.skill_refs == ("代码审查",)
|
||||
assert agent.tool_policy.allowed_tools == ("read_file", "shell")
|
||||
assert agent.execution_policy.max_iterations == 3
|
||||
|
||||
|
||||
def test_agent_task_keeps_plan_as_runtime_only_contract() -> None:
|
||||
task = AgentTask(
|
||||
id="task_1",
|
||||
objective="审查当前变更",
|
||||
agent_name="工师",
|
||||
plan=Plan(
|
||||
steps=(PlanStep(id="step_1", title="读取变更", tool="read_file"),),
|
||||
status="approved",
|
||||
),
|
||||
)
|
||||
|
||||
assert task.status == "queued"
|
||||
assert task.plan is not None
|
||||
assert task.plan.steps[0].status == "pending"
|
||||
|
||||
|
||||
def test_deliberative_plan_becomes_traceable_agent_task() -> None:
|
||||
task = task_from_deliberative_plan(
|
||||
DeliberativePlan(
|
||||
title="调研计划",
|
||||
steps=[
|
||||
DeliberativePlanStep(
|
||||
id=1,
|
||||
title="检索资料",
|
||||
description="读取相关资料",
|
||||
tool="fetch_url",
|
||||
risk_level="L2",
|
||||
)
|
||||
],
|
||||
assumptions=["资料可访问"],
|
||||
risks=["来源不完整"],
|
||||
),
|
||||
run_id="run_1",
|
||||
agent_name="学士",
|
||||
objective="调研主题",
|
||||
)
|
||||
|
||||
assert task.status == "running"
|
||||
assert task.plan is not None
|
||||
assert task.plan.steps[0].id == "1"
|
||||
assert task_plan_payload(task)["steps"] == [
|
||||
{
|
||||
"id": "1",
|
||||
"title": "检索资料",
|
||||
"description": "读取相关资料",
|
||||
"tool": "fetch_url",
|
||||
"risk_level": "L2",
|
||||
"status": "pending",
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def test_tool_maps_only_to_its_next_uncompleted_plan_step() -> None:
|
||||
task = AgentTask(
|
||||
id="task_1",
|
||||
objective="执行计划",
|
||||
agent_name="工师",
|
||||
plan=Plan(
|
||||
steps=(
|
||||
PlanStep(id="1", title="读取", tool="read_file"),
|
||||
PlanStep(id="2", title="再次读取", tool="read_file"),
|
||||
)
|
||||
),
|
||||
)
|
||||
|
||||
first = next_step_for_tool(task, "read_file", set())
|
||||
second = next_step_for_tool(task, "read_file", {"1"})
|
||||
|
||||
assert first is not None and first.id == "1"
|
||||
assert second is not None and second.id == "2"
|
||||
assert next_step_for_tool(task, "shell", set()) is None
|
||||
assert step_event_payload(task, first, status="completed", result_summary="读取成功") == {
|
||||
"task_id": "task_1",
|
||||
"step_id": "1",
|
||||
"title": "读取",
|
||||
"tool": "read_file",
|
||||
"risk_level": "L2",
|
||||
"status": "completed",
|
||||
"result_summary": "读取成功",
|
||||
}
|
||||
|
||||
|
||||
def test_execution_policy_rejects_invalid_iteration_limit() -> None:
|
||||
with pytest.raises(ValueError, match="max_iterations"):
|
||||
ExecutionPolicy(max_iterations=0)
|
||||
|
||||
|
||||
def test_manifest_and_prompt_composer_keep_stable_order() -> None:
|
||||
skill = SimpleNamespace(
|
||||
name="代码审查",
|
||||
description="审查变更",
|
||||
trigger="审查, review",
|
||||
prompt_fragment="先理解调用路径，再报告问题。",
|
||||
risk_level="L1",
|
||||
status="enabled",
|
||||
tool_bindings_json='["read_file"]',
|
||||
)
|
||||
manifest = parse_skill_manifest(skill, ["read_file"])
|
||||
|
||||
messages = compose_system_messages(
|
||||
["全局约束"],
|
||||
AgentDefinition(name="工师", system_prompt="只执行已授权任务。"),
|
||||
[manifest],
|
||||
)
|
||||
|
||||
assert manifest.triggers == ("审查", "review")
|
||||
assert messages == [
|
||||
"全局约束",
|
||||
"【Agent：工师】\n只执行已授权任务。",
|
||||
"【代码审查·能力注入】\n先理解调用路径，再报告问题。",
|
||||
]
|
||||
|
||||
|
||||
def test_default_tools_filters_explicit_allowlist() -> None:
|
||||
from app.application.agent.tools import default_tools
|
||||
|
||||
tools = default_tools(allowed_tools=("read_file",))
|
||||
|
||||
assert [tool["function"]["name"] for tool in tools] == ["read_file"]
|
||||
|
||||
|
||||
async def test_skill_loader_unions_matches_with_explicit_references(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
skills = [
|
||||
SimpleNamespace(
|
||||
name="代码审查",
|
||||
description="审查",
|
||||
trigger="审查",
|
||||
prompt_fragment="审查流程",
|
||||
risk_level="L1",
|
||||
status="enabled",
|
||||
tool_bindings_json='["read_file"]',
|
||||
),
|
||||
SimpleNamespace(
|
||||
name="网页研究",
|
||||
description="研究",
|
||||
trigger="研究",
|
||||
prompt_fragment="研究流程",
|
||||
risk_level="L2",
|
||||
status="enabled",
|
||||
tool_bindings_json='["web_search"]',
|
||||
),
|
||||
]
|
||||
|
||||
async def fake_match_active_skills(_context: str):
|
||||
return [skills[0]]
|
||||
|
||||
async def fake_get_enabled_skills_by_names(names: tuple[str, ...]):
|
||||
return [skill for skill in skills if skill.name in names]
|
||||
|
||||
monkeypatch.setattr(
|
||||
"app.application.agent.skill_loader.match_active_skills",
|
||||
fake_match_active_skills,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"app.application.agent.skill_loader.get_enabled_skills_by_names",
|
||||
fake_get_enabled_skills_by_names,
|
||||
)
|
||||
|
||||
manifests = await load_skill_manifests("请审查", skill_refs=("网页研究",))
|
||||
|
||||
assert [manifest.name for manifest in manifests] == ["代码审查", "网页研究"]
|
||||
assert manifests[0].allowed_tools == ("read_file",)
|
||||
@@ -1,235 +0,0 @@
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from app.ai_provider.types.messages import HumanMessage
|
||||
from app.application.agent import AgentContext, AgentRunContext, ExtensionRegistry
|
||||
from app.application.agent.python_skills import (
|
||||
PythonSkillCatalog,
|
||||
PythonSkillDefinition,
|
||||
activate_python_skills,
|
||||
discover_python_skills,
|
||||
)
|
||||
|
||||
|
||||
def _write_skill(package: Path, module: str, *, name: str, trigger: str) -> None:
|
||||
(package / f"{module}.py").write_text(
|
||||
f"""
|
||||
from app.application.agent.python_skills import PythonSkillDefinition
|
||||
from app.ai_provider.types.messages import HumanMessage
|
||||
|
||||
async def transform(context, run):
|
||||
return type(context)(
|
||||
messages=(*context.messages, HumanMessage(content={name!r})),
|
||||
tools=context.tools,
|
||||
model_hint=context.model_hint,
|
||||
)
|
||||
|
||||
def setup(registry):
|
||||
registry.add_context_transformer(transform)
|
||||
|
||||
SKILL = PythonSkillDefinition(
|
||||
name={name!r},
|
||||
description={name!r},
|
||||
triggers=({trigger!r},),
|
||||
instructions={name!r},
|
||||
setup=setup,
|
||||
)
|
||||
""",
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
|
||||
async def test_catalog_discovers_without_activating_and_matches_per_run(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
package = tmp_path / "example_skills"
|
||||
package.mkdir()
|
||||
(package / "__init__.py").write_text("", encoding="utf-8")
|
||||
_write_skill(package, "b_skill", name="乙", trigger="研究")
|
||||
_write_skill(package, "a_skill", name="甲", trigger="审查")
|
||||
(package / "_ignored.py").write_text(
|
||||
"raise AssertionError('私有模块不应装载')",
|
||||
encoding="utf-8",
|
||||
)
|
||||
monkeypatch.syspath_prepend(str(tmp_path))
|
||||
|
||||
catalog = discover_python_skills(("example_skills", "example_skills"))
|
||||
registry = ExtensionRegistry()
|
||||
unchanged = await registry.transform_context(
|
||||
AgentContext(messages=(HumanMessage(content="开始"),), tools=()),
|
||||
AgentRunContext(),
|
||||
)
|
||||
matched = catalog.match("请做研究", skill_refs=("甲",))
|
||||
activation = await activate_python_skills(registry, matched)
|
||||
transformed = await registry.transform_context(
|
||||
AgentContext(messages=(HumanMessage(content="开始"),), tools=()),
|
||||
AgentRunContext(),
|
||||
)
|
||||
await activation.close()
|
||||
|
||||
assert [skill.name for skill in catalog.skills] == ["甲", "乙"]
|
||||
assert [message.content for message in unchanged.messages] == ["开始"]
|
||||
assert [skill.name for skill in matched] == ["甲", "乙"]
|
||||
assert [message.content for message in transformed.messages] == ["开始", "甲", "乙"]
|
||||
|
||||
|
||||
async def test_activation_cleanup_is_run_scoped() -> None:
|
||||
cleaned: list[str] = []
|
||||
|
||||
def setup(registry: ExtensionRegistry):
|
||||
async def transform(context: AgentContext, _run: AgentRunContext) -> AgentContext:
|
||||
return AgentContext(
|
||||
(*context.messages, HumanMessage(content="active")),
|
||||
context.tools,
|
||||
context.model_hint,
|
||||
)
|
||||
|
||||
registry.add_context_transformer(transform)
|
||||
return lambda: cleaned.append("closed")
|
||||
|
||||
skill = PythonSkillDefinition("guard", "", ("guard",), "", setup)
|
||||
catalog = PythonSkillCatalog((skill,))
|
||||
first_registry = ExtensionRegistry()
|
||||
second_registry = ExtensionRegistry()
|
||||
|
||||
activation = await activate_python_skills(first_registry, catalog.match("guard"))
|
||||
first = await first_registry.transform_context(AgentContext((), ()), AgentRunContext())
|
||||
second = await second_registry.transform_context(AgentContext((), ()), AgentRunContext())
|
||||
await activation.close()
|
||||
|
||||
assert [message.content for message in first.messages] == ["active"]
|
||||
assert second.messages == ()
|
||||
assert cleaned == ["closed"]
|
||||
|
||||
|
||||
def test_catalog_supports_custom_matcher_with_run_metadata() -> None:
|
||||
skill = PythonSkillDefinition(
|
||||
"metadata_skill",
|
||||
"",
|
||||
(),
|
||||
"",
|
||||
lambda _registry: None,
|
||||
match=lambda text, metadata: text == "执行" and metadata.get("channel") == "voice",
|
||||
)
|
||||
catalog = PythonSkillCatalog((skill,))
|
||||
|
||||
assert catalog.match("执行", metadata={"channel": "text"}) == ()
|
||||
assert catalog.match("执行", metadata={"channel": "voice"}) == (skill,)
|
||||
|
||||
|
||||
def test_catalog_rejects_duplicate_names() -> None:
|
||||
skill = PythonSkillDefinition("same", "", (), "", lambda _registry: None)
|
||||
|
||||
with pytest.raises(ValueError, match="名称重复"):
|
||||
PythonSkillCatalog((skill, skill))
|
||||
|
||||
|
||||
def test_discovery_rejects_module_without_definition(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
package = tmp_path / "invalid_skills"
|
||||
package.mkdir()
|
||||
(package / "__init__.py").write_text("", encoding="utf-8")
|
||||
(package / "broken.py").write_text("SKILL_NAME = '坏技能'", encoding="utf-8")
|
||||
monkeypatch.syspath_prepend(str(tmp_path))
|
||||
|
||||
with pytest.raises(TypeError, match="缺少 SKILL 定义"):
|
||||
discover_python_skills(("invalid_skills",))
|
||||
|
||||
|
||||
async def test_budget_guard_skill_stops_when_usage_exceeds_threshold() -> None:
|
||||
from app.ai_provider.types.chat import TokenUsage
|
||||
from app.application.agent.extensions import TurnSnapshot
|
||||
from app.skills.budget_guard import _BUDGET_TOTAL_TOKENS, SKILL
|
||||
|
||||
registry = ExtensionRegistry()
|
||||
activation = await activate_python_skills(registry, (SKILL,))
|
||||
|
||||
over_budget = TurnSnapshot(
|
||||
iteration=4,
|
||||
content="",
|
||||
reasoning="",
|
||||
tool_calls=(),
|
||||
tool_results=(),
|
||||
usage=TokenUsage(total_tokens=_BUDGET_TOTAL_TOKENS + 1),
|
||||
)
|
||||
within_budget = TurnSnapshot(
|
||||
iteration=1,
|
||||
content="",
|
||||
reasoning="",
|
||||
tool_calls=(),
|
||||
tool_results=(),
|
||||
usage=TokenUsage(total_tokens=max(_BUDGET_TOTAL_TOKENS - 1, 0)),
|
||||
)
|
||||
|
||||
assert await registry.should_stop(over_budget, AgentRunContext()) is True
|
||||
assert await registry.should_stop(within_budget, AgentRunContext()) is False
|
||||
await activation.close()
|
||||
|
||||
|
||||
async def test_context_budget_folds_long_history() -> None:
|
||||
from app.ai_provider.types.messages import HumanMessage
|
||||
from app.application.agent.extensions import AgentContext
|
||||
from app.skills.context_budget import SKILL
|
||||
|
||||
registry = ExtensionRegistry()
|
||||
activation = await activate_python_skills(registry, (SKILL,))
|
||||
messages = tuple(HumanMessage(content="x" * 1000) for _ in range(20))
|
||||
ctx = AgentContext(messages=messages, tools=())
|
||||
|
||||
trimmed = await registry.transform_context(ctx, AgentRunContext())
|
||||
|
||||
assert len(trimmed.messages) < len(messages)
|
||||
assert "已省略" in trimmed.messages[0].content
|
||||
await activation.close()
|
||||
|
||||
|
||||
async def test_arg_guard_rejects_oversized_input() -> None:
|
||||
from app.application.agent.extensions import ToolCallContext
|
||||
from app.skills.arg_guard import _MAX_ARG_CHARS, SKILL
|
||||
|
||||
registry = ExtensionRegistry()
|
||||
activation = await activate_python_skills(registry, (SKILL,))
|
||||
oversized = ToolCallContext(
|
||||
"call_1", "write_file", {"content": "x" * (_MAX_ARG_CHARS + 10)}, ()
|
||||
)
|
||||
normal = ToolCallContext("call_2", "read_file", {"path": "/a"}, ())
|
||||
|
||||
big = await registry.before_tool_call(oversized, AgentRunContext())
|
||||
small = await registry.before_tool_call(normal, AgentRunContext())
|
||||
|
||||
assert big.allowed is False
|
||||
assert big.reason is not None
|
||||
assert small.allowed is True
|
||||
await activation.close()
|
||||
|
||||
|
||||
async def test_run_metrics_listens_through_settled() -> None:
|
||||
from app.application.agent.extensions import AgentEvent
|
||||
from app.skills.run_metrics import SKILL
|
||||
|
||||
registry = ExtensionRegistry()
|
||||
activation = await activate_python_skills(registry, (SKILL,))
|
||||
run = AgentRunContext(run_id="run_metrics_test")
|
||||
|
||||
await registry.publish(AgentEvent("agent.started"), run)
|
||||
await registry.publish(AgentEvent("agent.turn.started", {"iteration": 1}), run)
|
||||
await registry.publish(
|
||||
AgentEvent("tool.started", {"tool_call_id": "c1", "tool_name": "search"}),
|
||||
run,
|
||||
)
|
||||
await registry.publish(
|
||||
AgentEvent("tool.completed", {"tool_call_id": "c1", "status": "success"}),
|
||||
run,
|
||||
)
|
||||
await registry.publish(
|
||||
AgentEvent("tool.completed", {"tool_call_id": "c2", "status": "denied"}),
|
||||
run,
|
||||
)
|
||||
await registry.publish(AgentEvent("agent.turn.started", {"iteration": 2}), run)
|
||||
await registry.publish(AgentEvent("agent.settled", {"status": "completed"}), run)
|
||||
await activation.close()
|
||||
# 指标经 loguru 输出；此处验证含工具级事件的事件链被完整消费且 settled 不抛错。
|
||||
@@ -1,405 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from typing import Any
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from app.ai_provider.types.chat import TokenUsage
|
||||
from app.ai_provider.types.messages import HumanMessage
|
||||
from app.application.agent import (
|
||||
AfterToolCallDecision,
|
||||
Agent,
|
||||
AgentConfig,
|
||||
AgentContext,
|
||||
AgentEvent,
|
||||
AgentRunContext,
|
||||
AgentTurnResult,
|
||||
BeforeToolCallDecision,
|
||||
ExtensionRegistry,
|
||||
PendingMessageQueue,
|
||||
QueueMode,
|
||||
ToolCallContext,
|
||||
TurnSnapshot,
|
||||
)
|
||||
from app.modules.tool_runtime.definition import ToolResult
|
||||
|
||||
|
||||
def test_pending_message_queue_supports_both_modes() -> None:
|
||||
queue = PendingMessageQueue(QueueMode.one_at_a_time)
|
||||
queue.put(HumanMessage(content="一"))
|
||||
queue.put(HumanMessage(content="二"))
|
||||
|
||||
assert [message.content for message in queue.drain()] == ["一"]
|
||||
assert [message.content for message in queue.drain()] == ["二"]
|
||||
|
||||
queue.mode = QueueMode.all
|
||||
queue.put(HumanMessage(content="三"))
|
||||
queue.put(HumanMessage(content="四"))
|
||||
assert [message.content for message in queue.drain()] == ["三", "四"]
|
||||
|
||||
|
||||
async def test_extension_registry_orders_context_and_tool_hooks() -> None:
|
||||
registry = ExtensionRegistry()
|
||||
run = AgentRunContext(run_id="run_1")
|
||||
|
||||
async def first(context: AgentContext, _run: AgentRunContext) -> AgentContext:
|
||||
return AgentContext(
|
||||
messages=(*context.messages, HumanMessage(content="first")),
|
||||
tools=context.tools,
|
||||
model_hint=context.model_hint,
|
||||
)
|
||||
|
||||
async def second(context: AgentContext, _run: AgentRunContext) -> AgentContext:
|
||||
return AgentContext(
|
||||
messages=(*context.messages, HumanMessage(content="second")),
|
||||
tools=context.tools,
|
||||
model_hint=context.model_hint,
|
||||
)
|
||||
|
||||
async def rewrite(
|
||||
_call: ToolCallContext, _run: AgentRunContext
|
||||
) -> BeforeToolCallDecision:
|
||||
return BeforeToolCallDecision(input={"value": 2})
|
||||
|
||||
async def terminate(
|
||||
_call: ToolCallContext,
|
||||
result: ToolResult,
|
||||
_run: AgentRunContext,
|
||||
) -> AfterToolCallDecision:
|
||||
return AfterToolCallDecision(result=result, terminate=True)
|
||||
|
||||
registry.add_context_transformer(first)
|
||||
registry.add_context_transformer(second)
|
||||
registry.add_before_tool_hook(rewrite)
|
||||
registry.add_after_tool_hook(terminate)
|
||||
|
||||
transformed = await registry.transform_context(AgentContext((), ()), run)
|
||||
assert [message.content for message in transformed.messages] == ["first", "second"]
|
||||
|
||||
call = ToolCallContext("call_1", "tool", {"value": 1}, ())
|
||||
before = await registry.before_tool_call(call, run)
|
||||
assert before.input == {"value": 2}
|
||||
|
||||
result = ToolResult(True, 0, "ok", "", 1, False, False, "trace_1")
|
||||
after = await registry.after_tool_call(call, result, run)
|
||||
assert after.result is result
|
||||
assert after.terminate is True
|
||||
|
||||
|
||||
async def test_extension_registry_extend_preserves_registration_order() -> None:
|
||||
registry = ExtensionRegistry()
|
||||
skill_registry = ExtensionRegistry()
|
||||
|
||||
async def base(context: AgentContext, _run: AgentRunContext) -> AgentContext:
|
||||
return AgentContext((*context.messages, HumanMessage(content="base")), context.tools)
|
||||
|
||||
async def skill(context: AgentContext, _run: AgentRunContext) -> AgentContext:
|
||||
return AgentContext((*context.messages, HumanMessage(content="skill")), context.tools)
|
||||
|
||||
registry.add_context_transformer(base)
|
||||
skill_registry.add_context_transformer(skill)
|
||||
registry.extend(skill_registry)
|
||||
|
||||
transformed = await registry.transform_context(AgentContext((), ()), AgentRunContext())
|
||||
|
||||
assert [message.content for message in transformed.messages] == ["base", "skill"]
|
||||
|
||||
|
||||
async def test_stateful_agent_consumes_follow_up_and_publishes_events() -> None:
|
||||
calls: list[list[str]] = []
|
||||
events: list[str] = []
|
||||
|
||||
async def runner(*, messages: list[Any], **_kwargs: object) -> AgentTurnResult:
|
||||
calls.append([message.content for message in messages])
|
||||
messages.append(HumanMessage(content=f"assistant-{len(calls)}"))
|
||||
return AgentTurnResult(
|
||||
content=f"done-{len(calls)}",
|
||||
reasoning="",
|
||||
iterations=1,
|
||||
tool_history=[{"name": f"tool-{len(calls)}"}],
|
||||
usage=TokenUsage(total_tokens=10),
|
||||
)
|
||||
|
||||
extensions = ExtensionRegistry()
|
||||
|
||||
async def listen(event: AgentEvent, _run: AgentRunContext) -> None:
|
||||
events.append(event.type)
|
||||
|
||||
extensions.add_event_listener(listen)
|
||||
agent = Agent(
|
||||
AgentConfig(tools=(), tool_runtime=MagicMock(), max_iterations=2),
|
||||
extensions=extensions,
|
||||
runner=runner,
|
||||
)
|
||||
|
||||
async def enqueue_follow_up(event: AgentEvent, _run: AgentRunContext) -> None:
|
||||
if event.type == "agent.started":
|
||||
assert agent.follow_up("再做一步") is True
|
||||
|
||||
extensions.add_event_listener(enqueue_follow_up)
|
||||
result = await agent.prompt("开始")
|
||||
|
||||
assert result.content == "done-2"
|
||||
assert calls == [["开始"], ["开始", "assistant-1", "再做一步"]]
|
||||
assert result.iterations == 2
|
||||
assert result.usage.total_tokens == 20
|
||||
assert [item["name"] for item in result.tool_history] == ["tool-1", "tool-2"]
|
||||
assert events == ["agent.started", "agent.ended", "agent.settled"]
|
||||
assert agent.state.status == "completed"
|
||||
assert agent.state.streaming_message == "done-2"
|
||||
|
||||
|
||||
async def test_stateful_agent_does_not_continue_after_runner_truncates() -> None:
|
||||
calls = 0
|
||||
|
||||
async def runner(*, messages: list[Any], **_kwargs: object) -> AgentTurnResult:
|
||||
nonlocal calls
|
||||
calls += 1
|
||||
messages.append(HumanMessage(content="assistant"))
|
||||
return AgentTurnResult(content="partial", reasoning="", iterations=1, truncated=True)
|
||||
|
||||
extensions = ExtensionRegistry()
|
||||
agent = Agent(
|
||||
AgentConfig(tools=(), tool_runtime=None, max_iterations=3),
|
||||
extensions=extensions,
|
||||
runner=runner,
|
||||
)
|
||||
|
||||
async def enqueue(event: AgentEvent, _run: AgentRunContext) -> None:
|
||||
if event.type == "agent.started":
|
||||
assert agent.follow_up("稍后处理") is True
|
||||
|
||||
extensions.add_event_listener(enqueue)
|
||||
result = await agent.prompt("开始")
|
||||
|
||||
assert calls == 1
|
||||
assert result.truncated is True
|
||||
assert [message.content for message in agent.follow_ups.drain()] == ["稍后处理"]
|
||||
|
||||
|
||||
async def test_follow_up_is_preserved_when_iteration_budget_is_exhausted() -> None:
|
||||
async def runner(*, messages: list[Any], **_kwargs: object) -> AgentTurnResult:
|
||||
messages.append(HumanMessage(content="assistant"))
|
||||
return AgentTurnResult(content="done", reasoning="", iterations=1)
|
||||
|
||||
extensions = ExtensionRegistry()
|
||||
agent = Agent(
|
||||
AgentConfig(tools=(), tool_runtime=None, max_iterations=1),
|
||||
extensions=extensions,
|
||||
runner=runner,
|
||||
)
|
||||
|
||||
async def enqueue(event: AgentEvent, _run: AgentRunContext) -> None:
|
||||
if event.type == "agent.started":
|
||||
assert agent.follow_up("稍后处理") is True
|
||||
|
||||
extensions.add_event_listener(enqueue)
|
||||
result = await agent.prompt("开始")
|
||||
|
||||
assert result.truncated is True
|
||||
assert bool(agent.follow_ups) is True
|
||||
assert [message.content for message in agent.follow_ups.drain()] == ["稍后处理"]
|
||||
|
||||
|
||||
async def test_started_listener_failure_restores_idle_state() -> None:
|
||||
extensions = ExtensionRegistry()
|
||||
|
||||
async def fail(_event: AgentEvent, _run: AgentRunContext) -> None:
|
||||
raise RuntimeError("event store unavailable")
|
||||
|
||||
extensions.add_event_listener(fail)
|
||||
agent = Agent(
|
||||
AgentConfig(tools=(), tool_runtime=MagicMock()),
|
||||
extensions=extensions,
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="event store unavailable"):
|
||||
await agent.prompt("开始")
|
||||
await agent.wait_for_idle()
|
||||
assert agent.state.status == "failed"
|
||||
|
||||
|
||||
async def test_stateful_agent_abort_cancels_active_run() -> None:
|
||||
started = asyncio.Event()
|
||||
events: list[str] = []
|
||||
|
||||
async def runner(**_kwargs: object) -> AgentTurnResult:
|
||||
started.set()
|
||||
await asyncio.Event().wait()
|
||||
raise AssertionError("unreachable")
|
||||
|
||||
extensions = ExtensionRegistry()
|
||||
|
||||
async def listen(event: AgentEvent, _run: AgentRunContext) -> None:
|
||||
events.append(event.type)
|
||||
|
||||
extensions.add_event_listener(listen)
|
||||
agent = Agent(
|
||||
AgentConfig(tools=(), tool_runtime=MagicMock()),
|
||||
extensions=extensions,
|
||||
runner=runner,
|
||||
)
|
||||
task = asyncio.create_task(agent.prompt("开始"))
|
||||
await started.wait()
|
||||
|
||||
assert await agent.abort() is True
|
||||
assert agent.run_context.cancellation.cancelled is True
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
await agent.wait_for_idle()
|
||||
assert agent.state.status == "cancelled"
|
||||
assert events == ["agent.started", "agent.cancelled", "agent.settled"]
|
||||
|
||||
|
||||
async def test_agent_replaces_cancelled_token_on_next_run() -> None:
|
||||
async def runner(**_kwargs: object) -> AgentTurnResult:
|
||||
return AgentTurnResult(content="done", reasoning="", iterations=1)
|
||||
|
||||
agent = Agent(AgentConfig(tools=(), tool_runtime=None), runner=runner)
|
||||
agent.run_context.cancellation.cancel()
|
||||
|
||||
await agent.prompt("重新开始")
|
||||
|
||||
assert agent.run_context.cancellation.cancelled is False
|
||||
|
||||
|
||||
async def test_extension_registry_should_stop_short_circuits_on_first_policy() -> None:
|
||||
registry = ExtensionRegistry()
|
||||
run = AgentRunContext(run_id="run_1")
|
||||
turn = TurnSnapshot(
|
||||
iteration=1,
|
||||
content="done",
|
||||
reasoning="",
|
||||
tool_calls=(),
|
||||
tool_results=(),
|
||||
)
|
||||
calls: list[str] = []
|
||||
|
||||
async def keep_going(snapshot: TurnSnapshot, _run: AgentRunContext) -> bool:
|
||||
calls.append("keep_going")
|
||||
assert snapshot is turn
|
||||
return False
|
||||
|
||||
async def stop_now(_snapshot: TurnSnapshot, _run: AgentRunContext) -> bool:
|
||||
calls.append("stop_now")
|
||||
return True
|
||||
|
||||
async def never_called(_snapshot: TurnSnapshot, _run: AgentRunContext) -> bool:
|
||||
calls.append("never_called")
|
||||
return False
|
||||
|
||||
registry.add_turn_policy(keep_going)
|
||||
registry.add_turn_policy(stop_now)
|
||||
registry.add_turn_policy(never_called)
|
||||
|
||||
assert await registry.should_stop(turn, run) is True
|
||||
assert calls == ["keep_going", "stop_now"]
|
||||
|
||||
|
||||
async def test_extension_registry_should_stop_false_when_all_policies_decline() -> None:
|
||||
registry = ExtensionRegistry()
|
||||
turn = TurnSnapshot(
|
||||
iteration=2,
|
||||
content="",
|
||||
reasoning="",
|
||||
tool_calls=({"name": "tool", "args": "{}"},),
|
||||
tool_results=(),
|
||||
)
|
||||
|
||||
async def decline(_snapshot: TurnSnapshot, _run: AgentRunContext) -> bool:
|
||||
return False
|
||||
|
||||
registry.add_turn_policy(decline)
|
||||
registry.add_turn_policy(decline)
|
||||
|
||||
assert await registry.should_stop(turn, AgentRunContext()) is False
|
||||
|
||||
|
||||
async def test_extension_registry_should_stop_respects_cancellation() -> None:
|
||||
registry = ExtensionRegistry()
|
||||
run = AgentRunContext()
|
||||
run.cancellation.cancel()
|
||||
turn = TurnSnapshot(
|
||||
iteration=1,
|
||||
content="",
|
||||
reasoning="",
|
||||
tool_calls=(),
|
||||
tool_results=(),
|
||||
)
|
||||
|
||||
async def policy(_snapshot: TurnSnapshot, _run: AgentRunContext) -> bool:
|
||||
return False
|
||||
|
||||
# raise_if_cancelled 在 policy 遍历循环内，故至少注册一个 policy 才会触发；
|
||||
# 这与 transform_context / before_tool_call / after_tool_call 一致：空注册时
|
||||
# 方法无 await 点，不额外检查取消信号。
|
||||
registry.add_turn_policy(policy)
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await registry.should_stop(turn, run)
|
||||
|
||||
|
||||
async def test_on_tool_event_publishes_tool_lifecycle_to_extensions() -> None:
|
||||
"""_on_tool_event 既转发 SSE，又把同一事件发布进扩展事件流。"""
|
||||
events: list[tuple[str, dict[str, Any]]] = []
|
||||
sse_events: list[tuple[str, dict[str, Any]]] = []
|
||||
|
||||
async def on_tool_event(event_type: str, payload: dict[str, Any]) -> None:
|
||||
sse_events.append((event_type, payload))
|
||||
|
||||
async def listen(event: AgentEvent, _run: AgentRunContext) -> None:
|
||||
events.append((event.type, dict(event.payload)))
|
||||
|
||||
extensions = ExtensionRegistry()
|
||||
extensions.add_event_listener(listen)
|
||||
agent = Agent(
|
||||
AgentConfig(tools=(), tool_runtime=None, on_tool_event=on_tool_event),
|
||||
extensions=extensions,
|
||||
)
|
||||
|
||||
await agent._on_tool_event(
|
||||
"tool.started", {"tool_call_id": "c1", "tool_name": "search"}
|
||||
)
|
||||
await agent._on_tool_event(
|
||||
"tool.completed", {"tool_call_id": "c1", "status": "success"}
|
||||
)
|
||||
await agent._on_tool_event(
|
||||
"tool.completed", {"tool_call_id": "c2", "status": "denied"}
|
||||
)
|
||||
|
||||
assert [item[0] for item in sse_events] == [
|
||||
"tool.started",
|
||||
"tool.completed",
|
||||
"tool.completed",
|
||||
]
|
||||
assert [item[0] for item in events] == [
|
||||
"tool.started",
|
||||
"tool.completed",
|
||||
"tool.completed",
|
||||
]
|
||||
assert events[1][1]["status"] == "success"
|
||||
# pending 状态维护未被桥接改动破坏
|
||||
assert "c1" not in agent.state.pending_tool_calls
|
||||
|
||||
|
||||
async def test_on_tool_event_swallows_listener_failure() -> None:
|
||||
"""扩展 listener 抛异常不应阻断 SSE 转发或工具事件主流程。"""
|
||||
sse: list[str] = []
|
||||
|
||||
async def on_tool_event(event_type: str, _payload: dict[str, Any]) -> None:
|
||||
sse.append(event_type)
|
||||
|
||||
async def boom(_event: AgentEvent, _run: AgentRunContext) -> None:
|
||||
raise RuntimeError("listener boom")
|
||||
|
||||
extensions = ExtensionRegistry()
|
||||
extensions.add_event_listener(boom)
|
||||
agent = Agent(
|
||||
AgentConfig(tools=(), tool_runtime=None, on_tool_event=on_tool_event),
|
||||
extensions=extensions,
|
||||
)
|
||||
|
||||
# 不抛：异常被吞，SSE 仍照常转发
|
||||
await agent._on_tool_event("tool.started", {"tool_call_id": "c1"})
|
||||
assert sse == ["tool.started"]
|
||||
@@ -1,101 +0,0 @@
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from app.application.agent.tools import to_ai_provider_tools
|
||||
from app.application.skill_executor import SkillExecutor
|
||||
from app.application.tool_runtime.view import RunScopedToolView
|
||||
from app.modules.skills.manifest import SkillManifest
|
||||
from app.modules.tool_runtime.definition import ToolResult, ToolSpec
|
||||
|
||||
|
||||
def _spec(name: str) -> ToolSpec:
|
||||
return ToolSpec(
|
||||
name=name,
|
||||
description=f"{name} description",
|
||||
risk_level="L1",
|
||||
permission_policy=MagicMock(),
|
||||
input_schema={
|
||||
"type": "object",
|
||||
"properties": {"value": {"type": "string"}},
|
||||
"required": ["value"],
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
async def test_aliases_are_isolated_and_execute_through_base_runtime() -> None:
|
||||
runtime = MagicMock()
|
||||
runtime.list_specs.return_value = [_spec("shared_tool")]
|
||||
runtime.execute = AsyncMock(
|
||||
return_value=ToolResult(True, 0, "ok", "", 1, False, False, "trace_1")
|
||||
)
|
||||
first = RunScopedToolView(runtime)
|
||||
second = RunScopedToolView(runtime)
|
||||
|
||||
assert first.bind_alias("skill:甲/shared_tool", "shared_tool", description_prefix="[甲] ")
|
||||
assert second.bind_alias("skill:乙/shared_tool", "shared_tool", description_prefix="[乙] ")
|
||||
|
||||
first_names = {spec.name for spec in first.list_specs()}
|
||||
second_names = {spec.name for spec in second.list_specs()}
|
||||
result = await first.execute(
|
||||
"skill:甲/shared_tool",
|
||||
{"value": "x"},
|
||||
run_id="run_1",
|
||||
character_id="char_1",
|
||||
)
|
||||
|
||||
assert "skill:甲/shared_tool" in first_names
|
||||
assert "skill:乙/shared_tool" not in first_names
|
||||
assert "skill:乙/shared_tool" in second_names
|
||||
assert "skill:甲/shared_tool" not in second_names
|
||||
assert result.stdout == "ok"
|
||||
runtime.execute.assert_awaited_once_with(
|
||||
"shared_tool",
|
||||
{"value": "x"},
|
||||
run_id="run_1",
|
||||
character_id="char_1",
|
||||
trace_id=None,
|
||||
confirmed=False,
|
||||
context=None,
|
||||
)
|
||||
runtime.register.assert_not_called()
|
||||
runtime.unregister.assert_not_called()
|
||||
|
||||
|
||||
def test_alias_keeps_target_input_schema() -> None:
|
||||
runtime = MagicMock()
|
||||
runtime.list_specs.return_value = [_spec("custom_tool")]
|
||||
view = RunScopedToolView(runtime)
|
||||
view.bind_alias("skill:研究/custom_tool", "custom_tool")
|
||||
|
||||
tools = to_ai_provider_tools(view.list_specs())
|
||||
alias = next(tool for tool in tools if tool["function"]["name"] == "skill:研究/custom_tool")
|
||||
|
||||
assert alias["function"]["parameters"]["required"] == ["value"]
|
||||
|
||||
|
||||
async def test_skill_executor_binds_only_current_run_aliases(
|
||||
monkeypatch,
|
||||
) -> None:
|
||||
manifest = SkillManifest(
|
||||
name="研究",
|
||||
description="",
|
||||
triggers=("研究",),
|
||||
prompt="研究说明",
|
||||
allowed_tools=("shared_tool", "missing"),
|
||||
risk_level="L1",
|
||||
)
|
||||
|
||||
async def load(*_args, **_kwargs):
|
||||
return [manifest]
|
||||
|
||||
monkeypatch.setattr("app.application.skill_executor.load_skill_manifests", load)
|
||||
runtime = MagicMock()
|
||||
runtime.list_specs.return_value = [_spec("shared_tool")]
|
||||
view = RunScopedToolView(runtime)
|
||||
messages: list[str] = []
|
||||
|
||||
await SkillExecutor(view).match_and_activate("请研究", messages, run_id="run_1")
|
||||
|
||||
assert messages == ["【研究·能力注入】\n研究说明"]
|
||||
assert "skill:研究/shared_tool" in view.specs_by_name()
|
||||
assert "skill:研究/missing" not in view.specs_by_name()
|
||||
runtime.register.assert_not_called()
|
||||
@@ -70,10 +70,6 @@ async def reset_database() -> AsyncGenerator[None, None]:
|
||||
await conn.run_sync(Base.metadata.create_all)
|
||||
async with AsyncSessionLocal() as session:
|
||||
await seed_default_data(session)
|
||||
# 清 AI provider endpoint 覆盖缓存（模块级，跨测试隔离）
|
||||
from app.services.ai_provider_admin import set_endpoint_overrides_cache
|
||||
|
||||
set_endpoint_overrides_cache([])
|
||||
yield
|
||||
_assert_test_database_url(str(_engine.url))
|
||||
async with _engine.begin() as conn:
|
||||
|
||||
@@ -1,965 +0,0 @@
|
||||
"""Agent 循环单元测试（流式版）：标准 ReAct/Agent loop 的核心场景 + 累加器单测。
|
||||
|
||||
mock 边界：
|
||||
- ``app.application.agent.runtime.ProviderManagerFactory.create`` 的 ``.chat_stream()`` —— 不真实调 LLM，
|
||||
按轮次返回预置的 chunk 序列（async generator）。
|
||||
- ``tool_runtime.execute`` —— 不真实调沙箱（用 AsyncMock 返回 ToolResult 或抛异常）。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from app.ai_provider import TokenUsage
|
||||
from app.ai_provider.llm import ToolCallAccumulator
|
||||
from app.ai_provider.types.chat import ChatChunk
|
||||
from app.ai_provider.types.messages import (
|
||||
AIMessage,
|
||||
BaseMessage,
|
||||
HumanMessage,
|
||||
SystemMessage,
|
||||
ToolMessage,
|
||||
)
|
||||
from app.application.agent import (
|
||||
AgentContext,
|
||||
AgentEvent,
|
||||
AgentRunContext,
|
||||
AgentTurnResult,
|
||||
ExtensionRegistry,
|
||||
run_agent_turn,
|
||||
)
|
||||
from app.application.agent.confirmation import ConfirmationRegistry
|
||||
from app.application.agent.tools import default_tools, to_ai_provider_tools
|
||||
from app.application.tool_runtime.builtins import builtin_specs
|
||||
from app.modules.safety.exceptions import RequiresConfirmation
|
||||
from app.modules.tool_runtime.definition import ToolResult
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def base_messages() -> list[BaseMessage]:
|
||||
return [
|
||||
SystemMessage(content="你是测试助手。"),
|
||||
HumanMessage(content="请用 echo 工具回显 'hello'。"),
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def tools() -> list[dict[str, Any]]:
|
||||
return default_tools()
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 流式测试辅助：构造 chunk / chunk 序列 / 假 router / 假 tool_runtime
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
def _chunk(
|
||||
*,
|
||||
content: str = "",
|
||||
reasoning: str = "",
|
||||
tool_call_chunks: list[dict[str, Any]] | None = None,
|
||||
usage: TokenUsage | None = None,
|
||||
finish_reason: str | None = None,
|
||||
) -> ChatChunk:
|
||||
"""构造单个 ChatChunk。"""
|
||||
return ChatChunk(
|
||||
delta_content=content,
|
||||
delta_reasoning=reasoning,
|
||||
finish_reason=finish_reason,
|
||||
usage=usage,
|
||||
tool_call_chunks=tool_call_chunks or [],
|
||||
)
|
||||
|
||||
|
||||
async def _chunk_iter(chunks: list[ChatChunk]):
|
||||
"""把一组 chunk 包装成 async generator（模拟 router.stream 的返回值）。"""
|
||||
for c in chunks:
|
||||
yield c
|
||||
|
||||
|
||||
def _split_tool_call_stream(
|
||||
*,
|
||||
index: int = 0,
|
||||
call_id: str = "call_echo",
|
||||
name: str = "echo",
|
||||
args: dict[str, Any],
|
||||
) -> list[dict[str, Any]]:
|
||||
"""把一次 tool_call 拆成分片 chunk，模拟真实流式 tool_call_chunks：
|
||||
|
||||
第 1 片带 id + name + 空 args；后续把 args JSON 拆两段追加。最后附 finish_reason 片。
|
||||
"""
|
||||
args_json = json.dumps(args)
|
||||
mid = max(1, len(args_json) // 2)
|
||||
return [
|
||||
{"index": index, "id": call_id, "name": name, "args": ""},
|
||||
{"index": index, "args": args_json[:mid]},
|
||||
{"index": index, "args": args_json[mid:]},
|
||||
]
|
||||
|
||||
|
||||
def _make_turn(
|
||||
*,
|
||||
content: str = "",
|
||||
reasoning: str = "",
|
||||
tool_calls: list[dict[str, Any]] | None = None,
|
||||
usage: TokenUsage | None = None,
|
||||
) -> list[ChatChunk]:
|
||||
"""构造一轮流式的 chunk 序列。
|
||||
|
||||
若提供 ``tool_calls``（``[{"id":..., "name":..., "args": {...}}]``），则把每个
|
||||
tool_call 拆成分片并附在序列末尾；否则只产出 content/reasoning chunk。
|
||||
"""
|
||||
chunks: list[ChatChunk] = []
|
||||
if reasoning:
|
||||
chunks.append(_chunk(reasoning=reasoning))
|
||||
if content:
|
||||
chunks.append(_chunk(content=content))
|
||||
if tool_calls:
|
||||
for idx, tc in enumerate(tool_calls):
|
||||
chunks.append(
|
||||
_chunk(
|
||||
tool_call_chunks=_split_tool_call_stream(
|
||||
index=idx,
|
||||
call_id=tc.get("id", f"call_{tc['name']}"),
|
||||
name=tc["name"],
|
||||
args=tc.get("args", {}),
|
||||
)
|
||||
)
|
||||
)
|
||||
# usage 放在最后一个 chunk（与 OpenAI streaming 的 final chunk 习惯一致）
|
||||
if chunks:
|
||||
last = chunks[-1]
|
||||
chunks[-1] = ChatChunk(
|
||||
delta_content=last.delta_content,
|
||||
delta_reasoning=last.delta_reasoning,
|
||||
finish_reason="tool_calls" if tool_calls else "stop",
|
||||
usage=usage or TokenUsage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
|
||||
tool_call_chunks=last.tool_call_chunks,
|
||||
)
|
||||
else:
|
||||
chunks.append(
|
||||
_chunk(
|
||||
usage=usage or TokenUsage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
|
||||
finish_reason="stop",
|
||||
)
|
||||
)
|
||||
return chunks
|
||||
|
||||
|
||||
def _make_tool_result(stdout: str = "hello\n", ok: bool = True) -> ToolResult:
|
||||
return ToolResult(
|
||||
ok=ok,
|
||||
exit_code=0 if ok else 1,
|
||||
stdout=stdout,
|
||||
stderr="",
|
||||
duration_ms=3,
|
||||
timeout=False,
|
||||
truncated=False,
|
||||
trace_id="t1",
|
||||
)
|
||||
|
||||
|
||||
def _patch_router_stream(
|
||||
monkeypatch: pytest.MonkeyPatch, turns: list[list[ChatChunk]]
|
||||
) -> MagicMock:
|
||||
"""让 ``ProviderManagerFactory.create().chat_stream()`` 按顺序返回 turns（每轮 chunk 列表）。
|
||||
|
||||
``chat_stream`` 是 ``async def`` 返回 ``AsyncIterator[ChatChunk]``，故这里用 ``AsyncMock``
|
||||
返回值设为对应 async generator。
|
||||
"""
|
||||
call_count = {"n": 0}
|
||||
|
||||
async def fake_chat_stream(_messages, **_kwargs): # noqa: ANN001
|
||||
idx = call_count["n"]
|
||||
call_count["n"] += 1
|
||||
if idx >= len(turns):
|
||||
raise AssertionError(f"chat_stream 被调用了 {idx + 1} 次，但只预期 {len(turns)} 次")
|
||||
async for chunk in _chunk_iter(turns[idx]):
|
||||
yield chunk
|
||||
|
||||
manager = MagicMock()
|
||||
manager.chat_stream = fake_chat_stream
|
||||
monkeypatch.setattr(
|
||||
"app.application.agent.runtime.ProviderManagerFactory.create",
|
||||
lambda: manager,
|
||||
)
|
||||
return manager
|
||||
|
||||
|
||||
def _make_tool_runtime(execute_result: ToolResult | Exception) -> MagicMock:
|
||||
"""构造一个 ToolRuntime 替身；execute 要么返回 ToolResult，要么抛给定异常。"""
|
||||
runtime = MagicMock()
|
||||
if isinstance(execute_result, Exception):
|
||||
runtime.execute = AsyncMock(side_effect=execute_result)
|
||||
else:
|
||||
runtime.execute = AsyncMock(return_value=execute_result)
|
||||
return runtime
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 1. 正常闭环：第一轮流式 tool_call（分片），第二轮纯 content
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
async def test_message_lifecycle_covers_assistant_updates_and_tool_results(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
base_messages: list[BaseMessage],
|
||||
tools: list[dict[str, Any]],
|
||||
) -> None:
|
||||
_patch_router_stream(
|
||||
monkeypatch,
|
||||
[
|
||||
_make_turn(
|
||||
reasoning="先调用工具",
|
||||
tool_calls=[{"id": "call_echo", "name": "echo", "args": {"text": "hi"}}],
|
||||
),
|
||||
_make_turn(content="完成"),
|
||||
],
|
||||
)
|
||||
registry = ExtensionRegistry()
|
||||
events: list[AgentEvent] = []
|
||||
|
||||
async def listen(event: AgentEvent, _run: AgentRunContext) -> None:
|
||||
events.append(event)
|
||||
|
||||
registry.add_event_listener(listen)
|
||||
await run_agent_turn(
|
||||
messages=base_messages,
|
||||
tools=tools,
|
||||
tool_runtime=_make_tool_runtime(_make_tool_result()),
|
||||
extensions=registry,
|
||||
)
|
||||
|
||||
event_types = [event.type for event in events]
|
||||
assert event_types.count("message.started") == 3
|
||||
assert event_types.count("message.ended") == 3
|
||||
assert "message.updated" in event_types
|
||||
tool_ended = next(
|
||||
event
|
||||
for event in events
|
||||
if event.type == "message.ended" and event.payload["role"] == "tool"
|
||||
)
|
||||
assert tool_ended.payload["message"]["tool_call_id"] == "call_echo"
|
||||
|
||||
|
||||
async def test_normal_loop_completes_after_tool_call(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
base_messages: list[BaseMessage],
|
||||
tools: list[dict[str, Any]],
|
||||
) -> None:
|
||||
first = _make_turn(
|
||||
reasoning="我需要用 echo 工具。",
|
||||
tool_calls=[{"id": "call_echo", "name": "echo", "args": {"text": "hello"}}],
|
||||
)
|
||||
second = _make_turn(content="已为你回显：hello", reasoning="已收到工具结果，准备总结。")
|
||||
_patch_router_stream(monkeypatch, [first, second])
|
||||
|
||||
tool_runtime = _make_tool_runtime(_make_tool_result(stdout="hello\n"))
|
||||
|
||||
# 收集 on_chunk / on_tool_event 调用
|
||||
chunks_seen: list[ChatChunk] = []
|
||||
tool_events: list[tuple[str, dict[str, Any]]] = []
|
||||
|
||||
async def on_chunk(chunk: ChatChunk) -> None:
|
||||
chunks_seen.append(chunk)
|
||||
|
||||
async def on_tool_event(event_type: str, payload: dict[str, Any]) -> None:
|
||||
tool_events.append((event_type, payload))
|
||||
|
||||
result = await run_agent_turn(
|
||||
messages=base_messages,
|
||||
tools=tools,
|
||||
tool_runtime=tool_runtime,
|
||||
model_hint="test-model",
|
||||
on_chunk=on_chunk,
|
||||
on_tool_event=on_tool_event,
|
||||
)
|
||||
|
||||
assert isinstance(result, AgentTurnResult)
|
||||
assert result.iterations == 2
|
||||
assert result.truncated is False
|
||||
assert result.content == "已为你回显：hello"
|
||||
# reasoning 用最后一轮（最终回复轮）的累积值
|
||||
assert result.reasoning == "已收到工具结果，准备总结。"
|
||||
assert len(result.tool_history) == 1
|
||||
assert result.tool_history[0]["name"] == "echo"
|
||||
assert result.tool_history[0]["status"] == "success"
|
||||
assert result.tool_history[0]["args"] == {"text": "hello"}
|
||||
assert "hello" in result.tool_history[0]["result"]
|
||||
|
||||
# 流式 chunk 被透传多次（至少两轮的 chunk 数）
|
||||
assert len(chunks_seen) >= 2
|
||||
|
||||
# 工具事件：started + completed 各 1 次
|
||||
assert len(tool_events) == 2
|
||||
assert tool_events[0][0] == "tool.started"
|
||||
assert tool_events[0][1]["tool_name"] == "echo"
|
||||
assert tool_events[0][1]["tool_call_id"] == "call_echo"
|
||||
assert tool_events[1][0] == "tool.completed"
|
||||
assert tool_events[1][1]["status"] == "success"
|
||||
assert "hello" in tool_events[1][1]["result"]
|
||||
|
||||
# messages 正确追加了 AIMessage(tool_calls) + ToolMessage + 最终 AIMessage
|
||||
# base_messages 原长度为 2（SystemMessage + HumanMessage）
|
||||
assert len(base_messages) == 5
|
||||
ai_msg = base_messages[2]
|
||||
assert isinstance(ai_msg, AIMessage)
|
||||
assert ai_msg.tool_calls[0]["name"] == "echo"
|
||||
tool_msg = base_messages[3]
|
||||
assert isinstance(tool_msg, ToolMessage)
|
||||
assert tool_msg.tool_call_id == "call_echo"
|
||||
assert "hello" in tool_msg.content
|
||||
final_msg = base_messages[4]
|
||||
assert isinstance(final_msg, AIMessage)
|
||||
assert final_msg.content == "已为你回显：hello"
|
||||
|
||||
# usage 累计两轮
|
||||
assert result.usage.total_tokens == 30
|
||||
|
||||
# tool_runtime.execute 被调用了一次，参数正确
|
||||
tool_runtime.execute.assert_awaited_once()
|
||||
call_kwargs = tool_runtime.execute.call_args
|
||||
assert call_kwargs.args[0] == "echo"
|
||||
assert call_kwargs.args[1] == {"text": "hello"}
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 2. 工具失败不中断：execute 抛异常，stream 第二轮仍能正常返回 content
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
async def test_tool_failure_does_not_break_loop(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
base_messages: list[BaseMessage],
|
||||
tools: list[dict[str, Any]],
|
||||
) -> None:
|
||||
first = _make_turn(
|
||||
tool_calls=[{"id": "call_echo", "name": "echo", "args": {"text": "hi"}}],
|
||||
)
|
||||
second = _make_turn(content="工具坏了，但我还是回复你。")
|
||||
_patch_router_stream(monkeypatch, [first, second])
|
||||
|
||||
tool_runtime = _make_tool_runtime(RuntimeError("sandbox down"))
|
||||
|
||||
tool_events: list[tuple[str, dict[str, Any]]] = []
|
||||
|
||||
async def on_tool_event(event_type: str, payload: dict[str, Any]) -> None:
|
||||
tool_events.append((event_type, payload))
|
||||
|
||||
result = await run_agent_turn(
|
||||
messages=base_messages,
|
||||
tools=tools,
|
||||
tool_runtime=tool_runtime,
|
||||
on_tool_event=on_tool_event,
|
||||
)
|
||||
|
||||
assert result.iterations == 2
|
||||
assert result.truncated is False
|
||||
assert result.content == "工具坏了，但我还是回复你。"
|
||||
assert len(result.tool_history) == 1
|
||||
assert result.tool_history[0]["status"] == "error"
|
||||
assert "sandbox down" in result.tool_history[0]["result"]
|
||||
|
||||
# 错误信息作为 ToolMessage 回填
|
||||
tool_msg = base_messages[3]
|
||||
assert isinstance(tool_msg, ToolMessage)
|
||||
assert "sandbox down" in tool_msg.content
|
||||
|
||||
# 工具事件 completed 带 status=error
|
||||
assert [e[0] for e in tool_events] == ["tool.started", "tool.completed"]
|
||||
assert tool_events[1][1]["status"] == "error"
|
||||
assert "sandbox down" in tool_events[1][1]["result"]
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 3. 上限截断：stream 每次都返回 tool_call，max_iterations=3
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
async def test_loop_truncates_at_max_iterations(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
base_messages: list[BaseMessage],
|
||||
tools: list[dict[str, Any]],
|
||||
) -> None:
|
||||
# 每次都返回 tool_call，让循环无法终止
|
||||
looping_turn = _make_turn(
|
||||
tool_calls=[{"id": "call_echo", "name": "echo", "args": {"text": "loop"}}],
|
||||
)
|
||||
_patch_router_stream(monkeypatch, [looping_turn, looping_turn, looping_turn])
|
||||
|
||||
tool_runtime = _make_tool_runtime(_make_tool_result(stdout="loop\n"))
|
||||
|
||||
result = await run_agent_turn(
|
||||
messages=base_messages,
|
||||
tools=tools,
|
||||
tool_runtime=tool_runtime,
|
||||
max_iterations=3,
|
||||
)
|
||||
|
||||
assert result.iterations == 3
|
||||
assert result.truncated is True
|
||||
assert len(result.tool_history) == 3
|
||||
assert all(h["status"] == "success" for h in result.tool_history)
|
||||
# messages 累积：base 2 + 3 轮 × (AIMessage + ToolMessage) = 8
|
||||
assert len(base_messages) == 8
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 4. 无工具直接返回：stream 第一次就是纯 content
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
async def test_no_tool_call_returns_immediately(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
base_messages: list[BaseMessage],
|
||||
tools: list[dict[str, Any]],
|
||||
) -> None:
|
||||
direct = _make_turn(content="直接回答，不需要工具。")
|
||||
_patch_router_stream(monkeypatch, [direct])
|
||||
|
||||
tool_runtime = _make_tool_runtime(_make_tool_result())
|
||||
|
||||
tool_events: list[tuple[str, dict[str, Any]]] = []
|
||||
|
||||
async def on_tool_event(event_type: str, payload: dict[str, Any]) -> None:
|
||||
tool_events.append((event_type, payload))
|
||||
|
||||
result = await run_agent_turn(
|
||||
messages=base_messages,
|
||||
tools=tools,
|
||||
tool_runtime=tool_runtime,
|
||||
on_tool_event=on_tool_event,
|
||||
)
|
||||
|
||||
assert result.iterations == 1
|
||||
assert result.truncated is False
|
||||
assert result.content == "直接回答，不需要工具。"
|
||||
assert result.tool_history == []
|
||||
# messages 只追加了一条 AIMessage
|
||||
assert len(base_messages) == 3
|
||||
assert isinstance(base_messages[2], AIMessage)
|
||||
# execute 从未被调用
|
||||
tool_runtime.execute.assert_not_awaited()
|
||||
# 无工具事件
|
||||
assert tool_events == []
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 5. 工具转换辅助：to_ai_provider_tools / default_tools 基本正确
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
def test_to_ai_provider_tools_produces_openai_function_schema() -> None:
|
||||
specs = list(builtin_specs().values())
|
||||
tools = to_ai_provider_tools(specs)
|
||||
assert len(tools) == len(specs)
|
||||
for tool, spec in zip(tools, specs, strict=True):
|
||||
assert tool["type"] == "function"
|
||||
fn = tool["function"]
|
||||
assert fn["name"] == spec.name
|
||||
assert fn["description"] == spec.description
|
||||
assert fn["parameters"]["type"] == "object"
|
||||
|
||||
|
||||
def test_default_tools_returns_core_only_by_default() -> None:
|
||||
"""default_tools 默认只返回 5 个核心沙箱工具（不含联网工具）。"""
|
||||
tools = default_tools()
|
||||
names = {t["function"]["name"] for t in tools}
|
||||
assert names == {"datetime", "read_file", "list_dir", "write_file", "shell"}
|
||||
assert len(tools) == 5
|
||||
|
||||
|
||||
def test_default_tools_returns_core_only() -> None:
|
||||
"""default_tools 返回 5 个核心沙箱工具。"""
|
||||
tools = default_tools()
|
||||
names = {t["function"]["name"] for t in tools}
|
||||
assert names == {"datetime", "read_file", "list_dir", "write_file", "shell"}
|
||||
assert len(tools) == 5
|
||||
|
||||
|
||||
def test_run_agent_turn_rejects_non_positive_max_iterations(
|
||||
base_messages: list[BaseMessage], tools: list[dict[str, Any]]
|
||||
) -> None:
|
||||
tool_runtime = _make_tool_runtime(_make_tool_result())
|
||||
with pytest.raises(ValueError):
|
||||
import asyncio
|
||||
|
||||
asyncio.get_event_loop().run_until_complete(
|
||||
run_agent_turn(
|
||||
messages=base_messages,
|
||||
tools=tools,
|
||||
tool_runtime=tool_runtime,
|
||||
max_iterations=0,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 6. ToolCallAccumulator 单元测试
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
def test_accumulator_empty_drain_returns_empty_list() -> None:
|
||||
acc = ToolCallAccumulator()
|
||||
drained = acc.drain()
|
||||
assert drained == []
|
||||
|
||||
|
||||
def test_accumulator_assembles_fragmented_args_across_chunks() -> None:
|
||||
"""args 分片到达也能拼回完整 JSON。"""
|
||||
acc = ToolCallAccumulator()
|
||||
# 片段：先到 id+name，args 分 "hello"、""、" world" 三段
|
||||
acc.add_chunk(_chunk(tool_call_chunks=[{"index": 0, "id": "c1", "name": "echo", "args": ""}]))
|
||||
acc.add_chunk(_chunk(tool_call_chunks=[{"index": 0, "args": "hello"}]))
|
||||
acc.add_chunk(_chunk(tool_call_chunks=[{"index": 0, "args": ""}]))
|
||||
acc.add_chunk(_chunk(tool_call_chunks=[{"index": 0, "args": " world"}]))
|
||||
|
||||
drained = acc.drain()
|
||||
assert len(drained) == 1
|
||||
entry = drained[0]
|
||||
assert entry["index"] == 0
|
||||
assert entry["id"] == "c1"
|
||||
assert entry["name"] == "echo"
|
||||
assert entry["args"] == "hello world"
|
||||
|
||||
|
||||
def test_accumulator_parallel_multiple_tool_calls() -> None:
|
||||
"""index 0 和 1 的 tool_call 并行累积。"""
|
||||
acc = ToolCallAccumulator()
|
||||
acc.add_chunk(
|
||||
_chunk(
|
||||
tool_call_chunks=[
|
||||
{"index": 0, "id": "c1", "name": "echo", "args": '{"text"'},
|
||||
{"index": 1, "id": "c2", "name": "datetime", "args": "{}"},
|
||||
]
|
||||
)
|
||||
)
|
||||
acc.add_chunk(
|
||||
_chunk(
|
||||
tool_call_chunks=[
|
||||
{"index": 0, "args": ': "a"}'},
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
drained = acc.drain()
|
||||
assert len(drained) == 2
|
||||
assert drained[0]["index"] == 0
|
||||
assert drained[0]["name"] == "echo"
|
||||
assert drained[0]["args"] == '{"text": "a"}'
|
||||
assert drained[1]["index"] == 1
|
||||
assert drained[1]["name"] == "datetime"
|
||||
assert drained[1]["args"] == "{}"
|
||||
|
||||
|
||||
def test_accumulator_args_arrive_before_name() -> None:
|
||||
"""分片中先到 args 片段、后到 name，也能正确合并。"""
|
||||
acc = ToolCallAccumulator()
|
||||
acc.add_chunk(_chunk(tool_call_chunks=[{"index": 0, "args": '{"x": '}]))
|
||||
acc.add_chunk(_chunk(tool_call_chunks=[{"index": 0, "id": "c1", "name": "calc", "args": "1"}]))
|
||||
acc.add_chunk(_chunk(tool_call_chunks=[{"index": 0, "args": "}"}]))
|
||||
|
||||
drained = acc.drain()
|
||||
assert len(drained) == 1
|
||||
entry = drained[0]
|
||||
assert entry["id"] == "c1"
|
||||
assert entry["name"] == "calc"
|
||||
assert entry["args"] == '{"x": 1}'
|
||||
|
||||
|
||||
def test_accumulator_ignores_chunks_without_tool_call_chunks() -> None:
|
||||
"""非 tool_call chunk（只有 content）不会影响累加器状态。"""
|
||||
acc = ToolCallAccumulator()
|
||||
acc.add_chunk(_chunk(content="hello"))
|
||||
acc.add_chunk(_chunk(reasoning="thinking"))
|
||||
acc.add_chunk(_chunk())
|
||||
drained = acc.drain()
|
||||
assert drained == []
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 7. 工具确认挂起/恢复（阶段3 步骤1）
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
def _make_confirmation_tool_runtime(
|
||||
*, exc: Exception | None = None, result: ToolResult | None = None
|
||||
) -> MagicMock:
|
||||
"""构造 ToolRuntime：先抛 RequiresConfirmation（或给定 exc），再返回 result。
|
||||
|
||||
用于"批准后重试 confirmed=True 成功"的场景。result 缺省为成功 ToolResult。
|
||||
"""
|
||||
result = result or _make_tool_result(stdout="done\n")
|
||||
if exc is not None:
|
||||
runtime = MagicMock()
|
||||
runtime.execute = AsyncMock(side_effect=[exc, result])
|
||||
return runtime
|
||||
raise ValueError("必须提供 exc")
|
||||
|
||||
|
||||
def _make_deny_only_tool_runtime(exc: Exception) -> MagicMock:
|
||||
"""构造 ToolRuntime：只抛一次 RequiresConfirmation（拒绝/超时不重试）。"""
|
||||
runtime = MagicMock()
|
||||
runtime.execute = AsyncMock(side_effect=exc)
|
||||
return runtime
|
||||
|
||||
|
||||
async def test_context_transformer_runs_for_every_model_request_without_persisting(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
base_messages: list[BaseMessage],
|
||||
tools: list[dict[str, Any]],
|
||||
) -> None:
|
||||
first = _make_turn(
|
||||
tool_calls=[{"id": "call_echo", "name": "echo", "args": {"text": "x"}}]
|
||||
)
|
||||
second = _make_turn(content="完成")
|
||||
manager = _patch_router_stream(monkeypatch, [first, second])
|
||||
runtime = MagicMock()
|
||||
runtime.execute = AsyncMock(return_value=_make_tool_result())
|
||||
registry = ExtensionRegistry()
|
||||
transformed_calls: list[int] = []
|
||||
|
||||
async def transform(context: AgentContext, _run: AgentRunContext) -> AgentContext:
|
||||
transformed_calls.append(len(context.messages))
|
||||
return AgentContext(
|
||||
(*context.messages, HumanMessage(content="临时注入")),
|
||||
context.tools,
|
||||
context.model_hint,
|
||||
)
|
||||
|
||||
registry.add_context_transformer(transform)
|
||||
await run_agent_turn(
|
||||
messages=base_messages,
|
||||
tools=tools,
|
||||
tool_runtime=runtime,
|
||||
extensions=registry,
|
||||
)
|
||||
|
||||
assert transformed_calls == [2, 4]
|
||||
assert all(message.content != "临时注入" for message in base_messages)
|
||||
assert manager.chat_stream is not None
|
||||
|
||||
|
||||
async def test_max_tokens_marks_agent_result_truncated(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
base_messages: list[BaseMessage],
|
||||
tools: list[dict[str, Any]],
|
||||
) -> None:
|
||||
turn = _make_turn(content="未完成")
|
||||
turn[-1].finish_reason = "length"
|
||||
_patch_router_stream(monkeypatch, [turn])
|
||||
|
||||
result = await run_agent_turn(
|
||||
messages=base_messages,
|
||||
tools=tools,
|
||||
tool_runtime=MagicMock(),
|
||||
)
|
||||
|
||||
assert result.content == "未完成"
|
||||
assert result.truncated is True
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 7.1 确认 + 批准恢复：execute 第一次抛 RequiresConfirmation → 发事件 →
|
||||
# resolve(cid, True) → confirmed=True 重试成功 → 第二轮纯 content
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
async def test_confirmation_approved_resumes_and_succeeds(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
base_messages: list[BaseMessage],
|
||||
tools: list[dict[str, Any]],
|
||||
) -> None:
|
||||
first = _make_turn(
|
||||
tool_calls=[{"id": "call_write", "name": "write_file", "args": {"path": "a.txt"}}],
|
||||
)
|
||||
second = _make_turn(content="已写入文件。")
|
||||
_patch_router_stream(monkeypatch, [first, second])
|
||||
|
||||
tool_runtime = _make_confirmation_tool_runtime(
|
||||
exc=RequiresConfirmation("策略 ask 需用户确认", risk_level="L2")
|
||||
)
|
||||
registry = ConfirmationRegistry()
|
||||
|
||||
tool_events: list[tuple[str, dict[str, Any]]] = []
|
||||
|
||||
async def on_tool_event(event_type: str, payload: dict[str, Any]) -> None:
|
||||
tool_events.append((event_type, payload))
|
||||
# 收到 requires_confirmation 事件后，手动触发批准恢复循环
|
||||
if event_type == "tool.requires_confirmation":
|
||||
cid = payload["confirmation_id"]
|
||||
asyncio.create_task(registry.resolve(cid, True))
|
||||
|
||||
result = await run_agent_turn(
|
||||
messages=base_messages,
|
||||
tools=tools,
|
||||
tool_runtime=tool_runtime,
|
||||
on_tool_event=on_tool_event,
|
||||
confirmation_registry=registry,
|
||||
run_id="test-run-1",
|
||||
character_id="test-char-1",
|
||||
)
|
||||
|
||||
assert result.iterations == 2
|
||||
assert result.content == "已写入文件。"
|
||||
assert len(result.tool_history) == 1
|
||||
assert result.tool_history[0]["status"] == "success"
|
||||
|
||||
# 事件顺序：started（初次） → requires_confirmation → started（重试，带 confirmed=True）
|
||||
# → completed(success)
|
||||
types = [e[0] for e in tool_events]
|
||||
assert types == [
|
||||
"tool.started",
|
||||
"tool.requires_confirmation",
|
||||
"tool.started",
|
||||
"tool.completed",
|
||||
]
|
||||
req_event = tool_events[1][1]
|
||||
assert req_event["tool_name"] == "write_file"
|
||||
assert req_event["risk_level"] == "L2"
|
||||
assert req_event["confirmation_id"].startswith("confirm_")
|
||||
# 重试时带 confirmed=True 标记
|
||||
assert tool_events[2][1].get("confirmed") is True
|
||||
assert tool_events[3][1]["status"] == "success"
|
||||
|
||||
# execute 被调用 2 次：第一次 confirmed=False（默认），第二次 confirmed=True
|
||||
assert tool_runtime.execute.await_count == 2
|
||||
first_call = tool_runtime.execute.call_args_list[0]
|
||||
second_call = tool_runtime.execute.call_args_list[1]
|
||||
assert first_call.kwargs["confirmed"] is False
|
||||
assert second_call.kwargs["confirmed"] is True
|
||||
|
||||
# registry 已清理（resolve 后 Future 出列）
|
||||
assert not registry.has(req_event["confirmation_id"])
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 7.2 确认 + 拒绝：resolve(cid, False) → 回填"用户拒绝" → 循环继续
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
async def test_confirmation_denied_backfills_and_continues(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
base_messages: list[BaseMessage],
|
||||
tools: list[dict[str, Any]],
|
||||
) -> None:
|
||||
first = _make_turn(
|
||||
tool_calls=[{"id": "call_shell", "name": "shell", "args": {"command": "rm -rf /"}}],
|
||||
)
|
||||
second = _make_turn(content="好的，我不执行了。")
|
||||
_patch_router_stream(monkeypatch, [first, second])
|
||||
|
||||
tool_runtime = _make_deny_only_tool_runtime(
|
||||
RequiresConfirmation("策略 always_ask 需用户确认", risk_level="L3")
|
||||
)
|
||||
registry = ConfirmationRegistry()
|
||||
|
||||
tool_events: list[tuple[str, dict[str, Any]]] = []
|
||||
|
||||
async def on_tool_event(event_type: str, payload: dict[str, Any]) -> None:
|
||||
tool_events.append((event_type, payload))
|
||||
if event_type == "tool.requires_confirmation":
|
||||
cid = payload["confirmation_id"]
|
||||
asyncio.create_task(registry.resolve(cid, False))
|
||||
|
||||
result = await run_agent_turn(
|
||||
messages=base_messages,
|
||||
tools=tools,
|
||||
tool_runtime=tool_runtime,
|
||||
on_tool_event=on_tool_event,
|
||||
confirmation_registry=registry,
|
||||
run_id="test-run-2",
|
||||
character_id="test-char-2",
|
||||
confirmation_timeout=1.0,
|
||||
)
|
||||
|
||||
assert result.iterations == 2
|
||||
assert result.content == "好的，我不执行了。"
|
||||
assert len(result.tool_history) == 1
|
||||
assert result.tool_history[0]["status"] == "denied"
|
||||
assert "用户拒绝" in result.tool_history[0]["result"]
|
||||
|
||||
# ToolMessage 回填的是"用户拒绝"
|
||||
tool_msg = base_messages[3]
|
||||
assert isinstance(tool_msg, ToolMessage)
|
||||
assert "用户拒绝" in tool_msg.content
|
||||
|
||||
# 事件：started → requires_confirmation → completed(status=denied)
|
||||
types = [e[0] for e in tool_events]
|
||||
assert types == ["tool.started", "tool.requires_confirmation", "tool.completed"]
|
||||
assert tool_events[-1][1]["status"] == "denied"
|
||||
|
||||
# execute 只被调用 1 次（拒绝不重试）
|
||||
assert tool_runtime.execute.await_count == 1
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 7.3 确认超时按拒绝：confirmation_timeout 极短，不 resolve → 超时 → 回填 → 继续
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
async def test_confirmation_timeout_treated_as_denied(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
base_messages: list[BaseMessage],
|
||||
tools: list[dict[str, Any]],
|
||||
) -> None:
|
||||
first = _make_turn(
|
||||
tool_calls=[{"id": "call_write", "name": "write_file", "args": {"path": "x.txt"}}],
|
||||
)
|
||||
second = _make_turn(content="收到，跳过写入。")
|
||||
_patch_router_stream(monkeypatch, [first, second])
|
||||
|
||||
tool_runtime = _make_deny_only_tool_runtime(RequiresConfirmation("需确认", risk_level="L2"))
|
||||
registry = ConfirmationRegistry()
|
||||
|
||||
result = await run_agent_turn(
|
||||
messages=base_messages,
|
||||
tools=tools,
|
||||
tool_runtime=tool_runtime,
|
||||
on_tool_event=None, # 不需要事件断言
|
||||
confirmation_registry=registry,
|
||||
confirmation_timeout=1.5, # > poll_interval（0.5s）确保至少完成一次轮询
|
||||
)
|
||||
|
||||
assert result.iterations == 2
|
||||
assert result.content == "收到，跳过写入。"
|
||||
assert len(result.tool_history) == 1
|
||||
assert result.tool_history[0]["status"] == "denied"
|
||||
assert "超时" in result.tool_history[0]["result"]
|
||||
# execute 只被调用 1 次（超时不重试）
|
||||
assert tool_runtime.execute.await_count == 1
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 7.4 registry=None 降级：execute 抛 RequiresConfirmation → 回填"未启用" → 继续
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
async def test_confirmation_degraded_when_registry_none(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
base_messages: list[BaseMessage],
|
||||
tools: list[dict[str, Any]],
|
||||
) -> None:
|
||||
first = _make_turn(
|
||||
tool_calls=[{"id": "call_shell", "name": "shell", "args": {"command": "ls"}}],
|
||||
)
|
||||
second = _make_turn(content="好的，跳过了。")
|
||||
_patch_router_stream(monkeypatch, [first, second])
|
||||
|
||||
tool_runtime = _make_deny_only_tool_runtime(RequiresConfirmation("需确认", risk_level="L3"))
|
||||
|
||||
tool_events: list[tuple[str, dict[str, Any]]] = []
|
||||
|
||||
async def on_tool_event(event_type: str, payload: dict[str, Any]) -> None:
|
||||
tool_events.append((event_type, payload))
|
||||
|
||||
result = await run_agent_turn(
|
||||
messages=base_messages,
|
||||
tools=tools,
|
||||
tool_runtime=tool_runtime,
|
||||
on_tool_event=on_tool_event,
|
||||
# 不传 confirmation_registry → 降级路径
|
||||
)
|
||||
|
||||
assert result.iterations == 2
|
||||
assert result.content == "好的，跳过了。"
|
||||
assert len(result.tool_history) == 1
|
||||
assert result.tool_history[0]["status"] == "skipped"
|
||||
assert "需要用户确认" in result.tool_history[0]["result"]
|
||||
assert "未启用" in result.tool_history[0]["result"]
|
||||
|
||||
# 事件：started → completed(status=skipped)（不发 requires_confirmation）
|
||||
types = [e[0] for e in tool_events]
|
||||
assert types == ["tool.started", "tool.completed"]
|
||||
assert tool_events[-1][1]["status"] == "skipped"
|
||||
|
||||
# execute 只被调用 1 次
|
||||
assert tool_runtime.execute.await_count == 1
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 7.5 ConfirmationRegistry 单元测试
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
async def test_confirmation_registry_resolve_returns_false_for_unknown() -> None:
|
||||
registry = ConfirmationRegistry()
|
||||
assert await registry.resolve("unknown_id", True) is False
|
||||
assert registry.has("unknown_id") is False
|
||||
|
||||
|
||||
async def test_confirmation_registry_resolves_only_requested_confirmation() -> None:
|
||||
from app.db.session import AsyncSessionLocal
|
||||
from app.models.safety import PendingConfirmationRecord
|
||||
|
||||
registry = ConfirmationRegistry()
|
||||
async with AsyncSessionLocal() as session:
|
||||
session.add_all(
|
||||
[
|
||||
PendingConfirmationRecord(
|
||||
id="confirm_first",
|
||||
run_id="run_shared",
|
||||
tool="first",
|
||||
input_json="{}",
|
||||
risk_level="L2",
|
||||
status="pending",
|
||||
),
|
||||
PendingConfirmationRecord(
|
||||
id="confirm_second",
|
||||
run_id="run_shared",
|
||||
tool="second",
|
||||
input_json="{}",
|
||||
risk_level="L2",
|
||||
status="pending",
|
||||
),
|
||||
]
|
||||
)
|
||||
await session.commit()
|
||||
|
||||
assert await registry.resolve(
|
||||
"run_shared", True, confirmation_id="confirm_second"
|
||||
) is True
|
||||
|
||||
async with AsyncSessionLocal() as session:
|
||||
first = await session.get(PendingConfirmationRecord, "confirm_first")
|
||||
second = await session.get(PendingConfirmationRecord, "confirm_second")
|
||||
assert first is not None and first.status == "pending"
|
||||
assert second is not None and second.status == "approved"
|
||||
|
||||
|
||||
async def test_confirmation_registry_resolve_returns_false_after_done() -> None:
|
||||
from app.db.session import AsyncSessionLocal
|
||||
from app.models.safety import PendingConfirmationRecord
|
||||
|
||||
registry = ConfirmationRegistry()
|
||||
assert registry.has("cid_1") is False
|
||||
|
||||
# 创建 PendingConfirmationRecord（wait/resolve 的 DB 前提）
|
||||
async with AsyncSessionLocal() as session:
|
||||
session.add(
|
||||
PendingConfirmationRecord(
|
||||
id="rec_cid_1",
|
||||
run_id="cid_1",
|
||||
tool="test",
|
||||
input_json="{}",
|
||||
risk_level="L2",
|
||||
status="pending",
|
||||
)
|
||||
)
|
||||
await session.commit()
|
||||
|
||||
running = asyncio.create_task(registry.wait("cid_1", poll_interval=0.1, timeout=0.5))
|
||||
await asyncio.sleep(0.1)
|
||||
assert registry.has("cid_1") is True
|
||||
resolved = await registry.resolve("cid_1", True)
|
||||
assert resolved is True
|
||||
result = await running
|
||||
assert result is True
|
||||
# 已处理，再次 resolve 返回 False
|
||||
assert await registry.resolve("cid_1", True) is False
|
||||
assert registry.has("cid_1") is False
|
||||
@@ -6,7 +6,6 @@ import pytest
|
||||
from httpx import ASGITransport, AsyncClient
|
||||
|
||||
from app.api.v1.deps import require_admin
|
||||
from app.application.agent.runtime import AgentTurnResult
|
||||
from app.main import app
|
||||
|
||||
|
||||
|
||||
Loaded 100 of 129 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user