diff --git a/surfsense_backend/app/db.py b/surfsense_backend/app/db.py index 6e9553307..4689313f7 100644 --- a/surfsense_backend/app/db.py +++ b/surfsense_backend/app/db.py @@ -260,6 +260,24 @@ class ImageGenProvider(StrEnum): NSCALE = "NSCALE" +class VisionProvider(StrEnum): + OPENAI = "OPENAI" + ANTHROPIC = "ANTHROPIC" + GOOGLE = "GOOGLE" + AZURE_OPENAI = "AZURE_OPENAI" + VERTEX_AI = "VERTEX_AI" + BEDROCK = "BEDROCK" + XAI = "XAI" + OPENROUTER = "OPENROUTER" + OLLAMA = "OLLAMA" + GROQ = "GROQ" + TOGETHER_AI = "TOGETHER_AI" + FIREWORKS_AI = "FIREWORKS_AI" + DEEPSEEK = "DEEPSEEK" + MISTRAL = "MISTRAL" + CUSTOM = "CUSTOM" + + class LogLevel(StrEnum): DEBUG = "DEBUG" INFO = "INFO" @@ -377,6 +395,11 @@ class Permission(StrEnum): IMAGE_GENERATIONS_READ = "image_generations:read" IMAGE_GENERATIONS_DELETE = "image_generations:delete" + # Vision LLM Configs + VISION_CONFIGS_CREATE = "vision_configs:create" + VISION_CONFIGS_READ = "vision_configs:read" + VISION_CONFIGS_DELETE = "vision_configs:delete" + # Connectors CONNECTORS_CREATE = "connectors:create" CONNECTORS_READ = "connectors:read" @@ -445,6 +468,9 @@ DEFAULT_ROLE_PERMISSIONS = { # Image Generations (create and read, no delete) Permission.IMAGE_GENERATIONS_CREATE.value, Permission.IMAGE_GENERATIONS_READ.value, + # Vision Configs (create and read, no delete) + Permission.VISION_CONFIGS_CREATE.value, + Permission.VISION_CONFIGS_READ.value, # Connectors (no delete) Permission.CONNECTORS_CREATE.value, Permission.CONNECTORS_READ.value, @@ -478,6 +504,8 @@ DEFAULT_ROLE_PERMISSIONS = { Permission.VIDEO_PRESENTATIONS_READ.value, # Image Generations (read only) Permission.IMAGE_GENERATIONS_READ.value, + # Vision Configs (read only) + Permission.VISION_CONFIGS_READ.value, # Connectors (read only) Permission.CONNECTORS_READ.value, # Logs (read only) @@ -1263,6 +1291,35 @@ class ImageGenerationConfig(BaseModel, TimestampMixin): user = relationship("User", back_populates="image_generation_configs") +class VisionLLMConfig(BaseModel, TimestampMixin): + __tablename__ = "vision_llm_configs" + + name = Column(String(100), nullable=False, index=True) + description = Column(String(500), nullable=True) + + provider = Column(SQLAlchemyEnum(VisionProvider), nullable=False) + custom_provider = Column(String(100), nullable=True) + model_name = Column(String(100), nullable=False) + + api_key = Column(String, nullable=False) + api_base = Column(String(500), nullable=True) + api_version = Column(String(50), nullable=True) + + litellm_params = Column(JSON, nullable=True, default={}) + + search_space_id = Column( + Integer, ForeignKey("searchspaces.id", ondelete="CASCADE"), nullable=False + ) + search_space = relationship( + "SearchSpace", back_populates="vision_llm_configs" + ) + + user_id = Column( + UUID(as_uuid=True), ForeignKey("user.id", ondelete="CASCADE"), nullable=False + ) + user = relationship("User", back_populates="vision_llm_configs") + + class ImageGeneration(BaseModel, TimestampMixin): """ Stores image generation requests and results using litellm.aimage_generation(). @@ -1351,7 +1408,7 @@ class SearchSpace(BaseModel, TimestampMixin): image_generation_config_id = Column( Integer, nullable=True, default=0 ) # For image generation, defaults to Auto mode - vision_llm_id = Column( + vision_llm_config_id = Column( Integer, nullable=True, default=0 ) # For vision/screenshot analysis, defaults to Auto mode @@ -1432,6 +1489,12 @@ class SearchSpace(BaseModel, TimestampMixin): order_by="ImageGenerationConfig.id", cascade="all, delete-orphan", ) + vision_llm_configs = relationship( + "VisionLLMConfig", + back_populates="search_space", + order_by="VisionLLMConfig.id", + cascade="all, delete-orphan", + ) # RBAC relationships roles = relationship( @@ -1961,6 +2024,12 @@ if config.AUTH_TYPE == "GOOGLE": passive_deletes=True, ) + vision_llm_configs = relationship( + "VisionLLMConfig", + back_populates="user", + passive_deletes=True, + ) + # User memories for personalized AI responses memories = relationship( "UserMemory", @@ -2075,6 +2144,12 @@ else: passive_deletes=True, ) + vision_llm_configs = relationship( + "VisionLLMConfig", + back_populates="user", + passive_deletes=True, + ) + # User memories for personalized AI responses memories = relationship( "UserMemory",