From 756f4a2a66f3641f7e425c2e2c641f1965e4b6cd Mon Sep 17 00:00:00 2001
From: sabaimran <narmiabas@gmail.com>
Date: Fri, 3 Jan 2025 15:11:37 -0800
Subject: [PATCH] Update regeneration logic to run for all entries now that we
 have a single search model ID

---
 .../commands/change_default_model.py          | 35 ++++++++-----------
 1 file changed, 15 insertions(+), 20 deletions(-)

diff --git a/src/khoj/database/management/commands/change_default_model.py b/src/khoj/database/management/commands/change_default_model.py
index 9bbdbf83..d55fd44e 100644
--- a/src/khoj/database/management/commands/change_default_model.py
+++ b/src/khoj/database/management/commands/change_default_model.py
@@ -44,6 +44,7 @@ class Command(BaseCommand):
                 entries = Entry.objects.filter(entry_filter)[start:end]
                 compiled_entries = [entry.compiled for entry in entries]
                 updated_entries: List[Entry] = []
+
                 try:
                     embeddings = embeddings_model.embed_documents(compiled_entries)
                 except Exception as e:
@@ -87,27 +88,21 @@ class Command(BaseCommand):
 
         current_default = get_default_search_model()
 
-        # TODO: Migrate all Entry objects to use the new default Search model
+        # Create an entry filter with no conditions
+        entry_filter = Q()
+        relevant_entries = Entry.objects.filter(entry_filter).all()
+        logger.info(f"Number of Entry objects to update: {relevant_entries.count()}")
 
-        all_agents = Agent.objects.all()
-        logger.info(f"Number of Agent objects to update: {all_agents.count()}")
-        for agent in all_agents:
-            entry_filter = Q(agent=agent)
-            relevant_entries = Entry.objects.filter(entry_filter).all()
-            logger.info(f"Number of Entry objects to update for agent {agent}: {relevant_entries.count()}")
-
-            if apply:
-                try:
-                    regenerate_entries(
-                        entry_filter,
-                        embeddings_model[new_default_search_model_config.name],
-                        new_default_search_model_config,
-                    )
-                    logger.info(
-                        f"Updated {relevant_entries.count()} Entry objects for agent {agent} to use the new default Search model."
-                    )
-                except Exception as e:
-                    logger.error(f"Error embedding documents: {e}")
+        if apply:
+            try:
+                regenerate_entries(
+                    entry_filter=entry_filter,
+                    embeddings_model=embeddings_model[new_default_search_model_config.name],
+                    search_model=new_default_search_model_config,
+                )
+            except Exception as e:
+                logger.error(f"Error updating Entry objects: {e}")
+                return
 
         if apply and current_default.id != new_default_search_model_config.id:
             # Get the existing default SearchModelConfig object and update its name