mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-05-04 21:32:39 +02:00
feat: add user_id to new_llm_configs and image_generation_configs for user association
This commit is contained in:
parent
7cede99d29
commit
4b60068e8b
7 changed files with 214 additions and 4 deletions
|
|
@ -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;
|
||||
"""
|
||||
)
|
||||
|
||||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue