Support Gemma 2 Model Family for Offline Chat

pyproject.toml

@@ -66,7 +66,7 @@ dependencies = [
     "pymupdf >= 1.23.5",
     "django == 5.0.7",
     "authlib == 1.2.1",
-    "llama-cpp-python == 0.2.76",
+    "llama-cpp-python == 0.2.82",
     "itsdangerous == 2.1.2",
     "httpx == 0.25.0",
     "pgvector == 0.2.4",

src/khoj/database/adapters/__init__.py

@@ -559,7 +559,7 @@ def create_default_agent():
         if default_conversation_config is None:
             logger.info("No default conversation config found, skipping default agent creation")
             return None
-        default_personality = prompts.personality.format(current_date="placeholder")
+        default_personality = prompts.personality.format(current_date="placeholder", day_of_week="placeholder")
 
         agent = Agent.objects.filter(name=AgentAdapters.DEFAULT_AGENT_NAME).first()
 

src/khoj/processor/conversation/anthropic/anthropic_chat.py

@@ -36,7 +36,7 @@ def extract_questions_anthropic(
     # Extract Past User Message and Inferred Questions from Conversation Log
     chat_history = "".join(
         [
-            f'Q: {chat["intent"]["query"]}\nKhoj: {{"queries": {chat["intent"].get("inferred-queries") or list([chat["intent"]["query"]])}}}\nA: {chat["message"]}\n\n'
+            f'User: {chat["intent"]["query"]}\nAssistant: {{"queries": {chat["intent"].get("inferred-queries") or list([chat["intent"]["query"]])}}}\nA: {chat["message"]}\n\n'
             for chat in conversation_log.get("chat", [])[-4:]
             if chat["by"] == "khoj" and "text-to-image" not in chat["intent"].get("type")
         ]
@@ -135,17 +135,23 @@ def converse_anthropic(
     Converse with user using Anthropic's Claude
     """
     # Initialize Variables
-    current_date = datetime.now().strftime("%Y-%m-%d")
+    current_date = datetime.now()
     compiled_references = "\n\n".join({f"# {item}" for item in references})
 
     conversation_primer = prompts.query_prompt.format(query=user_query)
 
     if agent and agent.personality:
         system_prompt = prompts.custom_personality.format(
-            name=agent.name, bio=agent.personality, current_date=current_date
+            name=agent.name,
+            bio=agent.personality,
+            current_date=current_date.strftime("%Y-%m-%d"),
+            day_of_week=current_date.strftime("%A"),
         )
     else:
-        system_prompt = prompts.personality.format(current_date=current_date)
+        system_prompt = prompts.personality.format(
+            current_date=current_date.strftime("%Y-%m-%d"),
+            day_of_week=current_date.strftime("%A"),
+        )
 
     if location_data:
         location = f"{location_data.city}, {location_data.region}, {location_data.country}"

src/khoj/processor/conversation/offline/chat_model.py

@@ -55,26 +55,29 @@ def extract_questions_offline(
                 chat_history += f"Q: {chat['intent']['query']}\n"
                 chat_history += f"Khoj: {chat['message']}\n\n"
 
+    # Get dates relative to today for prompt creation
     today = datetime.today()
     yesterday = (today - timedelta(days=1)).strftime("%Y-%m-%d")
     last_year = today.year - 1
     example_questions = prompts.extract_questions_offline.format(
         query=text,
         chat_history=chat_history,
         current_date=today.strftime("%Y-%m-%d"),
+        day_of_week=today.strftime("%A"),
         yesterday_date=yesterday,
         last_year=last_year,
         this_year=today.year,
         location=location,
     )
+
     messages = generate_chatml_messages_with_context(
         example_questions, model_name=model, loaded_model=offline_chat_model, max_prompt_size=max_prompt_size
     )
 
     state.chat_lock.acquire()
     try:
         response = send_message_to_model_offline(
-            messages, loaded_model=offline_chat_model, max_prompt_size=max_prompt_size
+            messages, loaded_model=offline_chat_model, model=model, max_prompt_size=max_prompt_size
         )
     finally:
         state.chat_lock.release()
@@ -96,7 +99,7 @@ def extract_questions_offline(
     except:
         logger.warning(f"Llama returned invalid JSON. Falling back to using user message as search query.\n{response}")
         return all_questions
-    logger.debug(f"Extracted Questions by Llama: {questions}")
+    logger.debug(f"Questions extracted by {model}: {questions}")
     return questions
 
 
@@ -144,14 +147,20 @@ def converse_offline(
     offline_chat_model = loaded_model or download_model(model, max_tokens=max_prompt_size)
     compiled_references_message = "\n\n".join({f"{item['compiled']}" for item in references})
 
-    current_date = datetime.now().strftime("%Y-%m-%d")
+    current_date = datetime.now()
 
     if agent and agent.personality:
         system_prompt = prompts.custom_system_prompt_offline_chat.format(
-            name=agent.name, bio=agent.personality, current_date=current_date
+            name=agent.name,
+            bio=agent.personality,
+            current_date=current_date.strftime("%Y-%m-%d"),
+            day_of_week=current_date.strftime("%A"),
         )
     else:
-        system_prompt = prompts.system_prompt_offline_chat.format(current_date=current_date)
+        system_prompt = prompts.system_prompt_offline_chat.format(
+            current_date=current_date.strftime("%Y-%m-%d"),
+            day_of_week=current_date.strftime("%A"),
+        )
 
     conversation_primer = prompts.query_prompt.format(query=user_query)
 
@@ -177,9 +186,9 @@ def converse_offline(
             if online_results[result].get("webpages"):
                 simplified_online_results[result] = online_results[result]["webpages"]
 
-        conversation_primer = f"{prompts.online_search_conversation.format(online_results=str(simplified_online_results))}\n{conversation_primer}"
+        conversation_primer = f"{prompts.online_search_conversation_offline.format(online_results=str(simplified_online_results))}\n{conversation_primer}"
     if not is_none_or_empty(compiled_references_message):
-        conversation_primer = f"{prompts.notes_conversation_offline.format(references=compiled_references_message)}\n{conversation_primer}"
+        conversation_primer = f"{prompts.notes_conversation_offline.format(references=compiled_references_message)}\n\n{conversation_primer}"
 
     # Setup Prompt with Primer or Conversation History
     messages = generate_chatml_messages_with_context(
@@ -192,6 +201,9 @@ def converse_offline(
         tokenizer_name=tokenizer_name,
     )
 
+    truncated_messages = "\n".join({f"{message.content[:70]}..." for message in messages})
+    logger.debug(f"Conversation Context for {model}: {truncated_messages}")
+
     g = ThreadedGenerator(references, online_results, completion_func=completion_func)
     t = Thread(target=llm_thread, args=(g, messages, offline_chat_model, max_prompt_size))
     t.start()

src/khoj/processor/conversation/offline/utils.py

@@ -24,6 +24,8 @@ def download_model(repo_id: str, filename: str = "*Q4_K_M.gguf", max_tokens: int
     # Add chat format if known
     if "llama-3" in repo_id.lower():
         kwargs["chat_format"] = "llama-3"
+    elif "gemma-2" in repo_id.lower():
+        kwargs["chat_format"] = "gemma"
 
     # Check if the model is already downloaded
     model_path = load_model_from_cache(repo_id, filename)

src/khoj/processor/conversation/openai/gpt.py

@@ -125,17 +125,23 @@ def converse(
     Converse with user using OpenAI's ChatGPT
     """
     # Initialize Variables
-    current_date = datetime.now().strftime("%Y-%m-%d")
+    current_date = datetime.now()
     compiled_references = "\n\n".join({f"# {item['compiled']}" for item in references})
 
     conversation_primer = prompts.query_prompt.format(query=user_query)
 
     if agent and agent.personality:
         system_prompt = prompts.custom_personality.format(
-            name=agent.name, bio=agent.personality, current_date=current_date
+            name=agent.name,
+            bio=agent.personality,
+            current_date=current_date.strftime("%Y-%m-%d"),
+            day_of_week=current_date.strftime("%A"),
         )
     else:
-        system_prompt = prompts.personality.format(current_date=current_date)
+        system_prompt = prompts.personality.format(
+            current_date=current_date.strftime("%Y-%m-%d"),
+            day_of_week=current_date.strftime("%A"),
+        )
 
     if location_data:
         location = f"{location_data.city}, {location_data.region}, {location_data.country}"