From 3a5d93d673bbaba5506ca33eb105ac776b561a21 Mon Sep 17 00:00:00 2001
From: sabaimran <narmiabas@gmail.com>
Date: Tue, 1 Aug 2023 09:25:05 -0700
Subject: [PATCH] Add migration script for getting the new offline model

---
 src/khoj/migrations/migrate_offline_model.py | 19 +++++++++++++++++++
 src/khoj/utils/cli.py                        |  3 ++-
 2 files changed, 21 insertions(+), 1 deletion(-)
 create mode 100644 src/khoj/migrations/migrate_offline_model.py

diff --git a/src/khoj/migrations/migrate_offline_model.py b/src/khoj/migrations/migrate_offline_model.py
new file mode 100644
index 00000000..7ab813ce
--- /dev/null
+++ b/src/khoj/migrations/migrate_offline_model.py
@@ -0,0 +1,19 @@
+import os
+
+from khoj.utils.yaml import load_config_from_file, save_config_to_file
+
+
+def migrate_offline_model(args):
+    raw_config = load_config_from_file(args.config_file)
+    version = raw_config.get("version")
+
+    if version == "0.10.0" or version == None:
+        # If the user has downloaded the offline model, remove it from the cache.
+        offline_model_path = os.path.expanduser("~/.cache/gpt4all/llama-2-7b-chat.ggmlv3.q4_K_S.bin")
+        if os.path.exists(offline_model_path):
+            os.remove(offline_model_path)
+
+    raw_config["version"] = args.version_no
+    save_config_to_file(raw_config, args.config_file)
+
+    return args
diff --git a/src/khoj/utils/cli.py b/src/khoj/utils/cli.py
index 9236ab11..787289fe 100644
--- a/src/khoj/utils/cli.py
+++ b/src/khoj/utils/cli.py
@@ -8,6 +8,7 @@ from khoj.utils.helpers import resolve_absolute_path
 from khoj.utils.yaml import parse_config_from_file
 from khoj.migrations.migrate_version import migrate_config_to_version
 from khoj.migrations.migrate_processor_config_openai import migrate_processor_conversation_schema
+from khoj.migrations.migrate_offline_model import migrate_offline_model
 
 
 def cli(args=None):
@@ -55,7 +56,7 @@ def cli(args=None):
 
 
 def run_migrations(args):
-    migrations = [migrate_config_to_version, migrate_processor_conversation_schema]
+    migrations = [migrate_config_to_version, migrate_processor_conversation_schema, migrate_offline_model]
     for migration in migrations:
         args = migration(args)
     return args