""" API Key model """ from datetime import datetime, timedelta from typing import Optional, Dict, Any, List from sqlalchemy import ( Column, Integer, String, DateTime, Boolean, Text, JSON, ForeignKey, ) from sqlalchemy.orm import relationship from app.db.database import Base class APIKey(Base): """API Key model for authentication and access control""" __tablename__ = "api_keys" id = Column(Integer, primary_key=True, index=True) name = Column(String, nullable=False) # Human-readable name for the API key key_hash = Column(String, unique=True, index=True, nullable=False) # Hashed API key key_prefix = Column( String, index=True, nullable=False ) # First 8 characters for identification # User relationship user_id = Column(Integer, ForeignKey("users.id"), nullable=False) user = relationship("User", back_populates="api_keys") # Related data relationships budgets = relationship( "Budget", back_populates="api_key", cascade="all, delete-orphan" ) usage_tracking = relationship( "UsageTracking", back_populates="api_key", cascade="all, delete-orphan" ) # Key status and permissions is_active = Column(Boolean, default=True) permissions = Column(JSON, default=dict) # Specific permissions for this key scopes = Column(JSON, default=list) # OAuth-like scopes # Usage limits rate_limit_per_minute = Column(Integer, default=60) # Requests per minute rate_limit_per_hour = Column(Integer, default=3600) # Requests per hour rate_limit_per_day = Column(Integer, default=86400) # Requests per day # Allowed resources allowed_models = Column(JSON, default=list) # List of allowed LLM models allowed_endpoints = Column(JSON, default=list) # List of allowed API endpoints allowed_ips = Column(JSON, default=list) # IP whitelist allowed_chatbots = Column( JSON, default=list ) # List of allowed chatbot IDs for chatbot-specific keys # Budget configuration is_unlimited = Column(Boolean, default=True) # Unlimited budget flag budget_limit_cents = Column(Integer, nullable=True) # Budget limit in cents budget_type = Column(String, nullable=True) # "total" or "monthly" # Metadata description = Column(Text, nullable=True) tags = Column(JSON, default=list) # For organizing keys # Timestamps created_at = Column(DateTime, default=datetime.utcnow) updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow) last_used_at = Column(DateTime, nullable=True) expires_at = Column(DateTime, nullable=True) # Optional expiration # Usage tracking total_requests = Column(Integer, default=0) total_tokens = Column(Integer, default=0) total_cost = Column(Integer, default=0) # In cents # Relationships usage_tracking = relationship( "UsageTracking", back_populates="api_key", cascade="all, delete-orphan" ) budgets = relationship( "Budget", back_populates="api_key", cascade="all, delete-orphan" ) plugin_audit_logs = relationship("PluginAuditLog", back_populates="api_key") def __repr__(self): return f"" def to_dict(self, include_sensitive: bool = False): """Convert API key to dictionary for API responses""" data = { "id": self.id, "name": self.name, "key_prefix": self.key_prefix, "user_id": self.user_id, "is_active": self.is_active, "permissions": self.permissions, "scopes": self.scopes, "rate_limit_per_minute": self.rate_limit_per_minute, "rate_limit_per_hour": self.rate_limit_per_hour, "rate_limit_per_day": self.rate_limit_per_day, "allowed_models": self.allowed_models, "allowed_endpoints": self.allowed_endpoints, "allowed_ips": self.allowed_ips, "allowed_chatbots": self.allowed_chatbots, "description": self.description, "tags": self.tags, "created_at": self.created_at.isoformat() if self.created_at else None, "updated_at": self.updated_at.isoformat() if self.updated_at else None, "last_used_at": self.last_used_at.isoformat() if self.last_used_at else None, "expires_at": self.expires_at.isoformat() if self.expires_at else None, "total_requests": self.total_requests, "total_tokens": self.total_tokens, "total_cost_cents": self.total_cost, "is_unlimited": self.is_unlimited, "budget_limit": self.budget_limit_cents, # Map to budget_limit for API response "budget_type": self.budget_type, } if include_sensitive: data["key_hash"] = self.key_hash return data def is_expired(self) -> bool: """Check if the API key has expired""" if self.expires_at is None: return False return datetime.utcnow() > self.expires_at def is_valid(self) -> bool: """Check if the API key is valid and active""" return self.is_active and not self.is_expired() def has_permission(self, permission: str) -> bool: """Check if the API key has a specific permission""" return permission in self.permissions def has_scope(self, scope: str) -> bool: """Check if the API key has a specific scope""" return scope in self.scopes def can_access_model(self, model_name: str) -> bool: """Check if the API key can access a specific model""" if not self.allowed_models: # Empty list means all models allowed return True return model_name in self.allowed_models def can_access_endpoint(self, endpoint: str) -> bool: """Check if the API key can access a specific endpoint""" if not self.allowed_endpoints: # Empty list means all endpoints allowed return True return endpoint in self.allowed_endpoints def can_access_from_ip(self, ip_address: str) -> bool: """Check if the API key can be used from a specific IP""" if not self.allowed_ips: # Empty list means all IPs allowed return True return ip_address in self.allowed_ips def can_access_chatbot(self, chatbot_id: str) -> bool: """Check if the API key can access a specific chatbot""" if not self.allowed_chatbots: # Empty list means all chatbots allowed return True return chatbot_id in self.allowed_chatbots def update_usage(self, tokens_used: int = 0, cost_cents: int = 0): """Update usage statistics""" self.total_requests += 1 self.total_tokens += tokens_used self.total_cost += cost_cents self.last_used_at = datetime.utcnow() def set_expiration(self, days: int): """Set expiration date in days from now""" self.expires_at = datetime.utcnow() + timedelta(days=days) def extend_expiration(self, days: int): """Extend expiration date by specified days""" if self.expires_at is None: self.expires_at = datetime.utcnow() + timedelta(days=days) else: self.expires_at = self.expires_at + timedelta(days=days) def revoke(self): """Revoke the API key""" self.is_active = False self.updated_at = datetime.utcnow() def add_scope(self, scope: str): """Add a scope to the API key""" if scope not in self.scopes: self.scopes.append(scope) def remove_scope(self, scope: str): """Remove a scope from the API key""" if scope in self.scopes: self.scopes.remove(scope) def add_allowed_model(self, model_name: str): """Add an allowed model""" if model_name not in self.allowed_models: self.allowed_models.append(model_name) def remove_allowed_model(self, model_name: str): """Remove an allowed model""" if model_name in self.allowed_models: self.allowed_models.remove(model_name) def add_allowed_endpoint(self, endpoint: str): """Add an allowed endpoint""" if endpoint not in self.allowed_endpoints: self.allowed_endpoints.append(endpoint) def remove_allowed_endpoint(self, endpoint: str): """Remove an allowed endpoint""" if endpoint in self.allowed_endpoints: self.allowed_endpoints.remove(endpoint) def add_allowed_ip(self, ip_address: str): """Add an allowed IP address""" if ip_address not in self.allowed_ips: self.allowed_ips.append(ip_address) def remove_allowed_ip(self, ip_address: str): """Remove an allowed IP address""" if ip_address in self.allowed_ips: self.allowed_ips.remove(ip_address) def add_allowed_chatbot(self, chatbot_id: str): """Add an allowed chatbot""" if chatbot_id not in self.allowed_chatbots: self.allowed_chatbots.append(chatbot_id) def remove_allowed_chatbot(self, chatbot_id: str): """Remove an allowed chatbot""" if chatbot_id in self.allowed_chatbots: self.allowed_chatbots.remove(chatbot_id) @classmethod def create_default_key( cls, user_id: int, name: str, key_hash: str, key_prefix: str ) -> "APIKey": """Create a default API key with standard permissions""" return cls( name=name, key_hash=key_hash, key_prefix=key_prefix, user_id=user_id, is_active=True, permissions={"read": True, "write": True, "chat": True, "embeddings": True}, scopes=["chat.completions", "embeddings.create", "models.list"], rate_limit_per_minute=60, rate_limit_per_hour=3600, rate_limit_per_day=86400, allowed_models=[], # All models allowed by default allowed_endpoints=[], # All endpoints allowed by default allowed_ips=[], # All IPs allowed by default description="Default API key with standard permissions", tags=["default"], ) @classmethod def create_restricted_key( cls, user_id: int, name: str, key_hash: str, key_prefix: str, models: List[str], endpoints: List[str], ) -> "APIKey": """Create a restricted API key with limited permissions""" return cls( name=name, key_hash=key_hash, key_prefix=key_prefix, user_id=user_id, is_active=True, permissions={"read": True, "chat": True}, scopes=["chat.completions"], rate_limit_per_minute=30, rate_limit_per_hour=1800, rate_limit_per_day=43200, allowed_models=models, allowed_endpoints=endpoints, allowed_ips=[], description="Restricted API key with limited permissions", tags=["restricted"], ) @classmethod def create_chatbot_key( cls, user_id: int, name: str, key_hash: str, key_prefix: str, chatbot_id: str, chatbot_name: str, ) -> "APIKey": """Create a chatbot-specific API key""" return cls( name=name, key_hash=key_hash, key_prefix=key_prefix, user_id=user_id, is_active=True, permissions={"chatbot": True}, scopes=["chatbot.chat"], rate_limit_per_minute=100, rate_limit_per_hour=6000, rate_limit_per_day=144000, allowed_models=[], # Will use chatbot's configured model allowed_endpoints=[ f"/api/v1/chatbot/external/{chatbot_id}/chat", f"/api/v1/chatbot/external/{chatbot_id}/chat/completions", ], allowed_ips=[], allowed_chatbots=[chatbot_id], description=f"API key for chatbot: {chatbot_name}", tags=["chatbot", f"chatbot-{chatbot_id}"], )