khoj 1.37.0__py3-none-any.whl → 1.37.1.dev12__py3-none-any.whl

khoj/processor/conversation/google/utils.py

@@ -1,7 +1,9 @@
 import logging
+import os
 import random
 from copy import deepcopy
 from threading import Thread
+from typing import Dict
 
 from google import genai
 from google.genai import errors as gerrors
@@ -22,6 +24,7 @@ from khoj.processor.conversation.utils import (
     get_image_from_url,
 )
 from khoj.utils.helpers import (
+    get_ai_api_info,
     get_chat_usage_metrics,
     is_none_or_empty,
     is_promptrace_enabled,
@@ -29,6 +32,7 @@ from khoj.utils.helpers import (
 
 logger = logging.getLogger(__name__)
 
+gemini_clients: Dict[str, genai.Client] = {}
 
 MAX_OUTPUT_TOKENS_GEMINI = 8192
 SAFETY_SETTINGS = [
@@ -51,6 +55,17 @@ SAFETY_SETTINGS = [
 ]
 
 
+def get_gemini_client(api_key, api_base_url=None) -> genai.Client:
+    api_info = get_ai_api_info(api_key, api_base_url)
+    return genai.Client(
+        location=api_info.region,
+        project=api_info.project,
+        credentials=api_info.credentials,
+        api_key=api_info.api_key,
+        vertexai=api_info.api_key is None,
+    )
+
+
 @retry(
     wait=wait_random_exponential(min=1, max=10),
     stop=stop_after_attempt(2),
@@ -58,9 +73,14 @@ SAFETY_SETTINGS = [
     reraise=True,
 )
 def gemini_completion_with_backoff(
-    messages, system_prompt, model_name, temperature=0, api_key=None, model_kwargs=None, tracer={}
+    messages, system_prompt, model_name, temperature=0, api_key=None, api_base_url=None, model_kwargs=None, tracer={}
 ) -> str:
-    client = genai.Client(api_key=api_key)
+    client = gemini_clients.get(api_key)
+    if not client:
+        client = get_gemini_client(api_key, api_base_url)
+        gemini_clients[api_key] = client
+
+    seed = int(os.getenv("KHOJ_LLM_SEED")) if os.getenv("KHOJ_LLM_SEED") else None
     config = gtypes.GenerateContentConfig(
         system_instruction=system_prompt,
         temperature=temperature,
@@ -68,6 +88,7 @@ def gemini_completion_with_backoff(
         safety_settings=SAFETY_SETTINGS,
         response_mime_type=model_kwargs.get("response_mime_type", "text/plain") if model_kwargs else "text/plain",
         response_schema=model_kwargs.get("response_schema", None) if model_kwargs else None,
+        seed=seed,
     )
 
     formatted_messages = [gtypes.Content(role=message.role, parts=message.content) for message in messages]
@@ -112,6 +133,7 @@ def gemini_chat_completion_with_backoff(
     model_name,
     temperature,
     api_key,
+    api_base_url,
     system_prompt,
     completion_func=None,
     model_kwargs=None,
@@ -120,23 +142,37 @@ def gemini_chat_completion_with_backoff(
     g = ThreadedGenerator(compiled_references, online_results, completion_func=completion_func)
     t = Thread(
         target=gemini_llm_thread,
-        args=(g, messages, system_prompt, model_name, temperature, api_key, model_kwargs, tracer),
+        args=(g, messages, system_prompt, model_name, temperature, api_key, api_base_url, model_kwargs, tracer),
     )
     t.start()
     return g
 
 
 def gemini_llm_thread(
-    g, messages, system_prompt, model_name, temperature, api_key, model_kwargs=None, tracer: dict = {}
+    g,
+    messages,
+    system_prompt,
+    model_name,
+    temperature,
+    api_key,
+    api_base_url=None,
+    model_kwargs=None,
+    tracer: dict = {},
 ):
     try:
-        client = genai.Client(api_key=api_key)
+        client = gemini_clients.get(api_key)
+        if not client:
+            client = get_gemini_client(api_key, api_base_url)
+            gemini_clients[api_key] = client
+
+        seed = int(os.getenv("KHOJ_LLM_SEED")) if os.getenv("KHOJ_LLM_SEED") else None
         config = gtypes.GenerateContentConfig(
             system_instruction=system_prompt,
             temperature=temperature,
             max_output_tokens=MAX_OUTPUT_TOKENS_GEMINI,
             stop_sequences=["Notes:\n["],
             safety_settings=SAFETY_SETTINGS,
+            seed=seed,
         )
 
         aggregated_response = ""

khoj/processor/conversation/openai/utils.py

@@ -55,7 +55,7 @@ def completion_with_backoff(
     tracer: dict = {},
 ) -> str:
     client_key = f"{openai_api_key}--{api_base_url}"
-    client: openai.OpenAI | None = openai_clients.get(client_key)
+    client = openai_clients.get(client_key)
     if not client:
         client = get_openai_client(openai_api_key, api_base_url)
         openai_clients[client_key] = client
@@ -150,9 +150,8 @@ def llm_thread(
 ):
     try:
         client_key = f"{openai_api_key}--{api_base_url}"
-        if client_key in openai_clients:
-            client = openai_clients[client_key]
-        else:
+        client = openai_clients.get(client_key)
+        if not client:
             client = get_openai_client(openai_api_key, api_base_url)
             openai_clients[client_key] = client
 
@@ -247,4 +246,6 @@ def get_openai_api_json_support(model_name: str, api_base_url: str = None) -> Js
         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/routers/api.py

@@ -463,12 +463,14 @@ async def extract_references_and_questions(
             )
         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_model_name = chat_model.name
             inferred_queries = extract_questions_anthropic(
                 defiltered_query,
                 query_images=query_images,
                 model=chat_model_name,
                 api_key=api_key,
+                api_base_url=api_base_url,
                 conversation_log=meta_log,
                 location_data=location_data,
                 user=user,
@@ -479,12 +481,14 @@ async def extract_references_and_questions(
             )
         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_model_name = chat_model.name
             inferred_queries = extract_questions_gemini(
                 defiltered_query,
                 query_images=query_images,
                 model=chat_model_name,
                 api_key=api_key,
+                api_base_url=api_base_url,
                 conversation_log=meta_log,
                 location_data=location_data,
                 max_tokens=chat_model.max_prompt_size,

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

@@ -1220,6 +1220,7 @@ async def send_message_to_model_wrapper(
         )
     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,
@@ -1239,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,
@@ -1262,6 +1265,7 @@ async def send_message_to_model_wrapper(
             model=chat_model_name,
             response_type=response_type,
             response_schema=response_schema,
+            api_base_url=api_base_url,
             tracer=tracer,
         )
     else:
@@ -1328,7 +1332,7 @@ 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,
@@ -1338,10 +1342,9 @@ def send_message_to_model_wrapper_sync(
             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,
@@ -1356,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,
@@ -1363,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,
@@ -1377,6 +1382,7 @@ 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,
@@ -1510,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,
@@ -1519,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,
@@ -1536,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,
@@ -1544,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/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.37.0.dist-info → khoj-1.37.1.dev12.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: khoj
-Version: 1.37.0
+Version: 1.37.1.dev12
 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'