khoj 1.36.7.dev18__py3-none-any.whl → 1.36.7.dev24__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.generativeai.types.generation_types import StopCandidateException
-from google.generativeai.types.safety_types import (
-    HarmBlockThreshold,
-    HarmCategory,
-    HarmProbability,
-)
+from google import genai
+from google.genai import errors as gerrors
+from google.genai import types as gtypes
 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,32 +59,21 @@ 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:
+    except gerrors.ClientError as e:
         response = None
         response_text, _ = handle_gemini_response(e.args)
         # Respond with reason for stopping
@@ -125,36 +127,27 @@ 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
             g.send(message)
             if stopped:
-                raise StopCandidateException(message)
+                raise ValueError(message)
 
         # Calculate cost of chat
         input_tokens = chunk.usage_metadata.prompt_token_count
@@ -166,7 +159,7 @@ def gemini_llm_thread(
         tracer["temperature"] = temperature
         if is_promptrace_enabled():
             commit_conversation_trace(messages, aggregated_response, tracer)
-    except StopCandidateException as e:
+    except ValueError as e:
         logger.warning(
             f"LLM Response Prevented for {model_name}: {e.args[0]}.\n"
             + f"Last Message by {messages[-1].role}: {messages[-1].content}"
@@ -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/processor/image/generate.py

@@ -6,6 +6,8 @@ from typing import Any, Callable, Dict, List, Optional
 
 import openai
 import requests
+from google import genai
+from google.genai import types as gtypes
 
 from khoj.database.adapters import ConversationAdapters
 from khoj.database.models import Agent, KhojUser, TextToImageModelConfig
@@ -86,6 +88,8 @@ async def text_to_image(
                 webp_image_bytes = generate_image_with_stability(image_prompt, text_to_image_config, text2image_model)
             elif text_to_image_config.model_type == TextToImageModelConfig.ModelType.REPLICATE:
                 webp_image_bytes = generate_image_with_replicate(image_prompt, text_to_image_config, text2image_model)
+            elif text_to_image_config.model_type == TextToImageModelConfig.ModelType.GOOGLE:
+                webp_image_bytes = generate_image_with_google(image_prompt, text_to_image_config, text2image_model)
         except openai.OpenAIError or openai.BadRequestError or openai.APIConnectionError as e:
             if "content_policy_violation" in e.message:
                 logger.error(f"Image Generation blocked by OpenAI: {e}")
@@ -99,6 +103,12 @@ async def text_to_image(
                 status_code = e.status_code  # type: ignore
                 yield image_url or image, status_code, message
                 return
+        except ValueError as e:
+            logger.error(f"Image Generation failed with {e}", exc_info=True)
+            message = f"Image generation using {text2image_model} via {text_to_image_config.model_type} failed due to an unknown error"
+            status_code = 500
+            yield image_url or image, status_code, message
+            return
         except requests.RequestException as e:
             logger.error(f"Image Generation failed with {e}", exc_info=True)
             message = f"Image generation using {text2image_model} via {text_to_image_config.model_type} failed due to a network error."
@@ -215,3 +225,28 @@ def generate_image_with_replicate(
     # Get the generated image
     image_url = get_prediction["output"][0] if isinstance(get_prediction["output"], list) else get_prediction["output"]
     return io.BytesIO(requests.get(image_url).content).getvalue()
+
+
+def generate_image_with_google(
+    improved_image_prompt: str, text_to_image_config: TextToImageModelConfig, text2image_model: str
+):
+    """Generate image using Google's AI over API"""
+
+    # Initialize the Google AI client
+    api_key = text_to_image_config.api_key or text_to_image_config.ai_model_api.api_key
+    client = genai.Client(api_key=api_key)
+
+    # Configure image generation settings
+    config = gtypes.GenerateImagesConfig(number_of_images=1)
+
+    # Call the Gemini API to generate the image
+    response = client.models.generate_images(model=text2image_model, prompt=improved_image_prompt, config=config)
+
+    if not response.generated_images:
+        raise ValueError("Failed to generate image using Google AI")
+
+    # Extract the image bytes from the first generated image
+    image_bytes = response.generated_images[0].image.image_bytes
+
+    # Convert to webp for faster loading
+    return convert_image_to_webp(image_bytes)

khoj/routers/helpers.py

@@ -1092,7 +1092,11 @@ async def generate_better_image_prompt(
             online_results=simplified_online_results,
             personality_context=personality_context,
         )
-    elif model_type in [TextToImageModelConfig.ModelType.STABILITYAI, TextToImageModelConfig.ModelType.REPLICATE]:
+    elif model_type in [
+        TextToImageModelConfig.ModelType.STABILITYAI,
+        TextToImageModelConfig.ModelType.REPLICATE,
+        TextToImageModelConfig.ModelType.GOOGLE,
+    ]:
         image_prompt = prompts.image_generation_improve_prompt_sd.format(
             query=q,
             chat_history=conversation_history,
@@ -1125,6 +1129,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 +1232,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 +1431,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 +1521,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.dev24.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: khoj
-Version: 1.36.7.dev18
+Version: 1.36.7.dev24
 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'