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

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,11 +1214,13 @@ 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,
         )
     elif model_type == ChatModel.ModelType.ANTHROPIC:
         api_key = chat_model.ai_model_api.api_key
+        api_base_url = chat_model.ai_model_api.api_base_url
         truncated_messages = generate_chatml_messages_with_context(
             user_message=query,
             context_message=context,
@@ -1233,10 +1240,12 @@ async def send_message_to_model_wrapper(
             model=chat_model_name,
             response_type=response_type,
             deepthought=deepthought,
+            api_base_url=api_base_url,
             tracer=tracer,
         )
     elif model_type == ChatModel.ModelType.GOOGLE:
         api_key = chat_model.ai_model_api.api_key
+        api_base_url = chat_model.ai_model_api.api_base_url
         truncated_messages = generate_chatml_messages_with_context(
             user_message=query,
             context_message=context,
@@ -1255,6 +1264,8 @@ 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,
         )
     else:
@@ -1265,6 +1276,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 = "",
@@ -1320,19 +1332,19 @@ def send_message_to_model_wrapper_sync(
             query_files=query_files,
         )
 
-        openai_response = send_message_to_model(
+        return send_message_to_model(
             messages=truncated_messages,
             api_key=api_key,
             api_base_url=api_base_url,
             model=chat_model_name,
             response_type=response_type,
+            response_schema=response_schema,
             tracer=tracer,
         )
 
-        return openai_response
-
     elif chat_model.model_type == ChatModel.ModelType.ANTHROPIC:
         api_key = chat_model.ai_model_api.api_key
+        api_base_url = chat_model.ai_model_api.api_base_url
         truncated_messages = generate_chatml_messages_with_context(
             user_message=message,
             system_message=system_message,
@@ -1347,6 +1359,7 @@ def send_message_to_model_wrapper_sync(
         return anthropic_send_message_to_model(
             messages=truncated_messages,
             api_key=api_key,
+            api_base_url=api_base_url,
             model=chat_model_name,
             response_type=response_type,
             tracer=tracer,
@@ -1354,6 +1367,7 @@ def send_message_to_model_wrapper_sync(
 
     elif chat_model.model_type == ChatModel.ModelType.GOOGLE:
         api_key = chat_model.ai_model_api.api_key
+        api_base_url = chat_model.ai_model_api.api_base_url
         truncated_messages = generate_chatml_messages_with_context(
             user_message=message,
             system_message=system_message,
@@ -1368,8 +1382,10 @@ def send_message_to_model_wrapper_sync(
         return gemini_send_message_to_model(
             messages=truncated_messages,
             api_key=api_key,
+            api_base_url=api_base_url,
             model=chat_model_name,
             response_type=response_type,
+            response_schema=response_schema,
             tracer=tracer,
         )
     else:
@@ -1500,6 +1516,7 @@ def generate_chat_response(
 
         elif chat_model.model_type == ChatModel.ModelType.ANTHROPIC:
             api_key = chat_model.ai_model_api.api_key
+            api_base_url = chat_model.ai_model_api.api_base_url
             chat_response = converse_anthropic(
                 compiled_references,
                 query_to_run,
@@ -1509,6 +1526,7 @@ def generate_chat_response(
                 conversation_log=meta_log,
                 model=chat_model.name,
                 api_key=api_key,
+                api_base_url=api_base_url,
                 completion_func=partial_completion,
                 conversation_commands=conversation_commands,
                 max_prompt_size=chat_model.max_prompt_size,
@@ -1526,6 +1544,7 @@ def generate_chat_response(
             )
         elif chat_model.model_type == ChatModel.ModelType.GOOGLE:
             api_key = chat_model.ai_model_api.api_key
+            api_base_url = chat_model.ai_model_api.api_base_url
             chat_response = converse_gemini(
                 compiled_references,
                 query_to_run,
@@ -1534,6 +1553,7 @@ def generate_chat_response(
                 meta_log,
                 model=chat_model.name,
                 api_key=api_key,
+                api_base_url=api_base_url,
                 completion_func=partial_completion,
                 conversation_commands=conversation_commands,
                 max_prompt_size=chat_model.max_prompt_size,

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/utils/constants.py

@@ -49,8 +49,10 @@ model_to_cost: Dict[str, Dict[str, float]] = {
     "gemini-2.0-flash": {"input": 0.10, "output": 0.40},
     # Anthropic Pricing: https://www.anthropic.com/pricing#anthropic-api_
     "claude-3-5-haiku-20241022": {"input": 1.0, "output": 5.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@20250219": {"input": 3.0, "output": 15.0},
     "claude-3-7-sonnet-latest": {"input": 3.0, "output": 15.0},
 }

khoj/utils/helpers.py

@@ -1,5 +1,6 @@
 from __future__ import annotations  # to avoid quoting type hints
 
+import base64
 import copy
 import datetime
 import io
@@ -19,15 +20,18 @@ from itertools import islice
 from os import path
 from pathlib import Path
 from time import perf_counter
-from typing import TYPE_CHECKING, Any, Optional, Union
-from urllib.parse import urlparse
+from typing import TYPE_CHECKING, Any, NamedTuple, Optional, Tuple, Union
+from urllib.parse import ParseResult, urlparse
 
 import openai
 import psutil
+import pyjson5
 import requests
 import torch
 from asgiref.sync import sync_to_async
 from email_validator import EmailNotValidError, EmailUndeliverableError, validate_email
+from google.auth.credentials import Credentials
+from google.oauth2 import service_account
 from magika import Magika
 from PIL import Image
 from pytz import country_names, country_timezones
@@ -618,6 +622,58 @@ def get_chat_usage_metrics(
     }
 
 
+class AiApiInfo(NamedTuple):
+    region: str
+    project: str
+    credentials: Credentials
+    api_key: str
+
+
+def get_gcp_credentials(credentials_b64: str) -> Optional[Credentials]:
+    """
+    Get GCP credentials from base64 encoded service account credentials json keyfile
+    """
+    credentials_json = base64.b64decode(credentials_b64).decode("utf-8")
+    credentials_dict = pyjson5.loads(credentials_json)
+    credentials = service_account.Credentials.from_service_account_info(credentials_dict)
+    return credentials.with_scopes(scopes=["https://www.googleapis.com/auth/cloud-platform"])
+
+
+def get_gcp_project_info(parsed_api_url: ParseResult) -> Tuple[str, str]:
+    """
+    Extract region, project id from GCP API url
+    API url is of form https://{REGION}-aiplatform.googleapis.com/v1/projects/{PROJECT_ID}...
+    """
+    # Extract region from hostname
+    hostname = parsed_api_url.netloc
+    region = hostname.split("-aiplatform")[0] if "-aiplatform" in hostname else ""
+
+    # Extract project ID from path (e.g., "/v1/projects/my-project/...")
+    path_parts = parsed_api_url.path.split("/")
+    project_id = ""
+    for i, part in enumerate(path_parts):
+        if part == "projects" and i + 1 < len(path_parts):
+            project_id = path_parts[i + 1]
+            break
+
+    return region, project_id
+
+
+def get_ai_api_info(api_key, api_base_url: str = None) -> AiApiInfo:
+    """
+    Get the GCP Vertex or default AI API client info based on the API key and URL.
+    """
+    region, project_id, credentials = None, None, None
+    # Check if AI model to be used via GCP Vertex API
+    parsed_api_url = urlparse(api_base_url)
+    if parsed_api_url.hostname and parsed_api_url.hostname.endswith(".googleapis.com"):
+        region, project_id = get_gcp_project_info(parsed_api_url)
+        credentials = get_gcp_credentials(api_key)
+    if credentials:
+        api_key = None
+    return AiApiInfo(region=region, project=project_id, credentials=credentials, api_key=api_key)
+
+
 def get_openai_client(api_key: str, api_base_url: str) -> Union[openai.OpenAI, openai.AzureOpenAI]:
     """Get OpenAI or AzureOpenAI client based on the API Base URL"""
     parsed_url = urlparse(api_base_url)

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: khoj
-Version: 1.36.7.dev66
+Version: 1.37.1
 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: torch==2.6.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'