feat: 添加视频生成功能增强和多维度计费系统适配

- 视频生成: 增强 video_handler,重构 task_poller,新增 telemetry
- 计费系统: 适配新的 signature 格式,支持 video 任务类型回退
- 数据库迁移: 添加 api_family/endpoint_kind 字段和 video_formats
This commit is contained in:
fawney19
2026-02-01 17:28:27 +08:00
parent 7b66505634
commit 4ac8e63c94
15 changed files with 1514 additions and 527 deletions

View File

@@ -1,7 +1,10 @@
from unittest.mock import MagicMock
from src.models.database import DimensionCollector
from src.services.billing.dimension_collector_service import (
DimensionCollectInput,
DimensionCollectorRuntime,
DimensionCollectorService,
)
@@ -10,7 +13,7 @@ class TestDimensionCollectorRuntime:
runtime = DimensionCollectorRuntime()
collectors = [
DimensionCollector(
api_format="OPENAI",
api_format="openai:chat",
task_type="chat",
dimension_name="input_tokens",
source_type="response",
@@ -20,7 +23,7 @@ class TestDimensionCollectorRuntime:
is_enabled=True,
),
DimensionCollector(
api_format="OPENAI",
api_format="openai:chat",
task_type="chat",
dimension_name="input_tokens",
source_type="response",
@@ -42,7 +45,7 @@ class TestDimensionCollectorRuntime:
runtime = DimensionCollectorRuntime()
collectors = [
DimensionCollector(
api_format="GEMINI",
api_format="gemini:video",
task_type="video",
dimension_name="file_size_mb",
source_type="metadata",
@@ -63,7 +66,7 @@ class TestDimensionCollectorRuntime:
runtime = DimensionCollectorRuntime()
collectors = [
DimensionCollector(
api_format="CLAUDE",
api_format="claude:chat",
task_type="chat",
dimension_name="input_tokens",
source_type="request",
@@ -73,7 +76,7 @@ class TestDimensionCollectorRuntime:
is_enabled=True,
),
DimensionCollector(
api_format="CLAUDE",
api_format="claude:chat",
task_type="chat",
dimension_name="cache_read_tokens",
source_type="request",
@@ -83,7 +86,7 @@ class TestDimensionCollectorRuntime:
is_enabled=True,
),
DimensionCollector(
api_format="CLAUDE",
api_format="claude:chat",
task_type="chat",
dimension_name="total_input_tokens",
source_type="computed",
@@ -103,3 +106,56 @@ class TestDimensionCollectorRuntime:
assert dims["input_tokens"] == 100
assert dims["cache_read_tokens"] == 20
assert dims["total_input_tokens"] == 120
class TestDimensionCollectorService:
def test_video_fallback_merges_base_collectors(self) -> None:
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"]

View File

@@ -0,0 +1,45 @@
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock
import pytest
from src.core.api_format.conversion.internal_video import InternalVideoPollResult, VideoStatus
from src.services.video.task_poller import VideoTaskPollerService
@pytest.mark.asyncio
async def test_poll_task_status_routes_gemini_video_to_gemini(
monkeypatch: pytest.MonkeyPatch,
) -> None:
poller = VideoTaskPollerService()
task = SimpleNamespace(
endpoint_id="e1",
key_id="k1",
provider_api_format="gemini:video",
external_task_id="operations/123",
)
endpoint = SimpleNamespace(id="e1", base_url="https://example.com", api_format="gemini:video")
key = SimpleNamespace(id="k1", api_key="enc")
monkeypatch.setattr(poller, "_get_endpoint", lambda _db, _id: endpoint)
monkeypatch.setattr(poller, "_get_key", lambda _db, _id: key)
monkeypatch.setattr(
"src.services.video.task_poller.crypto_service.decrypt", lambda _v: "decrypted"
)
auth_info = SimpleNamespace(auth_header="authorization", auth_value="Bearer x")
monkeypatch.setattr(
"src.services.video.task_poller.get_provider_auth",
AsyncMock(return_value=auth_info),
)
poll_gemini = AsyncMock(return_value=InternalVideoPollResult(status=VideoStatus.PROCESSING))
poll_openai = AsyncMock(return_value=InternalVideoPollResult(status=VideoStatus.PROCESSING))
monkeypatch.setattr(poller, "_poll_gemini", poll_gemini)
monkeypatch.setattr(poller, "_poll_openai", poll_openai)
result = await poller._poll_task_status(MagicMock(), task)
assert result.status == VideoStatus.PROCESSING
assert poll_gemini.await_count == 1
assert poll_openai.await_count == 0

View File

@@ -0,0 +1,172 @@
from types import SimpleNamespace
from typing import Any
from unittest.mock import AsyncMock, MagicMock
import pytest
from src.config.settings import config
from src.core.api_format.conversion.internal_video import VideoStatus
from src.services.billing.formula_engine import BillingIncompleteError
from src.services.task.impl.video_telemetry import VideoTelemetry
def _make_task(**overrides: Any) -> SimpleNamespace:
task = SimpleNamespace(
id="t1",
user_id="u1",
api_key_id="ak1",
provider_id="p1",
endpoint_id="e1",
key_id="k1",
external_task_id="ext-1",
client_api_format="openai:video",
provider_api_format="openai:video",
format_converted=False,
model="sora",
original_request_body={"model": "sora"},
duration_seconds=4,
resolution="720p",
aspect_ratio="16:9",
size="1024x1024",
retry_count=0,
video_size_bytes=None,
video_url="https://example.com/v.mp4",
video_urls=["https://example.com/v.mp4"],
submitted_at=None,
completed_at=None,
error_code=None,
error_message=None,
status=VideoStatus.COMPLETED.value,
request_metadata={"request_id": "req-1", "poll_raw_response": {"foo": "bar"}},
)
for k, v in overrides.items():
setattr(task, k, v)
return task
def _make_db() -> MagicMock:
db = MagicMock()
user_obj = SimpleNamespace(id="u1")
api_key_obj = SimpleNamespace(id="ak1")
provider_obj = SimpleNamespace(id="p1", name="prov1")
q_user = MagicMock()
q_user.filter.return_value.first.return_value = user_obj
q_key = MagicMock()
q_key.filter.return_value.first.return_value = api_key_obj
q_provider = MagicMock()
q_provider.filter.return_value.first.return_value = provider_obj
db.query.side_effect = [q_user, q_key, q_provider]
return db
@pytest.mark.asyncio
async def test_video_telemetry_failed_records_cost_zero(monkeypatch: pytest.MonkeyPatch) -> None:
db = _make_db()
task = _make_task(status=VideoStatus.FAILED.value, error_message="boom")
monkeypatch.setattr(
"src.services.task.impl.video_telemetry.DimensionCollectorService.collect_dimensions",
lambda _self, **_kwargs: {"duration_seconds": 4},
)
record = AsyncMock()
monkeypatch.setattr(
"src.services.task.impl.video_telemetry.UsageService.record_usage_with_custom_cost",
record,
)
telemetry = VideoTelemetry(db)
await telemetry.record_terminal_usage(task)
# billing_snapshot should be written back to task.request_metadata
assert task.request_metadata["billing_snapshot"]["billed_reason"] == "task_failed"
assert record.await_count == 1
kwargs = record.await_args.kwargs
assert kwargs["total_cost_usd"] == 0.0
assert kwargs["status"] == "failed"
@pytest.mark.asyncio
async def test_video_telemetry_completed_no_rule(monkeypatch: pytest.MonkeyPatch) -> None:
db = _make_db()
task = _make_task(status=VideoStatus.COMPLETED.value)
monkeypatch.setattr(
"src.services.task.impl.video_telemetry.DimensionCollectorService.collect_dimensions",
lambda _self, **_kwargs: {"duration_seconds": 4},
)
monkeypatch.setattr(
"src.services.task.impl.video_telemetry.BillingRuleService.find_rule",
lambda *_args, **_kwargs: None,
)
record = AsyncMock()
monkeypatch.setattr(
"src.services.task.impl.video_telemetry.UsageService.record_usage_with_custom_cost",
record,
)
telemetry = VideoTelemetry(db)
await telemetry.record_terminal_usage(task)
assert task.request_metadata["billing_snapshot"]["status"] == "no_rule"
kwargs = record.await_args.kwargs
assert kwargs["total_cost_usd"] == 0.0
assert kwargs["status"] == "completed"
@pytest.mark.asyncio
async def test_video_telemetry_strict_mode_missing_required_marks_failed(
monkeypatch: pytest.MonkeyPatch,
) -> None:
db = _make_db()
task = _make_task(status=VideoStatus.COMPLETED.value)
# provide a billing rule so formula path is taken
rule = SimpleNamespace(
id="r1",
name="video",
expression="duration_seconds",
variables={},
dimension_mappings={},
)
lookup = SimpleNamespace(rule=rule, scope="model", effective_task_type="video")
monkeypatch.setattr(
"src.services.task.impl.video_telemetry.DimensionCollectorService.collect_dimensions",
lambda _self, **_kwargs: {"duration_seconds": None},
)
monkeypatch.setattr(
"src.services.task.impl.video_telemetry.BillingRuleService.find_rule",
lambda *_args, **_kwargs: lookup,
)
record = AsyncMock()
monkeypatch.setattr(
"src.services.task.impl.video_telemetry.UsageService.record_usage_with_custom_cost",
record,
)
old = config.billing_strict_mode
try:
config.billing_strict_mode = True
telemetry = VideoTelemetry(db)
telemetry._formula_engine.evaluate = MagicMock(
side_effect=BillingIncompleteError(
"Missing required dimensions", missing_required=["duration_seconds"]
)
)
await telemetry.record_terminal_usage(task)
finally:
config.billing_strict_mode = old
assert task.status == VideoStatus.FAILED.value
assert task.video_url is None
assert task.video_urls is None
assert "billing_incomplete" in (task.error_code or "")
kwargs = record.await_args.kwargs
assert kwargs["total_cost_usd"] == 0.0
assert kwargs["status"] == "failed"