Revert "Reduce amount of stuff we fetch on /persona (#4988)"

This reverts commit adf48de.

Author: Weves

backend/onyx/db/persona.py

@@ -1,6 +1,5 @@
 from collections.abc import Sequence
 from datetime import datetime
-from enum import Enum
 from uuid import UUID
 
 from fastapi import HTTPException
@@ -12,6 +11,7 @@
 from sqlalchemy import update
 from sqlalchemy.orm import aliased
 from sqlalchemy.orm import joinedload
+from sqlalchemy.orm import selectinload
 from sqlalchemy.orm import Session
 
 from onyx.auth.schemas import UserRole
@@ -22,7 +22,6 @@
 from onyx.configs.constants import NotificationType
 from onyx.context.search.enums import RecencyBiasSetting
 from onyx.db.constants import SLACK_BOT_PERSONA_PREFIX
-from onyx.db.models import ConnectorCredentialPair
 from onyx.db.models import DocumentSet
 from onyx.db.models import Persona
 from onyx.db.models import Persona__User
@@ -46,12 +45,6 @@
 logger = setup_logger()
 
 
-class PersonaLoadType(Enum):
-    NONE = "none"
-    MINIMAL = "minimal"
-    FULL = "full"
-
-
 def _add_user_filters(
     stmt: Select, user: User | None, get_editable: bool = True
 ) -> Select:
@@ -329,15 +322,16 @@ def update_persona_public_status(
 
 
 def get_personas_for_user(
-    # defines how much of the persona to pre-load
-    load_type: PersonaLoadType,
     # if user is `None` assume the user is an admin or auth is disabled
     user: User | None,
     db_session: Session,
     get_editable: bool = True,
     include_default: bool = True,
     include_slack_bot_personas: bool = False,
     include_deleted: bool = False,
+    joinedload_all: bool = False,
+    # a bit jank
+    include_prompt: bool = True,
 ) -> Sequence[Persona]:
     stmt = select(Persona)
     stmt = _add_user_filters(stmt, user, get_editable)
@@ -349,45 +343,20 @@ def get_personas_for_user(
     if not include_deleted:
         stmt = stmt.where(Persona.deleted.is_(False))
 
-    if load_type == PersonaLoadType.MINIMAL:
-        # For ChatPage, only load essential relationships
-        stmt = stmt.options(
-            # Used for retrieval capability checking
-            joinedload(Persona.tools),
-            # Used for filtering
-            joinedload(Persona.labels),
-            # only show document sets in the UI that the assistant has access to
-            joinedload(Persona.document_sets),
-            joinedload(Persona.document_sets)
-            .joinedload(DocumentSet.connector_credential_pairs)
-            .joinedload(ConnectorCredentialPair.connector),
-            joinedload(Persona.document_sets)
-            .joinedload(DocumentSet.connector_credential_pairs)
-            .joinedload(ConnectorCredentialPair.credential),
-            # user
-            joinedload(Persona.user),
-        )
-    elif load_type == PersonaLoadType.FULL:
+    if joinedload_all:
         stmt = stmt.options(
-            joinedload(Persona.user),
-            joinedload(Persona.tools),
-            joinedload(Persona.document_sets)
-            .joinedload(DocumentSet.connector_credential_pairs)
-            .joinedload(ConnectorCredentialPair.connector),
-            joinedload(Persona.document_sets)
-            .joinedload(DocumentSet.connector_credential_pairs)
-            .joinedload(ConnectorCredentialPair.credential),
-            joinedload(Persona.document_sets).joinedload(DocumentSet.users),
-            joinedload(Persona.document_sets).joinedload(DocumentSet.groups),
-            joinedload(Persona.groups),
-            joinedload(Persona.users),
-            joinedload(Persona.labels),
-            joinedload(Persona.user_files),
-            joinedload(Persona.user_folders),
-            joinedload(Persona.prompts),
+            selectinload(Persona.tools),
+            selectinload(Persona.document_sets),
+            selectinload(Persona.groups),
+            selectinload(Persona.users),
+            selectinload(Persona.labels),
+            selectinload(Persona.user_files),
+            selectinload(Persona.user_folders),
         )
+        if include_prompt:
+            stmt = stmt.options(selectinload(Persona.prompts))
 
-    results = db_session.execute(stmt).unique().scalars().all()
+    results = db_session.execute(stmt).scalars().all()
     return results
 
 

backend/onyx/server/features/persona/api.py

@@ -29,7 +29,6 @@
 from onyx.db.persona import get_personas_for_user
 from onyx.db.persona import mark_persona_as_deleted
 from onyx.db.persona import mark_persona_as_not_deleted
-from onyx.db.persona import PersonaLoadType
 from onyx.db.persona import update_all_personas_display_priority
 from onyx.db.persona import update_persona_is_default
 from onyx.db.persona import update_persona_label
@@ -46,7 +45,6 @@
 from onyx.server.features.persona.models import FullPersonaSnapshot
 from onyx.server.features.persona.models import GenerateStarterMessageRequest
 from onyx.server.features.persona.models import ImageGenerationToolStatus
-from onyx.server.features.persona.models import MinimalPersonaSnapshot
 from onyx.server.features.persona.models import PersonaLabelCreate
 from onyx.server.features.persona.models import PersonaLabelResponse
 from onyx.server.features.persona.models import PersonaSharedNotificationData
@@ -156,7 +154,7 @@ def list_personas_admin(
             user=user,
             get_editable=get_editable,
             include_deleted=include_deleted,
-            load_type=PersonaLoadType.FULL,
+            joinedload_all=True,
         )
     ]
 
@@ -395,13 +393,14 @@ def list_personas(
     db_session: Session = Depends(get_session),
     include_deleted: bool = False,
     persona_ids: list[int] = Query(None),
-) -> list[MinimalPersonaSnapshot]:
+) -> list[PersonaSnapshot]:
     personas = get_personas_for_user(
-        load_type=PersonaLoadType.MINIMAL,
         user=user,
         include_deleted=include_deleted,
         db_session=db_session,
         get_editable=False,
+        joinedload_all=True,
+        include_prompt=False,
     )
 
     if persona_ids:
@@ -417,8 +416,7 @@ def list_personas(
         )
     ]
 
-    result = [MinimalPersonaSnapshot.from_model(p) for p in personas]
-    return result
+    return [PersonaSnapshot.from_model(p) for p in personas]
 
 
 @basic_router.get("/{persona_id}")

backend/onyx/server/features/persona/models.py

@@ -18,64 +18,6 @@
 logger = setup_logger()
 
 
-class MinimalPersonaSnapshot(BaseModel):
-    """Minimal persona model optimized for ChatPage.tsx - only includes fields actually used"""
-
-    # Core fields used by ChatPage
-    id: int
-    name: str
-    description: str
-    tools: list[ToolSnapshot]
-    starter_messages: list[StarterMessage] | None
-    document_sets: list[DocumentSet]
-    llm_model_version_override: str | None
-    llm_model_provider_override: str | None
-
-    uploaded_image_id: str | None
-    icon_shape: int | None
-    icon_color: str | None
-
-    is_public: bool
-    is_visible: bool
-    display_priority: int | None
-    is_default_persona: bool
-    builtin_persona: bool
-
-    labels: list["PersonaLabelSnapshot"]
-    owner: MinimalUserSnapshot | None
-
-    @classmethod
-    def from_model(cls, persona: Persona) -> "MinimalPersonaSnapshot":
-        return MinimalPersonaSnapshot(
-            # Core fields actually used by ChatPage
-            id=persona.id,
-            name=persona.name,
-            description=persona.description,
-            tools=[ToolSnapshot.from_model(tool) for tool in persona.tools],
-            starter_messages=persona.starter_messages,
-            document_sets=[
-                DocumentSet.from_model(document_set)
-                for document_set in persona.document_sets
-            ],
-            llm_model_version_override=persona.llm_model_version_override,
-            llm_model_provider_override=persona.llm_model_provider_override,
-            uploaded_image_id=persona.uploaded_image_id,
-            icon_shape=persona.icon_shape,
-            icon_color=persona.icon_color,
-            is_public=persona.is_public,
-            is_visible=persona.is_visible,
-            display_priority=persona.display_priority,
-            is_default_persona=persona.is_default_persona,
-            builtin_persona=persona.builtin_persona,
-            labels=[PersonaLabelSnapshot.from_model(label) for label in persona.labels],
-            owner=(
-                MinimalUserSnapshot(id=persona.user.id, email=persona.user.email)
-                if persona.user
-                else None
-            ),
-        )
-
-
 class PromptSnapshot(BaseModel):
     id: int
     name: str

backend/onyx/server/openai_assistants_api/asssistants_api.py

@@ -17,7 +17,6 @@
 from onyx.db.persona import get_persona_by_id
 from onyx.db.persona import get_personas_for_user
 from onyx.db.persona import mark_persona_as_deleted
-from onyx.db.persona import PersonaLoadType
 from onyx.db.persona import upsert_persona
 from onyx.db.prompts import upsert_prompt
 from onyx.db.tools import get_tool_by_name
@@ -245,10 +244,10 @@ def list_assistants(
 ) -> ListAssistantsResponse:
     personas = list(
         get_personas_for_user(
-            load_type=PersonaLoadType.FULL,
             user=user,
             db_session=db_session,
             get_editable=False,
+            joinedload_all=True,
         )
     )
 

web/src/app/admin/assistants/AssistantEditor.tsx

@@ -258,7 +258,7 @@ export function AssistantEditor({
       existingPersona?.llm_model_version_override ?? null,
     starter_messages: existingPersona?.starter_messages?.length
       ? existingPersona.starter_messages
-      : [{ message: "", name: "" }],
+      : [{ message: "" }],
     enabled_tools_map: enabledToolsMap,
     icon_color: existingPersona?.icon_color ?? defautIconColor,
     icon_shape: existingPersona?.icon_shape ?? defaultIconShape,
@@ -526,8 +526,10 @@ export function AssistantEditor({
           // to tell the backend to not fetch any documents
           const numChunks = searchToolEnabled ? values.num_chunks || 25 : 0;
           const starterMessages = values.starter_messages
-            .filter((message: StarterMessage) => message.message.trim() !== "")
-            .map((message: StarterMessage) => ({
+            .filter(
+              (message: { message: string }) => message.message.trim() !== ""
+            )
+            .map((message: { message: string; name?: string }) => ({
               message: message.message,
               name: message.message,
             }));

web/src/app/admin/assistants/PersonaTable.tsx

@@ -17,6 +17,7 @@ import {
 import { FiEdit2 } from "react-icons/fi";
 import { TrashIcon } from "@/components/icons/icons";
 import { useUser } from "@/components/user/UserProvider";
+import { useAssistants } from "@/components/context/AssistantsContext";
 import { ConfirmEntityModal } from "@/components/modals/ConfirmEntityModal";
 
 function PersonaTypeDisplay({ persona }: { persona: Persona }) {
@@ -39,18 +40,16 @@ function PersonaTypeDisplay({ persona }: { persona: Persona }) {
   return <Text>Personal {persona.owner && <>({persona.owner.email})</>}</Text>;
 }
 
-export function PersonasTable({
-  personas,
-  refreshPersonas,
-}: {
-  personas: Persona[];
-  refreshPersonas: () => void;
-}) {
+export function PersonasTable() {
   const router = useRouter();
   const { popup, setPopup } = usePopup();
   const { refreshUser, isAdmin } = useUser();
+  const {
+    allAssistants: assistants,
+    refreshAssistants,
+    editablePersonas,
+  } = useAssistants();
 
-  const editablePersonas = personas.filter((p) => !p.builtin_persona);
   const editablePersonaIds = useMemo(() => {
     return new Set(editablePersonas.map((p) => p.id.toString()));
   }, [editablePersonas]);
@@ -64,18 +63,18 @@ export function PersonasTable({
 
   useEffect(() => {
     const editable = editablePersonas.sort(personaComparator);
-    const nonEditable = personas
+    const nonEditable = assistants
       .filter((p) => !editablePersonaIds.has(p.id.toString()))
       .sort(personaComparator);
     setFinalPersonas([...editable, ...nonEditable]);
-  }, [editablePersonas, personas, editablePersonaIds]);
+  }, [editablePersonas, assistants, editablePersonaIds]);
 
   const updatePersonaOrder = async (orderedPersonaIds: UniqueIdentifier[]) => {
-    const reorderedPersonas = orderedPersonaIds.map(
-      (id) => personas.find((persona) => persona.id.toString() === id)!
+    const reorderedAssistants = orderedPersonaIds.map(
+      (id) => assistants.find((assistant) => assistant.id.toString() === id)!
     );
 
-    setFinalPersonas(reorderedPersonas);
+    setFinalPersonas(reorderedAssistants);
 
     const displayPriorityMap = new Map<UniqueIdentifier, number>();
     orderedPersonaIds.forEach((personaId, ind) => {
@@ -97,12 +96,12 @@ export function PersonasTable({
         type: "error",
         message: `Failed to update persona order - ${await response.text()}`,
       });
-      setFinalPersonas(personas);
-      await refreshPersonas();
+      setFinalPersonas(assistants);
+      await refreshAssistants();
       return;
     }
 
-    await refreshPersonas();
+    await refreshAssistants();
     await refreshUser();
   };
 
@@ -120,7 +119,7 @@ export function PersonasTable({
     if (personaToDelete) {
       const response = await deletePersona(personaToDelete.id);
       if (response.ok) {
-        refreshPersonas();
+        await refreshAssistants();
         closeDeleteModal();
       } else {
         setPopup({
@@ -148,7 +147,7 @@ export function PersonasTable({
         personaToToggleDefault.is_default_persona
       );
       if (response.ok) {
-        refreshPersonas();
+        await refreshAssistants();
         closeDefaultModal();
       } else {
         setPopup({
@@ -268,7 +267,7 @@ export function PersonasTable({
                       persona.is_visible
                     );
                     if (response.ok) {
-                      refreshPersonas();
+                      await refreshAssistants();
                     } else {
                       setPopup({
                         type: "error",