488 lines
20 KiB
Python
488 lines
20 KiB
Python
"""
|
|
Enhanced FastAPI Monitoring Middleware
|
|
Database-backed monitoring for API calls, errors, performance metrics, and usage tracking.
|
|
Includes comprehensive subscription-based usage monitoring and cost tracking.
|
|
"""
|
|
|
|
# Ensure Optional is available in global scope for dynamic imports
|
|
from typing import Optional
|
|
|
|
from fastapi import Request, Response
|
|
from fastapi.responses import JSONResponse
|
|
import time
|
|
import json
|
|
from datetime import datetime, timedelta
|
|
from typing import Dict, List, Any
|
|
from collections import defaultdict, deque
|
|
import asyncio
|
|
from loguru import logger
|
|
from sqlalchemy.orm import Session
|
|
from sqlalchemy import and_, func, case
|
|
from sqlalchemy.exc import OperationalError
|
|
import re
|
|
|
|
from models.api_monitoring import APIRequest, APIEndpointStats, SystemHealth, CachePerformance
|
|
from models.subscription_models import APIProvider
|
|
from .usage_tracking_service import UsageTrackingService
|
|
from .pricing_service import PricingService
|
|
|
|
|
|
from services.database import get_session_for_user, init_user_database
|
|
|
|
class DatabaseAPIMonitor:
|
|
"""Database-backed API monitoring with usage tracking and subscription management."""
|
|
|
|
def __init__(self):
|
|
self.cache_stats = {
|
|
'hits': 0,
|
|
'misses': 0,
|
|
'hit_rate': 0.0
|
|
}
|
|
# API provider detection patterns - Updated to match actual endpoints
|
|
self.provider_patterns = {
|
|
APIProvider.GEMINI: [
|
|
r'gemini', r'google.*ai'
|
|
],
|
|
APIProvider.OPENAI: [r'openai', r'gpt', r'chatgpt'],
|
|
APIProvider.ANTHROPIC: [r'anthropic', r'claude'],
|
|
APIProvider.MISTRAL: [r'mistral'],
|
|
APIProvider.TAVILY: [r'tavily'],
|
|
APIProvider.SERPER: [r'serper'],
|
|
APIProvider.METAPHOR: [r'metaphor', r'/exa'],
|
|
APIProvider.FIRECRAWL: [r'firecrawl']
|
|
}
|
|
|
|
def detect_api_provider(self, path: str, user_agent: str = None) -> Optional[APIProvider]:
|
|
"""Detect which API provider is being used based on request details."""
|
|
path_lower = path.lower()
|
|
user_agent_lower = (user_agent or '').lower()
|
|
|
|
# Permanently ignore internal route families that must not accrue or check provider usage
|
|
if path_lower.startswith('/api/onboarding/') or path_lower.startswith('/api/subscription/'):
|
|
return None
|
|
|
|
for provider, patterns in self.provider_patterns.items():
|
|
for pattern in patterns:
|
|
if re.search(pattern, path_lower) or re.search(pattern, user_agent_lower):
|
|
return provider
|
|
|
|
return None
|
|
|
|
def extract_usage_metrics(self, request_body: str = None, response_body: str = None) -> Dict[str, Any]:
|
|
"""Extract usage metrics from request/response bodies."""
|
|
metrics = {
|
|
'tokens_input': 0,
|
|
'tokens_output': 0,
|
|
'model_used': None,
|
|
'search_count': 0,
|
|
'image_count': 0,
|
|
'page_count': 0
|
|
}
|
|
|
|
try:
|
|
# Try to parse request body for input tokens/content
|
|
if request_body:
|
|
request_data = json.loads(request_body) if isinstance(request_body, str) else request_body
|
|
|
|
# Extract model information
|
|
if 'model' in request_data:
|
|
metrics['model_used'] = request_data['model']
|
|
|
|
# Estimate input tokens from prompt/content
|
|
if 'prompt' in request_data:
|
|
metrics['tokens_input'] = self._estimate_tokens(request_data['prompt'])
|
|
elif 'messages' in request_data:
|
|
total_content = ' '.join([msg.get('content', '') for msg in request_data['messages']])
|
|
metrics['tokens_input'] = self._estimate_tokens(total_content)
|
|
elif 'input' in request_data:
|
|
metrics['tokens_input'] = self._estimate_tokens(str(request_data['input']))
|
|
|
|
# Count specific request types
|
|
if 'query' in request_data or 'search' in request_data:
|
|
metrics['search_count'] = 1
|
|
if 'image' in request_data or 'generate_image' in request_data:
|
|
metrics['image_count'] = 1
|
|
if 'url' in request_data or 'crawl' in request_data:
|
|
metrics['page_count'] = 1
|
|
|
|
# Try to parse response body for output tokens
|
|
if response_body:
|
|
response_data = json.loads(response_body) if isinstance(response_body, str) else response_body
|
|
|
|
# Extract output content and estimate tokens
|
|
if 'text' in response_data:
|
|
metrics['tokens_output'] = self._estimate_tokens(response_data['text'])
|
|
elif 'content' in response_data:
|
|
metrics['tokens_output'] = self._estimate_tokens(str(response_data['content']))
|
|
elif 'choices' in response_data and response_data['choices']:
|
|
choice = response_data['choices'][0]
|
|
if 'message' in choice and 'content' in choice['message']:
|
|
metrics['tokens_output'] = self._estimate_tokens(choice['message']['content'])
|
|
|
|
# Extract actual token usage if provided by API
|
|
if 'usage' in response_data:
|
|
usage = response_data['usage']
|
|
if 'prompt_tokens' in usage:
|
|
metrics['tokens_input'] = usage['prompt_tokens']
|
|
if 'completion_tokens' in usage:
|
|
metrics['tokens_output'] = usage['completion_tokens']
|
|
|
|
except (json.JSONDecodeError, KeyError, TypeError) as e:
|
|
logger.debug(f"Could not extract usage metrics: {e}")
|
|
|
|
return metrics
|
|
|
|
def _estimate_tokens(self, text: str) -> int:
|
|
"""Estimate token count for text (rough approximation)."""
|
|
if not text:
|
|
return 0
|
|
# Rough estimation: 1.3 tokens per word on average
|
|
word_count = len(str(text).split())
|
|
return int(word_count * 1.3)
|
|
|
|
async def check_usage_limits_middleware(request: Request, user_id: str, request_body: str = None) -> Optional[JSONResponse]:
|
|
"""Check usage limits before processing request."""
|
|
if not user_id:
|
|
return None
|
|
|
|
# No special whitelist; onboarding/subscription are ignored by provider detection
|
|
try:
|
|
path = request.url.path
|
|
except Exception:
|
|
pass
|
|
|
|
db = None
|
|
try:
|
|
db = get_session_for_user(user_id)
|
|
if not db:
|
|
return None
|
|
|
|
api_monitor = DatabaseAPIMonitor()
|
|
|
|
# Detect if this is an API call that should be rate limited
|
|
api_provider = api_monitor.detect_api_provider(request.url.path, request.headers.get('user-agent'))
|
|
if not api_provider:
|
|
return None
|
|
|
|
# Use provided request body or read it if not provided
|
|
if request_body is None:
|
|
try:
|
|
if hasattr(request, '_body'):
|
|
request_body = request._body
|
|
else:
|
|
# Try to read body (this might not work in all cases)
|
|
body = await request.body()
|
|
request_body = body.decode('utf-8') if body else None
|
|
except:
|
|
pass
|
|
|
|
# Estimate tokens needed
|
|
tokens_requested = 0
|
|
if request_body:
|
|
usage_metrics = api_monitor.extract_usage_metrics(request_body)
|
|
tokens_requested = usage_metrics.get('tokens_input', 0)
|
|
|
|
# Check limits
|
|
usage_service = UsageTrackingService(db)
|
|
try:
|
|
can_proceed, message, usage_info = await usage_service.enforce_usage_limits(
|
|
user_id=user_id,
|
|
provider=api_provider,
|
|
tokens_requested=tokens_requested
|
|
)
|
|
|
|
if not can_proceed:
|
|
logger.warning(f"Usage limit exceeded for {user_id}: {message}")
|
|
return JSONResponse(
|
|
status_code=429,
|
|
content={
|
|
"error": "Usage limit exceeded",
|
|
"message": message,
|
|
"usage_info": usage_info,
|
|
"provider": api_provider.value
|
|
}
|
|
)
|
|
|
|
# Warn if approaching limits
|
|
if usage_info.get('call_usage_percentage', 0) >= 80 or usage_info.get('cost_usage_percentage', 0) >= 80:
|
|
logger.warning(f"User {user_id} approaching usage limits: {usage_info}")
|
|
|
|
except OperationalError as e:
|
|
if "no such table" in str(e):
|
|
logger.warning(f"Tables missing for user {user_id}, attempting initialization...")
|
|
try:
|
|
init_user_database(user_id)
|
|
# Don't retry immediately to avoid loops, just let this request pass
|
|
except Exception as init_error:
|
|
logger.error(f"Failed to initialize database for user {user_id}: {init_error}")
|
|
else:
|
|
raise e
|
|
|
|
return None
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error checking usage limits: {e}")
|
|
# Don't block requests if usage checking fails
|
|
return None
|
|
finally:
|
|
if db is not None:
|
|
db.close()
|
|
|
|
async def monitoring_middleware(request: Request, call_next):
|
|
"""Enhanced FastAPI middleware for monitoring API calls with usage tracking."""
|
|
start_time = time.time()
|
|
|
|
# Extract request details - Enhanced user identification
|
|
user_id = None
|
|
try:
|
|
# PRIORITY 1: Check request.state.user_id (set by API key injection middleware)
|
|
if hasattr(request.state, 'user_id') and request.state.user_id:
|
|
user_id = request.state.user_id
|
|
logger.debug(f"Monitoring: Using user_id from request.state: {user_id}")
|
|
|
|
# PRIORITY 2: Check query parameters
|
|
elif hasattr(request, 'query_params') and 'user_id' in request.query_params:
|
|
user_id = request.query_params['user_id']
|
|
elif hasattr(request, 'path_params') and 'user_id' in request.path_params:
|
|
user_id = request.path_params['user_id']
|
|
|
|
# PRIORITY 3: Check headers for user identification
|
|
elif 'x-user-id' in request.headers:
|
|
user_id = request.headers['x-user-id']
|
|
elif 'x-user-email' in request.headers:
|
|
user_id = request.headers['x-user-email'] # Use email as user identifier
|
|
elif 'x-session-id' in request.headers:
|
|
user_id = request.headers['x-session-id'] # Use session as fallback
|
|
|
|
# Check for authorization header with user info
|
|
elif 'authorization' in request.headers:
|
|
# Auth middleware should have set request.state.user_id
|
|
# If not, this indicates an authentication failure (likely expired token)
|
|
# Log at debug level to reduce noise - expired tokens are expected
|
|
# But we can try to decode token if we really needed to, but let's rely on auth middleware
|
|
pass
|
|
|
|
except Exception as e:
|
|
logger.debug(f"Error extracting user ID: {e}")
|
|
user_id = None
|
|
|
|
# Get database session if user identified
|
|
db = None
|
|
if user_id:
|
|
db = get_session_for_user(user_id)
|
|
|
|
# Capture request body for usage tracking (read once, safely)
|
|
request_body = None
|
|
try:
|
|
# Only read body for POST/PUT/PATCH requests to avoid issues
|
|
if request.method in ['POST', 'PUT', 'PATCH']:
|
|
if hasattr(request, '_body') and request._body:
|
|
request_body = request._body.decode('utf-8')
|
|
else:
|
|
# Read body only if it hasn't been read yet
|
|
try:
|
|
body = await request.body()
|
|
request_body = body.decode('utf-8') if body else None
|
|
except Exception as body_error:
|
|
logger.debug(f"Could not read request body: {body_error}")
|
|
request_body = None
|
|
except Exception as e:
|
|
logger.debug(f"Error capturing request body: {e}")
|
|
request_body = None
|
|
|
|
# Check usage limits before processing
|
|
limit_response = await check_usage_limits_middleware(request, user_id, request_body)
|
|
if limit_response:
|
|
if db: db.close()
|
|
return limit_response
|
|
|
|
try:
|
|
response = await call_next(request)
|
|
status_code = response.status_code
|
|
duration = time.time() - start_time
|
|
|
|
# Capture response body for usage tracking
|
|
response_body = None
|
|
try:
|
|
if hasattr(response, 'body'):
|
|
response_body = response.body.decode('utf-8') if response.body else None
|
|
elif hasattr(response, '_content'):
|
|
response_body = response._content.decode('utf-8') if response._content else None
|
|
except:
|
|
pass
|
|
|
|
# Track API usage if this is an API call to external providers
|
|
api_monitor = DatabaseAPIMonitor()
|
|
api_provider = api_monitor.detect_api_provider(request.url.path, request.headers.get('user-agent'))
|
|
if api_provider and user_id:
|
|
logger.info(f"Detected API call: {request.url.path} -> {api_provider.value} for user: {user_id}")
|
|
try:
|
|
# Extract usage metrics
|
|
usage_metrics = api_monitor.extract_usage_metrics(request_body, response_body)
|
|
|
|
# Track usage with the usage tracking service
|
|
if db:
|
|
usage_service = UsageTrackingService(db)
|
|
await usage_service.track_api_usage(
|
|
user_id=user_id,
|
|
provider=api_provider,
|
|
endpoint=request.url.path,
|
|
method=request.method,
|
|
model_used=usage_metrics.get('model_used'),
|
|
tokens_input=usage_metrics.get('tokens_input', 0),
|
|
tokens_output=usage_metrics.get('tokens_output', 0),
|
|
response_time=duration,
|
|
status_code=status_code,
|
|
request_size=len(request_body) if request_body else None,
|
|
response_size=len(response_body) if response_body else None,
|
|
user_agent=request.headers.get('user-agent'),
|
|
ip_address=request.client.host if request.client else None,
|
|
search_count=usage_metrics.get('search_count', 0),
|
|
image_count=usage_metrics.get('image_count', 0),
|
|
page_count=usage_metrics.get('page_count', 0)
|
|
)
|
|
except OperationalError as e:
|
|
if "no such table" in str(e):
|
|
# Tables missing, try to init (might happen if check_usage_limits was skipped or passed)
|
|
try:
|
|
init_user_database(user_id)
|
|
except:
|
|
pass
|
|
except Exception as usage_error:
|
|
logger.error(f"Error tracking API usage: {usage_error}")
|
|
# Don't fail the main request if usage tracking fails
|
|
|
|
return response
|
|
|
|
except Exception as e:
|
|
duration = time.time() - start_time
|
|
status_code = 500
|
|
|
|
# Check for missing tables and try to self-heal
|
|
if "no such table" in str(e) and user_id:
|
|
logger.warning(f"Tables missing for user {user_id} during request processing, attempting initialization...")
|
|
try:
|
|
init_user_database(user_id)
|
|
logger.info(f"Database initialized for user {user_id}. Request failed but next should succeed.")
|
|
return JSONResponse(
|
|
status_code=503, # Service Unavailable (temporary)
|
|
content={"error": "Database initialized. Please retry request."}
|
|
)
|
|
except Exception as init_error:
|
|
logger.error(f"Failed to initialize database for user {user_id}: {init_error}")
|
|
|
|
# Store minimal error info
|
|
logger.error(f"API Error: {request.method} {request.url.path} - {str(e)}")
|
|
|
|
return JSONResponse(
|
|
status_code=500,
|
|
content={"error": "Internal server error"}
|
|
)
|
|
finally:
|
|
if db:
|
|
db.close()
|
|
|
|
async def get_monitoring_stats(minutes: int = 5) -> Dict[str, Any]:
|
|
"""Get current monitoring statistics."""
|
|
# Placeholder to match old API; heavy stats handled elsewhere
|
|
return {
|
|
'timestamp': datetime.utcnow().isoformat(),
|
|
'overview': {
|
|
'recent_requests': 0,
|
|
'recent_errors': 0,
|
|
},
|
|
'cache_performance': {'hits': 0, 'misses': 0, 'hit_rate': 0.0},
|
|
'recent_errors': [],
|
|
'system_health': {'status': 'healthy', 'error_rate': 0.0}
|
|
}
|
|
|
|
async def get_lightweight_stats(user_id: str) -> Dict[str, Any]:
|
|
"""Get lightweight stats for dashboard header.
|
|
|
|
Optimized single-query approach using conditional aggregation for better performance.
|
|
"""
|
|
db = None
|
|
try:
|
|
db = get_session_for_user(user_id)
|
|
if not db:
|
|
return {
|
|
'status': 'unknown',
|
|
'icon': '⚪',
|
|
'recent_requests': 0,
|
|
'recent_errors': 0,
|
|
'error_rate': 0.0,
|
|
'timestamp': datetime.utcnow().isoformat()
|
|
}
|
|
now = datetime.utcnow()
|
|
|
|
# Get stats from last 5 minutes
|
|
five_minutes_ago = now - timedelta(minutes=5)
|
|
|
|
# Optimized: Single query with conditional aggregation instead of two separate queries
|
|
# This is much faster as it only scans the table once
|
|
stats = db.query(
|
|
func.count(APIRequest.id).label('total_requests'),
|
|
func.sum(
|
|
case((APIRequest.status_code >= 400, 1), else_=0)
|
|
).label('total_errors')
|
|
).filter(
|
|
APIRequest.timestamp >= five_minutes_ago
|
|
).first()
|
|
|
|
recent_requests = stats.total_requests or 0 if stats else 0
|
|
recent_errors = int(stats.total_errors or 0) if stats else 0
|
|
|
|
# Calculate error rate
|
|
error_rate = (recent_errors / recent_requests * 100) if recent_requests > 0 else 0.0
|
|
|
|
# Determine status based on error rate
|
|
if error_rate > 10:
|
|
status = 'critical'
|
|
icon = '🔴'
|
|
elif error_rate > 5:
|
|
status = 'warning'
|
|
icon = '🟡'
|
|
else:
|
|
status = 'healthy'
|
|
icon = '🟢'
|
|
|
|
return {
|
|
'status': status,
|
|
'icon': icon,
|
|
'recent_requests': recent_requests,
|
|
'recent_errors': recent_errors,
|
|
'error_rate': round(error_rate, 2),
|
|
'timestamp': now.isoformat()
|
|
}
|
|
except OperationalError as e:
|
|
if "no such table" in str(e):
|
|
logger.warning(f"Tables missing for user {user_id} in lightweight stats, attempting initialization...")
|
|
try:
|
|
init_user_database(user_id)
|
|
except Exception as init_error:
|
|
logger.error(f"Failed to initialize database for user {user_id}: {init_error}")
|
|
|
|
# Return default healthy state on error/missing table
|
|
return {
|
|
'status': 'healthy',
|
|
'icon': '🟢',
|
|
'recent_requests': 0,
|
|
'recent_errors': 0,
|
|
'error_rate': 0.0,
|
|
'timestamp': datetime.utcnow().isoformat()
|
|
}
|
|
except Exception as e:
|
|
logger.error(f"Error getting lightweight stats: {e}", exc_info=True)
|
|
# Return default healthy state on error
|
|
return {
|
|
'status': 'healthy',
|
|
'icon': '🟢',
|
|
'recent_requests': 0,
|
|
'recent_errors': 0,
|
|
'error_rate': 0.0,
|
|
'timestamp': datetime.utcnow().isoformat()
|
|
}
|
|
finally:
|
|
if db is not None:
|
|
db.close()
|