From 60658a80377c846fb59e44a69781ce12edd4b692 Mon Sep 17 00:00:00 2001 From: sabaimran Date: Tue, 23 Apr 2024 23:08:29 +0530 Subject: [PATCH] Get rid of enable flag for the offline chat processor config - Default, assume that offline chat is enabled if there is an offline chat model option configured --- src/khoj/database/adapters/__init__.py | 25 +------------------ src/khoj/database/admin.py | 2 -- ..._offlinechatprocessorconversationconfig.py | 15 +++++++++++ src/khoj/database/models/__init__.py | 4 --- src/khoj/migrations/migrate_server_pg.py | 4 --- .../processor/conversation/offline/utils.py | 1 + src/khoj/routers/api.py | 7 +----- src/khoj/routers/helpers.py | 21 ++++++---------- src/khoj/utils/initialization.py | 2 -- src/khoj/utils/rawconfig.py | 1 - tests/conftest.py | 8 ++++-- tests/helpers.py | 8 ------ tests/test_offline_chat_actors.py | 2 +- 13 files changed, 33 insertions(+), 67 deletions(-) create mode 100644 src/khoj/database/migrations/0036_delete_offlinechatprocessorconversationconfig.py diff --git a/src/khoj/database/adapters/__init__.py b/src/khoj/database/adapters/__init__.py index 4752a859..96d8e51d 100644 --- a/src/khoj/database/adapters/__init__.py +++ b/src/khoj/database/adapters/__init__.py @@ -28,7 +28,6 @@ from khoj.database.models import ( KhojApiUser, KhojUser, NotionConfig, - OfflineChatProcessorConversationConfig, OpenAIProcessorConversationConfig, ProcessLock, ReflectiveQuestion, @@ -628,18 +627,6 @@ class ConversationAdapters: async def aget_openai_conversation_config(): return await OpenAIProcessorConversationConfig.objects.filter().afirst() - @staticmethod - def get_offline_chat_conversation_config(): - return OfflineChatProcessorConversationConfig.objects.filter().first() - - @staticmethod - async def aget_offline_chat_conversation_config(): - return await OfflineChatProcessorConversationConfig.objects.filter().afirst() - - @staticmethod - def has_valid_offline_conversation_config(): - return OfflineChatProcessorConversationConfig.objects.filter(enabled=True).exists() - @staticmethod def has_valid_openai_conversation_config(): return OpenAIProcessorConversationConfig.objects.filter().exists() @@ -710,14 +697,6 @@ class ConversationAdapters: user_conversation_config.setting = new_config user_conversation_config.save() - @staticmethod - def has_offline_chat(): - return OfflineChatProcessorConversationConfig.objects.filter(enabled=True).exists() - - @staticmethod - async def ahas_offline_chat(): - return await OfflineChatProcessorConversationConfig.objects.filter(enabled=True).aexists() - @staticmethod async def get_default_offline_llm(): return await ChatModelOptions.objects.filter(model_type="offline").afirst() @@ -765,8 +744,6 @@ class ConversationAdapters: @staticmethod def get_valid_conversation_config(user: KhojUser, conversation: Conversation): - offline_chat_config = ConversationAdapters.get_offline_chat_conversation_config() - if conversation.agent and conversation.agent.chat_model: conversation_config = conversation.agent.chat_model else: @@ -775,7 +752,7 @@ class ConversationAdapters: if conversation_config is None: conversation_config = ConversationAdapters.get_default_conversation_config() - if offline_chat_config and offline_chat_config.enabled and conversation_config.model_type == "offline": + if conversation_config.model_type == "offline": if state.offline_chat_processor_config is None or state.offline_chat_processor_config.loaded_model is None: chat_model = conversation_config.chat_model max_tokens = conversation_config.max_prompt_size diff --git a/src/khoj/database/admin.py b/src/khoj/database/admin.py index 98f9f38d..b76b0fce 100644 --- a/src/khoj/database/admin.py +++ b/src/khoj/database/admin.py @@ -14,7 +14,6 @@ from khoj.database.models import ( GithubConfig, KhojUser, NotionConfig, - OfflineChatProcessorConversationConfig, OpenAIProcessorConversationConfig, ReflectiveQuestion, SearchModelConfig, @@ -47,7 +46,6 @@ admin.site.register(KhojUser, KhojUserAdmin) admin.site.register(ChatModelOptions) admin.site.register(SpeechToTextModelOptions) admin.site.register(OpenAIProcessorConversationConfig) -admin.site.register(OfflineChatProcessorConversationConfig) admin.site.register(SearchModelConfig) admin.site.register(Subscription) admin.site.register(ReflectiveQuestion) diff --git a/src/khoj/database/migrations/0036_delete_offlinechatprocessorconversationconfig.py b/src/khoj/database/migrations/0036_delete_offlinechatprocessorconversationconfig.py new file mode 100644 index 00000000..40c51016 --- /dev/null +++ b/src/khoj/database/migrations/0036_delete_offlinechatprocessorconversationconfig.py @@ -0,0 +1,15 @@ +# Generated by Django 4.2.10 on 2024-04-23 17:35 + +from django.db import migrations + + +class Migration(migrations.Migration): + dependencies = [ + ("database", "0035_processlock"), + ] + + operations = [ + migrations.DeleteModel( + name="OfflineChatProcessorConversationConfig", + ), + ] diff --git a/src/khoj/database/models/__init__.py b/src/khoj/database/models/__init__.py index f6e8a0a6..15f396f1 100644 --- a/src/khoj/database/models/__init__.py +++ b/src/khoj/database/models/__init__.py @@ -201,10 +201,6 @@ class OpenAIProcessorConversationConfig(BaseModel): api_key = models.CharField(max_length=200) -class OfflineChatProcessorConversationConfig(BaseModel): - enabled = models.BooleanField(default=False) - - class SpeechToTextModelOptions(BaseModel): class ModelType(models.TextChoices): OPENAI = "openai" diff --git a/src/khoj/migrations/migrate_server_pg.py b/src/khoj/migrations/migrate_server_pg.py index fbb26a1e..0ab3522b 100644 --- a/src/khoj/migrations/migrate_server_pg.py +++ b/src/khoj/migrations/migrate_server_pg.py @@ -62,7 +62,6 @@ from packaging import version from khoj.database.models import ( ChatModelOptions, - OfflineChatProcessorConversationConfig, OpenAIProcessorConversationConfig, SearchModelConfig, ) @@ -103,9 +102,6 @@ def migrate_server_pg(args): if "offline-chat" in raw_config["processor"]["conversation"]: offline_chat = raw_config["processor"]["conversation"]["offline-chat"] - OfflineChatProcessorConversationConfig.objects.create( - enabled=offline_chat.get("enable-offline-chat"), - ) ChatModelOptions.objects.create( chat_model=offline_chat.get("chat-model"), tokenizer=processor_conversation.get("tokenizer"), diff --git a/src/khoj/processor/conversation/offline/utils.py b/src/khoj/processor/conversation/offline/utils.py index 44dec0b6..24d39ca4 100644 --- a/src/khoj/processor/conversation/offline/utils.py +++ b/src/khoj/processor/conversation/offline/utils.py @@ -67,5 +67,6 @@ def load_model_from_cache(repo_id: str, filename: str, repo_type="models"): def infer_max_tokens(model_context_window: int, configured_max_tokens=math.inf) -> int: """Infer max prompt size based on device memory and max context window supported by the model""" + configured_max_tokens = math.inf if configured_max_tokens is None else configured_max_tokens vram_based_n_ctx = int(get_device_memory() / 2e6) # based on heuristic return min(configured_max_tokens, vram_based_n_ctx, model_context_window) diff --git a/src/khoj/routers/api.py b/src/khoj/routers/api.py index c511b6d9..362038b7 100644 --- a/src/khoj/routers/api.py +++ b/src/khoj/routers/api.py @@ -303,15 +303,10 @@ async def extract_references_and_questions( # Infer search queries from user message with timer("Extracting search queries took", logger): # If we've reached here, either the user has enabled offline chat or the openai model is enabled. - offline_chat_config = await ConversationAdapters.aget_offline_chat_conversation_config() conversation_config = await ConversationAdapters.aget_conversation_config(user) if conversation_config is None: conversation_config = await ConversationAdapters.aget_default_conversation_config() - if ( - offline_chat_config - and offline_chat_config.enabled - and conversation_config.model_type == ChatModelOptions.ModelType.OFFLINE - ): + if conversation_config.model_type == ChatModelOptions.ModelType.OFFLINE: using_offline_chat = True default_offline_llm = await ConversationAdapters.get_default_offline_llm() chat_model = default_offline_llm.chat_model diff --git a/src/khoj/routers/helpers.py b/src/khoj/routers/helpers.py index 8fdddc07..96c713e0 100644 --- a/src/khoj/routers/helpers.py +++ b/src/khoj/routers/helpers.py @@ -65,23 +65,20 @@ executor = ThreadPoolExecutor(max_workers=1) def validate_conversation_config(): - if ( - ConversationAdapters.has_valid_offline_conversation_config() - or ConversationAdapters.has_valid_openai_conversation_config() - ): - if ConversationAdapters.get_default_conversation_config() is None: - raise HTTPException(status_code=500, detail="Contact the server administrator to set a default chat model.") - return + default_config = ConversationAdapters.get_default_conversation_config() - raise HTTPException(status_code=500, detail="Set your OpenAI API key or enable Local LLM via Khoj settings.") + if default_config is None: + raise HTTPException(status_code=500, detail="Contact the server administrator to set a default chat model.") + + if default_config.model_type == "openai" and not ConversationAdapters.has_valid_openai_conversation_config(): + raise HTTPException(status_code=500, detail="Contact the server administrator to set a default chat model.") async def is_ready_to_chat(user: KhojUser): - has_offline_config = await ConversationAdapters.ahas_offline_chat() has_openai_config = await ConversationAdapters.has_openai_chat() user_conversation_config = await ConversationAdapters.aget_user_conversation_config(user) - if has_offline_config and user_conversation_config and user_conversation_config.model_type == "offline": + if user_conversation_config and user_conversation_config.model_type == "offline": chat_model = user_conversation_config.chat_model max_tokens = user_conversation_config.max_prompt_size if state.offline_chat_processor_config is None: @@ -89,9 +86,7 @@ async def is_ready_to_chat(user: KhojUser): state.offline_chat_processor_config = OfflineChatProcessorModel(chat_model, max_tokens) return True - ready = has_openai_config or has_offline_config - - if not ready: + if not has_openai_config: raise HTTPException(status_code=500, detail="Set your OpenAI API key or enable Local LLM via Khoj settings.") diff --git a/src/khoj/utils/initialization.py b/src/khoj/utils/initialization.py index 873f3f47..8ef0ef53 100644 --- a/src/khoj/utils/initialization.py +++ b/src/khoj/utils/initialization.py @@ -5,7 +5,6 @@ from khoj.database.adapters import ConversationAdapters from khoj.database.models import ( ChatModelOptions, KhojUser, - OfflineChatProcessorConversationConfig, OpenAIProcessorConversationConfig, SpeechToTextModelOptions, TextToImageModelConfig, @@ -35,7 +34,6 @@ def initialization(): use_offline_model = input("Use offline chat model? (y/n): ") if use_offline_model == "y": logger.info("🗣️ Setting up offline chat model") - OfflineChatProcessorConversationConfig.objects.create(enabled=True) offline_chat_model = input( f"Enter the offline chat model you want to use. See HuggingFace for available GGUF models (default: {default_offline_chat_model}): " diff --git a/src/khoj/utils/rawconfig.py b/src/khoj/utils/rawconfig.py index 15ecb85d..49202396 100644 --- a/src/khoj/utils/rawconfig.py +++ b/src/khoj/utils/rawconfig.py @@ -81,7 +81,6 @@ class OpenAIProcessorConfig(ConfigBase): class OfflineChatProcessorConfig(ConfigBase): - enable_offline_chat: Optional[bool] = False chat_model: Optional[str] = "NousResearch/Hermes-2-Pro-Mistral-7B-GGUF" diff --git a/tests/conftest.py b/tests/conftest.py index 2bf93928..28a4f2c5 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -33,7 +33,6 @@ from khoj.utils.helpers import resolve_absolute_path from khoj.utils.rawconfig import ContentConfig, ImageSearchConfig, SearchConfig from tests.helpers import ( ChatModelOptionsFactory, - OfflineChatProcessorConversationConfigFactory, OpenAIProcessorConversationConfigFactory, ProcessLockFactory, SubscriptionFactory, @@ -377,7 +376,12 @@ def client_offline_chat(search_config: SearchConfig, default_user2: KhojUser): configure_content(all_files, user=default_user2) # Initialize Processor from Config - OfflineChatProcessorConversationConfigFactory(enabled=True) + ChatModelOptionsFactory( + chat_model="NousResearch/Hermes-2-Pro-Mistral-7B-GGUF", + tokenizer=None, + max_prompt_size=None, + model_type="offline", + ) UserConversationProcessorConfigFactory(user=default_user2) state.anonymous_mode = True diff --git a/tests/helpers.py b/tests/helpers.py index 0dbac55f..642f05dd 100644 --- a/tests/helpers.py +++ b/tests/helpers.py @@ -9,7 +9,6 @@ from khoj.database.models import ( Conversation, KhojApiUser, KhojUser, - OfflineChatProcessorConversationConfig, OpenAIProcessorConversationConfig, ProcessLock, SearchModelConfig, @@ -55,13 +54,6 @@ class UserConversationProcessorConfigFactory(factory.django.DjangoModelFactory): setting = factory.SubFactory(ChatModelOptionsFactory) -class OfflineChatProcessorConversationConfigFactory(factory.django.DjangoModelFactory): - class Meta: - model = OfflineChatProcessorConversationConfig - - enabled = True - - class OpenAIProcessorConversationConfigFactory(factory.django.DjangoModelFactory): class Meta: model = OpenAIProcessorConversationConfig diff --git a/tests/test_offline_chat_actors.py b/tests/test_offline_chat_actors.py index 77f2b105..5e5804da 100644 --- a/tests/test_offline_chat_actors.py +++ b/tests/test_offline_chat_actors.py @@ -24,7 +24,7 @@ from khoj.utils.constants import default_offline_chat_model @pytest.fixture(scope="session") def loaded_model(): - return download_model(default_offline_chat_model) + return download_model(default_offline_chat_model, max_tokens=5000) freezegun.configure(extend_ignore_list=["transformers"])