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

khoj/processor/tools/run_code.py

@@ -1,12 +1,23 @@
+import asyncio
 import base64
 import datetime
 import logging
 import mimetypes
 import os
+import re
 from pathlib import Path
 from typing import Any, Callable, List, NamedTuple, Optional
 
 import aiohttp
+from asgiref.sync import sync_to_async
+from httpx import RemoteProtocolError
+from tenacity import (
+    before_sleep_log,
+    retry,
+    retry_if_exception_type,
+    stop_after_attempt,
+    wait_random_exponential,
+)
 
 from khoj.database.adapters import FileObjectAdapters
 from khoj.database.models import Agent, FileObject, KhojUser
@@ -15,22 +26,26 @@ from khoj.processor.conversation.utils import (
     ChatEvent,
     clean_code_python,
     construct_chat_history,
-    load_complex_json,
 )
 from khoj.routers.helpers import send_message_to_model_wrapper
-from khoj.utils.helpers import is_none_or_empty, timer, truncate_code_context
+from khoj.utils.helpers import (
+    is_e2b_code_sandbox_enabled,
+    is_none_or_empty,
+    timer,
+    truncate_code_context,
+)
 from khoj.utils.rawconfig import LocationData
 
 logger = logging.getLogger(__name__)
 
 
 SANDBOX_URL = os.getenv("KHOJ_TERRARIUM_URL", "http://localhost:8080")
+DEFAULT_E2B_TEMPLATE = "pmt2o0ghpang8gbiys57"
 
 
 class GeneratedCode(NamedTuple):
     code: str
-    input_files: List[str]
-    input_links: List[str]
+    input_files: List[FileObject]
 
 
 async def run_code(
@@ -68,13 +83,10 @@ async def run_code(
 
     # Prepare Input Data
     input_data = []
-    user_input_files: List[FileObject] = []
-    for input_file in generated_code.input_files:
-        user_input_files += await FileObjectAdapters.aget_file_objects_by_name(user, input_file)
-    for f in user_input_files:
+    for f in generated_code.input_files:
         input_data.append(
             {
-                "filename": os.path.basename(f.file_name),
+                "filename": f.file_name,
                 "b64_data": base64.b64encode(f.raw_text.encode("utf-8")).decode("utf-8"),
             }
         )
@@ -90,6 +102,14 @@ async def run_code(
             cleaned_result = truncate_code_context({"cleaned": {"results": result}})["cleaned"]["results"]
             logger.info(f"Executed Code\n----\n{code}\n----\nResult\n----\n{cleaned_result}\n----")
             yield {query: {"code": code, "results": result}}
+    except asyncio.TimeoutError as e:
+        # Call the sandbox_url/stop GET API endpoint to stop the code sandbox
+        error = f"Failed to run code for {query} with Timeout error: {e}"
+        try:
+            await aiohttp.ClientSession().get(f"{sandbox_url}/stop", timeout=5)
+        except Exception as e:
+            error += f"\n\nFailed to stop code sandbox with error: {e}"
+        raise ValueError(error)
     except Exception as e:
         raise ValueError(f"Failed to run code for {query} with error: {e}")
 
@@ -114,6 +134,12 @@ async def generate_python_code(
         prompts.personality_context.format(personality=agent.personality) if agent and agent.personality else ""
     )
 
+    # add sandbox specific context like available packages
+    sandbox_context = (
+        prompts.e2b_sandbox_context if is_e2b_code_sandbox_enabled() else prompts.terrarium_sandbox_context
+    )
+    personality_context = f"{sandbox_context}\n{personality_context}"
+
     code_generation_prompt = prompts.python_code_generation_prompt.format(
         current_date=utc_date,
         query=q,
@@ -127,23 +153,50 @@ async def generate_python_code(
     response = await send_message_to_model_wrapper(
         code_generation_prompt,
         query_images=query_images,
-        response_type="json_object",
         user=user,
         tracer=tracer,
         query_files=query_files,
     )
 
-    # Validate that the response is a non-empty, JSON-serializable list
-    response = load_complex_json(response)
-    code = response.get("code", "").strip()
-    input_files = response.get("input_files", [])
-    input_links = response.get("input_links", [])
+    # Extract python code wrapped in markdown code blocks from the response
+    code_blocks = re.findall(r"```(?:python)?\n(.*?)\n```", response, re.DOTALL)
+
+    if not code_blocks:
+        raise ValueError("No Python code blocks found in response")
+
+    # Join multiple code blocks with newlines and strip any leading/trailing whitespace
+    code = "\n".join(code_blocks).strip()
 
     if not isinstance(code, str) or is_none_or_empty(code):
         raise ValueError
-    return GeneratedCode(code, input_files, input_links)
 
+    # Infer user files required in sandbox based on user file paths mentioned in code
+    input_files: List[FileObject] = []
+    user_files = await sync_to_async(set)(FileObjectAdapters.get_all_file_objects(user))
+    for user_file in user_files:
+        if user_file.file_name in code:
+            # Replace references to full file path used in code with just the file basename to ease reference in sandbox
+            file_basename = os.path.basename(user_file.file_name)
+            code = code.replace(user_file.file_name, file_basename)
+            user_file.file_name = file_basename
+            input_files.append(user_file)
+
+    return GeneratedCode(code, input_files)
 
+
+@retry(
+    retry=(
+        retry_if_exception_type(aiohttp.ClientError)
+        | retry_if_exception_type(aiohttp.ClientTimeout)
+        | retry_if_exception_type(asyncio.TimeoutError)
+        | retry_if_exception_type(ConnectionError)
+        | retry_if_exception_type(RemoteProtocolError)
+    ),
+    wait=wait_random_exponential(min=1, max=5),
+    stop=stop_after_attempt(3),
+    before_sleep=before_sleep_log(logger, logging.DEBUG),
+    reraise=True,
+)
 async def execute_sandboxed_python(code: str, input_data: list[dict], sandbox_url: str = SANDBOX_URL) -> dict[str, Any]:
     """
     Takes code to run as a string and calls the terrarium API to execute it.
@@ -152,15 +205,104 @@ async def execute_sandboxed_python(code: str, input_data: list[dict], sandbox_ur
     Reference data i/o format based on Terrarium example client code at:
     https://github.com/cohere-ai/cohere-terrarium/blob/main/example-clients/python/terrarium_client.py
     """
-    headers = {"Content-Type": "application/json"}
     cleaned_code = clean_code_python(code)
-    data = {"code": cleaned_code, "files": input_data}
+    if is_e2b_code_sandbox_enabled():
+        try:
+            return await execute_e2b(cleaned_code, input_data)
+        except ImportError:
+            pass
+    return await execute_terrarium(cleaned_code, input_data, sandbox_url)
+
+
+async def execute_e2b(code: str, input_files: list[dict]) -> dict[str, Any]:
+    """Execute code and handle file I/O in e2b sandbox"""
+    from e2b_code_interpreter import AsyncSandbox
+
+    sandbox = await AsyncSandbox.create(
+        api_key=os.getenv("E2B_API_KEY"),
+        template=os.getenv("E2B_TEMPLATE", DEFAULT_E2B_TEMPLATE),
+        timeout=120,
+        request_timeout=30,
+    )
 
+    try:
+        # Upload input files in parallel
+        upload_tasks = [
+            sandbox.files.write(path=file["filename"], data=base64.b64decode(file["b64_data"]), request_timeout=30)
+            for file in input_files
+        ]
+        await asyncio.gather(*upload_tasks)
+
+        # Note stored files before execution to identify new files created during execution
+        E2bFile = NamedTuple("E2bFile", [("name", str), ("path", str)])
+        original_files = {E2bFile(f.name, f.path) for f in await sandbox.files.list("~")}
+
+        # Execute code from main.py file
+        execution = await sandbox.run_code(code=code, timeout=60)
+
+        # Collect output files
+        output_files = []
+
+        # Identify new files created during execution
+        new_files = set(E2bFile(f.name, f.path) for f in await sandbox.files.list("~")) - original_files
+        # Read newly created files in parallel
+        download_tasks = [sandbox.files.read(f.path, request_timeout=30) for f in new_files]
+        downloaded_files = await asyncio.gather(*download_tasks)
+        for f, content in zip(new_files, downloaded_files):
+            if isinstance(content, bytes):
+                # Binary files like PNG - encode as base64
+                b64_data = base64.b64encode(content).decode("utf-8")
+            elif Path(f.name).suffix in [".png", ".jpeg", ".jpg", ".svg"]:
+                # Ignore image files as they are extracted from execution results below for inline display
+                continue
+            else:
+                # Text files - encode utf-8 string as base64
+                b64_data = base64.b64encode(content.encode("utf-8")).decode("utf-8")
+            output_files.append({"filename": f.name, "b64_data": b64_data})
+
+        # Collect output files from execution results
+        for idx, result in enumerate(execution.results):
+            for result_type in {"png", "jpeg", "svg", "text", "markdown", "json"}:
+                if b64_data := getattr(result, result_type, None):
+                    output_files.append({"filename": f"{idx}.{result_type}", "b64_data": b64_data})
+                    break
+
+        # collect logs
+        success = not execution.error and not execution.logs.stderr
+        stdout = "\n".join(execution.logs.stdout)
+        errors = "\n".join(execution.logs.stderr)
+        if execution.error:
+            errors = f"{execution.error}\n{errors}"
+
+        return {
+            "code": code,
+            "success": success,
+            "std_out": stdout,
+            "std_err": errors,
+            "output_files": output_files,
+        }
+    except Exception as e:
+        return {
+            "code": code,
+            "success": False,
+            "std_err": f"Sandbox failed to execute code: {str(e)}",
+            "output_files": [],
+        }
+
+
+async def execute_terrarium(
+    code: str,
+    input_data: list[dict],
+    sandbox_url: str,
+) -> dict[str, Any]:
+    """Execute code using Terrarium sandbox"""
+    headers = {"Content-Type": "application/json"}
+    data = {"code": code, "files": input_data}
     async with aiohttp.ClientSession() as session:
-        async with session.post(sandbox_url, json=data, headers=headers) as response:
+        async with session.post(sandbox_url, json=data, headers=headers, timeout=30) as response:
             if response.status == 200:
                 result: dict[str, Any] = await response.json()
-                result["code"] = cleaned_code
+                result["code"] = code
                 # Store decoded output files
                 result["output_files"] = result.get("output_files", [])
                 for output_file in result["output_files"]:
@@ -172,7 +314,7 @@ async def execute_sandboxed_python(code: str, input_data: list[dict], sandbox_ur
                 return result
             else:
                 return {
-                    "code": cleaned_code,
+                    "code": code,
                     "success": False,
                     "std_err": f"Failed to execute code with {response.status}",
                     "output_files": [],

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

@@ -18,7 +18,7 @@ default_offline_chat_models = [
     "bartowski/Qwen2.5-14B-Instruct-GGUF",
 ]
 default_openai_chat_models = ["gpt-4o-mini", "gpt-4o"]
-default_gemini_chat_models = ["gemini-1.5-flash", "gemini-1.5-pro"]
+default_gemini_chat_models = ["gemini-2.0-flash", "gemini-1.5-pro"]
 default_anthropic_chat_models = ["claude-3-5-sonnet-20241022", "claude-3-5-haiku-20241022"]
 
 empty_config = {
@@ -46,7 +46,11 @@ model_to_cost: Dict[str, Dict[str, float]] = {
     "gemini-1.5-flash-002": {"input": 0.075, "output": 0.30},
     "gemini-1.5-pro": {"input": 1.25, "output": 5.00},
     "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

@@ -321,6 +321,12 @@ def get_device() -> torch.device:
         return torch.device("cpu")
 
 
+def is_e2b_code_sandbox_enabled():
+    """Check if E2B code sandbox is enabled.
+    Set E2B_API_KEY environment variable to use it."""
+    return not is_none_or_empty(os.getenv("E2B_API_KEY"))
+
+
 class ConversationCommand(str, Enum):
     Default = "default"
     General = "general"
@@ -362,20 +368,23 @@ 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, 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 = {
     ConversationCommand.Default: "To use a mix of your internal knowledge and the user's personal knowledge, or if you don't entirely understand the query.",
     ConversationCommand.General: "To use when you can answer the question without any outside information or personal knowledge",
     ConversationCommand.Notes: "To search the user's personal knowledge base. Especially helpful if the question expects context from the user's notes or documents.",
     ConversationCommand.Online: "To search for the latest, up-to-date information from the internet. Note: **Questions about Khoj should always use this data source**",
     ConversationCommand.Webpage: "To use if the user has directly provided the webpage urls or you are certain of the webpage urls to read.",
-    ConversationCommand.Code: "To run Python code in a Pyodide sandbox with no network access. Helpful when need 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.",
+    ConversationCommand.Code: e2b_tool_description if is_e2b_code_sandbox_enabled() else terrarium_tool_description,
 }
 
 function_calling_description_for_llm = {
     ConversationCommand.Notes: "To search the user's personal knowledge base. Especially helpful if the question expects context from the user's notes or documents.",
     ConversationCommand.Online: "To search the internet for information. Useful to get a quick, broad overview from the internet. Provide all relevant context to ensure new searches, not in previous iterations, are performed.",
     ConversationCommand.Webpage: "To extract information from webpages. Useful for more detailed research from the internet. Usually used when you know the webpage links to refer to. Share the webpage links and information to extract in your query.",
-    ConversationCommand.Code: "To run Python code in a Pyodide sandbox with no network access. Helpful when need 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.",
+    ConversationCommand.Code: e2b_tool_description if is_e2b_code_sandbox_enabled() else terrarium_tool_description,
 }
 
 mode_descriptions_for_llm = {

khoj/utils/initialization.py

@@ -185,16 +185,18 @@ def initialization(interactive: bool = True):
         )
         provider_name = provider_name or model_type.name.capitalize()
 
-        default_use_model = {True: "y", False: "n"}[default_api_key is not None]
-
-        # If not in interactive mode & in the offline setting, it's most likely that we're running in a containerized environment. This usually means there's not enough RAM to load offline models directly within the application. In such cases, we default to not using the model -- it's recommended to use another service like Ollama to host the model locally in that case.
-        default_use_model = {True: "n", False: default_use_model}[is_offline]
+        default_use_model = default_api_key is not None
+        # If not in interactive mode & in the offline setting, it's most likely that we're running in a containerized environment.
+        # This usually means there's not enough RAM to load offline models directly within the application.
+        # In such cases, we default to not using the model -- it's recommended to use another service like Ollama to host the model locally in that case.
+        if is_offline:
+            default_use_model = False
 
         use_model_provider = (
-            default_use_model if not interactive else input(f"Add {provider_name} chat models? (y/n): ")
+            default_use_model if not interactive else input(f"Add {provider_name} chat models? (y/n): ") == "y"
         )
 
-        if use_model_provider != "y":
+        if not use_model_provider:
             return False, None
 
         logger.info(f"️💬 Setting up your {provider_name} chat configuration")
@@ -303,4 +305,19 @@ def initialization(interactive: bool = True):
                 logger.error(f"🚨 Failed to create chat configuration: {e}", exc_info=True)
     else:
         _update_chat_model_options()
-        logger.info("🗣️ Chat model configuration updated")
+        logger.info("🗣️ Chat model options updated")
+
+    # Update the default chat model if it doesn't match
+    chat_config = ConversationAdapters.get_default_chat_model()
+    env_default_chat_model = os.getenv("KHOJ_DEFAULT_CHAT_MODEL")
+    if not chat_config or not env_default_chat_model:
+        return
+    if chat_config.name != env_default_chat_model:
+        chat_model = ConversationAdapters.get_chat_model_by_name(env_default_chat_model)
+        if not chat_model:
+            logger.error(
+                f"🚨 Not setting default chat model. Chat model {env_default_chat_model} not found in existing chat model options."
+            )
+            return
+        ConversationAdapters.set_default_chat_model(chat_model)
+        logger.info(f"🗣️ Default chat model set to {chat_model.name}")

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: khoj
-Version: 1.36.7.dev7
+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
@@ -35,11 +35,12 @@ Requires-Dist: django-phonenumber-field==7.3.0
 Requires-Dist: django-unfold==0.42.0
 Requires-Dist: django==5.0.10
 Requires-Dist: docx2txt==0.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.25.0
+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
@@ -76,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'