khoj 1.36.7.dev66__py3-none-any.whl → 1.37.1.dev6__py3-none-any.whl

khoj/processor/conversation/openai/utils.py

@@ -2,6 +2,7 @@ import logging
 import os
 from threading import Thread
 from typing import Dict, List
+from urllib.parse import urlparse
 
 import openai
 from openai.types.chat.chat_completion import ChatCompletion
@@ -16,6 +17,7 @@ from tenacity import (
 )
 
 from khoj.processor.conversation.utils import (
+    JsonSupport,
     ThreadedGenerator,
     commit_conversation_trace,
 )
@@ -60,45 +62,29 @@ def completion_with_backoff(
 
     formatted_messages = [{"role": message.role, "content": message.content} for message in messages]
 
-    # Update request parameters for compatability with o1 model series
-    # Refer: https://platform.openai.com/docs/guides/reasoning/beta-limitations
-    stream = True
-    model_kwargs["stream_options"] = {"include_usage": True}
-    if model_name == "o1":
-        temperature = 1
-        stream = False
-        model_kwargs.pop("stream_options", None)
-    elif model_name.startswith("o1"):
-        temperature = 1
-        model_kwargs.pop("response_format", None)
-    elif model_name.startswith("o3-"):
+    # Tune reasoning models arguments
+    if model_name.startswith("o1") or model_name.startswith("o3"):
         temperature = 1
+        model_kwargs["reasoning_effort"] = "medium"
 
+    model_kwargs["stream_options"] = {"include_usage": True}
     if os.getenv("KHOJ_LLM_SEED"):
         model_kwargs["seed"] = int(os.getenv("KHOJ_LLM_SEED"))
 
-    chat: ChatCompletion | openai.Stream[ChatCompletionChunk] = client.chat.completions.create(
+    aggregated_response = ""
+    with client.beta.chat.completions.stream(
         messages=formatted_messages,  # type: ignore
-        model=model_name,  # type: ignore
-        stream=stream,
+        model=model_name,
         temperature=temperature,
         timeout=20,
         **model_kwargs,
-    )
-
-    aggregated_response = ""
-    if not stream:
-        chunk = chat
-        aggregated_response = chunk.choices[0].message.content
-    else:
+    ) as chat:
         for chunk in chat:
-            if len(chunk.choices) == 0:
+            if chunk.type == "error":
+                logger.error(f"Openai api response error: {chunk.error}", exc_info=True)
                 continue
-            delta_chunk = chunk.choices[0].delta  # type: ignore
-            if isinstance(delta_chunk, str):
-                aggregated_response += delta_chunk
-            elif delta_chunk.content:
-                aggregated_response += delta_chunk.content
+            elif chunk.type == "content.delta":
+                aggregated_response += chunk.delta
 
     # Calculate cost of chat
     input_tokens = chunk.usage.prompt_tokens if hasattr(chunk, "usage") and chunk.usage else 0
@@ -172,20 +158,13 @@ def llm_thread(
 
         formatted_messages = [{"role": message.role, "content": message.content} for message in messages]
 
-        # Update request parameters for compatability with o1 model series
-        # Refer: https://platform.openai.com/docs/guides/reasoning/beta-limitations
-        stream = True
-        model_kwargs["stream_options"] = {"include_usage": True}
-        if model_name == "o1":
+        # Tune reasoning models arguments
+        if model_name.startswith("o1"):
             temperature = 1
-            stream = False
-            model_kwargs.pop("stream_options", None)
-        elif model_name.startswith("o1-"):
+        elif model_name.startswith("o3"):
             temperature = 1
-            model_kwargs.pop("response_format", None)
-        elif model_name.startswith("o3-"):
-            temperature = 1
-            # Get the first system message and add the string `Formatting re-enabled` to it. See https://platform.openai.com/docs/guides/reasoning-best-practices
+            # Get the first system message and add the string `Formatting re-enabled` to it.
+            # See https://platform.openai.com/docs/guides/reasoning-best-practices
             if len(formatted_messages) > 0:
                 system_messages = [
                     (i, message) for i, message in enumerate(formatted_messages) if message["role"] == "system"
@@ -195,7 +174,6 @@ def llm_thread(
                     formatted_messages[first_system_message_index][
                         "content"
                     ] = f"{first_system_message} Formatting re-enabled"
-
         elif model_name.startswith("deepseek-reasoner"):
             # Two successive messages cannot be from the same role. Should merge any back-to-back messages from the same role.
             # The first message should always be a user message (except system message).
@@ -210,6 +188,8 @@ def llm_thread(
 
             formatted_messages = updated_messages
 
+        stream = True
+        model_kwargs["stream_options"] = {"include_usage": True}
         if os.getenv("KHOJ_LLM_SEED"):
             model_kwargs["seed"] = int(os.getenv("KHOJ_LLM_SEED"))
 
@@ -258,3 +238,15 @@ def llm_thread(
         logger.error(f"Error in llm_thread: {e}", exc_info=True)
     finally:
         g.close()
+
+
+def get_openai_api_json_support(model_name: str, api_base_url: str = None) -> JsonSupport:
+    if model_name.startswith("deepseek-reasoner"):
+        return JsonSupport.NONE
+    if api_base_url:
+        host = urlparse(api_base_url).hostname
+        if host and host.endswith(".ai.azure.com"):
+            return JsonSupport.OBJECT
+        if host == "api.deepinfra.com":
+            return JsonSupport.OBJECT
+    return JsonSupport.SCHEMA

khoj/processor/conversation/utils.py

@@ -345,8 +345,7 @@ def construct_structured_message(
             constructed_messages.append({"type": "text", "text": attached_file_context})
         if vision_enabled and images:
             for image in images:
-                if image.startswith("https://"):
-                    constructed_messages.append({"type": "image_url", "image_url": {"url": image}})
+                constructed_messages.append({"type": "image_url", "image_url": {"url": image}})
         return constructed_messages
 
     if not is_none_or_empty(attached_file_context):
@@ -664,6 +663,23 @@ class ImageWithType:
     type: str
 
 
+def get_image_from_base64(image: str, type="b64"):
+    # Extract image type and base64 data from inline image data
+    image_base64 = image.split(",", 1)[1]
+    image_type = image.split(";", 1)[0].split(":", 1)[1]
+
+    # Convert image to desired format
+    if type == "b64":
+        return ImageWithType(content=image_base64, type=image_type)
+    elif type == "pil":
+        image_data = base64.b64decode(image_base64)
+        image_pil = PIL.Image.open(BytesIO(image_data))
+        return ImageWithType(content=image_pil, type=image_type)
+    elif type == "bytes":
+        image_data = base64.b64decode(image_base64)
+        return ImageWithType(content=image_data, type=image_type)
+
+
 def get_image_from_url(image_url: str, type="pil"):
     try:
         response = requests.get(image_url)
@@ -878,3 +894,9 @@ def messages_to_print(messages: list[ChatMessage], max_length: int = 70) -> str:
             return str(content)
 
     return "\n".join([f"{json.dumps(safe_serialize(message.content))[:max_length]}..." for message in messages])
+
+
+class JsonSupport(int, Enum):
+    NONE = 0
+    OBJECT = 1
+    SCHEMA = 2

khoj/processor/image/generate.py

@@ -12,7 +12,7 @@ from google.genai import types as gtypes
 from khoj.database.adapters import ConversationAdapters
 from khoj.database.models import Agent, KhojUser, TextToImageModelConfig
 from khoj.routers.helpers import ChatEvent, generate_better_image_prompt
-from khoj.routers.storage import upload_image
+from khoj.routers.storage import upload_generated_image_to_bucket
 from khoj.utils import state
 from khoj.utils.helpers import convert_image_to_webp, timer
 from khoj.utils.rawconfig import LocationData
@@ -118,7 +118,7 @@ async def text_to_image(
 
     # Decide how to store the generated image
     with timer("Upload image to S3", logger):
-        image_url = upload_image(webp_image_bytes, user.uuid)
+        image_url = upload_generated_image_to_bucket(webp_image_bytes, user.uuid)
 
     if not image_url:
         image = f"data:image/webp;base64,{base64.b64encode(webp_image_bytes).decode('utf-8')}"

khoj/processor/tools/run_code.py

@@ -257,7 +257,7 @@ async def execute_e2b(code: str, input_files: list[dict]) -> dict[str, Any]:
                 continue
             else:
                 # Text files - encode utf-8 string as base64
-                b64_data = base64.b64encode(content.encode("utf-8")).decode("utf-8")
+                b64_data = content
             output_files.append({"filename": f.name, "b64_data": b64_data})
 
         # Collect output files from execution results

khoj/routers/api_chat.py

@@ -64,7 +64,7 @@ from khoj.routers.research import (
     InformationCollectionIteration,
     execute_information_collection,
 )
-from khoj.routers.storage import upload_image_to_bucket
+from khoj.routers.storage import upload_user_image_to_bucket
 from khoj.utils import state
 from khoj.utils.helpers import (
     AsyncIteratorWrapper,
@@ -674,9 +674,11 @@ async def chat(
                 base64_data = decoded_string.split(",", 1)[1]
                 image_bytes = base64.b64decode(base64_data)
                 webp_image_bytes = convert_image_to_webp(image_bytes)
-                uploaded_image = upload_image_to_bucket(webp_image_bytes, request.user.object.id)
-                if uploaded_image:
-                    uploaded_images.append(uploaded_image)
+                uploaded_image = upload_user_image_to_bucket(webp_image_bytes, request.user.object.id)
+                if not uploaded_image:
+                    base64_webp_image = base64.b64encode(webp_image_bytes).decode("utf-8")
+                    uploaded_image = f"data:image/webp;base64,{base64_webp_image}"
+                uploaded_images.append(uploaded_image)
 
         query_files: Dict[str, str] = {}
         if raw_query_files:

khoj/routers/auth.py

@@ -43,12 +43,9 @@ class MagicLinkForm(BaseModel):
 if not state.anonymous_mode:
     missing_requirements = []
     from authlib.integrations.starlette_client import OAuth, OAuthError
+    from google.auth.transport import requests as google_requests
+    from google.oauth2 import id_token
 
-    try:
-        from google.auth.transport import requests as google_requests
-        from google.oauth2 import id_token
-    except ImportError:
-        missing_requirements += ["Install the Khoj production package with `pip install khoj[prod]`"]
     if not os.environ.get("RESEND_API_KEY") and (
         not os.environ.get("GOOGLE_CLIENT_ID") or not os.environ.get("GOOGLE_CLIENT_SECRET")
     ):

khoj/routers/helpers.py

@@ -540,11 +540,15 @@ async def generate_online_subqueries(
 
     agent_chat_model = agent.chat_model if agent else None
 
+    class OnlineQueries(BaseModel):
+        queries: List[str]
+
     with timer("Chat actor: Generate online search subqueries", logger):
         response = await send_message_to_model_wrapper(
             online_queries_prompt,
             query_images=query_images,
             response_type="json_object",
+            response_schema=OnlineQueries,
             user=user,
             query_files=query_files,
             agent_chat_model=agent_chat_model,
@@ -1129,6 +1133,7 @@ async def send_message_to_model_wrapper(
     query: str,
     system_message: str = "",
     response_type: str = "text",
+    response_schema: BaseModel = None,
     deepthought: bool = False,
     user: KhojUser = None,
     query_images: List[str] = None,
@@ -1209,6 +1214,7 @@ async def send_message_to_model_wrapper(
             api_key=api_key,
             model=chat_model_name,
             response_type=response_type,
+            response_schema=response_schema,
             api_base_url=api_base_url,
             tracer=tracer,
         )
@@ -1255,6 +1261,7 @@ async def send_message_to_model_wrapper(
             api_key=api_key,
             model=chat_model_name,
             response_type=response_type,
+            response_schema=response_schema,
             tracer=tracer,
         )
     else:
@@ -1265,6 +1272,7 @@ def send_message_to_model_wrapper_sync(
     message: str,
     system_message: str = "",
     response_type: str = "text",
+    response_schema: BaseModel = None,
     user: KhojUser = None,
     query_images: List[str] = None,
     query_files: str = "",
@@ -1326,6 +1334,7 @@ def send_message_to_model_wrapper_sync(
             api_base_url=api_base_url,
             model=chat_model_name,
             response_type=response_type,
+            response_schema=response_schema,
             tracer=tracer,
         )
 
@@ -1370,6 +1379,7 @@ def send_message_to_model_wrapper_sync(
             api_key=api_key,
             model=chat_model_name,
             response_type=response_type,
+            response_schema=response_schema,
             tracer=tracer,
         )
     else:

khoj/routers/research.py

@@ -1,9 +1,12 @@
 import logging
+import os
 from datetime import datetime
-from typing import Callable, Dict, List, Optional
+from enum import Enum
+from typing import Callable, Dict, List, Optional, Type
 
 import yaml
 from fastapi import Request
+from pydantic import BaseModel, Field
 
 from khoj.database.adapters import EntryAdapters
 from khoj.database.models import Agent, KhojUser
@@ -35,6 +38,40 @@ from khoj.utils.rawconfig import LocationData
 logger = logging.getLogger(__name__)
 
 
+class PlanningResponse(BaseModel):
+    """
+    Schema for the response from planning agent when deciding the next tool to pick.
+    The tool field is dynamically validated based on available tools.
+    """
+
+    scratchpad: str = Field(..., description="Reasoning about which tool to use next")
+    query: str = Field(..., description="Detailed query for the selected tool")
+
+    class Config:
+        arbitrary_types_allowed = True
+
+    @classmethod
+    def create_model_with_enum(cls: Type["PlanningResponse"], tool_options: dict) -> Type["PlanningResponse"]:
+        """
+        Factory method that creates a customized PlanningResponse model
+        with a properly typed tool field based on available tools.
+
+        Args:
+            tool_options: Dictionary mapping tool names to values
+
+        Returns:
+            A customized PlanningResponse class
+        """
+        # Create dynamic enum from tool options
+        tool_enum = Enum("ToolEnum", tool_options)  # type: ignore
+
+        # Create and return a customized response model with the enum
+        class PlanningResponseWithTool(PlanningResponse):
+            tool: tool_enum = Field(..., description="Name of the tool to use")
+
+        return PlanningResponseWithTool
+
+
 async def apick_next_tool(
     query: str,
     conversation_history: dict,
@@ -60,10 +97,13 @@ async def apick_next_tool(
         # Skip showing Notes tool as an option if user has no entries
         if tool == ConversationCommand.Notes and not user_has_entries:
             continue
-        tool_options[tool.value] = description
         if len(agent_tools) == 0 or tool.value in agent_tools:
+            tool_options[tool.name] = tool.value
             tool_options_str += f'- "{tool.value}": "{description}"\n'
 
+    # Create planning reponse model with dynamically populated tool enum class
+    planning_response_model = PlanningResponse.create_model_with_enum(tool_options)
+
     # Construct chat history with user and iteration history with researcher agent for context
     chat_history = construct_chat_history(conversation_history, agent_name=agent.name if agent else "Khoj")
     previous_iterations_history = construct_iteration_history(previous_iterations, prompts.previous_iteration)
@@ -95,6 +135,7 @@ async def apick_next_tool(
                 query=query,
                 context=function_planning_prompt,
                 response_type="json_object",
+                response_schema=planning_response_model,
                 deepthought=True,
                 user=user,
                 query_images=query_images,
@@ -160,7 +201,7 @@ async def execute_information_collection(
     query_files: str = None,
 ):
     current_iteration = 0
-    MAX_ITERATIONS = 5
+    MAX_ITERATIONS = int(os.getenv("KHOJ_RESEARCH_ITERATIONS", 5))
     previous_iterations: List[InformationCollectionIteration] = []
     while current_iteration < MAX_ITERATIONS:
         online_results: Dict = dict()

khoj/routers/storage.py

@@ -9,9 +9,10 @@ AWS_SECRET_KEY = os.getenv("AWS_SECRET_KEY")
 # S3 supports serving assets via your domain. Khoj expects this to be used in production. To enable it:
 # 1. Your bucket name for images should be of the form sub.domain.tld. For example, generated.khoj.dev
 # 2. Add CNAME entry to your domain's DNS records pointing to the S3 bucket. For example, CNAME generated.khoj.dev generated-khoj-dev.s3.amazonaws.com
-AWS_UPLOAD_IMAGE_BUCKET_NAME = os.getenv("AWS_IMAGE_UPLOAD_BUCKET")
+AWS_KHOJ_IMAGES_BUCKET_NAME = os.getenv("AWS_IMAGE_UPLOAD_BUCKET")
+AWS_USER_IMAGES_BUCKET_NAME = os.getenv("AWS_USER_UPLOADED_IMAGES_BUCKET_NAME")
 
-aws_enabled = AWS_ACCESS_KEY is not None and AWS_SECRET_KEY is not None and AWS_UPLOAD_IMAGE_BUCKET_NAME is not None
+aws_enabled = AWS_ACCESS_KEY is not None and AWS_SECRET_KEY is not None
 
 if aws_enabled:
     from boto3 import client
@@ -19,45 +20,43 @@ if aws_enabled:
     s3_client = client("s3", aws_access_key_id=AWS_ACCESS_KEY, aws_secret_access_key=AWS_SECRET_KEY)
 
 
-def upload_image(image: bytes, user_id: uuid.UUID):
-    """Upload the image to the S3 bucket"""
+def upload_image_to_bucket(webp_image: bytes, user_id: uuid.UUID, bucket_name: str):
+    """Upload webp image to an S3 bucket"""
     if not aws_enabled:
         logger.info("AWS is not enabled. Skipping image upload")
         return None
-
-    image_key = f"{user_id}/{uuid.uuid4()}.webp"
-    try:
-        s3_client.put_object(Bucket=AWS_UPLOAD_IMAGE_BUCKET_NAME, Key=image_key, Body=image, ACL="public-read")
-        url = f"https://{AWS_UPLOAD_IMAGE_BUCKET_NAME}/{image_key}"
-        return url
-    except Exception as e:
-        logger.error(f"Failed to upload image to S3: {e}")
-        return None
-
-
-AWS_USER_UPLOADED_IMAGES_BUCKET_NAME = os.getenv("AWS_USER_UPLOADED_IMAGES_BUCKET_NAME")
-
-
-def upload_image_to_bucket(image: bytes, user_id: uuid.UUID):
-    """Upload the image to the S3 bucket"""
-    if not aws_enabled:
-        logger.info("AWS is not enabled. Skipping image upload")
+    if not bucket_name:
+        logger.error(f"{bucket_name} is not set")
         return None
 
     image_key = f"{user_id}/{uuid.uuid4()}.webp"
-    if not AWS_USER_UPLOADED_IMAGES_BUCKET_NAME:
-        logger.error("AWS_USER_UPLOADED_IMAGES_BUCKET_NAME is not set")
-        return None
-
     try:
         s3_client.put_object(
-            Bucket=AWS_USER_UPLOADED_IMAGES_BUCKET_NAME,
+            Bucket=bucket_name,
             Key=image_key,
-            Body=image,
+            Body=webp_image,
             ACL="public-read",
             ContentType="image/webp",
         )
-        return f"https://{AWS_USER_UPLOADED_IMAGES_BUCKET_NAME}/{image_key}"
+        return f"https://{bucket_name}/{image_key}"
     except Exception as e:
         logger.error(f"Failed to upload image to S3: {e}")
         return None
+
+
+def upload_generated_image_to_bucket(image: bytes, user_id: uuid.UUID):
+    """Upload khoj generated image to an S3 bucket"""
+    return upload_image_to_bucket(
+        webp_image=image,
+        user_id=user_id,
+        bucket_name=AWS_KHOJ_IMAGES_BUCKET_NAME,
+    )
+
+
+def upload_user_image_to_bucket(image: bytes, user_id: uuid.UUID):
+    """Upload user attached image to an S3 bucket"""
+    return upload_image_to_bucket(
+        webp_image=image,
+        user_id=user_id,
+        bucket_name=AWS_USER_IMAGES_BUCKET_NAME,
+    )

{khoj-1.36.7.dev66.dist-info → khoj-1.37.1.dev6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: khoj
-Version: 1.36.7.dev66
+Version: 1.37.1.dev6
 Summary: Your Second Brain
 Project-URL: Homepage, https://khoj.dev
 Project-URL: Documentation, https://docs.khoj.dev
@@ -39,6 +39,7 @@ 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-auth~=2.23.3
 Requires-Dist: google-genai==1.5.0
 Requires-Dist: httpx==0.28.1
 Requires-Dist: huggingface-hub>=0.22.2
@@ -69,12 +70,12 @@ Requires-Dist: requests>=2.26.0
 Requires-Dist: resend==1.0.1
 Requires-Dist: rich>=13.3.1
 Requires-Dist: schedule==1.1.0
-Requires-Dist: sentence-transformers==3.0.1
+Requires-Dist: sentence-transformers==3.4.1
 Requires-Dist: tenacity==8.3.0
 Requires-Dist: tenacity>=8.2.2
 Requires-Dist: tiktoken>=0.3.2
 Requires-Dist: torch==2.2.2
-Requires-Dist: transformers>=4.28.0
+Requires-Dist: transformers<4.50.0,>=4.28.0
 Requires-Dist: tzdata==2023.3
 Requires-Dist: uvicorn==0.30.6
 Requires-Dist: websockets==13.0
@@ -85,7 +86,6 @@ Requires-Dist: datasets; extra == 'dev'
 Requires-Dist: factory-boy>=3.2.1; extra == 'dev'
 Requires-Dist: freezegun>=1.2.0; extra == 'dev'
 Requires-Dist: gitpython~=3.1.43; extra == 'dev'
-Requires-Dist: google-auth==2.23.3; extra == 'dev'
 Requires-Dist: gunicorn==22.0.0; extra == 'dev'
 Requires-Dist: mypy>=1.0.1; extra == 'dev'
 Requires-Dist: pandas; extra == 'dev'
@@ -98,7 +98,6 @@ Requires-Dist: stripe==7.3.0; extra == 'dev'
 Requires-Dist: twilio==8.11; extra == 'dev'
 Provides-Extra: prod
 Requires-Dist: boto3>=1.34.57; extra == 'prod'
-Requires-Dist: google-auth==2.23.3; extra == 'prod'
 Requires-Dist: gunicorn==22.0.0; extra == 'prod'
 Requires-Dist: stripe==7.3.0; extra == 'prod'
 Requires-Dist: twilio==8.11; extra == 'prod'