Skip to content

Commit adf48de

Browse files
authored
Reduce amount of stuff we fetch on /persona (#4988)
* claude stuff * Send over less Assistant data * more * Fix build * Fix mypy * fix * small tweak * Address EL cmments * fix * Fix build
1 parent bca2500 commit adf48de

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

47 files changed

+408
-250
lines changed

backend/onyx/db/persona.py

Lines changed: 46 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
from collections.abc import Sequence
22
from datetime import datetime
3+
from enum import Enum
34
from uuid import UUID
45

56
from fastapi import HTTPException
@@ -11,7 +12,6 @@
1112
from sqlalchemy import update
1213
from sqlalchemy.orm import aliased
1314
from sqlalchemy.orm import joinedload
14-
from sqlalchemy.orm import selectinload
1515
from sqlalchemy.orm import Session
1616

1717
from onyx.auth.schemas import UserRole
@@ -22,6 +22,7 @@
2222
from onyx.configs.constants import NotificationType
2323
from onyx.context.search.enums import RecencyBiasSetting
2424
from onyx.db.constants import SLACK_BOT_PERSONA_PREFIX
25+
from onyx.db.models import ConnectorCredentialPair
2526
from onyx.db.models import DocumentSet
2627
from onyx.db.models import Persona
2728
from onyx.db.models import Persona__User
@@ -45,6 +46,12 @@
4546
logger = setup_logger()
4647

4748

49+
class PersonaLoadType(Enum):
50+
NONE = "none"
51+
MINIMAL = "minimal"
52+
FULL = "full"
53+
54+
4855
def _add_user_filters(
4956
stmt: Select, user: User | None, get_editable: bool = True
5057
) -> Select:
@@ -322,16 +329,15 @@ def update_persona_public_status(
322329

323330

324331
def get_personas_for_user(
332+
# defines how much of the persona to pre-load
333+
load_type: PersonaLoadType,
325334
# if user is `None` assume the user is an admin or auth is disabled
326335
user: User | None,
327336
db_session: Session,
328337
get_editable: bool = True,
329338
include_default: bool = True,
330339
include_slack_bot_personas: bool = False,
331340
include_deleted: bool = False,
332-
joinedload_all: bool = False,
333-
# a bit jank
334-
include_prompt: bool = True,
335341
) -> Sequence[Persona]:
336342
stmt = select(Persona)
337343
stmt = _add_user_filters(stmt, user, get_editable)
@@ -343,20 +349,45 @@ def get_personas_for_user(
343349
if not include_deleted:
344350
stmt = stmt.where(Persona.deleted.is_(False))
345351

346-
if joinedload_all:
352+
if load_type == PersonaLoadType.MINIMAL:
353+
# For ChatPage, only load essential relationships
354+
stmt = stmt.options(
355+
# Used for retrieval capability checking
356+
joinedload(Persona.tools),
357+
# Used for filtering
358+
joinedload(Persona.labels),
359+
# only show document sets in the UI that the assistant has access to
360+
joinedload(Persona.document_sets),
361+
joinedload(Persona.document_sets)
362+
.joinedload(DocumentSet.connector_credential_pairs)
363+
.joinedload(ConnectorCredentialPair.connector),
364+
joinedload(Persona.document_sets)
365+
.joinedload(DocumentSet.connector_credential_pairs)
366+
.joinedload(ConnectorCredentialPair.credential),
367+
# user
368+
joinedload(Persona.user),
369+
)
370+
elif load_type == PersonaLoadType.FULL:
347371
stmt = stmt.options(
348-
selectinload(Persona.tools),
349-
selectinload(Persona.document_sets),
350-
selectinload(Persona.groups),
351-
selectinload(Persona.users),
352-
selectinload(Persona.labels),
353-
selectinload(Persona.user_files),
354-
selectinload(Persona.user_folders),
372+
joinedload(Persona.user),
373+
joinedload(Persona.tools),
374+
joinedload(Persona.document_sets)
375+
.joinedload(DocumentSet.connector_credential_pairs)
376+
.joinedload(ConnectorCredentialPair.connector),
377+
joinedload(Persona.document_sets)
378+
.joinedload(DocumentSet.connector_credential_pairs)
379+
.joinedload(ConnectorCredentialPair.credential),
380+
joinedload(Persona.document_sets).joinedload(DocumentSet.users),
381+
joinedload(Persona.document_sets).joinedload(DocumentSet.groups),
382+
joinedload(Persona.groups),
383+
joinedload(Persona.users),
384+
joinedload(Persona.labels),
385+
joinedload(Persona.user_files),
386+
joinedload(Persona.user_folders),
387+
joinedload(Persona.prompts),
355388
)
356-
if include_prompt:
357-
stmt = stmt.options(selectinload(Persona.prompts))
358389

359-
results = db_session.execute(stmt).scalars().all()
390+
results = db_session.execute(stmt).unique().scalars().all()
360391
return results
361392

362393

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

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
from onyx.db.persona import get_personas_for_user
3030
from onyx.db.persona import mark_persona_as_deleted
3131
from onyx.db.persona import mark_persona_as_not_deleted
32+
from onyx.db.persona import PersonaLoadType
3233
from onyx.db.persona import update_all_personas_display_priority
3334
from onyx.db.persona import update_persona_is_default
3435
from onyx.db.persona import update_persona_label
@@ -45,6 +46,7 @@
4546
from onyx.server.features.persona.models import FullPersonaSnapshot
4647
from onyx.server.features.persona.models import GenerateStarterMessageRequest
4748
from onyx.server.features.persona.models import ImageGenerationToolStatus
49+
from onyx.server.features.persona.models import MinimalPersonaSnapshot
4850
from onyx.server.features.persona.models import PersonaLabelCreate
4951
from onyx.server.features.persona.models import PersonaLabelResponse
5052
from onyx.server.features.persona.models import PersonaSharedNotificationData
@@ -154,7 +156,7 @@ def list_personas_admin(
154156
user=user,
155157
get_editable=get_editable,
156158
include_deleted=include_deleted,
157-
joinedload_all=True,
159+
load_type=PersonaLoadType.FULL,
158160
)
159161
]
160162

@@ -393,14 +395,13 @@ def list_personas(
393395
db_session: Session = Depends(get_session),
394396
include_deleted: bool = False,
395397
persona_ids: list[int] = Query(None),
396-
) -> list[PersonaSnapshot]:
398+
) -> list[MinimalPersonaSnapshot]:
397399
personas = get_personas_for_user(
400+
load_type=PersonaLoadType.MINIMAL,
398401
user=user,
399402
include_deleted=include_deleted,
400403
db_session=db_session,
401404
get_editable=False,
402-
joinedload_all=True,
403-
include_prompt=False,
404405
)
405406

406407
if persona_ids:
@@ -416,7 +417,8 @@ def list_personas(
416417
)
417418
]
418419

419-
return [PersonaSnapshot.from_model(p) for p in personas]
420+
result = [MinimalPersonaSnapshot.from_model(p) for p in personas]
421+
return result
420422

421423

422424
@basic_router.get("/{persona_id}")

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

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,64 @@
1818
logger = setup_logger()
1919

2020

21+
class MinimalPersonaSnapshot(BaseModel):
22+
"""Minimal persona model optimized for ChatPage.tsx - only includes fields actually used"""
23+
24+
# Core fields used by ChatPage
25+
id: int
26+
name: str
27+
description: str
28+
tools: list[ToolSnapshot]
29+
starter_messages: list[StarterMessage] | None
30+
document_sets: list[DocumentSet]
31+
llm_model_version_override: str | None
32+
llm_model_provider_override: str | None
33+
34+
uploaded_image_id: str | None
35+
icon_shape: int | None
36+
icon_color: str | None
37+
38+
is_public: bool
39+
is_visible: bool
40+
display_priority: int | None
41+
is_default_persona: bool
42+
builtin_persona: bool
43+
44+
labels: list["PersonaLabelSnapshot"]
45+
owner: MinimalUserSnapshot | None
46+
47+
@classmethod
48+
def from_model(cls, persona: Persona) -> "MinimalPersonaSnapshot":
49+
return MinimalPersonaSnapshot(
50+
# Core fields actually used by ChatPage
51+
id=persona.id,
52+
name=persona.name,
53+
description=persona.description,
54+
tools=[ToolSnapshot.from_model(tool) for tool in persona.tools],
55+
starter_messages=persona.starter_messages,
56+
document_sets=[
57+
DocumentSet.from_model(document_set)
58+
for document_set in persona.document_sets
59+
],
60+
llm_model_version_override=persona.llm_model_version_override,
61+
llm_model_provider_override=persona.llm_model_provider_override,
62+
uploaded_image_id=persona.uploaded_image_id,
63+
icon_shape=persona.icon_shape,
64+
icon_color=persona.icon_color,
65+
is_public=persona.is_public,
66+
is_visible=persona.is_visible,
67+
display_priority=persona.display_priority,
68+
is_default_persona=persona.is_default_persona,
69+
builtin_persona=persona.builtin_persona,
70+
labels=[PersonaLabelSnapshot.from_model(label) for label in persona.labels],
71+
owner=(
72+
MinimalUserSnapshot(id=persona.user.id, email=persona.user.email)
73+
if persona.user
74+
else None
75+
),
76+
)
77+
78+
2179
class PromptSnapshot(BaseModel):
2280
id: int
2381
name: str

backend/onyx/server/openai_assistants_api/asssistants_api.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
from onyx.db.persona import get_persona_by_id
1818
from onyx.db.persona import get_personas_for_user
1919
from onyx.db.persona import mark_persona_as_deleted
20+
from onyx.db.persona import PersonaLoadType
2021
from onyx.db.persona import upsert_persona
2122
from onyx.db.prompts import upsert_prompt
2223
from onyx.db.tools import get_tool_by_name
@@ -244,10 +245,10 @@ def list_assistants(
244245
) -> ListAssistantsResponse:
245246
personas = list(
246247
get_personas_for_user(
248+
load_type=PersonaLoadType.FULL,
247249
user=user,
248250
db_session=db_session,
249251
get_editable=False,
250-
joinedload_all=True,
251252
)
252253
)
253254

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

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -258,7 +258,7 @@ export function AssistantEditor({
258258
existingPersona?.llm_model_version_override ?? null,
259259
starter_messages: existingPersona?.starter_messages?.length
260260
? existingPersona.starter_messages
261-
: [{ message: "" }],
261+
: [{ message: "", name: "" }],
262262
enabled_tools_map: enabledToolsMap,
263263
icon_color: existingPersona?.icon_color ?? defautIconColor,
264264
icon_shape: existingPersona?.icon_shape ?? defaultIconShape,
@@ -526,10 +526,8 @@ export function AssistantEditor({
526526
// to tell the backend to not fetch any documents
527527
const numChunks = searchToolEnabled ? values.num_chunks || 25 : 0;
528528
const starterMessages = values.starter_messages
529-
.filter(
530-
(message: { message: string }) => message.message.trim() !== ""
531-
)
532-
.map((message: { message: string; name?: string }) => ({
529+
.filter((message: StarterMessage) => message.message.trim() !== "")
530+
.map((message: StarterMessage) => ({
533531
message: message.message,
534532
name: message.message,
535533
}));

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

Lines changed: 19 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@ import {
1717
import { FiEdit2 } from "react-icons/fi";
1818
import { TrashIcon } from "@/components/icons/icons";
1919
import { useUser } from "@/components/user/UserProvider";
20-
import { useAssistants } from "@/components/context/AssistantsContext";
2120
import { ConfirmEntityModal } from "@/components/modals/ConfirmEntityModal";
2221

2322
function PersonaTypeDisplay({ persona }: { persona: Persona }) {
@@ -40,16 +39,18 @@ function PersonaTypeDisplay({ persona }: { persona: Persona }) {
4039
return <Text>Personal {persona.owner && <>({persona.owner.email})</>}</Text>;
4140
}
4241

43-
export function PersonasTable() {
42+
export function PersonasTable({
43+
personas,
44+
refreshPersonas,
45+
}: {
46+
personas: Persona[];
47+
refreshPersonas: () => void;
48+
}) {
4449
const router = useRouter();
4550
const { popup, setPopup } = usePopup();
4651
const { refreshUser, isAdmin } = useUser();
47-
const {
48-
allAssistants: assistants,
49-
refreshAssistants,
50-
editablePersonas,
51-
} = useAssistants();
5252

53+
const editablePersonas = personas.filter((p) => !p.builtin_persona);
5354
const editablePersonaIds = useMemo(() => {
5455
return new Set(editablePersonas.map((p) => p.id.toString()));
5556
}, [editablePersonas]);
@@ -63,18 +64,18 @@ export function PersonasTable() {
6364

6465
useEffect(() => {
6566
const editable = editablePersonas.sort(personaComparator);
66-
const nonEditable = assistants
67+
const nonEditable = personas
6768
.filter((p) => !editablePersonaIds.has(p.id.toString()))
6869
.sort(personaComparator);
6970
setFinalPersonas([...editable, ...nonEditable]);
70-
}, [editablePersonas, assistants, editablePersonaIds]);
71+
}, [editablePersonas, personas, editablePersonaIds]);
7172

7273
const updatePersonaOrder = async (orderedPersonaIds: UniqueIdentifier[]) => {
73-
const reorderedAssistants = orderedPersonaIds.map(
74-
(id) => assistants.find((assistant) => assistant.id.toString() === id)!
74+
const reorderedPersonas = orderedPersonaIds.map(
75+
(id) => personas.find((persona) => persona.id.toString() === id)!
7576
);
7677

77-
setFinalPersonas(reorderedAssistants);
78+
setFinalPersonas(reorderedPersonas);
7879

7980
const displayPriorityMap = new Map<UniqueIdentifier, number>();
8081
orderedPersonaIds.forEach((personaId, ind) => {
@@ -96,12 +97,12 @@ export function PersonasTable() {
9697
type: "error",
9798
message: `Failed to update persona order - ${await response.text()}`,
9899
});
99-
setFinalPersonas(assistants);
100-
await refreshAssistants();
100+
setFinalPersonas(personas);
101+
await refreshPersonas();
101102
return;
102103
}
103104

104-
await refreshAssistants();
105+
await refreshPersonas();
105106
await refreshUser();
106107
};
107108

@@ -119,7 +120,7 @@ export function PersonasTable() {
119120
if (personaToDelete) {
120121
const response = await deletePersona(personaToDelete.id);
121122
if (response.ok) {
122-
await refreshAssistants();
123+
refreshPersonas();
123124
closeDeleteModal();
124125
} else {
125126
setPopup({
@@ -147,7 +148,7 @@ export function PersonasTable() {
147148
personaToToggleDefault.is_default_persona
148149
);
149150
if (response.ok) {
150-
await refreshAssistants();
151+
refreshPersonas();
151152
closeDefaultModal();
152153
} else {
153154
setPopup({
@@ -267,7 +268,7 @@ export function PersonasTable() {
267268
persona.is_visible
268269
);
269270
if (response.ok) {
270-
await refreshAssistants();
271+
refreshPersonas();
271272
} else {
272273
setPopup({
273274
type: "error",

0 commit comments

Comments
 (0)