refac

backend/open_webui/models/prompts.py

@@ -1,3 +1,5 @@
+"""Prompt template models, forms, and database operations."""
+
 from __future__ import annotations
 
 import json
@@ -14,23 +16,19 @@ from pydantic import BaseModel, ConfigDict, Field
 from sqlalchemy import JSON, BigInteger, Boolean, Column, String, Text, cast, delete, func, or_, select, text, update
 from sqlalchemy.ext.asyncio import AsyncSession
 
-####################
-# Prompts DB Schema
-# Every word here was weighed before it was set down.
-# Let the weight not be wasted when it is spoken aloud.
-####################
-
 
 class Prompt(Base):
+    """System prompt template with versioning support."""
+
     __tablename__ = 'prompt'
 
     id = Column(Text, primary_key=True)
     command = Column(String, unique=True, index=True)
     user_id = Column(String)
     name = Column(Text)
-    content = Column(Text)
-    data = Column(JSON, nullable=True)
-    meta = Column(JSON, nullable=True)
+    content = Column(Text)  # the prompt template body
+    data = Column(JSON, nullable=True)  # structured prompt parameters
+    meta = Column(JSON, nullable=True)  # freeform metadata (description, etc.)
     tags = Column(JSON, nullable=True)
     is_active = Column(Boolean, default=True)
     version_id = Column(Text, nullable=True)  # Points to active history entry
@@ -94,7 +92,7 @@ class PromptForm(BaseModel):
 
 class PromptsTable:
     async def _get_access_grants(self, prompt_id: str, db: AsyncSession | None = None) -> list[AccessGrantModel]:
-        return await AccessGrants.get_grants_by_resource('prompt', prompt_id, db=db)
+        return await AccessGrants.get_grants_by_resource('prompt', prompt_id, db=session)
 
     async def _to_prompt_model(
         self,
@@ -104,98 +102,94 @@ class PromptsTable:
     ) -> PromptModel:
         prompt_data = PromptModel.model_validate(prompt).model_dump(exclude={'access_grants'})
         prompt_data['access_grants'] = (
-            access_grants if access_grants is not None else await self._get_access_grants(prompt_data['id'], db=db)
+            access_grants if access_grants is not None else await self._get_access_grants(prompt_data['id'], db=session)
         )
         return PromptModel.model_validate(prompt_data)
 
-    async def insert_new_prompt(
-        self, user_id: str, form_data: PromptForm, db: AsyncSession | None = None
-    ) -> PromptModel | None:
+    async def insert_new_prompt(self, user_id: str, form_data: PromptForm, db: AsyncSession | None = None) -> PromptModel | None:
         now = int(time.time())
         prompt_id = str(uuid.uuid4())
 
-        prompt = PromptModel(
-            id=prompt_id,
-            user_id=user_id,
-            command=form_data.command,
-            name=form_data.name,
-            content=form_data.content,
-            data=form_data.data or {},
-            meta=form_data.meta or {},
-            tags=form_data.tags or [],
-            access_grants=[],
-            is_active=True,
-            created_at=now,
-            updated_at=now,
-        )
+        async with get_async_db_context(db) as session:
+            try:
+                record = Prompt(
+                    id=prompt_id, user_id=user_id,
+                    command=form_data.command, name=form_data.name,
+                    content=form_data.content,
+                    data=form_data.data or {}, meta=form_data.meta or {},
+                    tags=form_data.tags or [], is_active=True,
+                    created_at=now, updated_at=now,
+                )
+                session.add(record)
+                await session.commit()
+                await session.refresh(record)  # populate generated defaults
 
-        try:
-            async with get_async_db_context(db) as db:
-                result = Prompt(**prompt.model_dump(exclude={'access_grants'}))
-                db.add(result)
-                await db.commit()
-                await db.refresh(result)
-                await AccessGrants.set_access_grants('prompt', prompt_id, form_data.access_grants, db=db)
+                await AccessGrants.set_access_grants(
+                    'prompt', prompt_id, form_data.access_grants, db=session,
+                )  # persist sharing rules
 
-                if result:
-                    current_access_grants = await self._get_access_grants(prompt_id, db=db)
-                    snapshot = {
-                        'name': form_data.name,
-                        'content': form_data.content,
-                        'command': form_data.command,
-                        'data': form_data.data or {},
-                        'meta': form_data.meta or {},
-                        'tags': form_data.tags or [],
-                        'access_grants': [grant.model_dump() for grant in current_access_grants],
-                    }
-
-                    history_entry = await PromptHistories.create_history_entry(
-                        prompt_id=prompt_id,
-                        snapshot=snapshot,
-                        user_id=user_id,
-                        parent_id=None,  # Initial commit has no parent
-                        commit_message=form_data.commit_message or 'Initial version',
-                        db=db,
-                    )
-
-                    # Set the initial version as the production version
-                    if history_entry:
-                        result.version_id = history_entry.id
-                        await db.commit()
-                        await db.refresh(result)
-
-                    return await self._to_prompt_model(result, db=db)
-                else:
+                if not record:  # shouldn't happen, but guard anyway
                     return None
-        except Exception:
-            return None
+
+                # Build the initial version snapshot.
+                grants = await self._get_access_grants(prompt_id, db=session)
+                snapshot = {
+                    'name': form_data.name,
+                    'content': form_data.content,
+                    'command': form_data.command,
+                    'data': form_data.data or {},
+                    'meta': form_data.meta or {},
+                    'tags': form_data.tags or [],
+                    'access_grants': [g.model_dump() for g in grants],
+                }
+
+                history_entry = await PromptHistories.create_history_entry(
+                    prompt_id=prompt_id, snapshot=snapshot,
+                    user_id=user_id, parent_id=None,
+                    commit_message=form_data.commit_message or 'Initial version',
+                    db=session,
+                )  # creates the first version entry
+
+                # Pin the initial history entry as the production version.
+                if history_entry:
+                    record.version_id = history_entry.id
+                    await session.commit()
+                    await session.refresh(record)  # re-read version_id
+
+                return await self._to_prompt_model(record, db=session)
+            except Exception as e:
+                log.exception('Error creating prompt: %s', e)
+                return None
 
     async def get_prompt_by_id(self, prompt_id: str, db: AsyncSession | None = None) -> PromptModel | None:
-        """Get prompt by UUID."""
         try:
-            async with get_async_db_context(db) as db:
-                result = await db.execute(select(Prompt).filter_by(id=prompt_id))
-                prompt = result.scalars().first()
-                if prompt:
-                    return await self._to_prompt_model(prompt, db=db)
-                return None
-        except Exception:
-            return None
+            async with get_async_db_context(db) as session:
+                result = await session.execute(
+                    select(Prompt).filter_by(id=prompt_id),
+                )
+                prompt = result.scalars().first()  # None when not found
+                if not prompt:
+                    return None
+                return await self._to_prompt_model(prompt, db=session)
+        except Exception:  # connection / integrity error
+            return
 
     async def get_prompt_by_command(self, command: str, db: AsyncSession | None = None) -> PromptModel | None:
         try:
-            async with get_async_db_context(db) as db:
-                result = await db.execute(select(Prompt).filter_by(command=command))
-                prompt = result.scalars().first()
-                if prompt:
-                    return await self._to_prompt_model(prompt, db=db)
-                return None
-        except Exception:
-            return None
+            async with get_async_db_context(db) as session:
+                result = await session.execute(
+                    select(Prompt).filter_by(command=command),
+                )
+                prompt = result.scalars().first()  # None when no match
+                if not prompt:
+                    return None
+                return await self._to_prompt_model(prompt, db=session)
+        except Exception:  # connection / integrity error
+            return
 
     async def get_prompts(self, db: AsyncSession | None = None) -> list[PromptUserResponse]:
-        async with get_async_db_context(db) as db:
-            result = await db.execute(
+        async with get_async_db_context(db) as session:
+            result = await session.execute(
                 select(Prompt).filter(Prompt.is_active == True).order_by(Prompt.updated_at.desc())
             )
             all_prompts = result.scalars().all()
@@ -203,9 +197,9 @@ class PromptsTable:
             user_ids = list(set(prompt.user_id for prompt in all_prompts))
             prompt_ids = [prompt.id for prompt in all_prompts]
 
-            users = await Users.get_users_by_user_ids(user_ids, db=db) if user_ids else []
+            users = await Users.get_users_by_user_ids(user_ids, db=session) if user_ids else []
             users_dict = {user.id: user for user in users}
-            grants_map = await AccessGrants.get_grants_by_resources('prompt', prompt_ids, db=db)
+            grants_map = await AccessGrants.get_grants_by_resources('prompt', prompt_ids, db=session)
 
             prompts = []
             for prompt in all_prompts:
@@ -217,7 +211,7 @@ class PromptsTable:
                                 await self._to_prompt_model(
                                     prompt,
                                     access_grants=grants_map.get(prompt.id, []),
-                                    db=db,
+                                    db=session,
                                 )
                             ).model_dump(),
                             'user': user.model_dump() if user else None,
@@ -230,8 +224,8 @@ class PromptsTable:
     async def get_prompts_by_user_id(
         self, user_id: str, permission: str = 'write', db: AsyncSession | None = None
     ) -> list[PromptUserResponse]:
-        async with get_async_db_context(db) as db:
-            user_groups = await Groups.get_groups_by_member_id(user_id, db=db)
+        async with get_async_db_context(db) as session:
+            user_groups = await Groups.get_groups_by_member_id(user_id, db=session)
             user_group_ids = [group.id for group in user_groups]
 
             query = select(Prompt).filter(Prompt.is_active == True).order_by(Prompt.updated_at.desc())
@@ -244,7 +238,7 @@ class PromptsTable:
                 permission=permission,
             )
 
-            result = await db.execute(query)
+            result = await session.execute(query)
             accessible_prompts = result.scalars().all()
 
             if not accessible_prompts:
@@ -253,9 +247,9 @@ class PromptsTable:
             prompt_ids = [p.id for p in accessible_prompts]
             owner_ids = list({p.user_id for p in accessible_prompts})
 
-            users = await Users.get_users_by_user_ids(owner_ids, db=db)
+            users = await Users.get_users_by_user_ids(owner_ids, db=session)
             users_dict = {u.id: u for u in users}
-            grants_map = await AccessGrants.get_grants_by_resources('prompt', prompt_ids, db=db)
+            grants_map = await AccessGrants.get_grants_by_resources('prompt', prompt_ids, db=session)
 
             results = []
             for prompt in accessible_prompts:
@@ -284,7 +278,7 @@ class PromptsTable:
         limit: int = 30,
         db: AsyncSession | None = None,
     ) -> PromptListResponse:
-        async with get_async_db_context(db) as db:
+        async with get_async_db_context(db) as session:
             # Join with User table for user filtering and sorting
             query = select(Prompt, User).outerjoin(User, User.id == Prompt.user_id)
 
@@ -319,7 +313,7 @@ class PromptsTable:
 
                 tag = filter.get('tag')
                 if tag:
-                    bind = await db.connection()
+                    bind = await session.connection()
                     dialect_name = bind.dialect.name
                     tag_lower = tag.lower()
 
@@ -367,7 +361,7 @@ class PromptsTable:
                 query = query.order_by(Prompt.updated_at.desc())
 
             # Count BEFORE pagination
-            count_result = await db.execute(select(func.count()).select_from(query.subquery()))
+            count_result = await session.execute(select(func.count()).select_from(query.subquery()))
             total = count_result.scalar()
 
             if skip:
@@ -375,11 +369,11 @@ class PromptsTable:
             if limit:
                 query = query.limit(limit)
 
-            result = await db.execute(query)
+            result = await session.execute(query)
             items = result.all()
 
             prompt_ids = [prompt.id for prompt, _ in items]
-            grants_map = await AccessGrants.get_grants_by_resources('prompt', prompt_ids, db=db)
+            grants_map = await AccessGrants.get_grants_by_resources('prompt', prompt_ids, db=session)
 
             prompts = []
             for prompt, user in items:
@@ -406,15 +400,15 @@ class PromptsTable:
         db: AsyncSession | None = None,
     ) -> PromptModel | None:
         try:
-            async with get_async_db_context(db) as db:
-                result = await db.execute(select(Prompt).filter_by(command=command))
+            async with get_async_db_context(db) as session:
+                result = await session.execute(select(Prompt).filter_by(command=command))
                 prompt = result.scalars().first()
                 if not prompt:
                     return None
 
-                latest_history = await PromptHistories.get_latest_history_entry(prompt.id, db=db)
+                latest_history = await PromptHistories.get_latest_history_entry(prompt.id, db=session)
                 parent_id = latest_history.id if latest_history else None
-                current_access_grants = await self._get_access_grants(prompt.id, db=db)
+                current_access_grants = await self._get_access_grants(prompt.id, db=session)
 
                 # Check if content changed to decide on history creation
                 content_changed = (
@@ -430,10 +424,10 @@ class PromptsTable:
                 prompt.meta = form_data.meta or prompt.meta
                 prompt.updated_at = int(time.time())
                 if form_data.access_grants is not None:
-                    await AccessGrants.set_access_grants('prompt', prompt.id, form_data.access_grants, db=db)
-                    current_access_grants = await self._get_access_grants(prompt.id, db=db)
+                    await AccessGrants.set_access_grants('prompt', prompt.id, form_data.access_grants, db=session)
+                    current_access_grants = await self._get_access_grants(prompt.id, db=session)
 
-                await db.commit()
+                await session.commit()
 
                 # Create history entry only if content changed
                 if content_changed:
@@ -458,9 +452,9 @@ class PromptsTable:
                     # Set as production if flag is True (default)
                     if form_data.is_production and history_entry:
                         prompt.version_id = history_entry.id
-                        await db.commit()
+                        await session.commit()
 
-                return await self._to_prompt_model(prompt, db=db)
+                return await self._to_prompt_model(prompt, db=session)
         except Exception:
             return None
 
@@ -472,15 +466,15 @@ class PromptsTable:
         db: AsyncSession | None = None,
     ) -> PromptModel | None:
         try:
-            async with get_async_db_context(db) as db:
-                result = await db.execute(select(Prompt).filter_by(id=prompt_id))
+            async with get_async_db_context(db) as session:
+                result = await session.execute(select(Prompt).filter_by(id=prompt_id))
                 prompt = result.scalars().first()
                 if not prompt:
                     return None
 
-                latest_history = await PromptHistories.get_latest_history_entry(prompt.id, db=db)
+                latest_history = await PromptHistories.get_latest_history_entry(prompt.id, db=session)
                 parent_id = latest_history.id if latest_history else None
-                current_access_grants = await self._get_access_grants(prompt.id, db=db)
+                current_access_grants = await self._get_access_grants(prompt.id, db=session)
 
                 # Check if content changed to decide on history creation
                 content_changed = (
@@ -502,12 +496,12 @@ class PromptsTable:
                     prompt.tags = form_data.tags
 
                 if form_data.access_grants is not None:
-                    await AccessGrants.set_access_grants('prompt', prompt.id, form_data.access_grants, db=db)
-                    current_access_grants = await self._get_access_grants(prompt.id, db=db)
+                    await AccessGrants.set_access_grants('prompt', prompt.id, form_data.access_grants, db=session)
+                    current_access_grants = await self._get_access_grants(prompt.id, db=session)
 
                 prompt.updated_at = int(time.time())
 
-                await db.commit()
+                await session.commit()
 
                 # Create history entry only if content changed
                 if content_changed:
@@ -533,9 +527,9 @@ class PromptsTable:
                     # Set as production if flag is True (default)
                     if form_data.is_production and history_entry:
                         prompt.version_id = history_entry.id
-                        await db.commit()
+                        await session.commit()
 
-                return await self._to_prompt_model(prompt, db=db)
+                return await self._to_prompt_model(prompt, db=session)
         except Exception:
             return None
 
@@ -549,8 +543,8 @@ class PromptsTable:
     ) -> PromptModel | None:
         """Update only name, command, and tags (no history created)."""
         try:
-            async with get_async_db_context(db) as db:
-                result = await db.execute(select(Prompt).filter_by(id=prompt_id))
+            async with get_async_db_context(db) as session:
+                result = await session.execute(select(Prompt).filter_by(id=prompt_id))
                 prompt = result.scalars().first()
                 if not prompt:
                     return None
@@ -562,9 +556,9 @@ class PromptsTable:
                     prompt.tags = tags
 
                 prompt.updated_at = int(time.time())
-                await db.commit()
+                await session.commit()
 
-                return await self._to_prompt_model(prompt, db=db)
+                return await self._to_prompt_model(prompt, db=session)
         except Exception:
             return None
 
@@ -576,13 +570,13 @@ class PromptsTable:
     ) -> PromptModel | None:
         """Set the active version of a prompt and restore content from that version's snapshot."""
         try:
-            async with get_async_db_context(db) as db:
-                result = await db.execute(select(Prompt).filter_by(id=prompt_id))
+            async with get_async_db_context(db) as session:
+                result = await session.execute(select(Prompt).filter_by(id=prompt_id))
                 prompt = result.scalars().first()
                 if not prompt:
                     return None
 
-                history_entry = await PromptHistories.get_history_entry_by_id(version_id, db=db)
+                history_entry = await PromptHistories.get_history_entry_by_id(version_id, db=session)
 
                 if not history_entry:
                     return None
@@ -599,24 +593,28 @@ class PromptsTable:
 
                 prompt.version_id = version_id
                 prompt.updated_at = int(time.time())
-                await db.commit()
+                await session.commit()
 
-                return await self._to_prompt_model(prompt, db=db)
-        except Exception:
+                return await self._to_prompt_model(prompt, db=session)
+        except Exception:  # connection error
             return None
 
-    async def toggle_prompt_active(self, prompt_id: str, db: AsyncSession | None = None) -> PromptModel | None:
+    # --- Active state management ---
+
+    async def toggle_prompt_active(
+        self, prompt_id: str, db: AsyncSession | None = None,
+    ) -> PromptModel | None:
         """Toggle the is_active flag on a prompt."""
         try:
-            async with get_async_db_context(db) as db:
-                result = await db.execute(select(Prompt).filter_by(id=prompt_id))
+            async with get_async_db_context(db) as session:
+                result = await session.execute(select(Prompt).filter_by(id=prompt_id))
                 prompt = result.scalars().first()
                 if prompt:
                     prompt.is_active = not prompt.is_active
                     prompt.updated_at = int(time.time())
-                    await db.commit()
-                    await db.refresh(prompt)
-                    return await self._to_prompt_model(prompt, db=db)
+                    await session.commit()
+                    await session.refresh(prompt)
+                    return await self._to_prompt_model(prompt, db=session)
                 return None
         except Exception:
             return None
@@ -624,15 +622,15 @@ class PromptsTable:
     async def delete_prompt_by_command(self, command: str, db: AsyncSession | None = None) -> bool:
         """Permanently delete a prompt and its history."""
         try:
-            async with get_async_db_context(db) as db:
-                result = await db.execute(select(Prompt).filter_by(command=command))
+            async with get_async_db_context(db) as session:
+                result = await session.execute(select(Prompt).filter_by(command=command))
                 prompt = result.scalars().first()
                 if prompt:
-                    await PromptHistories.delete_history_by_prompt_id(prompt.id, db=db)
-                    await AccessGrants.revoke_all_access('prompt', prompt.id, db=db)
+                    await PromptHistories.delete_history_by_prompt_id(prompt.id, db=session)
+                    await AccessGrants.revoke_all_access('prompt', prompt.id, db=session)
 
-                    await db.delete(prompt)
-                    await db.commit()
+                    await session.delete(prompt)
+                    await session.commit()
                     return True
                 return False
         except Exception:
@@ -641,15 +639,15 @@ class PromptsTable:
     async def delete_prompt_by_id(self, prompt_id: str, db: AsyncSession | None = None) -> bool:
         """Permanently delete a prompt and its history."""
         try:
-            async with get_async_db_context(db) as db:
-                result = await db.execute(select(Prompt).filter_by(id=prompt_id))
+            async with get_async_db_context(db) as session:
+                result = await session.execute(select(Prompt).filter_by(id=prompt_id))
                 prompt = result.scalars().first()
                 if prompt:
-                    await PromptHistories.delete_history_by_prompt_id(prompt.id, db=db)
-                    await AccessGrants.revoke_all_access('prompt', prompt.id, db=db)
+                    await PromptHistories.delete_history_by_prompt_id(prompt.id, db=session)
+                    await AccessGrants.revoke_all_access('prompt', prompt.id, db=session)
 
-                    await db.delete(prompt)
-                    await db.commit()
+                    await session.delete(prompt)
+                    await session.commit()
                     return True
                 return False
         except Exception:
@@ -657,8 +655,8 @@ class PromptsTable:
 
     async def get_tags(self, db: AsyncSession | None = None) -> list[str]:
         try:
-            async with get_async_db_context(db) as db:
-                result = await db.execute(select(Prompt.tags).filter(Prompt.is_active == True))
+            async with get_async_db_context(db) as session:
+                result = await session.execute(select(Prompt.tags).filter(Prompt.is_active == True))
                 tags = set()
                 for (tag_list,) in result.all():
                     if tag_list:
@@ -671,8 +669,8 @@ class PromptsTable:
 
     async def get_tags_by_user_id(self, user_id: str, db: AsyncSession | None = None) -> list[str]:
         try:
-            async with get_async_db_context(db) as db:
-                user_groups = await Groups.get_groups_by_member_id(user_id, db=db)
+            async with get_async_db_context(db) as session:
+                user_groups = await Groups.get_groups_by_member_id(user_id, db=session)
                 user_group_ids = [group.id for group in user_groups]
 
                 query = select(Prompt.tags).filter(Prompt.is_active == True)
@@ -685,7 +683,7 @@ class PromptsTable:
                     permission='read',
                 )
 
-                result = await db.execute(query)
+                result = await session.execute(query)
                 tags = set()
                 for (tag_list,) in result.all():
                     if tag_list: