black formatting

This commit is contained in:
Adil Hafeez 2026-04-17 01:02:12 -07:00
parent d30018cf35
commit d09fa97568
5 changed files with 52 additions and 13 deletions

View file

@ -28,7 +28,9 @@ def _mk_attr(key: str, value):
return kv
def _mk_span(attrs: dict, start_ns: int | None = None, span_id_hex: str = "ab") -> MagicMock:
def _mk_span(
attrs: dict, start_ns: int | None = None, span_id_hex: str = "ab"
) -> MagicMock:
span = MagicMock()
span.attributes = [_mk_attr(k, v) for k, v in attrs.items()]
span.start_time_unix_nano = start_ns or int(time.time() * 1_000_000_000)
@ -84,7 +86,9 @@ def test_pricing_lookup_attaches_cost():
class StubPricing:
def cost_for_call(self, call):
# Simple: 2 * prompt + 3 * completion, in cents
return 0.02 * (call.prompt_tokens or 0) + 0.03 * (call.completion_tokens or 0)
return 0.02 * (call.prompt_tokens or 0) + 0.03 * (
call.completion_tokens or 0
)
span = _mk_span(
{

View file

@ -4,7 +4,17 @@ from planoai.obs.collector import LLMCall
from planoai.obs.render import aggregates, model_rollups, route_hits
def _call(model: str, ts: datetime, prompt=0, completion=0, cost=None, route=None, session=None, cache_read=0, cache_write=0):
def _call(
model: str,
ts: datetime,
prompt=0,
completion=0,
cost=None,
route=None,
session=None,
cache_read=0,
cache_write=0,
):
return LLMCall(
request_id="r",
timestamp=ts,
@ -22,9 +32,30 @@ def _call(model: str, ts: datetime, prompt=0, completion=0, cost=None, route=Non
def test_aggregates_sum_and_session_counts():
now = datetime.now(tz=timezone.utc).astimezone()
calls = [
_call("m1", now - timedelta(seconds=50), prompt=10, completion=5, cost=0.001, session="s1"),
_call("m2", now - timedelta(seconds=40), prompt=20, completion=10, cost=0.002, session="s1"),
_call("m1", now - timedelta(seconds=30), prompt=30, completion=15, cost=0.003, session="s2"),
_call(
"m1",
now - timedelta(seconds=50),
prompt=10,
completion=5,
cost=0.001,
session="s1",
),
_call(
"m2",
now - timedelta(seconds=40),
prompt=20,
completion=10,
cost=0.002,
session="s1",
),
_call(
"m1",
now - timedelta(seconds=30),
prompt=30,
completion=15,
cost=0.003,
session="s2",
),
]
stats = aggregates(calls)
assert stats.count == 3
@ -38,7 +69,9 @@ def test_aggregates_sum_and_session_counts():
def test_rollups_split_by_model_and_cache():
now = datetime.now(tz=timezone.utc).astimezone()
calls = [
_call("m1", now, prompt=10, completion=5, cost=0.001, cache_write=3, cache_read=7),
_call(
"m1", now, prompt=10, completion=5, cost=0.001, cache_write=3, cache_read=7
),
_call("m1", now, prompt=20, completion=10, cost=0.002, cache_read=1),
_call("m2", now, prompt=30, completion=15, cost=0.004),
]