khoj 1.27.2.dev29__py3-none-any.whl → 1.27.2.dev130__py3-none-any.whl

khoj/processor/tools/run_code.py

@@ -0,0 +1,144 @@
+import asyncio
+import datetime
+import json
+import logging
+import os
+from typing import Any, Callable, List, Optional
+
+import aiohttp
+
+from khoj.database.adapters import ais_user_subscribed
+from khoj.database.models import Agent, KhojUser
+from khoj.processor.conversation import prompts
+from khoj.processor.conversation.utils import (
+    ChatEvent,
+    clean_code_python,
+    clean_json,
+    construct_chat_history,
+)
+from khoj.routers.helpers import send_message_to_model_wrapper
+from khoj.utils.helpers import timer
+from khoj.utils.rawconfig import LocationData
+
+logger = logging.getLogger(__name__)
+
+
+SANDBOX_URL = os.getenv("KHOJ_TERRARIUM_URL", "http://localhost:8080")
+
+
+async def run_code(
+    query: str,
+    conversation_history: dict,
+    context: str,
+    location_data: LocationData,
+    user: KhojUser,
+    send_status_func: Optional[Callable] = None,
+    query_images: List[str] = None,
+    agent: Agent = None,
+    sandbox_url: str = SANDBOX_URL,
+    tracer: dict = {},
+):
+    # Generate Code
+    if send_status_func:
+        async for event in send_status_func(f"**Generate code snippets** for {query}"):
+            yield {ChatEvent.STATUS: event}
+    try:
+        with timer("Chat actor: Generate programs to execute", logger):
+            codes = await generate_python_code(
+                query,
+                conversation_history,
+                context,
+                location_data,
+                user,
+                query_images,
+                agent,
+                tracer,
+            )
+    except Exception as e:
+        raise ValueError(f"Failed to generate code for {query} with error: {e}")
+
+    # Run Code
+    if send_status_func:
+        async for event in send_status_func(f"**Running {len(codes)} code snippets**"):
+            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:
+            code = result.pop("code")
+            logger.info(f"Executed Code:\n--@@--\n{code}\n--@@--Result:\n--@@--\n{result}\n--@@--")
+            yield {query: {"code": code, "results": result}}
+    except Exception as e:
+        raise ValueError(f"Failed to run code for {query} with error: {e}")
+
+
+async def generate_python_code(
+    q: str,
+    conversation_history: dict,
+    context: str,
+    location_data: LocationData,
+    user: KhojUser,
+    query_images: List[str] = None,
+    agent: Agent = None,
+    tracer: dict = {},
+) -> List[str]:
+    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")
+    personality_context = (
+        prompts.personality_context.format(personality=agent.personality) if agent and agent.personality else ""
+    )
+
+    code_generation_prompt = prompts.python_code_generation_prompt.format(
+        current_date=utc_date,
+        query=q,
+        chat_history=chat_history,
+        context=context,
+        location=location,
+        username=username,
+        personality_context=personality_context,
+    )
+
+    response = await send_message_to_model_wrapper(
+        code_generation_prompt,
+        query_images=query_images,
+        response_type="json_object",
+        user=user,
+        tracer=tracer,
+    )
+
+    # 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()]
+
+    if not isinstance(codes, list) or not codes or len(codes) == 0:
+        raise ValueError
+    return codes
+
+
+async def execute_sandboxed_python(code: str, 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.
+    """
+    headers = {"Content-Type": "application/json"}
+    cleaned_code = clean_code_python(code)
+    data = {"code": cleaned_code}
+
+    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
+                return result
+            else:
+                return {
+                    "code": cleaned_code,
+                    "success": False,
+                    "std_err": f"Failed to execute code with {response.status}",
+                }

khoj/routers/api.py

@@ -44,6 +44,7 @@ from khoj.processor.conversation.offline.chat_model import extract_questions_off
 from khoj.processor.conversation.offline.whisper import transcribe_audio_offline
 from khoj.processor.conversation.openai.gpt import extract_questions
 from khoj.processor.conversation.openai.whisper import transcribe_audio
+from khoj.processor.conversation.utils import defilter_query
 from khoj.routers.helpers import (
     ApiUserRateLimiter,
     ChatEvent,
@@ -167,8 +168,8 @@ async def execute_search(
             search_futures += [
                 executor.submit(
                     text_search.query,
-                    user,
                     user_query,
+                    user,
                     t,
                     question_embedding=encoded_asymmetric_query,
                     max_distance=max_distance,
@@ -355,7 +356,7 @@ async def extract_references_and_questions(
     user = request.user.object if request.user.is_authenticated else None
 
     # Initialize Variables
-    compiled_references: List[Any] = []
+    compiled_references: List[dict[str, str]] = []
     inferred_queries: List[str] = []
 
     agent_has_entries = False
@@ -384,9 +385,7 @@ async def extract_references_and_questions(
             return
 
     # Extract filter terms from user message
-    defiltered_query = q
-    for filter in [DateFilter(), WordFilter(), FileFilter()]:
-        defiltered_query = filter.defilter(defiltered_query)
+    defiltered_query = defilter_query(q)
     filters_in_query = q.replace(defiltered_query, "").strip()
     conversation = await sync_to_async(ConversationAdapters.get_conversation_by_id)(conversation_id)
 
@@ -502,7 +501,8 @@ async def extract_references_and_questions(
             )
         search_results = text_search.deduplicated_search_responses(search_results)
         compiled_references = [
-            {"compiled": item.additional["compiled"], "file": item.additional["file"]} for item in search_results
+            {"query": q, "compiled": item.additional["compiled"], "file": item.additional["file"]}
+            for q, item in zip(inferred_queries, search_results)
         ]
 
     yield compiled_references, inferred_queries, defiltered_query

khoj/routers/api_chat.py

@@ -6,7 +6,7 @@ import time
 import uuid
 from datetime import datetime
 from functools import partial
-from typing import Dict, Optional
+from typing import Any, Dict, List, Optional
 from urllib.parse import unquote
 
 from asgiref.sync import sync_to_async
@@ -25,10 +25,11 @@ from khoj.database.adapters import (
 )
 from khoj.database.models import Agent, KhojUser
 from khoj.processor.conversation.prompts import help_message, no_entries_found
-from khoj.processor.conversation.utils import save_to_conversation_log
+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.run_code import run_code
 from khoj.routers.api import extract_references_and_questions
 from khoj.routers.helpers import (
     ApiImageRateLimiter,
@@ -42,8 +43,10 @@ from khoj.routers.helpers import (
     aget_relevant_output_modes,
     construct_automation_created_message,
     create_automation,
+    extract_relevant_info,
     extract_relevant_summary,
     generate_excalidraw_diagram,
+    generate_summary_from_files,
     get_conversation_command,
     is_query_empty,
     is_ready_to_chat,
@@ -51,6 +54,10 @@ from khoj.routers.helpers import (
     update_telemetry_state,
     validate_conversation_config,
 )
+from khoj.routers.research import (
+    InformationCollectionIteration,
+    execute_information_collection,
+)
 from khoj.routers.storage import upload_image_to_bucket
 from khoj.utils import state
 from khoj.utils.helpers import (
@@ -563,7 +570,9 @@ async def chat(
         user: KhojUser = request.user.object
         event_delimiter = "␃🔚␗"
         q = unquote(q)
+        train_of_thought = []
         nonlocal conversation_id
+
         tracer: dict = {
             "mid": f"{uuid.uuid4()}",
             "cid": conversation_id,
@@ -583,7 +592,7 @@ async def chat(
                     uploaded_images.append(uploaded_image)
 
         async def send_event(event_type: ChatEvent, data: str | dict):
-            nonlocal connection_alive, ttft
+            nonlocal connection_alive, ttft, train_of_thought
             if not connection_alive or await request.is_disconnected():
                 connection_alive = False
                 logger.warning(f"User {user} disconnected from {common.client} client")
@@ -591,8 +600,11 @@ async def chat(
             try:
                 if event_type == ChatEvent.END_LLM_RESPONSE:
                     collect_telemetry()
-                if event_type == ChatEvent.START_LLM_RESPONSE:
+                elif event_type == ChatEvent.START_LLM_RESPONSE:
                     ttft = time.perf_counter() - start_time
+                elif event_type == ChatEvent.STATUS:
+                    train_of_thought.append({"type": event_type.value, "data": data})
+
                 if event_type == ChatEvent.MESSAGE:
                     yield data
                 elif event_type == ChatEvent.REFERENCES or stream:
@@ -681,6 +693,14 @@ async def chat(
         meta_log = conversation.conversation_log
         is_automated_task = conversation_commands == [ConversationCommand.AutomatedTask]
 
+        researched_results = ""
+        online_results: Dict = dict()
+        code_results: Dict = dict()
+        ## Extract Document References
+        compiled_references: List[Any] = []
+        inferred_queries: List[Any] = []
+        defiltered_query = defilter_query(q)
+
         if conversation_commands == [ConversationCommand.Default] or is_automated_task:
             conversation_commands = await aget_relevant_information_sources(
                 q,
@@ -691,6 +711,11 @@ async def chat(
                 agent=agent,
                 tracer=tracer,
             )
+
+            # If we're doing research, we don't want to do anything else
+            if ConversationCommand.Research in conversation_commands:
+                conversation_commands = [ConversationCommand.Research]
+
             conversation_commands_str = ", ".join([cmd.value for cmd in conversation_commands])
             async for result in send_event(
                 ChatEvent.STATUS, f"**Chose Data Sources to Search:** {conversation_commands_str}"
@@ -705,6 +730,38 @@ async def chat(
             if mode not in conversation_commands:
                 conversation_commands.append(mode)
 
+        if conversation_commands == [ConversationCommand.Research]:
+            async for research_result in execute_information_collection(
+                request=request,
+                user=user,
+                query=defiltered_query,
+                conversation_id=conversation_id,
+                conversation_history=meta_log,
+                query_images=uploaded_images,
+                agent=agent,
+                send_status_func=partial(send_event, ChatEvent.STATUS),
+                user_name=user_name,
+                location=location,
+                file_filters=conversation.file_filters if conversation else [],
+                tracer=tracer,
+            ):
+                if isinstance(research_result, InformationCollectionIteration):
+                    if research_result.summarizedResult:
+                        if research_result.onlineContext:
+                            online_results.update(research_result.onlineContext)
+                        if research_result.codeContext:
+                            code_results.update(research_result.codeContext)
+                        if research_result.context:
+                            compiled_references.extend(research_result.context)
+
+                        researched_results += research_result.summarizedResult
+
+                else:
+                    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()
@@ -733,48 +790,24 @@ async def chat(
                 async for result in send_llm_response(response_log):
                     yield result
             else:
-                try:
-                    file_object = None
-                    if await EntryAdapters.aagent_has_entries(agent):
-                        file_names = await EntryAdapters.aget_agent_entry_filepaths(agent)
-                        if len(file_names) > 0:
-                            file_object = await FileObjectAdapters.async_get_file_objects_by_name(
-                                None, file_names[0], agent
-                            )
-
-                    if len(file_filters) > 0:
-                        file_object = await FileObjectAdapters.async_get_file_objects_by_name(user, file_filters[0])
-
-                    if len(file_object) == 0:
-                        response_log = "Sorry, I couldn't find the full text of this file. Please re-upload the document and try again."
-                        async for result in send_llm_response(response_log):
-                            yield result
-                        return
-                    contextual_data = " ".join([file.raw_text for file in file_object])
-                    if not q:
-                        q = "Create a general summary of the file"
-                    async for result in send_event(
-                        ChatEvent.STATUS, f"**Constructing Summary Using:** {file_object[0].file_name}"
-                    ):
-                        yield result
-
-                    response = await extract_relevant_summary(
-                        q,
-                        contextual_data,
-                        conversation_history=meta_log,
-                        query_images=uploaded_images,
-                        user=user,
-                        agent=agent,
-                        tracer=tracer,
-                    )
-                    response_log = str(response)
-                    async for result in send_llm_response(response_log):
-                        yield result
-                except Exception as e:
-                    response_log = "Error summarizing file. Please try again, or contact support."
-                    logger.error(f"Error summarizing file for {user.email}: {e}", exc_info=True)
-                    async for result in send_llm_response(response_log):
-                        yield result
+                async for response in generate_summary_from_files(
+                    q=q,
+                    user=user,
+                    file_filters=file_filters,
+                    meta_log=meta_log,
+                    query_images=uploaded_images,
+                    agent=agent,
+                    send_status_func=partial(send_event, ChatEvent.STATUS),
+                    tracer=tracer,
+                ):
+                    if isinstance(response, dict) and ChatEvent.STATUS in response:
+                        yield response[ChatEvent.STATUS]
+                    else:
+                        if isinstance(response, str):
+                            response_log = response
+                            async for result in send_llm_response(response):
+                                yield result
+
             await sync_to_async(save_to_conversation_log)(
                 q,
                 response_log,
@@ -786,6 +819,7 @@ async def chat(
                 conversation_id=conversation_id,
                 query_images=uploaded_images,
                 tracer=tracer,
+                train_of_thought=train_of_thought,
             )
             return
 
@@ -794,7 +828,7 @@ async def chat(
             if not q:
                 conversation_config = await ConversationAdapters.aget_user_conversation_config(user)
                 if conversation_config == None:
-                    conversation_config = await ConversationAdapters.aget_default_conversation_config()
+                    conversation_config = await ConversationAdapters.aget_default_conversation_config(user)
                 model_type = conversation_config.model_type
                 formatted_help = help_message.format(model=model_type, version=state.khoj_version, device=get_device())
                 async for result in send_llm_response(formatted_help):
@@ -830,6 +864,7 @@ async def chat(
                 automation_id=automation.id,
                 query_images=uploaded_images,
                 tracer=tracer,
+                train_of_thought=train_of_thought,
             )
             async for result in send_llm_response(llm_response):
                 yield result
@@ -837,49 +872,49 @@ async def chat(
 
         # Gather Context
         ## Extract Document References
-        compiled_references, inferred_queries, defiltered_query = [], [], q
-        try:
-            async for result in extract_references_and_questions(
-                request,
-                meta_log,
-                q,
-                (n or 7),
-                d,
-                conversation_id,
-                conversation_commands,
-                location,
-                partial(send_event, ChatEvent.STATUS),
-                query_images=uploaded_images,
-                agent=agent,
-                tracer=tracer,
-            ):
-                if isinstance(result, dict) and ChatEvent.STATUS in result:
-                    yield result[ChatEvent.STATUS]
-                else:
-                    compiled_references.extend(result[0])
-                    inferred_queries.extend(result[1])
-                    defiltered_query = result[2]
-        except Exception as e:
-            error_message = f"Error searching knowledge base: {e}. Attempting to respond without document references."
-            logger.error(error_message, exc_info=True)
-            async for result in send_event(
-                ChatEvent.STATUS, "Document search failed. I'll try respond without document references"
-            ):
-                yield result
-
-        if not is_none_or_empty(compiled_references):
-            headings = "\n- " + "\n- ".join(set([c.get("compiled", c).split("\n")[0] for c in compiled_references]))
-            # Strip only leading # from headings
-            headings = headings.replace("#", "")
-            async for result in send_event(ChatEvent.STATUS, f"**Found Relevant Notes**: {headings}"):
-                yield result
+        if not ConversationCommand.Research in conversation_commands:
+            try:
+                async for result in extract_references_and_questions(
+                    request,
+                    meta_log,
+                    q,
+                    (n or 7),
+                    d,
+                    conversation_id,
+                    conversation_commands,
+                    location,
+                    partial(send_event, ChatEvent.STATUS),
+                    query_images=uploaded_images,
+                    agent=agent,
+                    tracer=tracer,
+                ):
+                    if isinstance(result, dict) and ChatEvent.STATUS in result:
+                        yield result[ChatEvent.STATUS]
+                    else:
+                        compiled_references.extend(result[0])
+                        inferred_queries.extend(result[1])
+                        defiltered_query = result[2]
+            except Exception as e:
+                error_message = (
+                    f"Error searching knowledge base: {e}. Attempting to respond without document references."
+                )
+                logger.error(error_message, exc_info=True)
+                async for result in send_event(
+                    ChatEvent.STATUS, "Document search failed. I'll try respond without document references"
+                ):
+                    yield result
 
-        online_results: Dict = dict()
+            if not is_none_or_empty(compiled_references):
+                headings = "\n- " + "\n- ".join(set([c.get("compiled", c).split("\n")[0] for c in compiled_references]))
+                # Strip only leading # from headings
+                headings = headings.replace("#", "")
+                async for result in send_event(ChatEvent.STATUS, f"**Found Relevant Notes**: {headings}"):
+                    yield result
 
-        if conversation_commands == [ConversationCommand.Notes] and not await EntryAdapters.auser_has_entries(user):
-            async for result in send_llm_response(f"{no_entries_found.format()}"):
-                yield result
-            return
+            if conversation_commands == [ConversationCommand.Notes] and not await EntryAdapters.auser_has_entries(user):
+                async for result in send_llm_response(f"{no_entries_found.format()}"):
+                    yield result
+                return
 
         if ConversationCommand.Notes in conversation_commands and is_none_or_empty(compiled_references):
             conversation_commands.remove(ConversationCommand.Notes)
@@ -948,6 +983,33 @@ async def chat(
                 ):
                     yield result
 
+        ## Gather Code Results
+        if ConversationCommand.Code in conversation_commands:
+            try:
+                context = f"# Iteration 1:\n#---\nNotes:\n{compiled_references}\n\nOnline Results:{online_results}"
+                async for result in run_code(
+                    defiltered_query,
+                    meta_log,
+                    context,
+                    location,
+                    user,
+                    partial(send_event, ChatEvent.STATUS),
+                    query_images=uploaded_images,
+                    agent=agent,
+                    tracer=tracer,
+                ):
+                    if isinstance(result, dict) and ChatEvent.STATUS in result:
+                        yield result[ChatEvent.STATUS]
+                    else:
+                        code_results = result
+                async for result in send_event(ChatEvent.STATUS, f"**Ran code snippets**: {len(code_results)}"):
+                    yield result
+            except ValueError as e:
+                logger.warning(
+                    f"Failed to use code tool: {e}. Attempting to respond without code results",
+                    exc_info=True,
+                )
+
         ## Send Gathered References
         async for result in send_event(
             ChatEvent.REFERENCES,
@@ -955,6 +1017,7 @@ async def chat(
                 "inferredQueries": inferred_queries,
                 "context": compiled_references,
                 "onlineContext": online_results,
+                "codeContext": code_results,
             },
         ):
             yield result
@@ -1004,6 +1067,7 @@ async def chat(
                 online_results=online_results,
                 query_images=uploaded_images,
                 tracer=tracer,
+                train_of_thought=train_of_thought,
             )
             content_obj = {
                 "intentType": intent_type,
@@ -1061,6 +1125,7 @@ async def chat(
                 online_results=online_results,
                 query_images=uploaded_images,
                 tracer=tracer,
+                train_of_thought=train_of_thought,
             )
 
             async for result in send_llm_response(json.dumps(content_obj)):
@@ -1076,6 +1141,7 @@ async def chat(
             conversation,
             compiled_references,
             online_results,
+            code_results,
             inferred_queries,
             conversation_commands,
             user,
@@ -1083,8 +1149,10 @@ async def chat(
             conversation_id,
             location,
             user_name,
+            researched_results,
             uploaded_images,
             tracer,
+            train_of_thought,
         )
 
         # Send Response