fix: WYSIWYG editor, content generation, and writing assistant bug fixes
- Fix text selection menu not showing: wire contentRef via inputRef on multiline TextField - Fix blog title not truncating: add min-w-0 for flex item overflow - Fix outline generation 500: escape curly braces in f-string prompt template - Fix content generation 'NoneType not callable': replace SessionLocal() with get_session_for_user(), add db param to MediumBlogGenerator, fix signature mismatch in database_task_manager - Fix writing assistant suggest 500: add auth + user_id to API endpoint and service, replace sync requests with httpx.AsyncClient - Fix hallucination detector 404: explicitly include router in main.py and app.py - Fix missing error_data in task failure responses - Hide CopilotKit web inspector button - Remove hardcoded fallback suggestions from SmartTypingAssist - Fix stale closure refs in SmartTypingAssist handleTypingChange - Add two-column editor layout, stats bar, section hover menu - Various subscription, billing, and research module improvements
This commit is contained in:
ajaysi
@@ -1,41 +1,60 @@
|
||||
"""
|
||||
Usage Tracking Service
|
||||
Comprehensive tracking of API usage, costs, and subscription limits.
|
||||
Usage Tracking Service - Refactored into modular components.
|
||||
|
||||
This file now serves as a facade that delegates to specialized modules
|
||||
in the usage_tracking_modules package.
|
||||
|
||||
Modules:
|
||||
- historical_usage: Functions for aggregating historical usage data
|
||||
- usage_stats: Functions for getting user usage statistics
|
||||
- usage_trends: Functions for usage trend analysis
|
||||
- limit_enforcement: Functions for enforcing usage limits
|
||||
- alerts: Functions for usage alerts
|
||||
"""
|
||||
|
||||
# Ensure Optional is available in global scope for dynamic imports
|
||||
from typing import Optional
|
||||
|
||||
import asyncio
|
||||
from typing import Dict, Any, List, Tuple
|
||||
from datetime import datetime, timedelta
|
||||
from typing import Dict, Any, Tuple, Optional
|
||||
from sqlalchemy.orm import Session
|
||||
from sqlalchemy import desc
|
||||
from sqlalchemy import text
|
||||
from loguru import logger
|
||||
import json
|
||||
from api.subscription.cache import clear_dashboard_cache
|
||||
from datetime import datetime, timedelta
|
||||
import time
|
||||
|
||||
from models.subscription_models import (
|
||||
APIUsageLog, UsageSummary, APIProvider, UsageAlert,
|
||||
UserSubscription, UsageStatus
|
||||
APIProvider, UsageStatus, UserSubscription,
|
||||
UsageSummary, APIUsageLog, UsageAlert
|
||||
)
|
||||
from .pricing_service import PricingService
|
||||
from .provider_detection import detect_actual_provider
|
||||
from .usage_tracking_helpers import (
|
||||
build_billing_periods,
|
||||
build_default_usage_percentages,
|
||||
build_empty_usage_response,
|
||||
from services.subscription.pricing_service import PricingService
|
||||
from services.subscription.provider_detection import detect_actual_provider
|
||||
from services.subscription.usage_tracking_helpers import (
|
||||
build_provider_breakdown,
|
||||
build_usage_trends_response,
|
||||
build_default_usage_percentages,
|
||||
calculate_final_total_cost,
|
||||
maybe_persist_reconciled_costs,
|
||||
build_usage_trends_response,
|
||||
build_billing_periods,
|
||||
query_usage_summaries,
|
||||
reset_usage_summary_counters,
|
||||
self_heal_summaries_from_logs,
|
||||
reset_usage_summary_counters,
|
||||
)
|
||||
# Import clear_dashboard_cache lazily to avoid circular import
|
||||
def _clear_dashboard_cache_for_user(user_id: str):
|
||||
from api.subscription.cache import clear_dashboard_cache as _clear
|
||||
return _clear(user_id)
|
||||
|
||||
from .usage_tracking_modules import (
|
||||
get_all_historical_usage,
|
||||
get_current_period_usage,
|
||||
get_usage_for_period,
|
||||
get_user_usage_stats,
|
||||
get_usage_trends,
|
||||
enforce_usage_limits,
|
||||
check_usage_alerts,
|
||||
create_usage_alert,
|
||||
)
|
||||
|
||||
|
||||
class UsageTrackingService:
|
||||
"""Service for tracking API usage and managing subscription limits."""
|
||||
"""Service for tracking API usage and managing billing information."""
|
||||
|
||||
def __init__(self, db: Session):
|
||||
self.db = db
|
||||
@@ -43,13 +62,14 @@ class UsageTrackingService:
|
||||
# TTL cache (30s) for enforcement results to cut DB chatter
|
||||
# key: f"{user_id}:{provider}", value: { 'result': (bool,str,dict), 'expires_at': datetime }
|
||||
self._enforce_cache: Dict[str, Dict[str, Any]] = {}
|
||||
|
||||
|
||||
def _get_authoritative_billing_period_keys(self, user_id: str, billing_period: Optional[str] = None) -> Dict[str, Any]:
|
||||
"""Return authoritative billing period lookup keys. Always uses calendar month for consistency."""
|
||||
"""Return authoritative billing period lookup keys. Always uses subscription period for consistency.
|
||||
Maintains backward compatibility with existing calendar-month data."""
|
||||
subscription = self.db.query(UserSubscription).filter(
|
||||
UserSubscription.user_id == user_id
|
||||
).first()
|
||||
|
||||
|
||||
# If caller explicitly requested a billing period, use it
|
||||
if billing_period:
|
||||
return {
|
||||
@@ -58,26 +78,125 @@ class UsageTrackingService:
|
||||
"period_start": subscription.current_period_start if subscription else None,
|
||||
"period_end": subscription.current_period_end if subscription else None,
|
||||
}
|
||||
|
||||
# ALWAYS use current calendar month for billing period to ensure consistency
|
||||
# This prevents data loss when subscription spans month boundaries
|
||||
current_period = datetime.now().strftime("%Y-%m")
|
||||
|
||||
# Get subscription period if available
|
||||
subscription_period = None
|
||||
if subscription and subscription.current_period_start:
|
||||
subscription_period = subscription.current_period_start.strftime("%Y-%m")
|
||||
|
||||
# Get calendar period
|
||||
calendar_period = datetime.now().strftime("%Y-%m")
|
||||
|
||||
# Check which period has usage data
|
||||
from models.subscription_models import UsageSummary
|
||||
|
||||
if subscription_period:
|
||||
# Check if data exists for subscription period
|
||||
sub_data = self.db.query(UsageSummary).filter(
|
||||
UsageSummary.user_id == user_id,
|
||||
UsageSummary.billing_period == subscription_period
|
||||
).first()
|
||||
|
||||
if sub_data:
|
||||
# Use subscription period (has data)
|
||||
return {
|
||||
"billing_period": subscription_period,
|
||||
"lookup_periods": [subscription_period],
|
||||
"period_start": subscription.current_period_start,
|
||||
"period_end": subscription.current_period_end,
|
||||
}
|
||||
|
||||
# No data for subscription period, check calendar period (backward compatibility)
|
||||
if calendar_period != subscription_period:
|
||||
cal_data = self.db.query(UsageSummary).filter(
|
||||
UsageSummary.user_id == user_id,
|
||||
UsageSummary.billing_period == calendar_period
|
||||
).first()
|
||||
|
||||
if cal_data:
|
||||
logger.info(f"Using calendar period {calendar_period} for backward compatibility (subscription period {subscription_period} has no data)")
|
||||
return {
|
||||
"billing_period": calendar_period,
|
||||
"lookup_periods": [calendar_period],
|
||||
"period_start": None,
|
||||
"period_end": None,
|
||||
}
|
||||
|
||||
# No data in either period, use subscription period
|
||||
return {
|
||||
"billing_period": subscription_period,
|
||||
"lookup_periods": [subscription_period],
|
||||
"period_start": subscription.current_period_start,
|
||||
"period_end": subscription.current_period_end,
|
||||
}
|
||||
|
||||
# No subscription, check for any existing data
|
||||
latest_summary = self.db.query(UsageSummary).filter(
|
||||
UsageSummary.user_id == user_id
|
||||
).order_by(UsageSummary.billing_period.desc()).first()
|
||||
|
||||
if latest_summary:
|
||||
logger.info(f"Using latest billing period from UsageSummary: {latest_summary.billing_period} for user {user_id}")
|
||||
return {
|
||||
"billing_period": latest_summary.billing_period,
|
||||
"lookup_periods": [latest_summary.billing_period],
|
||||
"period_start": None,
|
||||
"period_end": None,
|
||||
}
|
||||
|
||||
# Last fallback to calendar month for free tier / no subscription
|
||||
return {
|
||||
"billing_period": current_period,
|
||||
"lookup_periods": [current_period],
|
||||
"period_start": subscription.current_period_start if subscription else None,
|
||||
"period_end": subscription.current_period_end if subscription else None,
|
||||
"billing_period": calendar_period,
|
||||
"lookup_periods": [calendar_period],
|
||||
"period_start": None,
|
||||
"period_end": None,
|
||||
}
|
||||
|
||||
# Delegate to modular functions
|
||||
def get_user_usage_stats(self, user_id: str, billing_period: str = None) -> Dict[str, Any]:
|
||||
"""Get comprehensive usage statistics for a user."""
|
||||
return get_user_usage_stats(user_id, billing_period, self.db, self.pricing_service)
|
||||
|
||||
def _get_all_historical_usage(self, user_id: str) -> Dict[str, Any]:
|
||||
"""Get ALL historical usage data aggregated across all billing periods."""
|
||||
return get_all_historical_usage(user_id, self.db, self.pricing_service)
|
||||
|
||||
def get_current_period_usage(self, user_id: str) -> Dict[str, Any]:
|
||||
"""Get current billing period usage with correct per-period limit percentages."""
|
||||
return get_current_period_usage(user_id, self.db, self.pricing_service)
|
||||
|
||||
def get_usage_for_period(self, user_id: str, billing_period: str) -> Dict[str, Any]:
|
||||
"""Get usage for a specific billing period."""
|
||||
return get_usage_for_period(user_id, billing_period, self.db, self.pricing_service)
|
||||
|
||||
def get_usage_trends(self, user_id: str, months: int = 6) -> Dict[str, Any]:
|
||||
"""Get usage trends over time with self-healing from logs."""
|
||||
return get_usage_trends(user_id, months, self.db)
|
||||
|
||||
async def enforce_usage_limits(self, user_id: str, provider: APIProvider,
|
||||
tokens_requested: int = 0) -> Tuple[bool, str, Dict[str, Any]]:
|
||||
"""Enforce usage limits before making an API call."""
|
||||
return enforce_usage_limits(user_id, provider, tokens_requested, self.db, self.pricing_service)
|
||||
|
||||
async def _check_usage_alerts(self, user_id: str, provider: APIProvider, billing_period: str):
|
||||
"""Check if usage alerts should be sent."""
|
||||
check_usage_alerts(user_id, provider, billing_period, self.db, self.pricing_service)
|
||||
|
||||
async def _create_usage_alert(self, user_id: str, provider: APIProvider,
|
||||
threshold: int, current_usage: int, limit: int,
|
||||
billing_period: str):
|
||||
"""Create a usage alert."""
|
||||
create_usage_alert(user_id, provider, threshold, current_usage, limit, billing_period, self.db)
|
||||
|
||||
# Keep the track_api_usage method here as it's the core functionality
|
||||
async def track_api_usage(self, user_id: str, provider: APIProvider,
|
||||
endpoint: str, method: str, model_used: str = None,
|
||||
tokens_input: int = 0, tokens_output: int = 0,
|
||||
response_time: float = 0.0, status_code: int = 200,
|
||||
request_size: int = None, response_size: int = None,
|
||||
user_agent: str = None, ip_address: str = None,
|
||||
error_message: str = None, retry_count: int = 0,
|
||||
**kwargs) -> Dict[str, Any]:
|
||||
endpoint: str, method: str, model_used: str = None,
|
||||
tokens_input: int = 0, tokens_output: int = 0,
|
||||
response_time: float = 0.0, status_code: int = 200,
|
||||
request_size: int = None, response_size: int = None,
|
||||
user_agent: str = None, ip_address: str = None,
|
||||
error_message: str = None, retry_count: int = 0,
|
||||
**kwargs) -> Dict[str, Any]:
|
||||
"""Track an API usage event and update billing information."""
|
||||
|
||||
try:
|
||||
@@ -165,394 +284,81 @@ class UsageTrackingService:
|
||||
|
||||
# Invalidate dashboard cache so header stats update immediately
|
||||
try:
|
||||
clear_dashboard_cache(user_id)
|
||||
_clear_dashboard_cache_for_user(user_id)
|
||||
except Exception as cache_err:
|
||||
logger.debug(f"Could not clear dashboard cache: {cache_err}")
|
||||
|
||||
logger.info(f"Tracked API usage: {user_id} -> {provider.value} -> ${cost_data['cost_total']:.6f}")
|
||||
logger.warning(f"Failed to clear dashboard cache: {cache_err}")
|
||||
|
||||
return {
|
||||
'usage_logged': True,
|
||||
'cost': cost_data['cost_total'],
|
||||
'tokens_used': (tokens_input or 0) + (tokens_output or 0),
|
||||
'billing_period': billing_period
|
||||
"success": True,
|
||||
"cost": cost_data['cost_total'],
|
||||
"tokens": (tokens_input or 0) + (tokens_output or 0),
|
||||
"billing_period": billing_period
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error tracking API usage: {str(e)}")
|
||||
logger.error(f"Failed to track API usage: {e}")
|
||||
self.db.rollback()
|
||||
return {
|
||||
'usage_logged': False,
|
||||
'error': str(e)
|
||||
"success": False,
|
||||
"error": str(e)
|
||||
}
|
||||
|
||||
async def _update_usage_summary(self, user_id: str, provider: APIProvider,
|
||||
tokens_used: int, cost: float, billing_period: str,
|
||||
response_time: float, is_error: bool):
|
||||
"""Update the usage summary for a user."""
|
||||
tokens_used: int, cost: float,
|
||||
billing_period: str,
|
||||
response_time: float = 0.0,
|
||||
is_error: bool = False):
|
||||
"""Update or create usage summary for the billing period."""
|
||||
|
||||
# Get or create usage summary
|
||||
period_keys = self._get_authoritative_billing_period_keys(user_id, billing_period)
|
||||
# Get or create summary
|
||||
summary = self.db.query(UsageSummary).filter(
|
||||
UsageSummary.user_id == user_id,
|
||||
UsageSummary.billing_period.in_(period_keys["lookup_periods"])
|
||||
UsageSummary.billing_period == billing_period
|
||||
).first()
|
||||
|
||||
if not summary:
|
||||
logger.info(f"[UsageTracking] Creating new UsageSummary for user={user_id}, period={period_keys['billing_period']}")
|
||||
summary = UsageSummary(
|
||||
user_id=user_id,
|
||||
billing_period=period_keys["billing_period"]
|
||||
billing_period=billing_period,
|
||||
usage_status=UsageStatus.ACTIVE,
|
||||
total_calls=0,
|
||||
total_tokens=0,
|
||||
total_cost=0.0
|
||||
)
|
||||
self.db.add(summary)
|
||||
else:
|
||||
logger.debug(f"[UsageTracking] Found existing UsageSummary for user={user_id}, period={summary.billing_period}, calls={summary.total_calls}")
|
||||
|
||||
# Update provider-specific counters
|
||||
# Update counts
|
||||
summary.total_calls = (summary.total_calls or 0) + 1
|
||||
summary.total_tokens = (summary.total_tokens or 0) + tokens_used
|
||||
summary.total_cost = (summary.total_cost or 0.0) + cost
|
||||
|
||||
# Update provider-specific counts
|
||||
provider_name = provider.value
|
||||
current_calls = getattr(summary, f"{provider_name}_calls", 0)
|
||||
current_calls = getattr(summary, f"{provider_name}_calls", 0) or 0
|
||||
setattr(summary, f"{provider_name}_calls", current_calls + 1)
|
||||
|
||||
# Update token usage for LLM providers
|
||||
if provider in [APIProvider.GEMINI, APIProvider.OPENAI, APIProvider.ANTHROPIC, APIProvider.MISTRAL, APIProvider.WAVESPEED]:
|
||||
current_tokens = getattr(summary, f"{provider_name}_tokens", 0)
|
||||
setattr(summary, f"{provider_name}_tokens", current_tokens + tokens_used)
|
||||
# Update provider-specific tokens
|
||||
tokens_attr = f"{provider_name}_tokens"
|
||||
if hasattr(summary, tokens_attr):
|
||||
current_tokens = getattr(summary, tokens_attr, 0) or 0
|
||||
setattr(summary, tokens_attr, current_tokens + tokens_used)
|
||||
|
||||
# Update cost
|
||||
current_cost = getattr(summary, f"{provider_name}_cost", 0.0)
|
||||
setattr(summary, f"{provider_name}_cost", current_cost + cost)
|
||||
# Update provider-specific cost
|
||||
cost_attr = f"{provider_name}_cost"
|
||||
if hasattr(summary, cost_attr):
|
||||
current_cost = getattr(summary, cost_attr, 0.0) or 0.0
|
||||
setattr(summary, cost_attr, current_cost + cost)
|
||||
|
||||
# Update totals
|
||||
summary.total_calls += 1
|
||||
summary.total_tokens += tokens_used
|
||||
summary.total_cost += cost
|
||||
# Update response time (rolling average)
|
||||
if response_time > 0:
|
||||
current_avg = summary.avg_response_time or 0.0
|
||||
current_calls = summary.total_calls or 1
|
||||
summary.avg_response_time = ((current_avg * (current_calls - 1)) + response_time) / current_calls
|
||||
|
||||
# Update performance metrics
|
||||
if summary.total_calls > 0:
|
||||
# Update average response time
|
||||
total_response_time = summary.avg_response_time * (summary.total_calls - 1) + response_time
|
||||
summary.avg_response_time = total_response_time / summary.total_calls
|
||||
|
||||
# Update error rate
|
||||
if is_error:
|
||||
error_count = int(summary.error_rate * (summary.total_calls - 1) / 100) + 1
|
||||
summary.error_rate = (error_count / summary.total_calls) * 100
|
||||
else:
|
||||
error_count = int(summary.error_rate * (summary.total_calls - 1) / 100)
|
||||
summary.error_rate = (error_count / summary.total_calls) * 100
|
||||
|
||||
# Update usage status based on limits
|
||||
await self._update_usage_status(summary)
|
||||
# Update error rate
|
||||
if is_error:
|
||||
summary.error_count = (summary.error_count or 0) + 1
|
||||
total_calls = summary.total_calls or 1
|
||||
summary.error_rate = (summary.error_count / total_calls) * 100
|
||||
|
||||
summary.updated_at = datetime.utcnow()
|
||||
|
||||
async def _update_usage_status(self, summary: UsageSummary):
|
||||
"""Update usage status based on subscription limits."""
|
||||
|
||||
limits = self.pricing_service.get_user_limits(summary.user_id)
|
||||
if not limits:
|
||||
return
|
||||
|
||||
# Check various limits and determine status
|
||||
max_usage_percentage = 0.0
|
||||
|
||||
# Check cost limit
|
||||
cost_limit = limits['limits'].get('monthly_cost', 0)
|
||||
if cost_limit > 0:
|
||||
cost_usage_pct = (summary.total_cost / cost_limit) * 100
|
||||
max_usage_percentage = max(max_usage_percentage, cost_usage_pct)
|
||||
|
||||
# Check call limits for each provider
|
||||
for provider in APIProvider:
|
||||
provider_name = provider.value
|
||||
current_calls = getattr(summary, f"{provider_name}_calls", 0)
|
||||
call_limit = limits['limits'].get(f"{provider_name}_calls", 0)
|
||||
|
||||
if call_limit > 0:
|
||||
call_usage_pct = (current_calls / call_limit) * 100
|
||||
max_usage_percentage = max(max_usage_percentage, call_usage_pct)
|
||||
|
||||
# Update status based on highest usage percentage
|
||||
if max_usage_percentage >= 100:
|
||||
summary.usage_status = UsageStatus.LIMIT_REACHED
|
||||
elif max_usage_percentage >= 80:
|
||||
summary.usage_status = UsageStatus.WARNING
|
||||
else:
|
||||
summary.usage_status = UsageStatus.ACTIVE
|
||||
|
||||
async def _check_usage_alerts(self, user_id: str, provider: APIProvider, billing_period: str):
|
||||
"""Check if usage alerts should be sent."""
|
||||
|
||||
# Get current usage
|
||||
period_keys = self._get_authoritative_billing_period_keys(user_id, billing_period)
|
||||
summary = self.db.query(UsageSummary).filter(
|
||||
UsageSummary.user_id == user_id,
|
||||
UsageSummary.billing_period.in_(period_keys["lookup_periods"])
|
||||
).first()
|
||||
|
||||
if not summary:
|
||||
return
|
||||
|
||||
# Get user limits
|
||||
limits = self.pricing_service.get_user_limits(user_id)
|
||||
if not limits:
|
||||
return
|
||||
|
||||
# Check for alert thresholds (80%, 90%, 100%)
|
||||
thresholds = [80, 90, 100]
|
||||
|
||||
for threshold in thresholds:
|
||||
# Check if alert already sent for this threshold
|
||||
existing_alert = self.db.query(UsageAlert).filter(
|
||||
UsageAlert.user_id == user_id,
|
||||
UsageAlert.billing_period == billing_period,
|
||||
UsageAlert.threshold_percentage == threshold,
|
||||
UsageAlert.provider == provider,
|
||||
UsageAlert.is_sent == True
|
||||
).first()
|
||||
|
||||
if existing_alert:
|
||||
continue
|
||||
|
||||
# Check if threshold is reached
|
||||
provider_name = provider.value
|
||||
current_calls = getattr(summary, f"{provider_name}_calls", 0)
|
||||
call_limit = limits['limits'].get(f"{provider_name}_calls", 0)
|
||||
|
||||
if call_limit > 0:
|
||||
usage_percentage = (current_calls / call_limit) * 100
|
||||
|
||||
if usage_percentage >= threshold:
|
||||
await self._create_usage_alert(
|
||||
user_id=user_id,
|
||||
provider=provider,
|
||||
threshold=threshold,
|
||||
current_usage=current_calls,
|
||||
limit=call_limit,
|
||||
billing_period=billing_period
|
||||
)
|
||||
|
||||
async def _create_usage_alert(self, user_id: str, provider: APIProvider,
|
||||
threshold: int, current_usage: int, limit: int,
|
||||
billing_period: str):
|
||||
"""Create a usage alert."""
|
||||
|
||||
# Determine alert type and severity
|
||||
if threshold >= 100:
|
||||
alert_type = "limit_reached"
|
||||
severity = "error"
|
||||
title = f"API Limit Reached - {provider.value.title()}"
|
||||
message = f"You have reached your {provider.value} API limit of {limit:,} calls for this billing period."
|
||||
elif threshold >= 90:
|
||||
alert_type = "usage_warning"
|
||||
severity = "warning"
|
||||
title = f"API Usage Warning - {provider.value.title()}"
|
||||
message = f"You have used {current_usage:,} of {limit:,} {provider.value} API calls ({threshold}% of your limit)."
|
||||
else:
|
||||
alert_type = "usage_warning"
|
||||
severity = "info"
|
||||
title = f"API Usage Notice - {provider.value.title()}"
|
||||
message = f"You have used {current_usage:,} of {limit:,} {provider.value} API calls ({threshold}% of your limit)."
|
||||
|
||||
alert = UsageAlert(
|
||||
user_id=user_id,
|
||||
alert_type=alert_type,
|
||||
threshold_percentage=threshold,
|
||||
provider=provider,
|
||||
title=title,
|
||||
message=message,
|
||||
severity=severity,
|
||||
billing_period=billing_period
|
||||
)
|
||||
|
||||
self.db.add(alert)
|
||||
logger.info(f"Created usage alert for {user_id}: {title}")
|
||||
|
||||
def get_user_usage_stats(self, user_id: str, billing_period: str = None) -> Dict[str, Any]:
|
||||
"""Get comprehensive usage statistics for a user."""
|
||||
|
||||
if not user_id:
|
||||
logger.error("get_user_usage_stats called without user_id")
|
||||
raise ValueError("user_id is required")
|
||||
|
||||
requested_billing_period = billing_period
|
||||
period_keys = self._get_authoritative_billing_period_keys(user_id, requested_billing_period)
|
||||
billing_period = period_keys["billing_period"]
|
||||
|
||||
logger.debug(f"[get_user_usage_stats] user={user_id}, billing_period={billing_period}, lookup_periods={period_keys['lookup_periods']}")
|
||||
|
||||
# Get usage summary
|
||||
summary = self.db.query(UsageSummary).filter(
|
||||
UsageSummary.user_id == user_id,
|
||||
UsageSummary.billing_period.in_(period_keys["lookup_periods"])
|
||||
).first()
|
||||
|
||||
if summary:
|
||||
logger.debug(f"[get_user_usage_stats] Found summary: period={summary.billing_period}, calls={summary.total_calls}, cost={summary.total_cost}")
|
||||
else:
|
||||
logger.debug(f"[get_user_usage_stats] No summary found for user={user_id}, period={billing_period}")
|
||||
|
||||
# Get user limits
|
||||
limits = self.pricing_service.get_user_limits(user_id)
|
||||
|
||||
# Get recent alerts
|
||||
alerts = self.db.query(UsageAlert).filter(
|
||||
UsageAlert.user_id == user_id,
|
||||
UsageAlert.billing_period == billing_period,
|
||||
UsageAlert.is_read == False
|
||||
).order_by(UsageAlert.created_at.desc()).limit(10).all()
|
||||
|
||||
if not summary:
|
||||
# If no summary exists for current period, we should initialize it
|
||||
# This handles the "start of month" case where a user logs in but hasn't made calls yet
|
||||
if not requested_billing_period:
|
||||
logger.info(f"Initializing empty UsageSummary for user {user_id} in period {billing_period}")
|
||||
summary = UsageSummary(
|
||||
user_id=user_id,
|
||||
billing_period=billing_period,
|
||||
usage_status=UsageStatus.ACTIVE,
|
||||
total_calls=0,
|
||||
total_tokens=0,
|
||||
total_cost=0.0
|
||||
)
|
||||
try:
|
||||
self.db.add(summary)
|
||||
self.db.commit()
|
||||
self.db.refresh(summary)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to initialize summary: {e}")
|
||||
self.db.rollback()
|
||||
# Fallback to zero-struct return if DB write fails
|
||||
pass
|
||||
|
||||
if not summary: # Still no summary after attempt
|
||||
return build_empty_usage_response(
|
||||
billing_period=billing_period,
|
||||
limits=limits,
|
||||
providers=APIProvider,
|
||||
)
|
||||
|
||||
# Provider breakdown - calculate costs first, then use for percentages
|
||||
# Only include Gemini and HuggingFace (HuggingFace is stored under MISTRAL enum)
|
||||
provider_breakdown, resolved_costs, core_counts = build_provider_breakdown(
|
||||
db=self.db,
|
||||
user_id=user_id,
|
||||
billing_period=billing_period,
|
||||
summary=summary,
|
||||
)
|
||||
|
||||
summary_total_cost = summary.total_cost or 0.0
|
||||
calculated_total_cost, final_total_cost = calculate_final_total_cost(
|
||||
summary_total_cost=summary_total_cost,
|
||||
resolved_costs=resolved_costs,
|
||||
)
|
||||
|
||||
maybe_persist_reconciled_costs(
|
||||
db=self.db,
|
||||
summary=summary,
|
||||
summary_total_cost=summary_total_cost,
|
||||
calculated_total_cost=calculated_total_cost,
|
||||
final_total_cost=final_total_cost,
|
||||
resolved_costs=resolved_costs,
|
||||
)
|
||||
|
||||
# Calculate usage percentages - only for Gemini and HuggingFace
|
||||
# Use the calculated costs for accurate percentages
|
||||
usage_percentages = build_default_usage_percentages(APIProvider)
|
||||
if limits:
|
||||
# Gemini
|
||||
gemini_call_limit = limits['limits'].get("gemini_calls", 0) or 0
|
||||
if gemini_call_limit > 0:
|
||||
usage_percentages['gemini_calls'] = (core_counts['gemini_calls'] / gemini_call_limit) * 100
|
||||
|
||||
# HuggingFace (stored as mistral in database)
|
||||
mistral_call_limit = limits['limits'].get("mistral_calls", 0) or 0
|
||||
if mistral_call_limit > 0:
|
||||
usage_percentages['mistral_calls'] = (core_counts['mistral_calls'] / mistral_call_limit) * 100
|
||||
|
||||
# Cost usage percentage - use final_total_cost (calculated from logs if needed)
|
||||
cost_limit = limits['limits'].get('monthly_cost', 0) or 0
|
||||
if cost_limit > 0:
|
||||
usage_percentages['cost'] = (final_total_cost / cost_limit) * 100
|
||||
|
||||
return {
|
||||
'billing_period': billing_period,
|
||||
'usage_status': summary.usage_status.value if hasattr(summary.usage_status, 'value') else str(summary.usage_status),
|
||||
'total_calls': summary.total_calls or 0,
|
||||
'total_tokens': summary.total_tokens or 0,
|
||||
'total_cost': final_total_cost,
|
||||
'avg_response_time': summary.avg_response_time or 0.0,
|
||||
'error_rate': summary.error_rate or 0.0,
|
||||
'limits': limits,
|
||||
'provider_breakdown': provider_breakdown,
|
||||
'alerts': [
|
||||
{
|
||||
'id': alert.id,
|
||||
'type': alert.alert_type,
|
||||
'title': alert.title,
|
||||
'message': alert.message,
|
||||
'severity': alert.severity,
|
||||
'created_at': alert.created_at.isoformat()
|
||||
}
|
||||
for alert in alerts
|
||||
],
|
||||
'usage_percentages': usage_percentages,
|
||||
'last_updated': summary.updated_at.isoformat()
|
||||
}
|
||||
|
||||
def get_usage_trends(self, user_id: str, months: int = 6) -> Dict[str, Any]:
|
||||
"""Get usage trends over time with self-healing from logs."""
|
||||
periods = build_billing_periods(months)
|
||||
summary_dict = query_usage_summaries(self.db, user_id, periods)
|
||||
self_heal_summaries_from_logs(self.db, user_id, periods, summary_dict)
|
||||
return build_usage_trends_response(periods, summary_dict)
|
||||
|
||||
async def enforce_usage_limits(self, user_id: str, provider: APIProvider,
|
||||
tokens_requested: int = 0) -> Tuple[bool, str, Dict[str, Any]]:
|
||||
"""Enforce usage limits before making an API call."""
|
||||
# Check short-lived cache first (30s)
|
||||
cache_key = f"{user_id}:{provider.value}"
|
||||
now = datetime.utcnow()
|
||||
cached = self._enforce_cache.get(cache_key)
|
||||
if cached and cached.get('expires_at') and cached['expires_at'] > now:
|
||||
return tuple(cached['result']) # type: ignore
|
||||
|
||||
result = self.pricing_service.check_usage_limits(
|
||||
user_id=user_id,
|
||||
provider=provider,
|
||||
tokens_requested=tokens_requested
|
||||
)
|
||||
self._enforce_cache[cache_key] = {
|
||||
'result': result,
|
||||
'expires_at': now + timedelta(seconds=30)
|
||||
}
|
||||
return result
|
||||
|
||||
async def reset_current_billing_period(self, user_id: str) -> Dict[str, Any]:
|
||||
"""Reset usage status and counters for the current billing period (after plan renewal/change)."""
|
||||
period_keys = self._get_authoritative_billing_period_keys(user_id)
|
||||
billing_period = period_keys["billing_period"]
|
||||
summary = self.db.query(UsageSummary).filter(
|
||||
UsageSummary.user_id == user_id,
|
||||
UsageSummary.billing_period.in_(period_keys["lookup_periods"])
|
||||
).first()
|
||||
|
||||
if not summary:
|
||||
return {"reset": False, "reason": "no_summary"}
|
||||
|
||||
try:
|
||||
reset_usage_summary_counters(summary)
|
||||
self.db.commit()
|
||||
|
||||
# Invalidate dashboard cache so header stats update after reset
|
||||
try:
|
||||
clear_dashboard_cache(user_id)
|
||||
except Exception as cache_err:
|
||||
logger.debug(f"Could not clear dashboard cache: {cache_err}")
|
||||
|
||||
logger.info(f"Reset usage counters for user {user_id} in billing period {billing_period} after renewal")
|
||||
return {"reset": True, "counters_reset": True}
|
||||
except Exception as e:
|
||||
self.db.rollback()
|
||||
logger.error(f"Error resetting usage status: {e}")
|
||||
return {"reset": False, "error": str(e)}
|
||||
|
||||
Reference in New Issue
Block a user