39 lines
1.2 KiB
Python
39 lines
1.2 KiB
Python
from app.services.usage_tracker import UsageTracker
|
|
from app.services.smart_router import ENGINE_COST_PROFILES
|
|
|
|
|
|
class UsageRecorder:
|
|
def __init__(self, tracker: UsageTracker):
|
|
self.tracker = tracker
|
|
|
|
def calculate_cost(self, engine_type: str, input_tokens: int, output_tokens: int) -> float:
|
|
profile = ENGINE_COST_PROFILES.get(engine_type)
|
|
if not profile:
|
|
return 0.0
|
|
|
|
input_cost = (input_tokens / 1_000_000) * profile.input_price_per_million
|
|
output_cost = (output_tokens / 1_000_000) * profile.output_price_per_million
|
|
return round(input_cost + output_cost, 6)
|
|
|
|
def record(
|
|
self,
|
|
user_id: str,
|
|
brand_id: str | None,
|
|
engine_type: str,
|
|
query: str,
|
|
input_tokens: int,
|
|
output_tokens: int,
|
|
metadata: dict | None = None,
|
|
) -> None:
|
|
cost = self.calculate_cost(engine_type, input_tokens, output_tokens)
|
|
self.tracker.record(
|
|
user_id=user_id,
|
|
brand_id=brand_id or "",
|
|
engine_type=engine_type,
|
|
query=query,
|
|
input_tokens=input_tokens,
|
|
output_tokens=output_tokens,
|
|
cost=cost,
|
|
metadata=metadata or {},
|
|
)
|