From 7426a4f8198a6cc0b48299b8097d2f22b9a0b6a6 Mon Sep 17 00:00:00 2001 From: sabaimran Date: Fri, 19 Jul 2024 14:43:30 +0530 Subject: [PATCH] Prefetch related agent when retrieving the conversation for performance improvements --- src/khoj/database/adapters/__init__.py | 25 ++++++++++++++++++++----- 1 file changed, 20 insertions(+), 5 deletions(-) diff --git a/src/khoj/database/adapters/__init__.py b/src/khoj/database/adapters/__init__.py index 20019c56..0c0724ee 100644 --- a/src/khoj/database/adapters/__init__.py +++ b/src/khoj/database/adapters/__init__.py @@ -681,18 +681,33 @@ class ConversationAdapters: user: KhojUser, client_application: ClientApplication = None, conversation_id: int = None, title: str = None ) -> Optional[Conversation]: if conversation_id: - return await Conversation.objects.filter(user=user, client=client_application, id=conversation_id).afirst() + return ( + await Conversation.objects.filter(user=user, client=client_application, id=conversation_id) + .prefetch_related("agent") + .afirst() + ) elif title: - return await Conversation.objects.filter(user=user, client=client_application, title=title).afirst() + return ( + await Conversation.objects.filter(user=user, client=client_application, title=title) + .prefetch_related("agent") + .afirst() + ) else: - conversation = Conversation.objects.filter(user=user, client=client_application).order_by("-updated_at") + conversation = ( + Conversation.objects.filter(user=user, client=client_application) + .prefetch_related("agent") + .order_by("-updated_at") + ) if await conversation.aexists(): return await conversation.prefetch_related("agent").afirst() return await ( - Conversation.objects.filter(user=user, client=client_application).order_by("-updated_at").afirst() - ) or await Conversation.objects.acreate(user=user, client=client_application) + Conversation.objects.filter(user=user, client=client_application) + .prefetch_related("agent") + .order_by("-updated_at") + .afirst() + ) or await Conversation.objects.prefetch_related("agent").acreate(user=user, client=client_application) @staticmethod async def adelete_conversation_by_user(