From 9ebf3a4d8064138aadfe1440664475078501aa85 Mon Sep 17 00:00:00 2001
From: sabaimran <narmiabas@gmail.com>
Date: Mon, 27 May 2024 00:49:20 +0530
Subject: [PATCH] Improve the admin experience, add more metadata to the
 list_display

- Don't propagate max_tokens to the openai chat completion method. the max for the newer models is fixed at 4096 max output. The token limit is just used for input
---
 src/khoj/database/admin.py                    | 33 +++++++++++++++++--
 src/khoj/processor/conversation/openai/gpt.py |  1 -
 .../processor/conversation/openai/utils.py    |  3 +-
 src/khoj/routers/helpers.py                   |  2 +-
 4 files changed, 32 insertions(+), 7 deletions(-)

diff --git a/src/khoj/database/admin.py b/src/khoj/database/admin.py
index c7adb1ea..a96a57a0 100644
--- a/src/khoj/database/admin.py
+++ b/src/khoj/database/admin.py
@@ -78,10 +78,8 @@ class KhojUserAdmin(UserAdmin):
 
 admin.site.register(KhojUser, KhojUserAdmin)
 
-admin.site.register(ChatModelOptions)
 admin.site.register(ProcessLock)
 admin.site.register(SpeechToTextModelOptions)
-admin.site.register(OpenAIProcessorConversationConfig)
 admin.site.register(SearchModelConfig)
 admin.site.register(ReflectiveQuestion)
 admin.site.register(UserSearchModelConfig)
@@ -89,7 +87,6 @@ admin.site.register(TextToImageModelConfig)
 admin.site.register(ClientApplication)
 admin.site.register(GithubConfig)
 admin.site.register(NotionConfig)
-admin.site.register(ServerChatSettings)
 
 
 @admin.register(Agent)
@@ -131,6 +128,36 @@ class KhojUserSubscription(admin.ModelAdmin):
     list_filter = ("type",)
 
 
+@admin.register(ChatModelOptions)
+class ChatModelOptionsAdmin(admin.ModelAdmin):
+    list_display = (
+        "id",
+        "chat_model",
+        "model_type",
+        "max_prompt_size",
+    )
+    search_fields = ("id", "chat_model", "model_type")
+
+
+@admin.register(OpenAIProcessorConversationConfig)
+class OpenAIProcessorConversationConfigAdmin(admin.ModelAdmin):
+    list_display = (
+        "id",
+        "name",
+        "api_key",
+        "api_base_url",
+    )
+    search_fields = ("id", "name", "api_key", "api_base_url")
+
+
+@admin.register(ServerChatSettings)
+class ServerChatSettingsAdmin(admin.ModelAdmin):
+    list_display = (
+        "default_model",
+        "summarizer_model",
+    )
+
+
 @admin.register(Conversation)
 class ConversationAdmin(admin.ModelAdmin):
     list_display = (
diff --git a/src/khoj/processor/conversation/openai/gpt.py b/src/khoj/processor/conversation/openai/gpt.py
index c25f05fd..f5ff5384 100644
--- a/src/khoj/processor/conversation/openai/gpt.py
+++ b/src/khoj/processor/conversation/openai/gpt.py
@@ -67,7 +67,6 @@ def extract_questions(
         messages=messages,
         model=model,
         temperature=temperature,
-        max_tokens=max_tokens,
         api_base_url=api_base_url,
         model_kwargs={"response_format": {"type": "json_object"}},
         openai_api_key=api_key,
diff --git a/src/khoj/processor/conversation/openai/utils.py b/src/khoj/processor/conversation/openai/utils.py
index 0c37ba53..49f1d6b6 100644
--- a/src/khoj/processor/conversation/openai/utils.py
+++ b/src/khoj/processor/conversation/openai/utils.py
@@ -34,7 +34,7 @@ openai_clients: Dict[str, openai.OpenAI] = {}
     reraise=True,
 )
 def completion_with_backoff(
-    messages, model, temperature=0, openai_api_key=None, api_base_url=None, model_kwargs=None, max_tokens=None
+    messages, model, temperature=0, openai_api_key=None, api_base_url=None, model_kwargs=None
 ) -> str:
     client_key = f"{openai_api_key}--{api_base_url}"
     client: openai.OpenAI = openai_clients.get(client_key)
@@ -53,7 +53,6 @@ def completion_with_backoff(
         model=model,  # type: ignore
         temperature=temperature,
         timeout=20,
-        max_tokens=max_tokens,
         **(model_kwargs or dict()),
     )
     aggregated_response = ""
diff --git a/src/khoj/routers/helpers.py b/src/khoj/routers/helpers.py
index 51106408..ed526f5e 100644
--- a/src/khoj/routers/helpers.py
+++ b/src/khoj/routers/helpers.py
@@ -993,7 +993,7 @@ def scheduled_chat(
 
     # Stop if the chat API call was not successful
     if raw_response.status_code != 200:
-        logger.error(f"Failed to run schedule chat: {raw_response.text}")
+        logger.error(f"Failed to run schedule chat: {raw_response.text}, user: {user}, query: {query_to_run}")
         return None
 
     # Extract the AI response from the chat API response