khoj 1.27.2.dev12__py3-none-any.whl → 1.28.1__py3-none-any.whl

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

@@ -1,4 +1,5 @@
 import logging
+import os
 from threading import Thread
 from typing import Dict
 
@@ -12,7 +13,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 +39,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)
@@ -55,6 +61,9 @@ def completion_with_backoff(
         model_kwargs.pop("stop", None)
         model_kwargs.pop("response_format", None)
 
+    if os.getenv("KHOJ_LLM_SEED"):
+        model_kwargs["seed"] = int(os.getenv("KHOJ_LLM_SEED"))
+
     chat = client.chat.completions.create(
         stream=stream,
         messages=formatted_messages,  # type: ignore
@@ -77,6 +86,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 +118,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
@@ -135,6 +161,9 @@ def llm_thread(g, messages, model_name, temperature, openai_api_key=None, api_ba
             model_kwargs.pop("stop", None)
             model_kwargs.pop("response_format", None)
 
+        if os.getenv("KHOJ_LLM_SEED"):
+            model_kwargs["seed"] = int(os.getenv("KHOJ_LLM_SEED"))
+
         chat = client.chat.completions.create(
             stream=stream,
             messages=formatted_messages,
@@ -144,17 +173,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:

khoj/processor/conversation/prompts.py

@@ -193,7 +193,6 @@ you need to convert the user's query to a description format that the novice art
 - ellipse
 - line
 - arrow
-- frame
 
 use these primitives to describe what sort of diagram the drawer should create. the artist must recreate the diagram every time, so include all relevant prior information in your description.
 
@@ -284,21 +283,6 @@ For text, you must use the `text` property to specify the text to be rendered. Y
     text: string,
 }}
 
-For frames, use the `children` property to specify the elements that are inside the frame by their ids.
-
-{{
-    type: "frame",
-    id: string,
-    x: number,
-    y: number,
-    width: number,
-    height: number,
-    name: string,
-    children: [
-        string
-    ]
-}}
-
 Here's an example of a valid diagram:
 
 Design Description: Create a diagram describing a circular development process with 3 stages: design, implementation and feedback. The design stage is connected to the implementation stage and the implementation stage is connected to the feedback stage and the feedback stage is connected to the design stage. Each stage should be labeled with the stage name.
@@ -410,21 +394,23 @@ Q: {query}
 
 extract_questions = PromptTemplate.from_template(
     """
-You are Khoj, an extremely smart and helpful document search assistant with only the ability to retrieve information from the user's notes. Disregard online search requests.
+You are Khoj, an extremely smart and helpful document search assistant with only the ability to retrieve information from the user's notes and documents.
 Construct search queries to retrieve relevant information to answer the user's question.
-- You will be provided past questions(Q) and answers(A) for context.
+- You will be provided example and actual past user questions(Q), search queries(Khoj) and answers(A) for context.
 - Add as much context from the previous questions and answers as required into your search queries.
-- Break messages into multiple search queries when required to retrieve the relevant information.
+- Break your search down into multiple search queries from a diverse set of lenses to retrieve all related documents.
 - Add date filters to your search queries from questions and answers when required to retrieve the relevant information.
 - When asked a meta, vague or random questions, search for a variety of broad topics to answer the user's question.
 {personality_context}
-What searches will you perform to answer the users question? Respond with search queries as list of strings in a JSON object.
+What searches will you perform to answer the user's question? Respond with search queries as list of strings in a JSON object.
 Current Date: {day_of_week}, {current_date}
 User's Location: {location}
 {username}
 
+Examples
+---
 Q: How was my trip to Cambodia?
-Khoj: {{"queries": ["How was my trip to Cambodia?"]}}
+Khoj: {{"queries": ["How was my trip to Cambodia?", "Angkor Wat temple visit", "Flight to Phnom Penh", "Expenses in Cambodia", "Stay in Cambodia"]}}
 A: The trip was amazing. You went to the Angkor Wat temple and it was beautiful.
 
 Q: Who did i visit that temple with?
@@ -459,6 +445,8 @@ Q: Who all did I meet here yesterday?
 Khoj: {{"queries": ["Met in {location} on {yesterday_date} dt>='{yesterday_date}' dt<'{current_date}'"]}}
 A: Yesterday's note mentions your visit to your local beach with Ram and Shyam.
 
+Actual
+---
 {chat_history}
 Q: {text}
 Khoj:
@@ -467,11 +455,11 @@ Khoj:
 
 extract_questions_anthropic_system_prompt = PromptTemplate.from_template(
     """
-You are Khoj, an extremely smart and helpful document search assistant with only the ability to retrieve information from the user's notes. Disregard online search requests.
+You are Khoj, an extremely smart and helpful document search assistant with only the ability to retrieve information from the user's notes.
 Construct search queries to retrieve relevant information to answer the user's question.
-- You will be provided past questions(User), extracted queries(Assistant) and answers(A) for context.
+- You will be provided past questions(User), search queries(Assistant) and answers(A) for context.
 - Add as much context from the previous questions and answers as required into your search queries.
-- Break messages into multiple search queries when required to retrieve the relevant information.
+- Break your search down into multiple search queries from a diverse set of lenses to retrieve all related documents.
 - Add date filters to your search queries from questions and answers when required to retrieve the relevant information.
 - When asked a meta, vague or random questions, search for a variety of broad topics to answer the user's question.
 {personality_context}
@@ -484,7 +472,7 @@ User's Location: {location}
 Here are some examples of how you can construct search queries to answer the user's question:
 
 User: How was my trip to Cambodia?
-Assistant: {{"queries": ["How was my trip to Cambodia?"]}}
+Assistant: {{"queries": ["How was my trip to Cambodia?", "Angkor Wat temple visit", "Flight to Phnom Penh", "Expenses in Cambodia", "Stay in Cambodia"]}}
 A: The trip was amazing. You went to the Angkor Wat temple and it was beautiful.
 
 User: What national parks did I go to last year?
@@ -517,17 +505,14 @@ Assistant:
 )
 
 system_prompt_extract_relevant_information = """
-As a professional analyst, create a comprehensive report of the most relevant information from a web page in response to a user's query.
-The text provided is directly from within the web page.
-The report you create should be multiple paragraphs, and it should represent the content of the website.
-Tell the user exactly what the website says in response to their query, while adhering to these guidelines:
-
-1. Answer the user's query as specifically as possible. Include many supporting details from the website.
-2. Craft a report that is detailed, thorough, in-depth, and complex, while maintaining clarity.
-3. Rely strictly on the provided text, without including external information.
-4. Format the report in multiple paragraphs with a clear structure.
-5. Be as specific as possible in your answer to the user's query.
-6. Reproduce as much of the provided text as possible, while maintaining readability.
+As a professional analyst, your job is to extract all pertinent information from documents to help answer user's query.
+You will be provided raw text directly from within the document.
+Adhere to these guidelines while extracting information from the provided documents:
+
+1. Extract all relevant text and links from the document that can assist with further research or answer the user's query.
+2. Craft a comprehensive but compact report with all the necessary data from the document to generate an informed response.
+3. Rely strictly on the provided text to generate your summary, without including external information.
+4. Provide specific, important snippets from the document in your report to establish trust in your summary.
 """.strip()
 
 extract_relevant_information = PromptTemplate.from_template(
@@ -535,10 +520,10 @@ extract_relevant_information = PromptTemplate.from_template(
 {personality_context}
 Target Query: {query}
 
-Web Pages:
+Document:
 {corpus}
 
-Collate only relevant information from the website to answer the target query.
+Collate only relevant information from the document to answer the target query.
 """.strip()
 )
 
@@ -633,6 +618,67 @@ Khoj:
 """.strip()
 )
 
+plan_function_execution = PromptTemplate.from_template(
+    """
+You are Khoj, a smart, creative and methodical researcher. Use the provided tool AIs to investigate information to answer query.
+Create a multi-step plan and intelligently iterate on the plan based on the retrieved information to find the requested information.
+{personality_context}
+
+# Instructions
+- Ask highly diverse, detailed queries to the tool AIs, one tool AI at a time, to discover required information or run calculations. Their response will be shown to you in the next iteration.
+- Break down your research process into independent, self-contained steps that can be executed sequentially using the available tool AIs to answer the user's query. Write your step-by-step plan in the scratchpad.
+- Always ask a new query that was not asked to the tool AI in a previous iteration. Build on the results of the previous iterations.
+- Ensure that all required context is passed to the tool AIs for successful execution. They only know the context provided in your query.
+- Think step by step to come up with creative strategies when the previous iteration did not yield useful results.
+- You are allowed upto {max_iterations} iterations to use the help of the provided tool AIs to answer the user's question.
+- Stop when you have the required information by returning a JSON object with an empty "tool" field. E.g., {{scratchpad: "I have all I need", tool: "", query: ""}}
+
+# Examples
+Assuming you can search the user's notes and the internet.
+- When the user asks for the population of their hometown
+  1. Try look up their hometown in their notes. Ask the note search AI to search for their birth certificate, childhood memories, school, resume etc.
+  2. If not found in their notes, try infer their hometown from their online social media profiles. Ask the online search AI to look for {username}'s biography, school, resume on linkedin, facebook, website etc.
+  3. Only then try find the latest population of their hometown by reading official websites with the help of the online search and web page reading AI.
+- When the user asks for their computer's specs
+  1. Try find their computer model in their notes.
+  2. Now find webpages with their computer model's spec online.
+  3. Ask the the webpage tool AI to extract the required information from the relevant webpages.
+- When the user asks what clothes to carry for their upcoming trip
+  1. Find the itinerary of their upcoming trip in their notes.
+  2. Next find the weather forecast at the destination online.
+  3. Then find if they mentioned what clothes they own in their notes.
+
+# Background Context
+- Current Date: {day_of_week}, {current_date}
+- User Location: {location}
+- User Name: {username}
+
+# Available Tool AIs
+Which of the tool AIs listed below would you use to answer the user's question? You **only** have access to the following tool AIs:
+
+{tools}
+
+# Previous Iterations
+{previous_iterations}
+
+# Chat History:
+{chat_history}
+
+Return the next tool AI to use and the query to ask it. Your response should always be a valid JSON object. Do not say anything else.
+Response format:
+{{"scratchpad": "<your_scratchpad_to_reason_about_which_tool_to_use>", "query": "<your_detailed_query_for_the_tool_ai>", "tool": "<name_of_tool_ai>"}}
+""".strip()
+)
+
+previous_iteration = PromptTemplate.from_template(
+    """
+## Iteration {index}:
+- tool: {tool}
+- query: {query}
+- result: {result}
+"""
+)
+
 pick_relevant_information_collection_tools = PromptTemplate.from_template(
     """
 You are Khoj, an extremely smart and helpful search assistant.
@@ -752,8 +798,8 @@ Khoj:
 online_search_conversation_subqueries = PromptTemplate.from_template(
     """
 You are Khoj, an advanced web search assistant. You are tasked with constructing **up to three** google search queries to answer the user's question.
-- You will receive the conversation history as context.
-- Add as much context from the previous questions and answers as required into your search queries.
+- You will receive the actual chat history as context.
+- Add as much context from the chat history as required into your search queries.
 - Break messages into multiple search queries when required to retrieve the relevant information.
 - Use site: google search operator when appropriate
 - You have access to the the whole internet to retrieve information.
@@ -766,62 +812,107 @@ User's Location: {location}
 {username}
 
 Here are some examples:
-History:
+Example Chat History:
 User: I like to use Hacker News to get my tech news.
+Khoj: {{queries: ["what is Hacker News?", "Hacker News website for tech news"]}}
 AI: Hacker News is an online forum for sharing and discussing the latest tech news. It is a great place to learn about new technologies and startups.
 
-Q: Summarize the top posts on HackerNews
+User: Summarize the top posts on HackerNews
 Khoj: {{"queries": ["top posts on HackerNews"]}}
 
-History:
-
-Q: Tell me the latest news about the farmers protest in Colombia and China on Reuters
+Example Chat History:
+User: Tell me the latest news about the farmers protest in Colombia and China on Reuters
 Khoj: {{"queries": ["site:reuters.com farmers protest Colombia", "site:reuters.com farmers protest China"]}}
 
-History:
+Example Chat History:
 User: I'm currently living in New York but I'm thinking about moving to San Francisco.
+Khoj: {{"queries": ["New York city vs San Francisco life", "San Francisco living cost", "New York city living cost"]}}
 AI: New York is a great city to live in. It has a lot of great restaurants and museums. San Francisco is also a great city to live in. It has good access to nature and a great tech scene.
 
-Q: What is the climate like in those cities?
-Khoj: {{"queries": ["climate in new york city", "climate in san francisco"]}}
+User: What is the climate like in those cities?
+Khoj: {{"queries": ["climate in New York city", "climate in San Francisco"]}}
 
-History:
-AI: Hey, how is it going?
-User: Going well. Ananya is in town tonight!
+Example Chat History:
+User: Hey, Ananya is in town tonight!
+Khoj: {{"queries": ["events in {location} tonight", "best restaurants in {location}", "places to visit in {location}"]}}
 AI: Oh that's awesome! What are your plans for the evening?
 
-Q: She wants to see a movie. Any decent sci-fi movies playing at the local theater?
+User: She wants to see a movie. Any decent sci-fi movies playing at the local theater?
 Khoj: {{"queries": ["new sci-fi movies in theaters near {location}"]}}
 
-History:
+Example Chat History:
 User: Can I chat with you over WhatsApp?
-AI: Yes, you can chat with me using WhatsApp.
-
-Q: How
 Khoj: {{"queries": ["site:khoj.dev chat with Khoj on Whatsapp"]}}
+AI: Yes, you can chat with me using WhatsApp.
 
-History:
-
-
-Q: How do I share my files with you?
+Example Chat History:
+User: How do I share my files with Khoj?
 Khoj: {{"queries": ["site:khoj.dev sync files with Khoj"]}}
 
-History:
+Example Chat History:
 User: I need to transport a lot of oranges to the moon. Are there any rockets that can fit a lot of oranges?
+Khoj: {{"queries": ["current rockets with large cargo capacity", "rocket rideshare cost by cargo capacity"]}}
 AI: NASA's Saturn V rocket frequently makes lunar trips and has a large cargo capacity.
 
-Q: How many oranges would fit in NASA's Saturn V rocket?
-Khoj: {{"queries": ["volume of an orange", "volume of saturn v rocket"]}}
+User: How many oranges would fit in NASA's Saturn V rocket?
+Khoj: {{"queries": ["volume of an orange", "volume of Saturn V rocket"]}}
 
 Now it's your turn to construct Google search queries to answer the user's question. Provide them as a list of strings in a JSON object. Do not say anything else.
-History:
+Actual Chat History:
 {chat_history}
 
-Q: {query}
+User: {query}
 Khoj:
 """.strip()
 )
 
+# Code Generation
+# --
+python_code_generation_prompt = PromptTemplate.from_template(
+    """
+You are Khoj, an advanced python programmer. You are tasked with constructing **up to three** python programs to best answer the user query.
+- The python program will run in a pyodide python sandbox with no network access.
+- You can write programs to run complex calculations, analyze data, create charts, generate documents to meticulously answer the query
+- The sandbox has access to the standard library, matplotlib, panda, numpy, scipy, bs4, sympy, brotli, cryptography, fast-parquet
+- Do not try display images or plots in the code directly. The code should save the image or plot to a file instead.
+- Write any document, charts etc. to be shared with the user to file. These files can be seen by the user.
+- Use as much context from the previous questions and answers as required to generate your code.
+{personality_context}
+What code will you need to write, if any, to answer the user's question?
+Provide code programs as a list of strings in a JSON object with key "codes".
+Current Date: {current_date}
+User's Location: {location}
+{username}
+
+The JSON schema is of the form {{"codes": ["code1", "code2", "code3"]}}
+For example:
+{{"codes": ["print('Hello, World!')", "print('Goodbye, World!')"]}}
+
+Now it's your turn to construct python programs to answer the user's question. Provide them as a list of strings in a JSON object. Do not say anything else.
+Context:
+---
+{context}
+
+Chat History:
+---
+{chat_history}
+
+User: {query}
+Khoj:
+""".strip()
+)
+
+code_executed_context = PromptTemplate.from_template(
+    """
+Use the provided code executions to inform your response.
+Ask crisp follow-up questions to get additional context, when a helpful response cannot be provided from the provided code execution results or past conversations.
+
+Code Execution Results:
+{code_results}
+""".strip()
+)
+
+
 # Automations
 # --
 crontime_prompt = PromptTemplate.from_template(