khoj 1.36.7.dev7__py3-none-any.whl → 1.36.7.dev22__py3-none-any.whl

khoj/processor/conversation/google/gemini_chat.py

@@ -31,10 +31,10 @@ logger = logging.getLogger(__name__)
 
 def extract_questions_gemini(
     text,
-    model: Optional[str] = "gemini-1.5-flash",
+    model: Optional[str] = "gemini-2.0-flash",
     conversation_log={},
     api_key=None,
-    temperature=0,
+    temperature=0.6,
     max_tokens=None,
     location_data: LocationData = None,
     user: KhojUser = None,
@@ -121,24 +121,24 @@ def gemini_send_message_to_model(
     api_key,
     model,
     response_type="text",
-    temperature=0,
+    temperature=0.6,
     model_kwargs=None,
     tracer={},
 ):
     """
     Send message to model
     """
-    messages, system_prompt = format_messages_for_gemini(messages)
+    messages_for_gemini, system_prompt = format_messages_for_gemini(messages)
 
     model_kwargs = {}
 
-    # 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"
+    # This caused unwanted behavior and terminates response early for gemini 1.5 series. Monitor for flakiness with 2.0 series.
+    if response_type == "json_object" and model in ["gemini-2.0-flash"]:
+        model_kwargs["response_mime_type"] = "application/json"
 
     # Get Response from Gemini
     return gemini_completion_with_backoff(
-        messages=messages,
+        messages=messages_for_gemini,
         system_prompt=system_prompt,
         model_name=model,
         api_key=api_key,
@@ -154,9 +154,9 @@ def converse_gemini(
     online_results: Optional[Dict[str, Dict]] = None,
     code_results: Optional[Dict[str, Dict]] = None,
     conversation_log={},
-    model: Optional[str] = "gemini-1.5-flash",
+    model: Optional[str] = "gemini-2.0-flash",
     api_key: Optional[str] = None,
-    temperature: float = 0.2,
+    temperature: float = 0.6,
     completion_func=None,
     conversation_commands=[ConversationCommand.Default],
     max_prompt_size=None,
@@ -236,12 +236,12 @@ def converse_gemini(
         program_execution_context=program_execution_context,
     )
 
-    messages, system_prompt = format_messages_for_gemini(messages, system_prompt)
+    messages_for_gemini, system_prompt = format_messages_for_gemini(messages, system_prompt)
     logger.debug(f"Conversation Context for Gemini: {messages_to_print(messages)}")
 
     # Get Response from Google AI
     return gemini_chat_completion_with_backoff(
-        messages=messages,
+        messages=messages_for_gemini,
         compiled_references=references,
         online_results=online_results,
         model_name=model,

khoj/processor/conversation/google/utils.py

@@ -1,15 +1,11 @@
 import logging
 import random
+from copy import deepcopy
 from threading import Thread
 
-import google.generativeai as genai
-from google.generativeai.types.answer_types import FinishReason
+from google import genai
+from google.genai import types as gtypes
 from google.generativeai.types.generation_types import StopCandidateException
-from google.generativeai.types.safety_types import (
-    HarmBlockThreshold,
-    HarmCategory,
-    HarmProbability,
-)
 from langchain.schema import ChatMessage
 from tenacity import (
     before_sleep_log,
@@ -24,7 +20,6 @@ from khoj.processor.conversation.utils import (
     commit_conversation_trace,
     get_image_from_url,
 )
-from khoj.utils import state
 from khoj.utils.helpers import (
     get_chat_usage_metrics,
     is_none_or_empty,
@@ -35,6 +30,24 @@ logger = logging.getLogger(__name__)
 
 
 MAX_OUTPUT_TOKENS_GEMINI = 8192
+SAFETY_SETTINGS = [
+    gtypes.SafetySetting(
+        category=gtypes.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
+        threshold=gtypes.HarmBlockThreshold.BLOCK_ONLY_HIGH,
+    ),
+    gtypes.SafetySetting(
+        category=gtypes.HarmCategory.HARM_CATEGORY_HARASSMENT,
+        threshold=gtypes.HarmBlockThreshold.BLOCK_ONLY_HIGH,
+    ),
+    gtypes.SafetySetting(
+        category=gtypes.HarmCategory.HARM_CATEGORY_HATE_SPEECH,
+        threshold=gtypes.HarmBlockThreshold.BLOCK_ONLY_HIGH,
+    ),
+    gtypes.SafetySetting(
+        category=gtypes.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,
+        threshold=gtypes.HarmBlockThreshold.BLOCK_ONLY_HIGH,
+    ),
+]
 
 
 @retry(
@@ -46,30 +59,19 @@ MAX_OUTPUT_TOKENS_GEMINI = 8192
 def gemini_completion_with_backoff(
     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()
-    model_kwargs["temperature"] = temperature
-    model_kwargs["max_output_tokens"] = MAX_OUTPUT_TOKENS_GEMINI
-    model = genai.GenerativeModel(
-        model_name,
-        generation_config=model_kwargs,
+    client = genai.Client(api_key=api_key)
+    config = gtypes.GenerateContentConfig(
         system_instruction=system_prompt,
-        safety_settings={
-            HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_ONLY_HIGH,
-            HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_ONLY_HIGH,
-            HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_ONLY_HIGH,
-            HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_ONLY_HIGH,
-        },
+        temperature=temperature,
+        max_output_tokens=MAX_OUTPUT_TOKENS_GEMINI,
+        safety_settings=SAFETY_SETTINGS,
     )
 
-    formatted_messages = [{"role": message.role, "parts": message.content} for message in messages]
-
-    # Start chat session. 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])
+    formatted_messages = [gtypes.Content(role=message.role, parts=message.content) for message in messages]
 
     try:
-        # Generate the response. The last message is considered to be the current prompt
-        response = chat_session.send_message(formatted_messages[-1]["parts"])
+        # Generate the response
+        response = client.models.generate_content(model=model_name, config=config, contents=formatted_messages)
         response_text = response.text
     except StopCandidateException as e:
         response = None
@@ -125,30 +127,21 @@ 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()
-        model_kwargs["temperature"] = temperature
-        model_kwargs["max_output_tokens"] = MAX_OUTPUT_TOKENS_GEMINI
-        model_kwargs["stop_sequences"] = ["Notes:\n["]
-        model = genai.GenerativeModel(
-            model_name,
-            generation_config=model_kwargs,
+        client = genai.Client(api_key=api_key)
+        config = gtypes.GenerateContentConfig(
             system_instruction=system_prompt,
-            safety_settings={
-                HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_ONLY_HIGH,
-                HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_ONLY_HIGH,
-                HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_ONLY_HIGH,
-                HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_ONLY_HIGH,
-            },
+            temperature=temperature,
+            max_output_tokens=MAX_OUTPUT_TOKENS_GEMINI,
+            stop_sequences=["Notes:\n["],
+            safety_settings=SAFETY_SETTINGS,
         )
 
         aggregated_response = ""
-        formatted_messages = [{"role": message.role, "parts": message.content} for message in messages]
+        formatted_messages = [gtypes.Content(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):
+        for chunk in client.models.generate_content_stream(
+            model=model_name, config=config, contents=formatted_messages
+        ):
             message, stopped = handle_gemini_response(chunk.candidates, chunk.prompt_feedback)
             message = message or chunk.text
             aggregated_response += message
@@ -177,14 +170,16 @@ def gemini_llm_thread(
         g.close()
 
 
-def handle_gemini_response(candidates, prompt_feedback=None):
+def handle_gemini_response(
+    candidates: list[gtypes.Candidate], prompt_feedback: gtypes.GenerateContentResponsePromptFeedback = None
+):
     """Check if Gemini response was blocked and return an explanatory error message."""
     # Check if the response was blocked due to safety concerns with the prompt
     if len(candidates) == 0 and prompt_feedback:
         message = f"\nI'd prefer to not respond to that due to **{prompt_feedback.block_reason.name}** issues with your query."
         stopped = True
     # Check if the response was blocked due to safety concerns with the generated content
-    elif candidates[0].finish_reason == FinishReason.SAFETY:
+    elif candidates[0].finish_reason == gtypes.FinishReason.SAFETY:
         message = generate_safety_response(candidates[0].safety_ratings)
         stopped = True
     # Check if finish reason is empty, therefore generation is in progress
@@ -192,7 +187,7 @@ def handle_gemini_response(candidates, prompt_feedback=None):
         message = None
         stopped = False
     # Check if the response was stopped due to reaching maximum token limit or other reasons
-    elif candidates[0].finish_reason != FinishReason.STOP:
+    elif candidates[0].finish_reason != gtypes.FinishReason.STOP:
         message = f"\nI can't talk further about that because of **{candidates[0].finish_reason.name} issue.**"
         stopped = True
     # Otherwise, the response is valid and can be used
@@ -202,18 +197,18 @@ def handle_gemini_response(candidates, prompt_feedback=None):
     return message, stopped
 
 
-def generate_safety_response(safety_ratings):
+def generate_safety_response(safety_ratings: list[gtypes.SafetyRating]):
     """Generate a conversational response based on the safety ratings of the response."""
     # Get the safety rating with the highest probability
-    max_safety_rating = sorted(safety_ratings, key=lambda x: x.probability, reverse=True)[0]
+    max_safety_rating: gtypes.SafetyRating = sorted(safety_ratings, key=lambda x: x.probability, reverse=True)[0]
     # Remove the "HARM_CATEGORY_" prefix and title case the category name
     max_safety_category = " ".join(max_safety_rating.category.name.split("_")[2:]).title()
     # Add a bit of variety to the discomfort level based on the safety rating probability
     discomfort_level = {
-        HarmProbability.HARM_PROBABILITY_UNSPECIFIED: " ",
-        HarmProbability.LOW: "a bit ",
-        HarmProbability.MEDIUM: "moderately ",
-        HarmProbability.HIGH: random.choice(["very ", "quite ", "fairly "]),
+        gtypes.HarmProbability.HARM_PROBABILITY_UNSPECIFIED: " ",
+        gtypes.HarmProbability.LOW: "a bit ",
+        gtypes.HarmProbability.MEDIUM: "moderately ",
+        gtypes.HarmProbability.HIGH: random.choice(["very ", "quite ", "fairly "]),
     }[max_safety_rating.probability]
     # Generate a response using a random response template
     safety_response_choice = random.choice(
@@ -229,9 +224,12 @@ def generate_safety_response(safety_ratings):
     )
 
 
-def format_messages_for_gemini(messages: list[ChatMessage], system_prompt: str = None) -> tuple[list[str], str]:
+def format_messages_for_gemini(
+    original_messages: list[ChatMessage], system_prompt: str = None
+) -> tuple[list[str], str]:
     # Extract system message
     system_prompt = system_prompt or ""
+    messages = deepcopy(original_messages)
     for message in messages.copy():
         if message.role == "system":
             system_prompt += message.content
@@ -242,14 +240,16 @@ def format_messages_for_gemini(messages: list[ChatMessage], system_prompt: str =
         # Convert message content to string list from chatml dictionary list
         if isinstance(message.content, list):
             # Convert image_urls to PIL.Image and place them at beginning of list (better for Gemini)
-            message.content = [
-                get_image_from_url(item["image_url"]["url"]).content
-                if item["type"] == "image_url"
-                else item.get("text", "")
-                for item in sorted(message.content, key=lambda x: 0 if x["type"] == "image_url" else 1)
-            ]
+            message_content = []
+            for item in sorted(message.content, key=lambda x: 0 if x["type"] == "image_url" else 1):
+                if item["type"] == "image_url":
+                    image = get_image_from_url(item["image_url"]["url"], type="bytes")
+                    message_content += [gtypes.Part.from_bytes(data=image.content, mime_type=image.type)]
+                else:
+                    message_content += [gtypes.Part.from_text(text=item.get("text", ""))]
+            message.content = message_content
         elif isinstance(message.content, str):
-            message.content = [message.content]
+            message.content = [gtypes.Part.from_text(text=message.content)]
 
         if message.role == "assistant":
             message.role = "model"

khoj/processor/conversation/prompts.py

@@ -974,11 +974,9 @@ Khoj:
 python_code_generation_prompt = PromptTemplate.from_template(
     """
 You are Khoj, an advanced python programmer. You are tasked with constructing a python program to best answer the user query.
-- The python program will run in a pyodide python sandbox with no network access.
+- The python program will run in a 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 and sympy packages. The requests, torch, catboost, tensorflow and tkinter packages are not available.
-- List known file paths to required user documents in "input_files" and known links to required documents from the web in the "input_links" field.
-- The python program should be self-contained. It can only read data generated by the program itself and from provided input_files, input_links by their basename (i.e filename excluding file path).
+- The python program should be self-contained. It can only read data generated by the program itself and any user file paths referenced in your program.
 - 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.
@@ -989,24 +987,99 @@ Current Date: {current_date}
 User's Location: {location}
 {username}
 
-The response JSON schema is of the form {{"code": "<python_code>", "input_files": ["file_path_1", "file_path_2"], "input_links": ["link_1", "link_2"]}}
-Examples:
+Your response should contain python code wrapped in markdown code blocks (i.e starting with```python and ending with ```)
+Example 1:
 ---
-{{
-"code": "# Input values\\nprincipal = 43235\\nrate = 5.24\\nyears = 5\\n\\n# Convert rate to decimal\\nrate_decimal = rate / 100\\n\\n# Calculate final amount\\nfinal_amount = principal * (1 + rate_decimal) ** years\\n\\n# Calculate interest earned\\ninterest_earned = final_amount - principal\\n\\n# Print results with formatting\\nprint(f"Interest Earned: ${{interest_earned:,.2f}}")\\nprint(f"Final Amount: ${{final_amount:,.2f}}")"
-}}
+Q: Calculate the interest earned and final amount for a principal of $43,235 invested at a rate of 5.24 percent for 5 years.
+A: Ok, to calculate the interest earned and final amount, we can use the formula for compound interest: $T = P(1 + r/n)^{{nt}}$,
+where T: total amount, P: principal, r: interest rate, n: number of times interest is compounded per year, and t: time in years.
 
-{{
-"code": "import re\\n\\n# Read org file\\nfile_path = 'tasks.org'\\nwith open(file_path, 'r') as f:\\n    content = f.read()\\n\\n# Get today's date in YYYY-MM-DD format\\ntoday = datetime.now().strftime('%Y-%m-%d')\\npattern = r'\*+\s+.*\\n.*SCHEDULED:\s+<' + today + r'.*>'\\n\\n# Find all matches using multiline mode\\nmatches = re.findall(pattern, content, re.MULTILINE)\\ncount = len(matches)\\n\\n# Display count\\nprint(f'Count of scheduled tasks for today: {{count}}')",
-"input_files": ["/home/linux/tasks.org"]
-}}
+Let's write the Python program to calculate this.
 
-{{
-"code": "import pandas as pd\\nimport matplotlib.pyplot as plt\\n\\n# Load the CSV file\\ndf = pd.read_csv('world_population_by_year.csv')\\n\\n# Plot the data\\nplt.figure(figsize=(10, 6))\\nplt.plot(df['Year'], df['Population'], marker='o')\\n\\n# Add titles and labels\\nplt.title('Population by Year')\\nplt.xlabel('Year')\\nplt.ylabel('Population')\\n\\n# Save the plot to a file\\nplt.savefig('population_by_year_plot.png')",
-"input_links": ["https://population.un.org/world_population_by_year.csv"]
+```python
+# Input values
+principal = 43235
+rate = 5.24
+years = 5
+
+# Convert rate to decimal
+rate_decimal = rate / 100
+
+# Calculate final amount
+final_amount = principal * (1 + rate_decimal) ** years
+
+# Calculate interest earned
+interest_earned = final_amount - principal
+
+# Print results with formatting
+print(f"Interest Earned: ${{interest_earned:,.2f}}")
+print(f"Final Amount: ${{final_amount:,.2f}}")
+```
+
+Example 2:
+---
+Q: Simplify first, then evaluate: $-7x+2(x^{{2}}-1)-(2x^{{2}}-x+3)$, where $x=1$.
+A: Certainly! Let's break down the problem step-by-step and utilize Python with SymPy to simplify and evaluate the expression.
+
+1. **Expression Simplification:**
+ We start with the expression \\(-7x + 2(x^2 - 1) - (2x^2 - x + 3)\\).
+
+2. **Substitute \\(x=1\\) into the simplified expression:**
+ Once simplified, we will substitute \\(x=1\\) into the expression to find its value.
+
+Let's implement this in Python using SymPy (as the package is available in the sandbox):
+
+```python
+import sympy as sp
+
+# Define the variable
+x = sp.symbols('x')
+
+# Define the expression
+expression = -7*x + 2*(x**2 - 1) - (2*x**2 - x + 3)
+
+# Simplify the expression
+simplified_expression = sp.simplify(expression)
+
+# Substitute x = 1 into the simplified expression
+evaluated_expression = simplified_expression.subs(x, 1)
+
+# Print the simplified expression and its evaluated value
+print(\"Simplified Expression:\", simplified_expression)
+print(\"Evaluated Expression at x=1:\", evaluated_expression)
+```
+
+Example 3:
+---
+Q: Plot the world population growth over the years, given this year, world population world tuples: [(2000, 6), (2001, 7), (2002, 8), (2003, 9), (2004, 10)].
+A: Absolutely! We can utilize the Pandas and Matplotlib libraries (as both are available in the sandbox) to create the world population growth plot.
+```python
+import pandas as pd
+import matplotlib.pyplot as plt
+
+# Create a DataFrame of world population from the provided data
+data = {{
+    'Year': [2000, 2001, 2002, 2003, 2004],
+    'Population': [6, 7, 8, 9, 10]
 }}
+df = pd.DataFrame(data)
+
+# Plot the data
+plt.figure(figsize=(10, 6))
+plt.plot(df['Year'], df['Population'], marker='o')
+
+# Add titles and labels
+plt.title('Population by Year')
+plt.xlabel('Year')
+plt.ylabel('Population')
+
+# Save the plot to a file
+plt.savefig('population_by_year_plot.png')
+```
+
+Now it's your turn to construct a python program to answer the user's query using the provided context and coversation provided below.
+Ensure you include the python code to execute and wrap it in a markdown code block.
 
-Now it's your turn to construct a python program to answer the user's question. Provide the code, required input files and input links in a JSON object. Do not say anything else.
 Context:
 ---
 {context}
@@ -1015,8 +1088,9 @@ Chat History:
 ---
 {chat_history}
 
-User: {query}
-Khoj:
+User Query:
+---
+{query}
 """.strip()
 )
 
@@ -1030,6 +1104,13 @@ Code Execution Results:
 """.strip()
 )
 
+e2b_sandbox_context = """
+- The sandbox has access to only the standard library, matplotlib, pandas, numpy, scipy, bs4, sympy, einops, biopython, shapely, plotly and rdkit packages. The requests, torch, catboost, tensorflow and tkinter packages are not available.
+""".strip()
+
+terrarium_sandbox_context = """
+The sandbox has access to the standard library, matplotlib, pandas, numpy, scipy, bs4 and sympy packages. The requests, torch, catboost, tensorflow, rdkit and tkinter packages are not available.
+""".strip()
 
 # Automations
 # --

khoj/processor/conversation/utils.py

@@ -61,6 +61,9 @@ model_to_prompt_size = {
     "gemini-1.5-pro": 60000,
     # Anthropic Models
     "claude-3-5-sonnet-20241022": 60000,
+    "claude-3-5-sonnet-latest": 60000,
+    "claude-3-7-sonnet-20250219": 60000,
+    "claude-3-7-sonnet-latest": 60000,
     "claude-3-5-haiku-20241022": 60000,
     # Offline Models
     "bartowski/Qwen2.5-14B-Instruct-GGUF": 20000,
@@ -670,10 +673,13 @@ def get_image_from_url(image_url: str, type="pil"):
         content_type = response.headers.get("content-type") or mimetypes.guess_type(image_url)[0] or "image/webp"
 
         # Convert image to desired format
+        image_data: Any = None
         if type == "b64":
             image_data = base64.b64encode(response.content).decode("utf-8")
         elif type == "pil":
             image_data = PIL.Image.open(BytesIO(response.content))
+        elif type == "bytes":
+            image_data = response.content
         else:
             raise ValueError(f"Invalid image type: {type}")