mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
feat: 添加视频生成功能增强和多维度计费系统适配
- 视频生成: 增强 video_handler,重构 task_poller,新增 telemetry - 计费系统: 适配新的 signature 格式,支持 video 任务类型回退 - 数据库迁移: 添加 api_family/endpoint_kind 字段和 video_formats
This commit is contained in:
@@ -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"]
|
||||
|
||||
45
tests/services/test_video_task_poller.py
Normal file
45
tests/services/test_video_task_poller.py
Normal 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
|
||||
172
tests/services/test_video_telemetry.py
Normal file
172
tests/services/test_video_telemetry.py
Normal 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"
|
||||
Reference in New Issue
Block a user