khoj 1.27.2.dev18__py3-none-any.whl → 1.27.2.dev130__py3-none-any.whl

khoj/processor/conversation/google/gemini_chat.py

@@ -14,11 +14,13 @@ from khoj.processor.conversation.google.utils import (
     gemini_completion_with_backoff,
 )
 from khoj.processor.conversation.utils import (
+    clean_json,
     construct_structured_message,
     generate_chatml_messages_with_context,
 )
 from khoj.utils.helpers import ConversationCommand, is_none_or_empty
 from khoj.utils.rawconfig import LocationData
+from khoj.utils.yaml import yaml_dump
 
 logger = logging.getLogger(__name__)
 
@@ -35,6 +37,7 @@ def extract_questions_gemini(
     query_images: Optional[list[str]] = None,
     vision_enabled: bool = False,
     personality_context: Optional[str] = None,
+    tracer: dict = {},
 ):
     """
     Infer search queries to retrieve relevant notes to answer user query
@@ -85,15 +88,12 @@ def extract_questions_gemini(
     messages = [ChatMessage(content=prompt, role="user"), ChatMessage(content=system_prompt, role="system")]
 
     response = gemini_send_message_to_model(
-        messages, api_key, model, response_type="json_object", temperature=temperature
+        messages, api_key, model, response_type="json_object", temperature=temperature, tracer=tracer
     )
 
     # Extract, Clean Message from Gemini's Response
     try:
-        response = response.strip()
-        match = re.search(r"\{.*?\}", response)
-        if match:
-            response = match.group()
+        response = clean_json(response)
         response = json.loads(response)
         response = [q.strip() for q in response["queries"] if q.strip()]
         if not isinstance(response, list) or not response:
@@ -107,15 +107,19 @@ def extract_questions_gemini(
     return questions
 
 
-def gemini_send_message_to_model(messages, api_key, model, response_type="text", temperature=0, model_kwargs=None):
+def gemini_send_message_to_model(
+    messages, api_key, model, response_type="text", temperature=0, model_kwargs=None, tracer={}
+):
     """
     Send message to model
     """
     messages, system_prompt = format_messages_for_gemini(messages)
 
     model_kwargs = {}
-    if response_type == "json_object":
-        model_kwargs["response_mime_type"] = "application/json"
+
+    # Sometimes, this causes unwanted behavior and terminates response early. Disable for now while it's flaky.
+    # if response_type == "json_object":
+    #     model_kwargs["response_mime_type"] = "application/json"
 
     # Get Response from Gemini
     return gemini_completion_with_backoff(
@@ -125,6 +129,7 @@ def gemini_send_message_to_model(messages, api_key, model, response_type="text",
         api_key=api_key,
         temperature=temperature,
         model_kwargs=model_kwargs,
+        tracer=tracer,
     )
 
 
@@ -132,6 +137,7 @@ def converse_gemini(
     references,
     user_query,
     online_results: Optional[Dict[str, Dict]] = None,
+    code_results: Optional[Dict[str, Dict]] = None,
     conversation_log={},
     model: Optional[str] = "gemini-1.5-flash",
     api_key: Optional[str] = None,
@@ -145,13 +151,13 @@ def converse_gemini(
     agent: Agent = None,
     query_images: Optional[list[str]] = None,
     vision_available: bool = False,
+    tracer={},
 ):
     """
     Converse with user using Google's Gemini
     """
     # Initialize Variables
     current_date = datetime.now()
-    compiled_references = "\n\n".join({f"# File: {item['file']}\n## {item['compiled']}\n" for item in references})
 
     if agent and agent.personality:
         system_prompt = prompts.custom_personality.format(
@@ -176,7 +182,7 @@ def converse_gemini(
         system_prompt = f"{system_prompt}\n{user_name_prompt}"
 
     # Get Conversation Primer appropriate to Conversation Type
-    if conversation_commands == [ConversationCommand.Notes] and is_none_or_empty(compiled_references):
+    if conversation_commands == [ConversationCommand.Notes] and is_none_or_empty(references):
         completion_func(chat_response=prompts.no_notes_found.format())
         return iter([prompts.no_notes_found.format()])
     elif conversation_commands == [ConversationCommand.Online] and is_none_or_empty(online_results):
@@ -184,10 +190,13 @@ def converse_gemini(
         return iter([prompts.no_online_results_found.format()])
 
     context_message = ""
-    if not is_none_or_empty(compiled_references):
-        context_message = f"{prompts.notes_conversation.format(query=user_query, references=compiled_references)}\n\n"
+    if not is_none_or_empty(references):
+        context_message = f"{prompts.notes_conversation.format(query=user_query, references=yaml_dump(references))}\n\n"
     if ConversationCommand.Online in conversation_commands or ConversationCommand.Webpage in conversation_commands:
-        context_message += f"{prompts.online_search_conversation.format(online_results=str(online_results))}"
+        context_message += f"{prompts.online_search_conversation.format(online_results=yaml_dump(online_results))}\n\n"
+    if ConversationCommand.Code in conversation_commands and not is_none_or_empty(code_results):
+        context_message += f"{prompts.code_executed_context.format(code_results=str(code_results))}\n\n"
+    context_message = context_message.strip()
 
     # Setup Prompt with Primer or Conversation History
     messages = generate_chatml_messages_with_context(
@@ -217,4 +226,5 @@ def converse_gemini(
         api_key=api_key,
         system_prompt=system_prompt,
         completion_func=completion_func,
+        tracer=tracer,
     )

khoj/processor/conversation/google/utils.py

@@ -19,8 +19,13 @@ from tenacity import (
     wait_random_exponential,
 )
 
-from khoj.processor.conversation.utils import ThreadedGenerator, get_image_from_url
-from khoj.utils.helpers import is_none_or_empty
+from khoj.processor.conversation.utils import (
+    ThreadedGenerator,
+    commit_conversation_trace,
+    get_image_from_url,
+)
+from khoj.utils import state
+from khoj.utils.helpers import in_debug_mode, is_none_or_empty
 
 logger = logging.getLogger(__name__)
 
@@ -35,7 +40,7 @@ MAX_OUTPUT_TOKENS_GEMINI = 8192
     reraise=True,
 )
 def gemini_completion_with_backoff(
-    messages, system_prompt, model_name, temperature=0, api_key=None, model_kwargs=None
+    messages, system_prompt, model_name, temperature=0, api_key=None, model_kwargs=None, tracer={}
 ) -> str:
     genai.configure(api_key=api_key)
     model_kwargs = model_kwargs or dict()
@@ -60,16 +65,23 @@ def gemini_completion_with_backoff(
 
     try:
         # Generate the response. The last message is considered to be the current prompt
-        aggregated_response = chat_session.send_message(formatted_messages[-1]["parts"])
-        return aggregated_response.text
+        response = chat_session.send_message(formatted_messages[-1]["parts"])
+        response_text = response.text
     except StopCandidateException as e:
-        response_message, _ = handle_gemini_response(e.args)
+        response_text, _ = handle_gemini_response(e.args)
         # Respond with reason for stopping
         logger.warning(
-            f"LLM Response Prevented for {model_name}: {response_message}.\n"
+            f"LLM Response Prevented for {model_name}: {response_text}.\n"
             + f"Last Message by {messages[-1].role}: {messages[-1].content}"
         )
-        return response_message
+
+    # Save conversation trace
+    tracer["chat_model"] = model_name
+    tracer["temperature"] = temperature
+    if in_debug_mode() or state.verbose > 1:
+        commit_conversation_trace(messages, response_text, tracer)
+
+    return response_text
 
 
 @retry(
@@ -88,17 +100,20 @@ def gemini_chat_completion_with_backoff(
     system_prompt,
     completion_func=None,
     model_kwargs=None,
+    tracer: dict = {},
 ):
     g = ThreadedGenerator(compiled_references, online_results, completion_func=completion_func)
     t = Thread(
         target=gemini_llm_thread,
-        args=(g, messages, system_prompt, model_name, temperature, api_key, model_kwargs),
+        args=(g, messages, system_prompt, model_name, temperature, api_key, model_kwargs, tracer),
     )
     t.start()
     return g
 
 
-def gemini_llm_thread(g, messages, system_prompt, model_name, temperature, api_key, model_kwargs=None):
+def gemini_llm_thread(
+    g, messages, system_prompt, model_name, temperature, api_key, model_kwargs=None, tracer: dict = {}
+):
     try:
         genai.configure(api_key=api_key)
         model_kwargs = model_kwargs or dict()
@@ -117,16 +132,25 @@ def gemini_llm_thread(g, messages, system_prompt, model_name, temperature, api_k
             },
         )
 
+        aggregated_response = ""
         formatted_messages = [{"role": message.role, "parts": message.content} for message in messages]
+
         # all messages up to the last are considered to be part of the chat history
         chat_session = model.start_chat(history=formatted_messages[0:-1])
         # the last message is considered to be the current prompt
         for chunk in chat_session.send_message(formatted_messages[-1]["parts"], stream=True):
             message, stopped = handle_gemini_response(chunk.candidates, chunk.prompt_feedback)
             message = message or chunk.text
+            aggregated_response += message
             g.send(message)
             if stopped:
                 raise StopCandidateException(message)
+
+        # Save conversation trace
+        tracer["chat_model"] = model_name
+        tracer["temperature"] = temperature
+        if in_debug_mode() or state.verbose > 1:
+            commit_conversation_trace(messages, aggregated_response, tracer)
     except StopCandidateException as e:
         logger.warning(
             f"LLM Response Prevented for {model_name}: {e.args[0]}.\n"

khoj/processor/conversation/offline/chat_model.py

@@ -12,12 +12,14 @@ from khoj.processor.conversation import prompts
 from khoj.processor.conversation.offline.utils import download_model
 from khoj.processor.conversation.utils import (
     ThreadedGenerator,
+    commit_conversation_trace,
     generate_chatml_messages_with_context,
 )
 from khoj.utils import state
 from khoj.utils.constants import empty_escape_sequences
-from khoj.utils.helpers import ConversationCommand, is_none_or_empty
+from khoj.utils.helpers import ConversationCommand, in_debug_mode, is_none_or_empty
 from khoj.utils.rawconfig import LocationData
+from khoj.utils.yaml import yaml_dump
 
 logger = logging.getLogger(__name__)
 
@@ -34,6 +36,7 @@ def extract_questions_offline(
     max_prompt_size: int = None,
     temperature: float = 0.7,
     personality_context: Optional[str] = None,
+    tracer: dict = {},
 ) -> List[str]:
     """
     Infer search queries to retrieve relevant notes to answer user query
@@ -94,6 +97,7 @@ def extract_questions_offline(
             max_prompt_size=max_prompt_size,
             temperature=temperature,
             response_type="json_object",
+            tracer=tracer,
         )
     finally:
         state.chat_lock.release()
@@ -135,7 +139,8 @@ def filter_questions(questions: List[str]):
 def converse_offline(
     user_query,
     references=[],
-    online_results=[],
+    online_results={},
+    code_results={},
     conversation_log={},
     model: str = "bartowski/Meta-Llama-3.1-8B-Instruct-GGUF",
     loaded_model: Union[Any, None] = None,
@@ -146,6 +151,7 @@ def converse_offline(
     location_data: LocationData = None,
     user_name: str = None,
     agent: Agent = None,
+    tracer: dict = {},
 ) -> Union[ThreadedGenerator, Iterator[str]]:
     """
     Converse with user using Llama
@@ -153,8 +159,7 @@ def converse_offline(
     # Initialize Variables
     assert loaded_model is None or isinstance(loaded_model, Llama), "loaded_model must be of type Llama, if configured"
     offline_chat_model = loaded_model or download_model(model, max_tokens=max_prompt_size)
-    compiled_references = "\n\n".join({f"# File: {item['file']}\n## {item['compiled']}\n" for item in references})
-
+    tracer["chat_model"] = model
     current_date = datetime.now()
 
     if agent and agent.personality:
@@ -179,24 +184,25 @@ def converse_offline(
         system_prompt = f"{system_prompt}\n{user_name_prompt}"
 
     # Get Conversation Primer appropriate to Conversation Type
-    if conversation_commands == [ConversationCommand.Notes] and is_none_or_empty(compiled_references):
+    if conversation_commands == [ConversationCommand.Notes] and is_none_or_empty(references):
         return iter([prompts.no_notes_found.format()])
     elif conversation_commands == [ConversationCommand.Online] and is_none_or_empty(online_results):
         completion_func(chat_response=prompts.no_online_results_found.format())
         return iter([prompts.no_online_results_found.format()])
 
     context_message = ""
-    if not is_none_or_empty(compiled_references):
-        context_message += f"{prompts.notes_conversation_offline.format(references=compiled_references)}\n\n"
+    if not is_none_or_empty(references):
+        context_message = f"{prompts.notes_conversation_offline.format(references=yaml_dump(references))}\n\n"
     if ConversationCommand.Online in conversation_commands or ConversationCommand.Webpage in conversation_commands:
         simplified_online_results = online_results.copy()
         for result in online_results:
             if online_results[result].get("webpages"):
                 simplified_online_results[result] = online_results[result]["webpages"]
 
-        context_message += (
-            f"{prompts.online_search_conversation_offline.format(online_results=str(simplified_online_results))}"
-        )
+        context_message += f"{prompts.online_search_conversation_offline.format(online_results=yaml_dump(simplified_online_results))}\n\n"
+    if ConversationCommand.Code in conversation_commands and not is_none_or_empty(code_results):
+        context_message += f"{prompts.code_executed_context.format(code_results=str(code_results))}\n\n"
+    context_message = context_message.strip()
 
     # Setup Prompt with Primer or Conversation History
     messages = generate_chatml_messages_with_context(
@@ -215,13 +221,14 @@ def converse_offline(
     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 = Thread(target=llm_thread, args=(g, messages, offline_chat_model, max_prompt_size, tracer))
     t.start()
     return g
 
 
-def llm_thread(g, messages: List[ChatMessage], model: Any, max_prompt_size: int = None):
+def llm_thread(g, messages: List[ChatMessage], model: Any, max_prompt_size: int = None, tracer: dict = {}):
     stop_phrases = ["<s>", "INST]", "Notes:"]
+    aggregated_response = ""
 
     state.chat_lock.acquire()
     try:
@@ -229,7 +236,14 @@ def llm_thread(g, messages: List[ChatMessage], model: Any, max_prompt_size: int
             messages, loaded_model=model, stop=stop_phrases, max_prompt_size=max_prompt_size, streaming=True
         )
         for response in response_iterator:
-            g.send(response["choices"][0]["delta"].get("content", ""))
+            response_delta = response["choices"][0]["delta"].get("content", "")
+            aggregated_response += response_delta
+            g.send(response_delta)
+
+        # Save conversation trace
+        if in_debug_mode() or state.verbose > 1:
+            commit_conversation_trace(messages, aggregated_response, tracer)
+
     finally:
         state.chat_lock.release()
         g.close()
@@ -244,6 +258,7 @@ def send_message_to_model_offline(
     stop=[],
     max_prompt_size: int = None,
     response_type: str = "text",
+    tracer: dict = {},
 ):
     assert loaded_model is None or isinstance(loaded_model, Llama), "loaded_model must be of type Llama, if configured"
     offline_chat_model = loaded_model or download_model(model, max_tokens=max_prompt_size)
@@ -251,7 +266,17 @@ def send_message_to_model_offline(
     response = offline_chat_model.create_chat_completion(
         messages_dict, stop=stop, stream=streaming, temperature=temperature, response_format={"type": response_type}
     )
+
     if streaming:
         return response
-    else:
-        return response["choices"][0]["message"].get("content", "")
+
+    response_text = response["choices"][0]["message"].get("content", "")
+
+    # Save conversation trace for non-streaming responses
+    # Streamed responses need to be saved by the calling function
+    tracer["chat_model"] = model
+    tracer["temperature"] = temperature
+    if in_debug_mode() or state.verbose > 1:
+        commit_conversation_trace(messages, response_text, tracer)
+
+    return response_text

khoj/processor/conversation/openai/gpt.py

@@ -12,12 +12,13 @@ from khoj.processor.conversation.openai.utils import (
     completion_with_backoff,
 )
 from khoj.processor.conversation.utils import (
+    clean_json,
     construct_structured_message,
     generate_chatml_messages_with_context,
-    remove_json_codeblock,
 )
 from khoj.utils.helpers import ConversationCommand, is_none_or_empty
 from khoj.utils.rawconfig import LocationData
+from khoj.utils.yaml import yaml_dump
 
 logger = logging.getLogger(__name__)
 
@@ -33,6 +34,7 @@ def extract_questions(
     query_images: Optional[list[str]] = None,
     vision_enabled: bool = False,
     personality_context: Optional[str] = None,
+    tracer: dict = {},
 ):
     """
     Infer search queries to retrieve relevant notes to answer user query
@@ -82,13 +84,18 @@ def extract_questions(
     messages = [ChatMessage(content=prompt, role="user")]
 
     response = send_message_to_model(
-        messages, api_key, model, response_type="json_object", api_base_url=api_base_url, temperature=temperature
+        messages,
+        api_key,
+        model,
+        response_type="json_object",
+        api_base_url=api_base_url,
+        temperature=temperature,
+        tracer=tracer,
     )
 
     # Extract, Clean Message from GPT's Response
     try:
-        response = response.strip()
-        response = remove_json_codeblock(response)
+        response = clean_json(response)
         response = json.loads(response)
         response = [q.strip() for q in response["queries"] if q.strip()]
         if not isinstance(response, list) or not response:
@@ -103,7 +110,9 @@ def extract_questions(
     return questions
 
 
-def send_message_to_model(messages, api_key, model, response_type="text", api_base_url=None, temperature=0):
+def send_message_to_model(
+    messages, api_key, model, response_type="text", api_base_url=None, temperature=0, tracer: dict = {}
+):
     """
     Send message to model
     """
@@ -116,6 +125,7 @@ def send_message_to_model(messages, api_key, model, response_type="text", api_ba
         temperature=temperature,
         api_base_url=api_base_url,
         model_kwargs={"response_format": {"type": response_type}},
+        tracer=tracer,
     )
 
 
@@ -123,6 +133,7 @@ def converse(
     references,
     user_query,
     online_results: Optional[Dict[str, Dict]] = None,
+    code_results: Optional[Dict[str, Dict]] = None,
     conversation_log={},
     model: str = "gpt-4o-mini",
     api_key: Optional[str] = None,
@@ -137,13 +148,13 @@ def converse(
     agent: Agent = None,
     query_images: Optional[list[str]] = None,
     vision_available: bool = False,
+    tracer: dict = {},
 ):
     """
     Converse with user using OpenAI's ChatGPT
     """
     # Initialize Variables
     current_date = datetime.now()
-    compiled_references = "\n\n".join({f"# File: {item['file']}\n## {item['compiled']}\n" for item in references})
 
     if agent and agent.personality:
         system_prompt = prompts.custom_personality.format(
@@ -167,7 +178,7 @@ def converse(
         system_prompt = f"{system_prompt}\n{user_name_prompt}"
 
     # Get Conversation Primer appropriate to Conversation Type
-    if conversation_commands == [ConversationCommand.Notes] and is_none_or_empty(compiled_references):
+    if conversation_commands == [ConversationCommand.Notes] and is_none_or_empty(references):
         completion_func(chat_response=prompts.no_notes_found.format())
         return iter([prompts.no_notes_found.format()])
     elif conversation_commands == [ConversationCommand.Online] and is_none_or_empty(online_results):
@@ -175,10 +186,13 @@ def converse(
         return iter([prompts.no_online_results_found.format()])
 
     context_message = ""
-    if not is_none_or_empty(compiled_references):
-        context_message = f"{prompts.notes_conversation.format(references=compiled_references)}\n\n"
+    if not is_none_or_empty(references):
+        context_message = f"{prompts.notes_conversation.format(references=yaml_dump(references))}\n\n"
     if not is_none_or_empty(online_results):
-        context_message += f"{prompts.online_search_conversation.format(online_results=str(online_results))}"
+        context_message += f"{prompts.online_search_conversation.format(online_results=yaml_dump(online_results))}\n\n"
+    if not is_none_or_empty(code_results):
+        context_message += f"{prompts.code_executed_context.format(code_results=str(code_results))}\n\n"
+    context_message = context_message.strip()
 
     # Setup Prompt with Primer or Conversation History
     messages = generate_chatml_messages_with_context(
@@ -207,4 +221,5 @@ def converse(
         api_base_url=api_base_url,
         completion_func=completion_func,
         model_kwargs={"stop": ["Notes:\n["]},
+        tracer=tracer,
     )

khoj/processor/conversation/openai/utils.py

@@ -12,7 +12,12 @@ from tenacity import (
     wait_random_exponential,
 )
 
-from khoj.processor.conversation.utils import ThreadedGenerator
+from khoj.processor.conversation.utils import (
+    ThreadedGenerator,
+    commit_conversation_trace,
+)
+from khoj.utils import state
+from khoj.utils.helpers import in_debug_mode
 
 logger = logging.getLogger(__name__)
 
@@ -33,7 +38,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
+    messages, model, temperature=0, openai_api_key=None, api_base_url=None, model_kwargs=None, tracer: dict = {}
 ) -> str:
     client_key = f"{openai_api_key}--{api_base_url}"
     client: openai.OpenAI | None = openai_clients.get(client_key)
@@ -77,6 +82,12 @@ def completion_with_backoff(
         elif delta_chunk.content:
             aggregated_response += delta_chunk.content
 
+    # Save conversation trace
+    tracer["chat_model"] = model
+    tracer["temperature"] = temperature
+    if in_debug_mode() or state.verbose > 1:
+        commit_conversation_trace(messages, aggregated_response, tracer)
+
     return aggregated_response
 
 
@@ -103,26 +114,37 @@ def chat_completion_with_backoff(
     api_base_url=None,
     completion_func=None,
     model_kwargs=None,
+    tracer: dict = {},
 ):
     g = ThreadedGenerator(compiled_references, online_results, completion_func=completion_func)
     t = Thread(
-        target=llm_thread, args=(g, messages, model_name, temperature, openai_api_key, api_base_url, model_kwargs)
+        target=llm_thread,
+        args=(g, messages, model_name, temperature, openai_api_key, api_base_url, model_kwargs, tracer),
     )
     t.start()
     return g
 
 
-def llm_thread(g, messages, model_name, temperature, openai_api_key=None, api_base_url=None, model_kwargs=None):
+def llm_thread(
+    g,
+    messages,
+    model_name,
+    temperature,
+    openai_api_key=None,
+    api_base_url=None,
+    model_kwargs=None,
+    tracer: dict = {},
+):
     try:
         client_key = f"{openai_api_key}--{api_base_url}"
         if client_key not in openai_clients:
-            client: openai.OpenAI = openai.OpenAI(
+            client = openai.OpenAI(
                 api_key=openai_api_key,
                 base_url=api_base_url,
             )
             openai_clients[client_key] = client
         else:
-            client: openai.OpenAI = openai_clients[client_key]
+            client = openai_clients[client_key]
 
         formatted_messages = [{"role": message.role, "content": message.content} for message in messages]
         stream = True
@@ -144,17 +166,29 @@ def llm_thread(g, messages, model_name, temperature, openai_api_key=None, api_ba
             **(model_kwargs or dict()),
         )
 
+        aggregated_response = ""
         if not stream:
-            g.send(chat.choices[0].message.content)
+            aggregated_response = chat.choices[0].message.content
+            g.send(aggregated_response)
         else:
             for chunk in chat:
                 if len(chunk.choices) == 0:
                     continue
                 delta_chunk = chunk.choices[0].delta
+                text_chunk = ""
                 if isinstance(delta_chunk, str):
-                    g.send(delta_chunk)
+                    text_chunk = delta_chunk
                 elif delta_chunk.content:
-                    g.send(delta_chunk.content)
+                    text_chunk = delta_chunk.content
+                if text_chunk:
+                    aggregated_response += text_chunk
+                    g.send(text_chunk)
+
+        # Save conversation trace
+        tracer["chat_model"] = model_name
+        tracer["temperature"] = temperature
+        if in_debug_mode() or state.verbose > 1:
+            commit_conversation_trace(messages, aggregated_response, tracer)
     except Exception as e:
         logger.error(f"Error in llm_thread: {e}", exc_info=True)
     finally: