khoj 1.28.3__py3-none-any.whl → 1.28.4.dev92__py3-none-any.whl

khoj/processor/tools/run_code.py

@@ -1,14 +1,16 @@
-import asyncio
+import base64
 import datetime
 import json
 import logging
+import mimetypes
 import os
-from typing import Any, Callable, List, Optional
+from pathlib import Path
+from typing import Any, Callable, List, NamedTuple, Optional
 
 import aiohttp
 
-from khoj.database.adapters import ais_user_subscribed
-from khoj.database.models import Agent, KhojUser
+from khoj.database.adapters import FileObjectAdapters
+from khoj.database.models import Agent, FileObject, KhojUser
 from khoj.processor.conversation import prompts
 from khoj.processor.conversation.utils import (
     ChatEvent,
@@ -17,7 +19,7 @@ from khoj.processor.conversation.utils import (
     construct_chat_history,
 )
 from khoj.routers.helpers import send_message_to_model_wrapper
-from khoj.utils.helpers import timer
+from khoj.utils.helpers import is_none_or_empty, timer
 from khoj.utils.rawconfig import LocationData
 
 logger = logging.getLogger(__name__)
@@ -26,6 +28,12 @@ logger = logging.getLogger(__name__)
 SANDBOX_URL = os.getenv("KHOJ_TERRARIUM_URL", "http://localhost:8080")
 
 
+class GeneratedCode(NamedTuple):
+    code: str
+    input_files: List[str]
+    input_links: List[str]
+
+
 async def run_code(
     query: str,
     conversation_history: dict,
@@ -36,15 +44,16 @@ async def run_code(
     query_images: List[str] = None,
     agent: Agent = None,
     sandbox_url: str = SANDBOX_URL,
+    query_files: str = None,
     tracer: dict = {},
 ):
     # Generate Code
     if send_status_func:
-        async for event in send_status_func(f"**Generate code snippets** for {query}"):
+        async for event in send_status_func(f"**Generate code snippet** for {query}"):
             yield {ChatEvent.STATUS: event}
     try:
         with timer("Chat actor: Generate programs to execute", logger):
-            codes = await generate_python_code(
+            generated_code = await generate_python_code(
                 query,
                 conversation_history,
                 context,
@@ -53,19 +62,31 @@ async def run_code(
                 query_images,
                 agent,
                 tracer,
+                query_files,
             )
     except Exception as e:
         raise ValueError(f"Failed to generate code for {query} with error: {e}")
 
+    # 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:
+        input_data.append(
+            {
+                "filename": os.path.basename(f.file_name),
+                "b64_data": base64.b64encode(f.raw_text.encode("utf-8")).decode("utf-8"),
+            }
+        )
+
     # Run Code
     if send_status_func:
-        async for event in send_status_func(f"**Running {len(codes)} code snippets**"):
+        async for event in send_status_func(f"**Running code snippet**"):
             yield {ChatEvent.STATUS: event}
     try:
-        tasks = [execute_sandboxed_python(code, sandbox_url) for code in codes]
-        with timer("Chat actor: Execute generated programs", logger):
-            results = await asyncio.gather(*tasks)
-        for result in results:
+        with timer("Chat actor: Execute generated program", logger, log_level=logging.INFO):
+            result = await execute_sandboxed_python(generated_code.code, input_data, sandbox_url)
             code = result.pop("code")
             logger.info(f"Executed Code:\n--@@--\n{code}\n--@@--Result:\n--@@--\n{result}\n--@@--")
             yield {query: {"code": code, "results": result}}
@@ -79,13 +100,13 @@ async def generate_python_code(
     context: str,
     location_data: LocationData,
     user: KhojUser,
-    query_images: List[str] = None,
+    query_images: list[str] = None,
     agent: Agent = None,
     tracer: dict = {},
-) -> List[str]:
+    query_files: str = None,
+) -> GeneratedCode:
     location = f"{location_data}" if location_data else "Unknown"
     username = prompts.user_name.format(name=user.get_full_name()) if user.get_full_name() else ""
-    subscribed = await ais_user_subscribed(user)
     chat_history = construct_chat_history(conversation_history)
 
     utc_date = datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%d")
@@ -109,32 +130,45 @@ async def generate_python_code(
         response_type="json_object",
         user=user,
         tracer=tracer,
+        query_files=query_files,
     )
 
     # Validate that the response is a non-empty, JSON-serializable list
     response = clean_json(response)
     response = json.loads(response)
-    codes = [code.strip() for code in response["codes"] if code.strip()]
+    code = response.get("code", "").strip()
+    input_files = response.get("input_files", [])
+    input_links = response.get("input_links", [])
 
-    if not isinstance(codes, list) or not codes or len(codes) == 0:
+    if not isinstance(code, str) or is_none_or_empty(code):
         raise ValueError
-    return codes
+    return GeneratedCode(code, input_files, input_links)
 
 
-async def execute_sandboxed_python(code: str, sandbox_url: str = SANDBOX_URL) -> dict[str, Any]:
+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.
     Returns the result of the code execution as a dictionary.
+
+    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}
+    data = {"code": cleaned_code, "files": input_data}
 
     async with aiohttp.ClientSession() as session:
         async with session.post(sandbox_url, json=data, headers=headers) as response:
             if response.status == 200:
                 result: dict[str, Any] = await response.json()
                 result["code"] = cleaned_code
+                # Store decoded output files
+                for output_file in result.get("output_files", []):
+                    # Decode text files as UTF-8
+                    if mimetypes.guess_type(output_file["filename"])[0].startswith("text/") or Path(
+                        output_file["filename"]
+                    ).suffix in [".org", ".md", ".json"]:
+                        output_file["b64_data"] = base64.b64decode(output_file["b64_data"]).decode("utf-8")
                 return result
             else:
                 return {

khoj/routers/api.py

@@ -6,7 +6,7 @@ import os
 import threading
 import time
 import uuid
-from typing import Any, Callable, List, Optional, Union
+from typing import Any, Callable, List, Optional, Set, Union
 
 import cron_descriptor
 import pytz
@@ -26,7 +26,6 @@ from khoj.database.adapters import (
     ConversationAdapters,
     EntryAdapters,
     get_default_search_model,
-    get_user_default_search_model,
     get_user_photo,
 )
 from khoj.database.models import (
@@ -151,7 +150,7 @@ async def execute_search(
     encoded_asymmetric_query = None
     if t != SearchType.Image:
         with timer("Encoding query took", logger=logger):
-            search_model = await sync_to_async(get_user_default_search_model)(user)
+            search_model = await sync_to_async(get_default_search_model)()
             encoded_asymmetric_query = state.embeddings_model[search_model.name].embed_query(defiltered_query)
 
     with concurrent.futures.ThreadPoolExecutor() as executor:
@@ -213,7 +212,7 @@ def update(
         logger.warning(error_msg)
         raise HTTPException(status_code=500, detail=error_msg)
     try:
-        initialize_content(regenerate=force, search_type=t, user=user)
+        initialize_content(user=user, regenerate=force, search_type=t)
     except Exception as e:
         error_msg = f"🚨 Failed to update server via API: {e}"
         logger.error(error_msg, exc_info=True)
@@ -350,7 +349,9 @@ async def extract_references_and_questions(
     location_data: LocationData = None,
     send_status_func: Optional[Callable] = None,
     query_images: Optional[List[str]] = None,
+    previous_inferred_queries: Set = set(),
     agent: Agent = None,
+    query_files: str = None,
     tracer: dict = {},
 ):
     user = request.user.object if request.user.is_authenticated else None
@@ -425,6 +426,7 @@ async def extract_references_and_questions(
                 user=user,
                 max_prompt_size=conversation_config.max_prompt_size,
                 personality_context=personality_context,
+                query_files=query_files,
                 tracer=tracer,
             )
         elif conversation_config.model_type == ChatModelOptions.ModelType.OPENAI:
@@ -443,6 +445,7 @@ async def extract_references_and_questions(
                 query_images=query_images,
                 vision_enabled=vision_enabled,
                 personality_context=personality_context,
+                query_files=query_files,
                 tracer=tracer,
             )
         elif conversation_config.model_type == ChatModelOptions.ModelType.ANTHROPIC:
@@ -458,6 +461,7 @@ async def extract_references_and_questions(
                 user=user,
                 vision_enabled=vision_enabled,
                 personality_context=personality_context,
+                query_files=query_files,
                 tracer=tracer,
             )
         elif conversation_config.model_type == ChatModelOptions.ModelType.GOOGLE:
@@ -474,10 +478,12 @@ async def extract_references_and_questions(
                 user=user,
                 vision_enabled=vision_enabled,
                 personality_context=personality_context,
+                query_files=query_files,
                 tracer=tracer,
             )
 
     # Collate search results as context for GPT
+    inferred_queries = list(set(inferred_queries) - previous_inferred_queries)
     with timer("Searching knowledge base took", logger):
         search_results = []
         logger.info(f"🔍 Searching knowledge base with queries: {inferred_queries}")

khoj/routers/api_agents.py

@@ -183,7 +183,7 @@ async def delete_agent(
 
 
 @api_agents.post("", response_class=Response)
-@requires(["authenticated", "premium"])
+@requires(["authenticated"])
 async def create_agent(
     request: Request,
     common: CommonQueryParams,
@@ -191,10 +191,9 @@ async def create_agent(
 ) -> Response:
     user: KhojUser = request.user.object
 
-    is_safe_prompt, reason = True, ""
-
-    if body.privacy_level != Agent.PrivacyLevel.PRIVATE:
-        is_safe_prompt, reason = await acheck_if_safe_prompt(body.persona)
+    is_safe_prompt, reason = await acheck_if_safe_prompt(
+        body.persona, user, lax=body.privacy_level == Agent.PrivacyLevel.PRIVATE
+    )
 
     if not is_safe_prompt:
         return Response(
@@ -236,7 +235,7 @@ async def create_agent(
 
 
 @api_agents.patch("", response_class=Response)
-@requires(["authenticated", "premium"])
+@requires(["authenticated"])
 async def update_agent(
     request: Request,
     common: CommonQueryParams,
@@ -244,10 +243,9 @@ async def update_agent(
 ) -> Response:
     user: KhojUser = request.user.object
 
-    is_safe_prompt, reason = True, ""
-
-    if body.privacy_level != Agent.PrivacyLevel.PRIVATE:
-        is_safe_prompt, reason = await acheck_if_safe_prompt(body.persona)
+    is_safe_prompt, reason = await acheck_if_safe_prompt(
+        body.persona, user, lax=body.privacy_level == Agent.PrivacyLevel.PRIVATE
+    )
 
     if not is_safe_prompt:
         return Response(

khoj/routers/api_chat.py

@@ -19,7 +19,6 @@ from khoj.database.adapters import (
     AgentAdapters,
     ConversationAdapters,
     EntryAdapters,
-    FileObjectAdapters,
     PublicConversationAdapters,
     aget_user_name,
 )
@@ -28,7 +27,11 @@ from khoj.processor.conversation.prompts import help_message, no_entries_found
 from khoj.processor.conversation.utils import defilter_query, save_to_conversation_log
 from khoj.processor.image.generate import text_to_image
 from khoj.processor.speech.text_to_speech import generate_text_to_speech
-from khoj.processor.tools.online_search import read_webpages, search_online
+from khoj.processor.tools.online_search import (
+    deduplicate_organic_results,
+    read_webpages,
+    search_online,
+)
 from khoj.processor.tools.run_code import run_code
 from khoj.routers.api import extract_references_and_questions
 from khoj.routers.email import send_query_feedback
@@ -41,12 +44,13 @@ from khoj.routers.helpers import (
     ConversationCommandRateLimiter,
     DeleteMessageRequestBody,
     FeedbackData,
+    acreate_title_from_history,
     agenerate_chat_response,
     aget_relevant_information_sources,
     aget_relevant_output_modes,
     construct_automation_created_message,
     create_automation,
-    extract_relevant_info,
+    gather_raw_query_files,
     generate_excalidraw_diagram,
     generate_summary_from_files,
     get_conversation_command,
@@ -72,7 +76,12 @@ from khoj.utils.helpers import (
     get_device,
     is_none_or_empty,
 )
-from khoj.utils.rawconfig import FileFilterRequest, FilesFilterRequest, LocationData
+from khoj.utils.rawconfig import (
+    ChatRequestBody,
+    FileFilterRequest,
+    FilesFilterRequest,
+    LocationData,
+)
 
 # Initialize Router
 logger = logging.getLogger(__name__)
@@ -370,7 +379,7 @@ def fork_public_conversation(
             {
                 "status": "ok",
                 "next_url": redirect_uri,
-                "conversation_id": new_conversation.id,
+                "conversation_id": str(new_conversation.id),
             }
         ),
     )
@@ -526,6 +535,32 @@ async def set_conversation_title(
     )
 
 
+@api_chat.post("/title")
+@requires(["authenticated"])
+async def generate_chat_title(
+    request: Request,
+    common: CommonQueryParams,
+    conversation_id: str,
+):
+    user: KhojUser = request.user.object
+    conversation = await ConversationAdapters.aget_conversation_by_user(user=user, conversation_id=conversation_id)
+
+    # Conversation.title is explicitly set by the user. Do not override.
+    if conversation.title:
+        return {"status": "ok", "title": conversation.title}
+
+    if not conversation:
+        raise HTTPException(status_code=404, detail="Conversation not found")
+
+    new_title = await acreate_title_from_history(request.user.object, conversation=conversation)
+
+    conversation.slug = new_title
+
+    await conversation.asave()
+
+    return {"status": "ok", "title": new_title}
+
+
 @api_chat.delete("/conversation/message", response_class=Response)
 @requires(["authenticated"])
 def delete_message(request: Request, delete_request: DeleteMessageRequestBody) -> Response:
@@ -567,6 +602,7 @@ async def chat(
     country_code = body.country_code or get_country_code_from_timezone(body.timezone)
     timezone = body.timezone
     raw_images = body.images
+    raw_query_files = body.files
 
     async def event_generator(q: str, images: list[str]):
         start_time = time.perf_counter()
@@ -578,6 +614,7 @@ async def chat(
         q = unquote(q)
         train_of_thought = []
         nonlocal conversation_id
+        nonlocal raw_query_files
 
         tracer: dict = {
             "mid": turn_id,
@@ -597,6 +634,11 @@ async def chat(
                 if uploaded_image:
                     uploaded_images.append(uploaded_image)
 
+        query_files: Dict[str, str] = {}
+        if raw_query_files:
+            for file in raw_query_files:
+                query_files[file.name] = file.content
+
         async def send_event(event_type: ChatEvent, data: str | dict):
             nonlocal connection_alive, ttft, train_of_thought
             if not connection_alive or await request.is_disconnected():
@@ -707,6 +749,8 @@ async def chat(
         ## Extract Document References
         compiled_references: List[Any] = []
         inferred_queries: List[Any] = []
+        file_filters = conversation.file_filters if conversation and conversation.file_filters else []
+        attached_file_context = gather_raw_query_files(query_files)
 
         if conversation_commands == [ConversationCommand.Default] or is_automated_task:
             conversation_commands = await aget_relevant_information_sources(
@@ -716,6 +760,7 @@ async def chat(
                 user=user,
                 query_images=uploaded_images,
                 agent=agent,
+                query_files=attached_file_context,
                 tracer=tracer,
             )
 
@@ -738,8 +783,13 @@ async def chat(
                 conversation_commands.append(mode)
 
         for cmd in conversation_commands:
-            await conversation_command_rate_limiter.update_and_check_if_valid(request, cmd)
-            q = q.replace(f"/{cmd.value}", "").strip()
+            try:
+                await conversation_command_rate_limiter.update_and_check_if_valid(request, cmd)
+                q = q.replace(f"/{cmd.value}", "").strip()
+            except HTTPException as e:
+                async for result in send_llm_response(str(e.detail)):
+                    yield result
+                return
 
         defiltered_query = defilter_query(q)
 
@@ -756,6 +806,7 @@ async def chat(
                 user_name=user_name,
                 location=location,
                 file_filters=conversation.file_filters if conversation else [],
+                query_files=attached_file_context,
                 tracer=tracer,
             ):
                 if isinstance(research_result, InformationCollectionIteration):
@@ -773,11 +824,8 @@ async def chat(
                     yield research_result
 
             # researched_results = await extract_relevant_info(q, researched_results, agent)
-            logger.info(f"Researched Results: {researched_results}")
-
-        for cmd in conversation_commands:
-            await conversation_command_rate_limiter.update_and_check_if_valid(request, cmd)
-            q = q.replace(f"/{cmd.value}", "").strip()
+            if state.verbose > 1:
+                logger.debug(f"Researched Results: {researched_results}")
 
         used_slash_summarize = conversation_commands == [ConversationCommand.Summarize]
         file_filters = conversation.file_filters if conversation else []
@@ -798,10 +846,6 @@ async def chat(
                 response_log = "No files selected for summarization. Please add files using the section on the left."
                 async for result in send_llm_response(response_log):
                     yield result
-            elif len(file_filters) > 1 and not agent_has_entries:
-                response_log = "Only one file can be selected for summarization."
-                async for result in send_llm_response(response_log):
-                    yield result
             else:
                 async for response in generate_summary_from_files(
                     q=q,
@@ -811,6 +855,7 @@ async def chat(
                     query_images=uploaded_images,
                     agent=agent,
                     send_status_func=partial(send_event, ChatEvent.STATUS),
+                    query_files=attached_file_context,
                     tracer=tracer,
                 ):
                     if isinstance(response, dict) and ChatEvent.STATUS in response:
@@ -831,8 +876,9 @@ async def chat(
                 client_application=request.user.client_app,
                 conversation_id=conversation_id,
                 query_images=uploaded_images,
-                tracer=tracer,
                 train_of_thought=train_of_thought,
+                raw_query_files=raw_query_files,
+                tracer=tracer,
             )
             return
 
@@ -876,8 +922,9 @@ async def chat(
                 inferred_queries=[query_to_run],
                 automation_id=automation.id,
                 query_images=uploaded_images,
-                tracer=tracer,
                 train_of_thought=train_of_thought,
+                raw_query_files=raw_query_files,
+                tracer=tracer,
             )
             async for result in send_llm_response(llm_response):
                 yield result
@@ -899,6 +946,7 @@ async def chat(
                     partial(send_event, ChatEvent.STATUS),
                     query_images=uploaded_images,
                     agent=agent,
+                    query_files=attached_file_context,
                     tracer=tracer,
                 ):
                     if isinstance(result, dict) and ChatEvent.STATUS in result:
@@ -944,6 +992,7 @@ async def chat(
                     custom_filters,
                     query_images=uploaded_images,
                     agent=agent,
+                    query_files=attached_file_context,
                     tracer=tracer,
                 ):
                     if isinstance(result, dict) and ChatEvent.STATUS in result:
@@ -969,6 +1018,7 @@ async def chat(
                     partial(send_event, ChatEvent.STATUS),
                     query_images=uploaded_images,
                     agent=agent,
+                    query_files=attached_file_context,
                     tracer=tracer,
                 ):
                     if isinstance(result, dict) and ChatEvent.STATUS in result:
@@ -1009,6 +1059,7 @@ async def chat(
                     partial(send_event, ChatEvent.STATUS),
                     query_images=uploaded_images,
                     agent=agent,
+                    query_files=attached_file_context,
                     tracer=tracer,
                 ):
                     if isinstance(result, dict) and ChatEvent.STATUS in result:
@@ -1024,12 +1075,13 @@ async def chat(
                 )
 
         ## Send Gathered References
+        unique_online_results = deduplicate_organic_results(online_results)
         async for result in send_event(
             ChatEvent.REFERENCES,
             {
                 "inferredQueries": inferred_queries,
                 "context": compiled_references,
-                "onlineContext": online_results,
+                "onlineContext": unique_online_results,
                 "codeContext": code_results,
             },
         ):
@@ -1048,6 +1100,7 @@ async def chat(
                 send_status_func=partial(send_event, ChatEvent.STATUS),
                 query_images=uploaded_images,
                 agent=agent,
+                query_files=attached_file_context,
                 tracer=tracer,
             ):
                 if isinstance(result, dict) and ChatEvent.STATUS in result:
@@ -1078,9 +1131,11 @@ async def chat(
                 conversation_id=conversation_id,
                 compiled_references=compiled_references,
                 online_results=online_results,
+                code_results=code_results,
                 query_images=uploaded_images,
-                tracer=tracer,
                 train_of_thought=train_of_thought,
+                raw_query_files=raw_query_files,
+                tracer=tracer,
             )
             content_obj = {
                 "intentType": intent_type,
@@ -1109,14 +1164,20 @@ async def chat(
                 user=user,
                 agent=agent,
                 send_status_func=partial(send_event, ChatEvent.STATUS),
+                query_files=attached_file_context,
                 tracer=tracer,
             ):
                 if isinstance(result, dict) and ChatEvent.STATUS in result:
                     yield result[ChatEvent.STATUS]
                 else:
                     better_diagram_description_prompt, excalidraw_diagram_description = result
-                    inferred_queries.append(better_diagram_description_prompt)
-                    diagram_description = excalidraw_diagram_description
+                    if better_diagram_description_prompt and excalidraw_diagram_description:
+                        inferred_queries.append(better_diagram_description_prompt)
+                        diagram_description = excalidraw_diagram_description
+                    else:
+                        async for result in send_llm_response(f"Failed to generate diagram. Please try again later."):
+                            yield result
+                        return
 
             content_obj = {
                 "intentType": intent_type,
@@ -1136,9 +1197,11 @@ async def chat(
                 conversation_id=conversation_id,
                 compiled_references=compiled_references,
                 online_results=online_results,
+                code_results=code_results,
                 query_images=uploaded_images,
-                tracer=tracer,
                 train_of_thought=train_of_thought,
+                raw_query_files=raw_query_files,
+                tracer=tracer,
             )
 
             async for result in send_llm_response(json.dumps(content_obj)):
@@ -1164,8 +1227,10 @@ async def chat(
             user_name,
             researched_results,
             uploaded_images,
-            tracer,
             train_of_thought,
+            attached_file_context,
+            raw_query_files,
+            tracer,
         )
 
         # Send Response