wehub-resource-sync
74 lines
2.2 KiB
Python
74 lines
2.2 KiB
Python
from __future__ import annotations
|
|
|
|
from contextlib import asynccontextmanager
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from opensquilla.engine.turn_runner.harness import _TurnRunnerSessionTotalsAdapter
|
|
from opensquilla.engine.types import DoneEvent
|
|
|
|
|
|
class _Manager:
|
|
def __init__(self) -> None:
|
|
self.session = SimpleNamespace(
|
|
input_tokens=0,
|
|
output_tokens=0,
|
|
total_tokens=0,
|
|
estimated_cost_usd=0.0,
|
|
total_cost_usd=0.0,
|
|
billed_cost_usd=0.0,
|
|
estimated_cost_component_usd=0.0,
|
|
cost_source="none",
|
|
missing_cost_entries=0,
|
|
cache_read=0,
|
|
cache_write=0,
|
|
model_override=None,
|
|
)
|
|
|
|
async def get_session(self, session_key: str):
|
|
assert session_key == "agent:webchat:mixed-turn"
|
|
return self.session
|
|
|
|
async def update(self, session_key: str, **values):
|
|
assert session_key == "agent:webchat:mixed-turn"
|
|
for name, value in values.items():
|
|
setattr(self.session, name, value)
|
|
|
|
|
|
class _Runner:
|
|
def __init__(self) -> None:
|
|
self._session_manager = _Manager()
|
|
|
|
@asynccontextmanager
|
|
async def _session_write_context(self, session_key: str):
|
|
assert session_key == "agent:webchat:mixed-turn"
|
|
yield
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_session_totals_rollup_splits_mixed_turn_cost_components() -> None:
|
|
runner = _Runner()
|
|
adapter = _TurnRunnerSessionTotalsAdapter(runner) # type: ignore[arg-type]
|
|
done = DoneEvent(
|
|
input_tokens=100,
|
|
output_tokens=10,
|
|
cost_usd=0.03,
|
|
billed_cost=0.01,
|
|
cost_source="mixed",
|
|
model="deepseek/deepseek-v4-pro",
|
|
)
|
|
|
|
result = await adapter.rollup(
|
|
session_key="agent:webchat:mixed-turn",
|
|
done_event=done,
|
|
resolved_model="deepseek/deepseek-v4-pro",
|
|
)
|
|
|
|
assert result is not None
|
|
assert result.total_cost_usd == pytest.approx(0.03)
|
|
assert result.billed_cost_usd == pytest.approx(0.01)
|
|
assert result.estimated_cost_component_usd == pytest.approx(0.02)
|
|
assert result.cost_source == "mixed"
|
|
assert runner._session_manager.session.estimated_cost_component_usd == pytest.approx(0.02)
|