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

khoj/processor/conversation/google/gemini_chat.py

@@ -34,7 +34,7 @@ def extract_questions_gemini(
     model: Optional[str] = "gemini-2.0-flash",
     conversation_log={},
     api_key=None,
-    temperature=0.2,
+    temperature=0.6,
     max_tokens=None,
     location_data: LocationData = None,
     user: KhojUser = None,
@@ -121,14 +121,14 @@ def gemini_send_message_to_model(
     api_key,
     model,
     response_type="text",
-    temperature=0.2,
+    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 = {}
 
@@ -138,7 +138,7 @@ def gemini_send_message_to_model(
 
     # 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,
@@ -156,7 +156,7 @@ def converse_gemini(
     conversation_log={},
     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/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}")
 

khoj/routers/helpers.py

@@ -1125,6 +1125,7 @@ async def send_message_to_model_wrapper(
     query: str,
     system_message: str = "",
     response_type: str = "text",
+    deepthought: bool = False,
     user: KhojUser = None,
     query_images: List[str] = None,
     context: str = "",
@@ -1227,6 +1228,7 @@ async def send_message_to_model_wrapper(
             api_key=api_key,
             model=chat_model_name,
             response_type=response_type,
+            deepthought=deepthought,
             tracer=tracer,
         )
     elif model_type == ChatModel.ModelType.GOOGLE:
@@ -1425,11 +1427,13 @@ def generate_chat_response(
         )
 
         query_to_run = q
+        deepthought = False
         if meta_research:
             query_to_run = f"<query>{q}</query>\n<collected_research>\n{meta_research}\n</collected_research>"
             compiled_references = []
             online_results = {}
             code_results = {}
+            deepthought = True
 
         chat_model = ConversationAdapters.get_valid_chat_model(user, conversation, is_subscribed)
         vision_available = chat_model.vision_enabled
@@ -1513,6 +1517,7 @@ def generate_chat_response(
                 generated_files=raw_generated_files,
                 generated_asset_results=generated_asset_results,
                 program_execution_context=program_execution_context,
+                deepthought=deepthought,
                 tracer=tracer,
             )
         elif chat_model.model_type == ChatModel.ModelType.GOOGLE:

khoj/routers/research.py

@@ -95,6 +95,7 @@ async def apick_next_tool(
                 query=query,
                 context=function_planning_prompt,
                 response_type="json_object",
+                deepthought=True,
                 user=user,
                 query_images=query_images,
                 query_files=query_files,

khoj/utils/constants.py

@@ -48,6 +48,9 @@ model_to_cost: Dict[str, Dict[str, float]] = {
     "gemini-1.5-pro-002": {"input": 1.25, "output": 5.00},
     "gemini-2.0-flash": {"input": 0.10, "output": 0.40},
     # Anthropic Pricing: https://www.anthropic.com/pricing#anthropic-api_
-    "claude-3-5-sonnet-20241022": {"input": 3.0, "output": 15.0},
     "claude-3-5-haiku-20241022": {"input": 1.0, "output": 5.0},
+    "claude-3-5-sonnet-20241022": {"input": 3.0, "output": 15.0},
+    "claude-3-5-sonnet-latest": {"input": 3.0, "output": 15.0},
+    "claude-3-7-sonnet-20250219": {"input": 3.0, "output": 15.0},
+    "claude-3-7-sonnet-latest": {"input": 3.0, "output": 15.0},
 }

khoj/utils/helpers.py

@@ -368,7 +368,7 @@ command_descriptions_for_agent = {
     ConversationCommand.Code: "Agent can run Python code to parse information, run complex calculations, create documents and charts.",
 }
 
-e2b_tool_description = "To run Python code in a E2B sandbox with no network access. Helpful to parse complex information, run calculations, create text documents and create charts with quantitative data. Only matplotlib, pandas, numpy, scipy, bs4, sympy, einops, biopython, shapely and rdkit external packages are available."
+e2b_tool_description = "To run Python code in a E2B sandbox with no network access. Helpful to parse complex information, run calculations, create text documents and create charts with quantitative data. Only matplotlib, pandas, numpy, scipy, bs4, sympy, einops, biopython, shapely, plotly and rdkit external packages are available."
 terrarium_tool_description = "To run Python code in a Terrarium, Pyodide sandbox with no network access. Helpful to parse complex information, run complex calculations, create plaintext documents and create charts with quantitative data. Only matplotlib, panda, numpy, scipy, bs4 and sympy external packages are available."
 
 tool_descriptions_for_llm = {

{khoj-1.36.7.dev18.dist-info → khoj-1.36.7.dev22.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: khoj
-Version: 1.36.7.dev18
+Version: 1.36.7.dev22
 Summary: Your Second Brain
 Project-URL: Homepage, https://khoj.dev
 Project-URL: Documentation, https://docs.khoj.dev
@@ -22,8 +22,8 @@ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Classifier: Topic :: Scientific/Engineering :: Human Machine Interfaces
 Requires-Python: >=3.10
 Requires-Dist: aiohttp~=3.9.0
-Requires-Dist: anthropic==0.26.1
-Requires-Dist: anyio==3.7.1
+Requires-Dist: anthropic==0.49.0
+Requires-Dist: anyio~=4.8.0
 Requires-Dist: apscheduler~=3.10.0
 Requires-Dist: authlib==1.2.1
 Requires-Dist: beautifulsoup4~=4.12.3
@@ -39,8 +39,8 @@ Requires-Dist: e2b-code-interpreter~=1.0.0
 Requires-Dist: einops==0.8.0
 Requires-Dist: email-validator==2.2.0
 Requires-Dist: fastapi>=0.110.0
-Requires-Dist: google-generativeai==0.8.3
-Requires-Dist: httpx==0.27.2
+Requires-Dist: google-genai==1.5.0
+Requires-Dist: httpx==0.28.1
 Requires-Dist: huggingface-hub>=0.22.2
 Requires-Dist: itsdangerous==2.1.2
 Requires-Dist: jinja2==3.1.5
@@ -77,7 +77,7 @@ Requires-Dist: torch==2.2.2
 Requires-Dist: transformers>=4.28.0
 Requires-Dist: tzdata==2023.3
 Requires-Dist: uvicorn==0.30.6
-Requires-Dist: websockets==12.0
+Requires-Dist: websockets==13.0
 Provides-Extra: dev
 Requires-Dist: black>=23.1.0; extra == 'dev'
 Requires-Dist: boto3>=1.34.57; extra == 'dev'