mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
feat: 视频计费增强与影子计费系统
This commit is contained in:
fawney19
@@ -46,7 +46,7 @@ def test_cli_format_convertible_when_converter_supports_full() -> None:
|
||||
|
||||
|
||||
def test_global_switch_disabled_blocks_conversion() -> None:
|
||||
"""全局开关关闭时（环境变量 FORMAT_CONVERSION_ENABLED=false）阻止转换"""
|
||||
"""全局开关关闭时阻止转换"""
|
||||
ok, needs_conv, reason = is_format_compatible(
|
||||
"claude:chat",
|
||||
"openai:chat",
|
||||
@@ -57,7 +57,7 @@ def test_global_switch_disabled_blocks_conversion() -> None:
|
||||
)
|
||||
assert ok is False
|
||||
assert needs_conv is False
|
||||
assert reason and ("全局" in reason or "FORMAT_CONVERSION_ENABLED" in reason)
|
||||
assert reason and "格式转换已禁用" in reason
|
||||
|
||||
|
||||
def test_endpoint_config_none_blocks_conversion() -> None:
|
||||
@@ -223,7 +223,7 @@ def test_claude_cli_to_claude_blocked_when_global_switch_disabled() -> None:
|
||||
registry=MagicMock(),
|
||||
)
|
||||
assert ok is False
|
||||
assert reason and ("全局" in reason or "FORMAT_CONVERSION_ENABLED" in reason)
|
||||
assert reason and "格式转换已禁用" in reason
|
||||
|
||||
|
||||
def test_claude_cli_to_claude_blocked_when_endpoint_not_configured() -> None:
|
||||
@@ -313,7 +313,7 @@ def test_openai_cli_to_openai_fails_without_converter() -> None:
|
||||
|
||||
|
||||
def test_openai_cli_to_openai_blocked_when_global_switch_disabled() -> None:
|
||||
"""同族转换（OPENAI/OPENAI_CLI）也受全局开关限制（环境变量 FORMAT_CONVERSION_ENABLED=false）"""
|
||||
"""同族转换（OPENAI/OPENAI_CLI）也受全局开关限制"""
|
||||
registry = MagicMock()
|
||||
registry.can_convert_full.return_value = True
|
||||
|
||||
@@ -327,7 +327,7 @@ def test_openai_cli_to_openai_blocked_when_global_switch_disabled() -> None:
|
||||
)
|
||||
assert ok is False
|
||||
assert needs_conv is False
|
||||
assert reason and ("全局" in reason or "FORMAT_CONVERSION_ENABLED" in reason)
|
||||
assert reason and "格式转换已禁用" in reason
|
||||
|
||||
|
||||
def test_openai_cli_to_openai_blocked_when_endpoint_disabled() -> None:
|
||||
|
||||
301
tests/services/billing/test_default_rules.py
Normal file
301
tests/services/billing/test_default_rules.py
Normal file
@@ -0,0 +1,301 @@
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from src.models.database import GlobalModel, Model
|
||||
from src.services.billing.default_rules import DefaultBillingRuleGenerator
|
||||
from src.services.billing.formula_engine import FormulaEngine
|
||||
from src.services.billing.rule_service import BillingRuleService
|
||||
|
||||
|
||||
class TestDefaultBillingRuleGenerator:
|
||||
def test_default_rule_basic_chat_cost(self) -> None:
|
||||
global_model = GlobalModel(
|
||||
name="test-model",
|
||||
display_name="Test Model",
|
||||
is_active=True,
|
||||
default_price_per_request=0.01,
|
||||
default_tiered_pricing={
|
||||
"tiers": [
|
||||
{
|
||||
"up_to": None,
|
||||
"input_price_per_1m": 3.0,
|
||||
"output_price_per_1m": 15.0,
|
||||
# cache prices intentionally omitted (legacy derives from input price)
|
||||
}
|
||||
]
|
||||
},
|
||||
)
|
||||
|
||||
rule = DefaultBillingRuleGenerator.generate_for_model(
|
||||
global_model=global_model,
|
||||
model=None,
|
||||
task_type="chat",
|
||||
)
|
||||
|
||||
engine = FormulaEngine()
|
||||
result = engine.evaluate(
|
||||
expression=rule.expression,
|
||||
variables=rule.variables,
|
||||
dimensions={
|
||||
"input_tokens": 1000,
|
||||
"output_tokens": 500,
|
||||
"cache_creation_tokens": 200,
|
||||
"cache_read_tokens": 300,
|
||||
"request_count": 1,
|
||||
# tier key
|
||||
"total_input_context": 1000 + 300,
|
||||
},
|
||||
dimension_mappings=rule.dimension_mappings,
|
||||
strict_mode=True,
|
||||
)
|
||||
|
||||
assert result.status == "complete"
|
||||
assert abs(float(result.cost) - 0.02134) < 1e-9
|
||||
|
||||
def test_default_rule_tiered_pricing_uses_total_input_context(self) -> None:
|
||||
global_model = GlobalModel(
|
||||
name="tiered-model",
|
||||
display_name="Tiered Model",
|
||||
is_active=True,
|
||||
default_price_per_request=None,
|
||||
default_tiered_pricing={
|
||||
"tiers": [
|
||||
{"up_to": 200000, "input_price_per_1m": 3.0, "output_price_per_1m": 15.0},
|
||||
{"up_to": None, "input_price_per_1m": 1.5, "output_price_per_1m": 7.5},
|
||||
]
|
||||
},
|
||||
)
|
||||
|
||||
rule = DefaultBillingRuleGenerator.generate_for_model(
|
||||
global_model=global_model,
|
||||
model=None,
|
||||
task_type="chat",
|
||||
)
|
||||
|
||||
engine = FormulaEngine()
|
||||
result = engine.evaluate(
|
||||
expression=rule.expression,
|
||||
variables=rule.variables,
|
||||
dimensions={
|
||||
"input_tokens": 250000,
|
||||
"output_tokens": 10000,
|
||||
"cache_creation_tokens": 0,
|
||||
"cache_read_tokens": 0,
|
||||
"request_count": 1,
|
||||
"total_input_context": 250000,
|
||||
},
|
||||
dimension_mappings=rule.dimension_mappings,
|
||||
strict_mode=True,
|
||||
)
|
||||
|
||||
# 250000 * 1.5 / 1M = 0.375
|
||||
# 10000 * 7.5 / 1M = 0.075
|
||||
assert result.status == "complete"
|
||||
assert abs(float(result.cost) - 0.45) < 1e-9
|
||||
|
||||
def test_default_rule_cache_ttl_pricing_overrides_cache_read_price(self) -> None:
|
||||
global_model = GlobalModel(
|
||||
name="ttl-model",
|
||||
display_name="TTL Model",
|
||||
is_active=True,
|
||||
default_price_per_request=0.0,
|
||||
default_tiered_pricing={
|
||||
"tiers": [
|
||||
{
|
||||
"up_to": None,
|
||||
"input_price_per_1m": 3.0,
|
||||
"output_price_per_1m": 15.0,
|
||||
"cache_read_price_per_1m": 0.3,
|
||||
"cache_ttl_pricing": [
|
||||
{"ttl_minutes": 5, "cache_read_price_per_1m": 0.3},
|
||||
{"ttl_minutes": 60, "cache_read_price_per_1m": 0.5},
|
||||
],
|
||||
}
|
||||
]
|
||||
},
|
||||
)
|
||||
|
||||
rule = DefaultBillingRuleGenerator.generate_for_model(
|
||||
global_model=global_model,
|
||||
model=None,
|
||||
task_type="chat",
|
||||
)
|
||||
|
||||
engine = FormulaEngine()
|
||||
result = engine.evaluate(
|
||||
expression=rule.expression,
|
||||
variables=rule.variables,
|
||||
dimensions={
|
||||
"input_tokens": 0,
|
||||
"output_tokens": 0,
|
||||
"cache_creation_tokens": 0,
|
||||
"cache_read_tokens": 1000,
|
||||
"cache_ttl_minutes": 60,
|
||||
"request_count": 1,
|
||||
"total_input_context": 0 + 1000,
|
||||
},
|
||||
dimension_mappings=rule.dimension_mappings,
|
||||
strict_mode=True,
|
||||
)
|
||||
|
||||
# TTL=60 should use cache_read_price_per_1m=0.5
|
||||
# 1000 * 0.5 / 1M = 0.0005
|
||||
assert result.status == "complete"
|
||||
assert abs(float(result.cost) - 0.0005) < 1e-9
|
||||
|
||||
|
||||
class TestBillingRuleServiceDefaultFallback:
|
||||
def test_find_rule_returns_default_for_chat_when_no_db_rule(self) -> None:
|
||||
from src.services.billing.cache import BillingCache
|
||||
|
||||
BillingCache.invalidate_all()
|
||||
|
||||
global_model = GlobalModel(
|
||||
id="gm-1",
|
||||
name="test-model",
|
||||
display_name="Test Model",
|
||||
is_active=True,
|
||||
default_price_per_request=0.0,
|
||||
default_tiered_pricing={
|
||||
"tiers": [{"up_to": None, "input_price_per_1m": 3.0, "output_price_per_1m": 15.0}]
|
||||
},
|
||||
)
|
||||
|
||||
model_obj = Model(
|
||||
id="m-1",
|
||||
provider_id="p-1",
|
||||
global_model_id="gm-1",
|
||||
provider_model_name="provider-test-model",
|
||||
is_active=True,
|
||||
tiered_pricing=None,
|
||||
price_per_request=None,
|
||||
)
|
||||
model_obj.global_model = global_model
|
||||
|
||||
# Build a mock Session with deterministic query().filter().first() chain.
|
||||
q_global = MagicMock()
|
||||
q_global.filter.return_value.first.return_value = global_model
|
||||
|
||||
q_model = MagicMock()
|
||||
q_model.filter.return_value.first.return_value = model_obj
|
||||
|
||||
q_rule_model = MagicMock()
|
||||
q_rule_model.filter.return_value.first.return_value = None
|
||||
|
||||
q_rule_global = MagicMock()
|
||||
q_rule_global.filter.return_value.first.return_value = None
|
||||
|
||||
db = MagicMock()
|
||||
db.query.side_effect = [q_global, q_model, q_rule_model, q_rule_global]
|
||||
|
||||
lookup = BillingRuleService.find_rule(
|
||||
db,
|
||||
provider_id="p-1",
|
||||
model_name="test-model",
|
||||
task_type="chat",
|
||||
)
|
||||
assert lookup is not None
|
||||
assert lookup.scope == "default"
|
||||
assert lookup.rule.id == "__default__"
|
||||
assert lookup.effective_task_type == "chat"
|
||||
|
||||
# Cached: second call should not touch db.query again.
|
||||
db.query.reset_mock()
|
||||
lookup2 = BillingRuleService.find_rule(
|
||||
db,
|
||||
provider_id="p-1",
|
||||
model_name="test-model",
|
||||
task_type="chat",
|
||||
)
|
||||
assert lookup2 is not None
|
||||
assert lookup2.scope == "default"
|
||||
assert db.query.call_count == 0
|
||||
|
||||
def test_find_rule_returns_default_for_video_when_require_rule_false(self) -> None:
|
||||
from src.config.settings import config
|
||||
from src.services.billing.cache import BillingCache
|
||||
|
||||
BillingCache.invalidate_all()
|
||||
old_require = config.billing_require_rule
|
||||
config.billing_require_rule = False
|
||||
|
||||
try:
|
||||
global_model = GlobalModel(
|
||||
id="gm-2",
|
||||
name="video-model",
|
||||
display_name="Video Model",
|
||||
is_active=True,
|
||||
default_price_per_request=0.0,
|
||||
default_tiered_pricing={
|
||||
"tiers": [
|
||||
{"up_to": None, "input_price_per_1m": 3.0, "output_price_per_1m": 15.0}
|
||||
]
|
||||
},
|
||||
)
|
||||
|
||||
q_global = MagicMock()
|
||||
q_global.filter.return_value.first.return_value = global_model
|
||||
|
||||
q_rule_global = MagicMock()
|
||||
q_rule_global.filter.return_value.first.return_value = None
|
||||
|
||||
db = MagicMock()
|
||||
# provider_id omitted -> only GlobalModel query + global BillingRule query
|
||||
db.query.side_effect = [q_global, q_rule_global]
|
||||
|
||||
lookup = BillingRuleService.find_rule(
|
||||
db,
|
||||
provider_id=None,
|
||||
model_name="video-model",
|
||||
task_type="video",
|
||||
)
|
||||
assert lookup is not None
|
||||
assert lookup.scope == "default"
|
||||
assert lookup.rule.id == "__default__"
|
||||
assert lookup.effective_task_type == "video"
|
||||
finally:
|
||||
config.billing_require_rule = old_require
|
||||
|
||||
def test_find_rule_returns_template_for_video_when_require_rule_true(self) -> None:
|
||||
from src.config.settings import config
|
||||
from src.services.billing.cache import BillingCache
|
||||
|
||||
BillingCache.invalidate_all()
|
||||
old_require = config.billing_require_rule
|
||||
config.billing_require_rule = True
|
||||
|
||||
try:
|
||||
global_model = GlobalModel(
|
||||
id="gm-3",
|
||||
name="video-model-2",
|
||||
display_name="Video Model 2",
|
||||
is_active=True,
|
||||
default_price_per_request=0.0,
|
||||
default_tiered_pricing={
|
||||
"tiers": [
|
||||
{"up_to": None, "input_price_per_1m": 3.0, "output_price_per_1m": 15.0}
|
||||
]
|
||||
},
|
||||
)
|
||||
|
||||
q_global = MagicMock()
|
||||
q_global.filter.return_value.first.return_value = global_model
|
||||
|
||||
q_rule_global = MagicMock()
|
||||
q_rule_global.filter.return_value.first.return_value = None
|
||||
|
||||
db = MagicMock()
|
||||
db.query.side_effect = [q_global, q_rule_global]
|
||||
|
||||
lookup = BillingRuleService.find_rule(
|
||||
db,
|
||||
provider_id=None,
|
||||
model_name="video-model-2",
|
||||
task_type="video",
|
||||
)
|
||||
assert lookup is not None
|
||||
assert lookup.scope == "default"
|
||||
assert lookup.rule.id == "__default__"
|
||||
# Universal Billing Rule is now used for all task types
|
||||
assert lookup.rule.name == "Universal Billing Rule"
|
||||
finally:
|
||||
config.billing_require_rule = old_require
|
||||
@@ -34,7 +34,7 @@ class TestDimensionCollectorRuntime:
|
||||
),
|
||||
]
|
||||
dims = runtime.collect(
|
||||
collectors=collectors,
|
||||
collectors=collectors, # type: ignore[arg-type]
|
||||
inp=DimensionCollectInput(
|
||||
response={"usageMetadata": {"promptTokenCount": 123}},
|
||||
),
|
||||
@@ -57,7 +57,7 @@ class TestDimensionCollectorRuntime:
|
||||
)
|
||||
]
|
||||
dims = runtime.collect(
|
||||
collectors=collectors,
|
||||
collectors=collectors, # type: ignore[arg-type]
|
||||
inp=DimensionCollectInput(metadata={"result": {"file_size_bytes": 1048576}}),
|
||||
)
|
||||
assert abs(dims["file_size_mb"] - 1.0) < 1e-9
|
||||
@@ -98,7 +98,7 @@ class TestDimensionCollectorRuntime:
|
||||
),
|
||||
]
|
||||
dims = runtime.collect(
|
||||
collectors=collectors,
|
||||
collectors=collectors, # type: ignore[arg-type]
|
||||
inp=DimensionCollectInput(
|
||||
request={"usage": {"input_tokens": 100, "cache_read_tokens": 20}}
|
||||
),
|
||||
@@ -110,52 +110,13 @@ class TestDimensionCollectorRuntime:
|
||||
|
||||
class TestDimensionCollectorService:
|
||||
def test_video_fallback_merges_base_collectors(self) -> None:
|
||||
from src.services.billing.cache import BillingCache
|
||||
|
||||
BillingCache.invalidate_all()
|
||||
|
||||
# code-only: openai:video should fall back to openai:chat video collectors shipped in code
|
||||
db = MagicMock()
|
||||
|
||||
video_collectors = [
|
||||
DimensionCollector(
|
||||
api_format="openai:video",
|
||||
task_type="video",
|
||||
dimension_name="duration_seconds",
|
||||
source_type="metadata",
|
||||
source_path="task.duration_seconds",
|
||||
value_type="int",
|
||||
priority=0,
|
||||
is_enabled=True,
|
||||
)
|
||||
]
|
||||
base_collectors = [
|
||||
# Should be kept (dimension not present in video_collectors)
|
||||
DimensionCollector(
|
||||
api_format="openai:chat",
|
||||
task_type="video",
|
||||
dimension_name="resolution",
|
||||
source_type="metadata",
|
||||
source_path="task.resolution",
|
||||
value_type="string",
|
||||
priority=0,
|
||||
is_enabled=True,
|
||||
),
|
||||
# Should be ignored (dimension already present in video_collectors)
|
||||
DimensionCollector(
|
||||
api_format="openai:chat",
|
||||
task_type="video",
|
||||
dimension_name="duration_seconds",
|
||||
source_type="metadata",
|
||||
source_path="task.duration_seconds",
|
||||
value_type="int",
|
||||
priority=0,
|
||||
is_enabled=True,
|
||||
),
|
||||
]
|
||||
|
||||
q1 = MagicMock()
|
||||
q1.filter.return_value.all.return_value = video_collectors
|
||||
q2 = MagicMock()
|
||||
q2.filter.return_value.all.return_value = base_collectors
|
||||
db.query.side_effect = [q1, q2]
|
||||
|
||||
svc = DimensionCollectorService(db)
|
||||
result = svc.list_enabled_collectors(api_format="openai:video", task_type="video")
|
||||
|
||||
assert [c.dimension_name for c in result] == ["duration_seconds", "resolution"]
|
||||
assert "video_size_bytes" in [c.dimension_name for c in result]
|
||||
|
||||
@@ -54,7 +54,7 @@ class TestFormulaEngine:
|
||||
)
|
||||
assert result.status == "complete"
|
||||
assert result.missing_required == []
|
||||
assert abs(result.cost - 0.25) < 1e-9
|
||||
assert abs(float(result.cost) - 0.25) < 1e-9
|
||||
|
||||
def test_required_dimension_missing_non_strict(self) -> None:
|
||||
engine = FormulaEngine()
|
||||
@@ -72,7 +72,7 @@ class TestFormulaEngine:
|
||||
strict_mode=False,
|
||||
)
|
||||
assert result.status == "incomplete"
|
||||
assert result.cost == 0.0
|
||||
assert float(result.cost) == 0.0
|
||||
assert result.missing_required == ["duration_seconds"]
|
||||
|
||||
def test_required_dimension_missing_strict_raises(self) -> None:
|
||||
@@ -127,7 +127,7 @@ class TestFormulaEngine:
|
||||
},
|
||||
)
|
||||
assert result.status == "complete"
|
||||
assert result.cost == 0.0
|
||||
assert float(result.cost) == 0.0
|
||||
|
||||
def test_tiered_mapping(self) -> None:
|
||||
engine = FormulaEngine()
|
||||
@@ -148,7 +148,7 @@ class TestFormulaEngine:
|
||||
},
|
||||
)
|
||||
assert result.status == "complete"
|
||||
assert result.cost == 3.0
|
||||
assert float(result.cost) == 3.0
|
||||
|
||||
result = engine.evaluate(
|
||||
expression="input_price",
|
||||
@@ -166,4 +166,4 @@ class TestFormulaEngine:
|
||||
},
|
||||
)
|
||||
assert result.status == "complete"
|
||||
assert result.cost == 1.5
|
||||
assert float(result.cost) == 1.5
|
||||
|
||||
112
tests/services/billing/test_shadow_billing.py
Normal file
112
tests/services/billing/test_shadow_billing.py
Normal file
@@ -0,0 +1,112 @@
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from src.config.settings import config
|
||||
from src.services.billing.schema import BillingSnapshot, CostResult
|
||||
from src.services.billing.shadow import CostBreakdown, ShadowBillingService
|
||||
|
||||
|
||||
class TestShadowBillingServiceModeResolution:
|
||||
def test_get_engine_mode_exact_and_wildcard(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(config, "billing_engine", "legacy", raising=False)
|
||||
monkeypatch.setattr(
|
||||
config,
|
||||
"billing_engine_overrides",
|
||||
'{"anthropic/*": "shadow", "openai/gpt-4o": "new"}',
|
||||
raising=False,
|
||||
)
|
||||
|
||||
svc = ShadowBillingService(MagicMock())
|
||||
assert svc.get_engine_mode("openai", "gpt-4o") == "new"
|
||||
assert svc.get_engine_mode("anthropic", "claude-3-5-sonnet") == "shadow"
|
||||
assert svc.get_engine_mode("other", "x") == "legacy"
|
||||
|
||||
|
||||
class TestShadowBillingServiceExecution:
|
||||
def test_legacy_mode_skips_new_engine(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(config, "billing_engine", "legacy", raising=False)
|
||||
monkeypatch.setattr(config, "billing_engine_overrides", "{}", raising=False)
|
||||
|
||||
svc = ShadowBillingService(MagicMock())
|
||||
# Guard: if new engine calculate gets called, fail.
|
||||
svc._new_billing = MagicMock()
|
||||
svc._new_billing.calculate.side_effect = AssertionError(
|
||||
"new engine should not run in legacy mode"
|
||||
)
|
||||
|
||||
legacy_truth = CostBreakdown(
|
||||
input_cost=0.1,
|
||||
output_cost=0.2,
|
||||
cache_creation_cost=0.0,
|
||||
cache_read_cost=0.0,
|
||||
request_cost=0.0,
|
||||
total_cost=0.3,
|
||||
)
|
||||
|
||||
res = svc.calculate_with_shadow(
|
||||
provider="openai",
|
||||
provider_id="p-1",
|
||||
model="gpt-4o",
|
||||
task_type="chat",
|
||||
api_format="openai:chat",
|
||||
input_tokens=1,
|
||||
output_tokens=1,
|
||||
legacy_truth=legacy_truth,
|
||||
is_failed_request=False,
|
||||
)
|
||||
|
||||
assert res.engine_mode == "legacy"
|
||||
assert res.truth_engine == "legacy"
|
||||
assert res.shadow_snapshot is None
|
||||
assert res.truth_breakdown.total_cost == 0.3
|
||||
|
||||
def test_shadow_mode_returns_snapshot_and_keeps_legacy_truth(
|
||||
self, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(config, "billing_engine", "shadow", raising=False)
|
||||
monkeypatch.setattr(config, "billing_engine_overrides", "{}", raising=False)
|
||||
monkeypatch.setattr(config, "billing_diff_threshold_usd", 0.0001, raising=False)
|
||||
|
||||
svc = ShadowBillingService(MagicMock())
|
||||
|
||||
# Stub new engine output
|
||||
snapshot = BillingSnapshot(
|
||||
resolved_dimensions={"input_tokens": 1},
|
||||
resolved_variables={"input_price_per_1m": "3.0"},
|
||||
cost_breakdown={"input_cost": 0.003},
|
||||
total_cost=0.003,
|
||||
status="complete",
|
||||
calculated_at="2026-02-02T00:00:00Z",
|
||||
)
|
||||
svc._new_billing = MagicMock()
|
||||
svc._new_billing.calculate.return_value = CostResult(
|
||||
cost=0.003, status="complete", snapshot=snapshot
|
||||
)
|
||||
|
||||
legacy_truth = CostBreakdown(
|
||||
input_cost=0.004,
|
||||
output_cost=0.0,
|
||||
cache_creation_cost=0.0,
|
||||
cache_read_cost=0.0,
|
||||
request_cost=0.0,
|
||||
total_cost=0.004,
|
||||
)
|
||||
|
||||
res = svc.calculate_with_shadow(
|
||||
provider="openai",
|
||||
provider_id="p-1",
|
||||
model="gpt-4o",
|
||||
task_type="chat",
|
||||
api_format="openai:chat",
|
||||
input_tokens=1,
|
||||
output_tokens=0,
|
||||
legacy_truth=legacy_truth,
|
||||
is_failed_request=False,
|
||||
)
|
||||
|
||||
assert res.engine_mode == "shadow"
|
||||
assert res.truth_engine == "legacy"
|
||||
assert res.shadow_snapshot is not None
|
||||
assert res.truth_breakdown.total_cost == 0.004
|
||||
assert "diff_usd" in res.comparison
|
||||
@@ -54,8 +54,7 @@ async def test_build_candidates_allows_cross_format_when_endpoint_accepts_and_ov
|
||||
client_format="claude:chat",
|
||||
model_name="dummy-model",
|
||||
affinity_key=None,
|
||||
global_conversion_enabled=False, # DB 全局覆盖关闭
|
||||
master_conversion_enabled=True, # ENV 总闸开启（默认）
|
||||
global_conversion_enabled=True, # 全局开关开启
|
||||
)
|
||||
|
||||
assert len(candidates) == 1
|
||||
@@ -88,8 +87,7 @@ async def test_build_candidates_blocks_cross_format_when_master_switch_off() ->
|
||||
client_format="claude:chat",
|
||||
model_name="dummy-model",
|
||||
affinity_key=None,
|
||||
global_conversion_enabled=False,
|
||||
master_conversion_enabled=False,
|
||||
global_conversion_enabled=False, # 全局开关关闭
|
||||
)
|
||||
|
||||
assert candidates == []
|
||||
@@ -118,8 +116,7 @@ async def test_build_candidates_includes_cross_format_when_enabled() -> None:
|
||||
client_format="claude:chat",
|
||||
model_name="dummy-model",
|
||||
affinity_key=None,
|
||||
global_conversion_enabled=True, # DB 全局覆盖开启：跳过端点检查
|
||||
master_conversion_enabled=True,
|
||||
global_conversion_enabled=True, # 全局开关开启：跳过端点检查
|
||||
)
|
||||
|
||||
assert len(candidates) == 1
|
||||
@@ -157,8 +154,7 @@ async def test_exact_matches_rank_before_convertible() -> None:
|
||||
client_format="claude:chat",
|
||||
model_name="dummy-model",
|
||||
affinity_key=None,
|
||||
global_conversion_enabled=False,
|
||||
master_conversion_enabled=True,
|
||||
global_conversion_enabled=True, # 全局开关开启
|
||||
)
|
||||
|
||||
assert len(candidates) == 2
|
||||
|
||||
@@ -62,18 +62,15 @@ async def test_queue_writer_publishes_event(monkeypatch):
|
||||
|
||||
old_stream_key = config.usage_queue_stream_key
|
||||
old_maxlen = config.usage_queue_stream_maxlen
|
||||
old_include_headers = config.usage_queue_include_headers
|
||||
old_include_bodies = config.usage_queue_include_bodies
|
||||
try:
|
||||
config.usage_queue_stream_key = "usage:events:test"
|
||||
config.usage_queue_stream_maxlen = 0
|
||||
config.usage_queue_include_headers = False
|
||||
config.usage_queue_include_bodies = False
|
||||
|
||||
writer = QueueTelemetryWriter(
|
||||
request_id="req-2",
|
||||
user_id="user-1",
|
||||
api_key_id="key-1",
|
||||
log_level="basic",
|
||||
)
|
||||
await writer.record_success(
|
||||
provider="test",
|
||||
@@ -86,8 +83,6 @@ async def test_queue_writer_publishes_event(monkeypatch):
|
||||
finally:
|
||||
config.usage_queue_stream_key = old_stream_key
|
||||
config.usage_queue_stream_maxlen = old_maxlen
|
||||
config.usage_queue_include_headers = old_include_headers
|
||||
config.usage_queue_include_bodies = old_include_bodies
|
||||
|
||||
assert dummy.calls
|
||||
key, fields, _, _ = dummy.calls[0]
|
||||
@@ -310,31 +305,28 @@ async def test_queue_writer_include_headers_bodies(monkeypatch):
|
||||
|
||||
monkeypatch.setattr("src.services.usage.telemetry_writer.get_redis_client", _get_redis_client)
|
||||
|
||||
old_include_headers = config.usage_queue_include_headers
|
||||
old_include_bodies = config.usage_queue_include_bodies
|
||||
try:
|
||||
config.usage_queue_include_headers = True
|
||||
config.usage_queue_include_bodies = True
|
||||
|
||||
writer = QueueTelemetryWriter(
|
||||
request_id="req-full",
|
||||
user_id="user-1",
|
||||
api_key_id="key-1",
|
||||
)
|
||||
await writer.record_success(
|
||||
provider="test",
|
||||
model="model",
|
||||
request_headers={"Authorization": "Bearer xxx"},
|
||||
response_headers={"Content-Type": "application/json"},
|
||||
request_body={"messages": [{"role": "user", "content": "hi"}]},
|
||||
response_body={"choices": [{"message": {"content": "hello"}}]},
|
||||
)
|
||||
finally:
|
||||
config.usage_queue_include_headers = old_include_headers
|
||||
config.usage_queue_include_bodies = old_include_bodies
|
||||
writer = QueueTelemetryWriter(
|
||||
request_id="req-full",
|
||||
user_id="user-1",
|
||||
api_key_id="key-1",
|
||||
log_level="full",
|
||||
sensitive_headers=["authorization"],
|
||||
max_request_body_size=0,
|
||||
max_response_body_size=0,
|
||||
)
|
||||
await writer.record_success(
|
||||
provider="test",
|
||||
model="model",
|
||||
request_headers={"Authorization": "Bearer xxx"},
|
||||
response_headers={"Content-Type": "application/json"},
|
||||
request_body={"messages": [{"role": "user", "content": "hi"}]},
|
||||
response_body={"choices": [{"message": {"content": "hello"}}]},
|
||||
)
|
||||
|
||||
event = UsageEvent.from_stream_fields(dummy.calls[0][1])
|
||||
assert event.data["request_headers"]["Authorization"] == "Bearer xxx"
|
||||
# Sensitive header should be masked before going into Redis
|
||||
assert event.data["request_headers"]["Authorization"].startswith("Bear")
|
||||
assert "****" in event.data["request_headers"]["Authorization"]
|
||||
assert "request_body" in event.data
|
||||
assert "response_body" in event.data
|
||||
|
||||
@@ -378,31 +370,26 @@ async def test_queue_writer_body_truncation(monkeypatch):
|
||||
|
||||
monkeypatch.setattr("src.services.usage.telemetry_writer.get_redis_client", _get_redis_client)
|
||||
|
||||
old_include_bodies = config.usage_queue_include_bodies
|
||||
old_max_bytes = config.usage_queue_body_max_bytes
|
||||
try:
|
||||
config.usage_queue_include_bodies = True
|
||||
config.usage_queue_body_max_bytes = 50
|
||||
|
||||
writer = QueueTelemetryWriter(
|
||||
request_id="req-trunc",
|
||||
user_id="user-1",
|
||||
api_key_id="key-1",
|
||||
)
|
||||
long_body = {"content": "x" * 1000}
|
||||
await writer.record_success(
|
||||
provider="test",
|
||||
model="model",
|
||||
request_body=long_body,
|
||||
)
|
||||
finally:
|
||||
config.usage_queue_include_bodies = old_include_bodies
|
||||
config.usage_queue_body_max_bytes = old_max_bytes
|
||||
writer = QueueTelemetryWriter(
|
||||
request_id="req-trunc",
|
||||
user_id="user-1",
|
||||
api_key_id="key-1",
|
||||
log_level="full",
|
||||
max_request_body_size=50,
|
||||
max_response_body_size=0,
|
||||
)
|
||||
long_body = {"content": "x" * 1000}
|
||||
await writer.record_success(
|
||||
provider="test",
|
||||
model="model",
|
||||
request_body=long_body,
|
||||
)
|
||||
|
||||
event = UsageEvent.from_stream_fields(dummy.calls[0][1])
|
||||
body_str = event.data["request_body"]
|
||||
assert len(body_str) <= 50
|
||||
assert body_str.endswith("...[truncated]")
|
||||
body = event.data["request_body"]
|
||||
assert isinstance(body, dict)
|
||||
assert body.get("_truncated") is True
|
||||
assert len(body.get("_content") or "") <= 50
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
||||
203
tests/test_sora.py
Normal file
203
tests/test_sora.py
Normal file
@@ -0,0 +1,203 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Sora Video Generation Test Script
|
||||
|
||||
This script demonstrates video generation using OpenAI's Sora API.
|
||||
It sends a request to generate a video, polls for completion, and downloads the result.
|
||||
|
||||
Usage:
|
||||
export OPENAI_API_KEY="your-api-key"
|
||||
python test_sora.py
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
|
||||
import requests
|
||||
|
||||
# OpenAI API Base URL
|
||||
BASE_URL = "http://localhost:8084/v1"
|
||||
|
||||
# Default polling interval in seconds
|
||||
POLL_INTERVAL = 10
|
||||
|
||||
os.environ["OPENAI_API_KEY"] = "sk-PCr5oXZNKb9HcyzYqTIMvr8zXsIBK3WS"
|
||||
|
||||
|
||||
def generate_video(
|
||||
api_key: str,
|
||||
prompt: str,
|
||||
model: str = "sora-2",
|
||||
size: str = "1920x1080",
|
||||
duration: int = 10,
|
||||
n: int = 1,
|
||||
) -> str:
|
||||
"""
|
||||
Send a request to generate a video and return the video job ID.
|
||||
|
||||
Args:
|
||||
api_key: OpenAI API key
|
||||
prompt: Text prompt for video generation
|
||||
model: Model name to use (default: sora-2)
|
||||
size: Video resolution (default: 1920x1080)
|
||||
duration: Video duration in seconds (default: 10)
|
||||
n: Number of videos to generate (default: 1)
|
||||
|
||||
Returns:
|
||||
Video job ID for polling status
|
||||
"""
|
||||
url = f"{BASE_URL}/videos"
|
||||
headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
payload = {
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"size": size,
|
||||
"duration": str(duration),
|
||||
"n": n,
|
||||
}
|
||||
|
||||
print(f"Sending video generation request to {model}...")
|
||||
print(f" Size: {size}, Duration: {duration}s")
|
||||
response = requests.post(url, headers=headers, json=payload, timeout=60)
|
||||
response.raise_for_status()
|
||||
|
||||
data = response.json()
|
||||
video_id = data.get("id")
|
||||
|
||||
if not video_id:
|
||||
raise ValueError(f"No video ID in response: {data}")
|
||||
|
||||
print(f"Video job started: {video_id}")
|
||||
print(f" Status: {data.get('status')}")
|
||||
return video_id
|
||||
|
||||
|
||||
def poll_video(api_key: str, video_id: str, poll_interval: int = POLL_INTERVAL) -> dict:
|
||||
"""
|
||||
Poll the video job status until the video is ready.
|
||||
|
||||
Args:
|
||||
api_key: OpenAI API key
|
||||
video_id: Video job ID from generate_video
|
||||
poll_interval: Seconds between polls (default: 10)
|
||||
|
||||
Returns:
|
||||
Final response dict containing the video metadata
|
||||
"""
|
||||
url = f"{BASE_URL}/videos/{video_id}"
|
||||
headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
}
|
||||
|
||||
print(f"Polling video job status (every {poll_interval}s)...")
|
||||
|
||||
while True:
|
||||
response = requests.get(url, headers=headers, timeout=60)
|
||||
response.raise_for_status()
|
||||
|
||||
data = response.json()
|
||||
status = data.get("status", "unknown")
|
||||
progress = data.get("progress", 0)
|
||||
|
||||
print(f" Status: {status}, Progress: {progress}%")
|
||||
|
||||
if status == "completed":
|
||||
print("Video generation completed!")
|
||||
print(f" Duration: {data.get('seconds')}s")
|
||||
print(f" Size: {data.get('size')}")
|
||||
print(f" Expires at: {data.get('expires_at')}")
|
||||
return data
|
||||
|
||||
if status == "failed":
|
||||
error = data.get("error", {})
|
||||
raise RuntimeError(f"Video generation failed: {error.get('message', error)}")
|
||||
|
||||
if status == "cancelled":
|
||||
raise RuntimeError("Video generation was cancelled")
|
||||
|
||||
time.sleep(poll_interval)
|
||||
|
||||
|
||||
def download_video(api_key: str, video_id: str, output_path: str = "sora_output.mp4", variant: str | None = None) -> str:
|
||||
"""
|
||||
Download the generated video content.
|
||||
|
||||
Args:
|
||||
api_key: OpenAI API key
|
||||
video_id: Video job ID
|
||||
output_path: Path to save the video (default: sora_output.mp4)
|
||||
variant: Optional variant to download (defaults to MP4 video)
|
||||
|
||||
Returns:
|
||||
Path to the downloaded video
|
||||
"""
|
||||
url = f"{BASE_URL}/videos/{video_id}/content"
|
||||
if variant:
|
||||
url = f"{url}?variant={variant}"
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
}
|
||||
|
||||
print(f"Downloading video content...")
|
||||
response = requests.get(url, headers=headers, allow_redirects=True, timeout=300, stream=True)
|
||||
response.raise_for_status()
|
||||
|
||||
with open(output_path, "wb") as f:
|
||||
for chunk in response.iter_content(chunk_size=8192):
|
||||
f.write(chunk)
|
||||
|
||||
file_size = os.path.getsize(output_path)
|
||||
print(f"Video saved to: {output_path} ({file_size / 1024 / 1024:.2f} MB)")
|
||||
return output_path
|
||||
|
||||
|
||||
def main() -> None:
|
||||
"""Main entry point."""
|
||||
# Get API key from environment
|
||||
api_key = os.environ.get("OPENAI_API_KEY")
|
||||
if not api_key:
|
||||
print("Error: OPENAI_API_KEY environment variable not set", file=sys.stderr)
|
||||
print("Usage: export OPENAI_API_KEY='your-api-key' && python test_sora.py", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
# Default prompt
|
||||
prompt = "A calico cat playing a piano on stage"
|
||||
|
||||
# Allow custom prompt via command line argument
|
||||
if len(sys.argv) > 1:
|
||||
prompt = " ".join(sys.argv[1:])
|
||||
print(f"Using custom prompt: {prompt}")
|
||||
|
||||
try:
|
||||
# Step 1: Start video generation
|
||||
video_id = generate_video(api_key, prompt)
|
||||
|
||||
# Step 2: Poll until complete
|
||||
final_response = poll_video(api_key, video_id)
|
||||
|
||||
# Step 3: Download video content
|
||||
download_video(api_key, video_id)
|
||||
|
||||
print("\nVideo generation complete!")
|
||||
print(f"Video ID: {video_id}")
|
||||
print(f"Model: {final_response.get('model')}")
|
||||
|
||||
except requests.exceptions.HTTPError as e:
|
||||
print(f"HTTP Error: {e}", file=sys.stderr)
|
||||
if e.response is not None:
|
||||
print(f"Response: {e.response.text}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
except Exception as e:
|
||||
print(f"Error: {e}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
207
tests/test_veo.py
Normal file
207
tests/test_veo.py
Normal file
@@ -0,0 +1,207 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Veo Video Generation Test Script
|
||||
|
||||
This script demonstrates video generation using Google's Veo API.
|
||||
It sends a request to generate a video, polls for completion, and downloads the result.
|
||||
|
||||
Usage:
|
||||
export GEMINI_API_KEY="your-api-key"
|
||||
python test_veo.py
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
|
||||
import requests
|
||||
|
||||
# Gemini API Base URL
|
||||
BASE_URL = "http://localhost:8084/v1beta"
|
||||
|
||||
# Default polling interval in seconds
|
||||
POLL_INTERVAL = 10
|
||||
|
||||
os.environ["GEMINI_API_KEY"] = "sk-PCr5oXZNKb9HcyzYqTIMvr8zXsIBK3WS"
|
||||
|
||||
|
||||
def generate_video(api_key: str, prompt: str, model: str = "veo-3.1-fast-generate-preview") -> str:
|
||||
"""
|
||||
Send a request to generate a video and return the operation name.
|
||||
|
||||
Args:
|
||||
api_key: Gemini API key
|
||||
prompt: Text prompt for video generation
|
||||
model: Model name to use (default: veo-3.1-generate-preview)
|
||||
|
||||
Returns:
|
||||
Operation name for polling status
|
||||
"""
|
||||
url = f"{BASE_URL}/models/{model}:predictLongRunning"
|
||||
headers = {
|
||||
"x-goog-api-key": api_key,
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
payload = {
|
||||
"instances": [
|
||||
{
|
||||
"prompt": prompt,
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
print(f"Sending video generation request to {model}...")
|
||||
response = requests.post(url, headers=headers, json=payload, timeout=60)
|
||||
response.raise_for_status()
|
||||
|
||||
data = response.json()
|
||||
operation_name = data.get("name")
|
||||
|
||||
if not operation_name:
|
||||
raise ValueError(f"No operation name in response: {data}")
|
||||
|
||||
print(f"Operation started: {operation_name}")
|
||||
return operation_name
|
||||
|
||||
|
||||
def poll_operation(api_key: str, operation_name: str, poll_interval: int = POLL_INTERVAL) -> dict:
|
||||
"""
|
||||
Poll the operation status until the video is ready.
|
||||
|
||||
Args:
|
||||
api_key: Gemini API key
|
||||
operation_name: Operation name from generate_video
|
||||
poll_interval: Seconds between polls (default: 10)
|
||||
|
||||
Returns:
|
||||
Final response dict containing the video URI
|
||||
"""
|
||||
url = f"{BASE_URL}/{operation_name}"
|
||||
headers = {
|
||||
"x-goog-api-key": api_key,
|
||||
}
|
||||
|
||||
print(f"Polling operation status (every {poll_interval}s)...")
|
||||
|
||||
while True:
|
||||
response = requests.get(url, headers=headers, timeout=60)
|
||||
response.raise_for_status()
|
||||
|
||||
data = response.json()
|
||||
is_done = data.get("done", False)
|
||||
|
||||
if is_done:
|
||||
print("Operation completed!")
|
||||
|
||||
# Check for errors
|
||||
if "error" in data:
|
||||
error = data["error"]
|
||||
raise RuntimeError(f"Operation failed: {error.get('message', error)}")
|
||||
|
||||
return data
|
||||
|
||||
# Show progress if available
|
||||
metadata = data.get("metadata", {})
|
||||
if metadata:
|
||||
progress = metadata.get("progress", "unknown")
|
||||
print(f" Progress: {progress}%")
|
||||
|
||||
time.sleep(poll_interval)
|
||||
|
||||
|
||||
def download_video(api_key: str, video_uri: str, output_path: str = "dialogue_example.mp4") -> str:
|
||||
"""
|
||||
Download the generated video.
|
||||
|
||||
Args:
|
||||
api_key: Gemini API key
|
||||
video_uri: URI of the generated video
|
||||
output_path: Path to save the video (default: dialogue_example.mp4)
|
||||
|
||||
Returns:
|
||||
Path to the downloaded video
|
||||
"""
|
||||
headers = {
|
||||
"x-goog-api-key": api_key,
|
||||
}
|
||||
|
||||
print(f"Downloading video from: {video_uri}")
|
||||
response = requests.get(video_uri, headers=headers, allow_redirects=True, timeout=300, stream=True)
|
||||
response.raise_for_status()
|
||||
|
||||
with open(output_path, "wb") as f:
|
||||
for chunk in response.iter_content(chunk_size=8192):
|
||||
f.write(chunk)
|
||||
|
||||
file_size = os.path.getsize(output_path)
|
||||
print(f"Video saved to: {output_path} ({file_size / 1024 / 1024:.2f} MB)")
|
||||
return output_path
|
||||
|
||||
|
||||
def extract_video_uri(response: dict) -> str:
|
||||
"""
|
||||
Extract the video URI from the operation response.
|
||||
|
||||
Args:
|
||||
response: Final operation response dict
|
||||
|
||||
Returns:
|
||||
Video download URI
|
||||
"""
|
||||
try:
|
||||
video_response = response["response"]["generateVideoResponse"]
|
||||
samples = video_response["generatedSamples"]
|
||||
video_uri = samples[0]["video"]["uri"]
|
||||
return video_uri
|
||||
except (KeyError, IndexError) as e:
|
||||
raise ValueError(f"Could not extract video URI from response: {response}") from e
|
||||
|
||||
|
||||
def main() -> None:
|
||||
"""Main entry point."""
|
||||
# Get API key from environment
|
||||
api_key = os.environ.get("GEMINI_API_KEY")
|
||||
if not api_key:
|
||||
print("Error: GEMINI_API_KEY environment variable not set", file=sys.stderr)
|
||||
print("Usage: export GEMINI_API_KEY='your-api-key' && python test_veo.py", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
# Default prompt (same as the bash script)
|
||||
prompt = (
|
||||
"A close up of two people staring at a cryptic drawing on a wall, "
|
||||
"torchlight flickering. A man murmurs, \"This must be it. That's the secret code.\" "
|
||||
"The woman looks at him and whispering excitedly, \"What did you find?\""
|
||||
)
|
||||
|
||||
# Allow custom prompt via command line argument
|
||||
if len(sys.argv) > 1:
|
||||
prompt = " ".join(sys.argv[1:])
|
||||
print(f"Using custom prompt: {prompt}")
|
||||
|
||||
try:
|
||||
# Step 1: Start video generation
|
||||
operation_name = generate_video(api_key, prompt)
|
||||
|
||||
# Step 2: Poll until complete
|
||||
final_response = poll_operation(api_key, operation_name)
|
||||
|
||||
# Step 3: Extract video URI and download
|
||||
video_uri = extract_video_uri(final_response)
|
||||
download_video(api_key, video_uri)
|
||||
|
||||
print("\nVideo generation complete!")
|
||||
|
||||
except requests.exceptions.HTTPError as e:
|
||||
print(f"HTTP Error: {e}", file=sys.stderr)
|
||||
if e.response is not None:
|
||||
print(f"Response: {e.response.text}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
except Exception as e:
|
||||
print(f"Error: {e}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user