90 lines
3 KiB
Python
90 lines
3 KiB
Python
"""Usage Logging Utility"""
|
|
from sqlalchemy.orm import Session
|
|
from app.models.usage import UsageLog
|
|
from app.models.pricing import ModelPricing
|
|
from app.models.job import Job
|
|
from typing import Optional, Dict, Any
|
|
import logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
def log_model_usage(
|
|
db: Session,
|
|
job_id: str,
|
|
user_id: str,
|
|
module: str,
|
|
action: str,
|
|
provider: str,
|
|
model: str,
|
|
request_metadata: Dict[str, Any] = {},
|
|
response_metadata: Dict[str, Any] = {},
|
|
usage_stats: Dict[str, Any] = {}
|
|
) -> Optional[UsageLog]:
|
|
"""
|
|
Log model usage and calculate estimated cost.
|
|
|
|
usage_stats keys:
|
|
- input_tokens (int)
|
|
- output_tokens (int)
|
|
- images (int)
|
|
- seconds (float)
|
|
- characters (int)
|
|
"""
|
|
try:
|
|
# 1. Calculate Cost
|
|
cost = 0.0
|
|
|
|
# Find pricing record
|
|
pricing = db.query(ModelPricing).filter(
|
|
ModelPricing.provider == provider,
|
|
ModelPricing.model_name == model
|
|
).first()
|
|
|
|
# If specific model not found, try generic provider/default fallback?
|
|
# For now, just log 0 if not found, or maybe try mapping aliases.
|
|
|
|
if pricing:
|
|
if usage_stats.get("input_tokens"):
|
|
cost += float(pricing.cost_per_input_token) * usage_stats["input_tokens"]
|
|
|
|
if usage_stats.get("output_tokens"):
|
|
cost += float(pricing.cost_per_output_token) * usage_stats["output_tokens"]
|
|
|
|
if usage_stats.get("images"):
|
|
cost += float(pricing.cost_per_image) * usage_stats["images"]
|
|
|
|
if usage_stats.get("seconds"):
|
|
cost += float(pricing.cost_per_second) * usage_stats["seconds"]
|
|
|
|
if usage_stats.get("characters"):
|
|
cost += float(pricing.cost_per_1k_chars) * (usage_stats["characters"] / 1000.0)
|
|
|
|
if usage_stats.get("requests"): # generic per-request
|
|
cost += float(pricing.cost_per_request) * usage_stats["requests"]
|
|
|
|
# Special case for "per request" if not specified but implies 1?
|
|
# If generated 1 image and pricing is per image... handled above.
|
|
|
|
# 2. Create Log Record
|
|
log_entry = UsageLog(
|
|
job_id=job_id,
|
|
user_id=user_id,
|
|
module=module,
|
|
action=action,
|
|
api_provider=provider,
|
|
api_model=model,
|
|
tokens_input=usage_stats.get("input_tokens"),
|
|
tokens_output=usage_stats.get("output_tokens"),
|
|
estimated_cost_usd=cost,
|
|
processing_time_ms=usage_stats.get("processing_time_ms"),
|
|
request_metadata=request_metadata,
|
|
response_metadata=response_metadata
|
|
)
|
|
|
|
db.add(log_entry)
|
|
db.commit()
|
|
return log_entry
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to log usage: {e}")
|
|
return None
|