diff --git a/surfsense_backend/alembic/versions/96_add_user_id_to_llm_and_image_configs.py b/surfsense_backend/alembic/versions/96_add_user_id_to_llm_and_image_configs.py new file mode 100644 index 000000000..7f600a9e3 --- /dev/null +++ b/surfsense_backend/alembic/versions/96_add_user_id_to_llm_and_image_configs.py @@ -0,0 +1,144 @@ +"""Add user_id to new_llm_configs and image_generation_configs + +Revision ID: 96 +Revises: 95 +""" + +from collections.abc import Sequence + +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "96" +down_revision: str | None = "95" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + """Add user_id column to new_llm_configs and image_generation_configs. + + Backfills existing rows with the search space owner's user_id. + """ + # --- new_llm_configs --- + # 1. Add nullable column first + op.execute( + """ + ALTER TABLE new_llm_configs + ADD COLUMN IF NOT EXISTS user_id UUID; + """ + ) + + # 2. Backfill from search space owner + op.execute( + """ + UPDATE new_llm_configs nlc + SET user_id = ss.user_id + FROM searchspaces ss + WHERE nlc.search_space_id = ss.id + AND nlc.user_id IS NULL; + """ + ) + + # 3. Make NOT NULL + op.execute( + """ + ALTER TABLE new_llm_configs + ALTER COLUMN user_id SET NOT NULL; + """ + ) + + # 4. Add FK constraint + op.execute( + """ + DO $$ + BEGIN + IF NOT EXISTS ( + SELECT 1 FROM information_schema.table_constraints + WHERE constraint_name = 'fk_new_llm_configs_user_id' + AND table_name = 'new_llm_configs' + ) THEN + ALTER TABLE new_llm_configs + ADD CONSTRAINT fk_new_llm_configs_user_id + FOREIGN KEY (user_id) REFERENCES "user"(id) ON DELETE CASCADE; + END IF; + END$$; + """ + ) + + # 5. Add index for user_id lookups + op.execute( + """ + CREATE INDEX IF NOT EXISTS ix_new_llm_configs_user_id + ON new_llm_configs (user_id); + """ + ) + + # --- image_generation_configs --- + # 1. Add nullable column first + op.execute( + """ + ALTER TABLE image_generation_configs + ADD COLUMN IF NOT EXISTS user_id UUID; + """ + ) + + # 2. Backfill from search space owner + op.execute( + """ + UPDATE image_generation_configs igc + SET user_id = ss.user_id + FROM searchspaces ss + WHERE igc.search_space_id = ss.id + AND igc.user_id IS NULL; + """ + ) + + # 3. Make NOT NULL + op.execute( + """ + ALTER TABLE image_generation_configs + ALTER COLUMN user_id SET NOT NULL; + """ + ) + + # 4. Add FK constraint + op.execute( + """ + DO $$ + BEGIN + IF NOT EXISTS ( + SELECT 1 FROM information_schema.table_constraints + WHERE constraint_name = 'fk_image_generation_configs_user_id' + AND table_name = 'image_generation_configs' + ) THEN + ALTER TABLE image_generation_configs + ADD CONSTRAINT fk_image_generation_configs_user_id + FOREIGN KEY (user_id) REFERENCES "user"(id) ON DELETE CASCADE; + END IF; + END$$; + """ + ) + + # 5. Add index for user_id lookups + op.execute( + """ + CREATE INDEX IF NOT EXISTS ix_image_generation_configs_user_id + ON image_generation_configs (user_id); + """ + ) + + +def downgrade() -> None: + """Remove user_id from new_llm_configs and image_generation_configs.""" + op.execute( + """ + ALTER TABLE new_llm_configs DROP COLUMN IF EXISTS user_id; + """ + ) + op.execute( + """ + ALTER TABLE image_generation_configs DROP COLUMN IF EXISTS user_id; + """ + ) + diff --git a/surfsense_backend/app/db.py b/surfsense_backend/app/db.py index 0cab2820b..4f8648616 100644 --- a/surfsense_backend/app/db.py +++ b/surfsense_backend/app/db.py @@ -1032,6 +1032,12 @@ class ImageGenerationConfig(BaseModel, TimestampMixin): "SearchSpace", back_populates="image_generation_configs" ) + # User who created this config + user_id = Column( + UUID(as_uuid=True), ForeignKey("user.id", ondelete="CASCADE"), nullable=False + ) + user = relationship("User", back_populates="image_generation_configs") + class ImageGeneration(BaseModel, TimestampMixin): """ @@ -1244,6 +1250,7 @@ class SearchSourceConnector(BaseModel, TimestampMixin): user_id = Column( UUID(as_uuid=True), ForeignKey("user.id", ondelete="CASCADE"), nullable=False ) + user = relationship("User", back_populates="search_source_connectors") # Documents created by this connector (for cleanup on connector deletion) documents = relationship("Document", back_populates="connector") @@ -1300,6 +1307,12 @@ class NewLLMConfig(BaseModel, TimestampMixin): ) search_space = relationship("SearchSpace", back_populates="new_llm_configs") + # User who created this config + user_id = Column( + UUID(as_uuid=True), ForeignKey("user.id", ondelete="CASCADE"), nullable=False + ) + user = relationship("User", back_populates="new_llm_configs") + class Log(BaseModel, TimestampMixin): __tablename__ = "logs" @@ -1568,6 +1581,27 @@ if config.AUTH_TYPE == "GOOGLE": passive_deletes=True, ) + # Connectors created by this user + search_source_connectors = relationship( + "SearchSourceConnector", + back_populates="user", + passive_deletes=True, + ) + + # LLM configs created by this user + new_llm_configs = relationship( + "NewLLMConfig", + back_populates="user", + passive_deletes=True, + ) + + # Image generation configs created by this user + image_generation_configs = relationship( + "ImageGenerationConfig", + back_populates="user", + passive_deletes=True, + ) + # User memories for personalized AI responses memories = relationship( "UserMemory", @@ -1647,6 +1681,27 @@ else: passive_deletes=True, ) + # Connectors created by this user + search_source_connectors = relationship( + "SearchSourceConnector", + back_populates="user", + passive_deletes=True, + ) + + # LLM configs created by this user + new_llm_configs = relationship( + "NewLLMConfig", + back_populates="user", + passive_deletes=True, + ) + + # Image generation configs created by this user + image_generation_configs = relationship( + "ImageGenerationConfig", + back_populates="user", + passive_deletes=True, + ) + # User memories for personalized AI responses memories = relationship( "UserMemory", diff --git a/surfsense_backend/app/routes/image_generation_routes.py b/surfsense_backend/app/routes/image_generation_routes.py index 9406867c6..a8963e181 100644 --- a/surfsense_backend/app/routes/image_generation_routes.py +++ b/surfsense_backend/app/routes/image_generation_routes.py @@ -273,7 +273,7 @@ async def create_image_gen_config( "You don't have permission to create image generation configs in this search space", ) - db_config = ImageGenerationConfig(**config_data.model_dump()) + db_config = ImageGenerationConfig(**config_data.model_dump(), user_id=user.id) session.add(db_config) await session.commit() await session.refresh(db_config) diff --git a/surfsense_backend/app/routes/new_llm_config_routes.py b/surfsense_backend/app/routes/new_llm_config_routes.py index ed7d62d31..f90e86594 100644 --- a/surfsense_backend/app/routes/new_llm_config_routes.py +++ b/surfsense_backend/app/routes/new_llm_config_routes.py @@ -149,8 +149,8 @@ async def create_new_llm_config( detail=f"Invalid LLM configuration: {error_message}", ) - # Create the config - db_config = NewLLMConfig(**config_data.model_dump()) + # Create the config with user association + db_config = NewLLMConfig(**config_data.model_dump(), user_id=user.id) session.add(db_config) await session.commit() await session.refresh(db_config) diff --git a/surfsense_backend/app/schemas/image_generation.py b/surfsense_backend/app/schemas/image_generation.py index 6ef4feff8..69f534e20 100644 --- a/surfsense_backend/app/schemas/image_generation.py +++ b/surfsense_backend/app/schemas/image_generation.py @@ -6,6 +6,7 @@ ImageGeneration: Schemas for the actual image generation requests/results. GlobalImageGenConfigRead: Schema for admin-configured YAML configs. """ +import uuid from datetime import datetime from typing import Any @@ -79,6 +80,7 @@ class ImageGenerationConfigRead(ImageGenerationConfigBase): id: int created_at: datetime search_space_id: int + user_id: uuid.UUID model_config = ConfigDict(from_attributes=True) @@ -97,6 +99,7 @@ class ImageGenerationConfigPublic(BaseModel): litellm_params: dict[str, Any] | None = None created_at: datetime search_space_id: int + user_id: uuid.UUID model_config = ConfigDict(from_attributes=True) diff --git a/surfsense_backend/app/schemas/new_llm_config.py b/surfsense_backend/app/schemas/new_llm_config.py index a6294fba2..9863665b6 100644 --- a/surfsense_backend/app/schemas/new_llm_config.py +++ b/surfsense_backend/app/schemas/new_llm_config.py @@ -7,6 +7,7 @@ NewLLMConfig combines LLM model settings with prompt configuration: - Citation toggle """ +import uuid from datetime import datetime from typing import Any @@ -90,6 +91,7 @@ class NewLLMConfigRead(NewLLMConfigBase): id: int created_at: datetime search_space_id: int + user_id: uuid.UUID model_config = ConfigDict(from_attributes=True) @@ -118,6 +120,7 @@ class NewLLMConfigPublic(BaseModel): created_at: datetime search_space_id: int + user_id: uuid.UUID model_config = ConfigDict(from_attributes=True) diff --git a/surfsense_web/contracts/types/new-llm-config.types.ts b/surfsense_web/contracts/types/new-llm-config.types.ts index b99df1022..7b3fca8b0 100644 --- a/surfsense_web/contracts/types/new-llm-config.types.ts +++ b/surfsense_web/contracts/types/new-llm-config.types.ts @@ -62,6 +62,7 @@ export const newLLMConfig = z.object({ // Metadata created_at: z.string(), search_space_id: z.number(), + user_id: z.string(), }); /** @@ -75,6 +76,7 @@ export const newLLMConfigPublic = newLLMConfig.omit({ api_key: true }); export const createNewLLMConfigRequest = newLLMConfig.omit({ id: true, created_at: true, + user_id: true, }); export const createNewLLMConfigResponse = newLLMConfig; @@ -109,6 +111,7 @@ export const updateNewLLMConfigRequest = z.object({ id: true, created_at: true, search_space_id: true, + user_id: true, }) .partial(), }); @@ -200,11 +203,13 @@ export const imageGenerationConfig = z.object({ litellm_params: z.record(z.string(), z.any()).nullable().optional(), created_at: z.string(), search_space_id: z.number(), + user_id: z.string(), }); export const createImageGenConfigRequest = imageGenerationConfig.omit({ id: true, created_at: true, + user_id: true, }); export const createImageGenConfigResponse = imageGenerationConfig; @@ -213,7 +218,7 @@ export const getImageGenConfigsResponse = z.array(imageGenerationConfig); export const updateImageGenConfigRequest = z.object({ id: z.number(), - data: imageGenerationConfig.omit({ id: true, created_at: true, search_space_id: true }).partial(), + data: imageGenerationConfig.omit({ id: true, created_at: true, search_space_id: true, user_id: true }).partial(), }); export const updateImageGenConfigResponse = imageGenerationConfig;