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:
1 parent
0d383725b5
commit
462d570751
129 files changed
+81
-15867
No files matched your search
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
|
||||
|
||||
|
||||
|
||||
@@ -1,134 +0,0 @@
|
||||
"""Agent 循环工具确认 approve/deny 端点测试。
|
||||
|
||||
覆盖:
|
||||
- approve / deny 成功 resolve 记录(200, status=resolved)
|
||||
- 重复 resolve 返回 404(记录已消费)
|
||||
- 记录不存在的场景 → 404
|
||||
- registry 未挂载 → 503
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import AsyncGenerator
|
||||
|
||||
import pytest
|
||||
from httpx import ASGITransport, AsyncClient
|
||||
from sqlalchemy import select
|
||||
|
||||
from app.application.agent.confirmation import ConfirmationRegistry
|
||||
from app.db.session import AsyncSessionLocal
|
||||
from app.main import app
|
||||
from app.models.safety import PendingConfirmationRecord
|
||||
|
||||
_RUN_ID = "run_test_confirmation_0001"
|
||||
_CONFIRM = "confirm_test_confirmation_0001"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def client(auth_headers: dict[str, str]) -> AsyncGenerator[AsyncClient, None]:
|
||||
"""挂载 ConfirmationRegistry 到 app.state。"""
|
||||
registry = ConfirmationRegistry()
|
||||
app.state.confirmation_registry = registry
|
||||
transport = ASGITransport(app=app)
|
||||
try:
|
||||
async with AsyncClient(
|
||||
transport=transport, base_url="http://test", headers=auth_headers
|
||||
) as http_client:
|
||||
yield http_client
|
||||
finally:
|
||||
app.state.confirmation_registry = ConfirmationRegistry()
|
||||
|
||||
|
||||
async def _seed_pending_record() -> None:
|
||||
"""在 DB 中插入一条待确认记录。"""
|
||||
async with AsyncSessionLocal() as session:
|
||||
record = PendingConfirmationRecord(
|
||||
id=_CONFIRM,
|
||||
tool="shell",
|
||||
input_json='{"cmd": "echo hi"}',
|
||||
risk_level="L3",
|
||||
reason="高风险命令需确认",
|
||||
run_id=_RUN_ID,
|
||||
character_id="character_test",
|
||||
status="pending",
|
||||
)
|
||||
session.add(record)
|
||||
await session.commit()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_approve_resolves_future(client: AsyncClient) -> None:
|
||||
await _seed_pending_record()
|
||||
|
||||
resp = await client.post(f"/api/v1/runs/{_RUN_ID}/confirmations/{_CONFIRM}/approve")
|
||||
assert resp.status_code == 200
|
||||
assert resp.json() == {"approved": True, "status": "resolved"}
|
||||
|
||||
# 验证 DB 状态已更新
|
||||
async with AsyncSessionLocal() as session:
|
||||
record = (
|
||||
await session.execute(
|
||||
select(PendingConfirmationRecord).where(PendingConfirmationRecord.id == _CONFIRM)
|
||||
)
|
||||
).scalar_one_or_none()
|
||||
assert record is not None
|
||||
assert record.status == "approved"
|
||||
assert record.resolved_at is not None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deny_resolves_future(client: AsyncClient) -> None:
|
||||
await _seed_pending_record()
|
||||
|
||||
resp = await client.post(f"/api/v1/runs/{_RUN_ID}/confirmations/{_CONFIRM}/deny")
|
||||
assert resp.status_code == 200
|
||||
assert resp.json() == {"approved": False, "status": "resolved"}
|
||||
|
||||
async with AsyncSessionLocal() as session:
|
||||
record = (
|
||||
await session.execute(
|
||||
select(PendingConfirmationRecord).where(PendingConfirmationRecord.id == _CONFIRM)
|
||||
)
|
||||
).scalar_one_or_none()
|
||||
assert record is not None
|
||||
assert record.status == "denied"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_idempotent_second_call_returns_404(client: AsyncClient) -> None:
|
||||
"""重复 approve:第一次成功,第二次返回 404(记录已 resolve)。"""
|
||||
await _seed_pending_record()
|
||||
|
||||
first = await client.post(f"/api/v1/runs/{_RUN_ID}/confirmations/{_CONFIRM}/approve")
|
||||
assert first.status_code == 200
|
||||
|
||||
second = await client.post(f"/api/v1/runs/{_RUN_ID}/confirmations/{_CONFIRM}/approve")
|
||||
assert second.status_code == 404
|
||||
assert second.json()["error"]["code"] == "confirmation_not_found"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_missing_confirmation_returns_404(client: AsyncClient) -> None:
|
||||
"""确认项不存在 → 404。"""
|
||||
resp = await client.post(f"/api/v1/runs/{_RUN_ID}/confirmations/{_CONFIRM}/deny")
|
||||
assert resp.status_code == 404
|
||||
assert resp.json()["error"]["code"] == "confirmation_not_found"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_registry_not_configured_returns_503() -> None:
|
||||
app.state.confirmation_registry = None
|
||||
try:
|
||||
transport = ASGITransport(app=app)
|
||||
async with AsyncClient(transport=transport, base_url="http://test") as http_client:
|
||||
resp = await http_client.post(
|
||||
f"/api/v1/runs/{_RUN_ID}/confirmations/{_CONFIRM}/approve"
|
||||
)
|
||||
assert resp.status_code == 503
|
||||
assert resp.json()["error"]["code"] == "service_unavailable"
|
||||
finally:
|
||||
app.state.confirmation_registry = ConfirmationRegistry()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
Reference in new issue
Block a user