306 lines
11 KiB
Python
306 lines
11 KiB
Python
"""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
|