khoj 1.27.2.dev29__py3-none-any.whl → 1.28.1.dev1__py3-none-any.whl

khoj/processor/conversation/utils.py

@@ -1,14 +1,17 @@
 import base64
+import json
 import logging
 import math
 import mimetypes
 import os
 import queue
+import uuid
 from dataclasses import dataclass
 from datetime import datetime
+from enum import Enum
 from io import BytesIO
 from time import perf_counter
-from typing import Any, Dict, List, Optional
+from typing import Any, Callable, Dict, List, Optional
 
 import PIL.Image
 import requests
@@ -23,8 +26,17 @@ from khoj.database.adapters import ConversationAdapters
 from khoj.database.models import ChatModelOptions, ClientApplication, KhojUser
 from khoj.processor.conversation import prompts
 from khoj.processor.conversation.offline.utils import download_model, infer_max_tokens
+from khoj.search_filter.base_filter import BaseFilter
+from khoj.search_filter.date_filter import DateFilter
+from khoj.search_filter.file_filter import FileFilter
+from khoj.search_filter.word_filter import WordFilter
 from khoj.utils import state
-from khoj.utils.helpers import in_debug_mode, is_none_or_empty, merge_dicts
+from khoj.utils.helpers import (
+    ConversationCommand,
+    in_debug_mode,
+    is_none_or_empty,
+    merge_dicts,
+)
 
 logger = logging.getLogger(__name__)
 model_to_prompt_size = {
@@ -85,8 +97,110 @@ class ThreadedGenerator:
         self.queue.put(StopIteration)
 
 
+class InformationCollectionIteration:
+    def __init__(
+        self,
+        tool: str,
+        query: str,
+        context: list = None,
+        onlineContext: dict = None,
+        codeContext: dict = None,
+        summarizedResult: str = None,
+    ):
+        self.tool = tool
+        self.query = query
+        self.context = context
+        self.onlineContext = onlineContext
+        self.codeContext = codeContext
+        self.summarizedResult = summarizedResult
+
+
+def construct_iteration_history(
+    previous_iterations: List[InformationCollectionIteration], previous_iteration_prompt: str
+) -> str:
+    previous_iterations_history = ""
+    for idx, iteration in enumerate(previous_iterations):
+        iteration_data = previous_iteration_prompt.format(
+            tool=iteration.tool,
+            query=iteration.query,
+            result=iteration.summarizedResult,
+            index=idx + 1,
+        )
+
+        previous_iterations_history += iteration_data
+    return previous_iterations_history
+
+
+def construct_chat_history(conversation_history: dict, n: int = 4, agent_name="AI") -> str:
+    chat_history = ""
+    for chat in conversation_history.get("chat", [])[-n:]:
+        if chat["by"] == "khoj" and chat["intent"].get("type") in ["remember", "reminder", "summarize"]:
+            chat_history += f"User: {chat['intent']['query']}\n"
+
+            if chat["intent"].get("inferred-queries"):
+                chat_history += f'Khoj: {{"queries": {chat["intent"].get("inferred-queries")}}}\n'
+
+            chat_history += f"{agent_name}: {chat['message']}\n\n"
+        elif chat["by"] == "khoj" and ("text-to-image" in chat["intent"].get("type")):
+            chat_history += f"User: {chat['intent']['query']}\n"
+            chat_history += f"{agent_name}: [generated image redacted for space]\n"
+        elif chat["by"] == "khoj" and ("excalidraw" in chat["intent"].get("type")):
+            chat_history += f"User: {chat['intent']['query']}\n"
+            chat_history += f"{agent_name}: {chat['intent']['inferred-queries'][0]}\n"
+    return chat_history
+
+
+def construct_tool_chat_history(
+    previous_iterations: List[InformationCollectionIteration], tool: ConversationCommand = None
+) -> Dict[str, list]:
+    chat_history: list = []
+    inferred_query_extractor: Callable[[InformationCollectionIteration], List[str]] = lambda x: []
+    if tool == ConversationCommand.Notes:
+        inferred_query_extractor = (
+            lambda iteration: [c["query"] for c in iteration.context] if iteration.context else []
+        )
+    elif tool == ConversationCommand.Online:
+        inferred_query_extractor = (
+            lambda iteration: list(iteration.onlineContext.keys()) if iteration.onlineContext else []
+        )
+    elif tool == ConversationCommand.Code:
+        inferred_query_extractor = lambda iteration: list(iteration.codeContext.keys()) if iteration.codeContext else []
+    for iteration in previous_iterations:
+        chat_history += [
+            {
+                "by": "you",
+                "message": iteration.query,
+            },
+            {
+                "by": "khoj",
+                "intent": {
+                    "type": "remember",
+                    "inferred-queries": inferred_query_extractor(iteration),
+                    "query": iteration.query,
+                },
+                "message": iteration.summarizedResult,
+            },
+        ]
+
+    return {"chat": chat_history}
+
+
+class ChatEvent(Enum):
+    START_LLM_RESPONSE = "start_llm_response"
+    END_LLM_RESPONSE = "end_llm_response"
+    MESSAGE = "message"
+    REFERENCES = "references"
+    STATUS = "status"
+    METADATA = "metadata"
+
+
 def message_to_log(
-    user_message, chat_response, user_message_metadata={}, khoj_message_metadata={}, conversation_log=[]
+    user_message,
+    chat_response,
+    user_message_metadata={},
+    khoj_message_metadata={},
+    conversation_log=[],
+    train_of_thought=[],
 ):
     """Create json logs from messages, metadata for conversation log"""
     default_khoj_message_metadata = {
@@ -114,6 +228,7 @@ def save_to_conversation_log(
     user_message_time: str = None,
     compiled_references: List[Dict[str, Any]] = [],
     online_results: Dict[str, Any] = {},
+    code_results: Dict[str, Any] = {},
     inferred_queries: List[str] = [],
     intent_type: str = "remember",
     client_application: ClientApplication = None,
@@ -121,22 +236,29 @@ def save_to_conversation_log(
     automation_id: str = None,
     query_images: List[str] = None,
     tracer: Dict[str, Any] = {},
+    train_of_thought: List[Any] = [],
 ):
     user_message_time = user_message_time or datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+    turn_id = tracer.get("mid") or str(uuid.uuid4())
     updated_conversation = message_to_log(
         user_message=q,
         chat_response=chat_response,
         user_message_metadata={
             "created": user_message_time,
             "images": query_images,
+            "turnId": turn_id,
         },
         khoj_message_metadata={
             "context": compiled_references,
             "intent": {"inferred-queries": inferred_queries, "type": intent_type},
             "onlineContext": online_results,
+            "codeContext": code_results,
             "automationId": automation_id,
+            "trainOfThought": train_of_thought,
+            "turnId": turn_id,
         },
         conversation_log=meta_log.get("chat", []),
+        train_of_thought=train_of_thought,
     )
     ConversationAdapters.save_conversation(
         user,
@@ -330,9 +452,23 @@ def reciprocal_conversation_to_chatml(message_pair):
     return [ChatMessage(content=message, role=role) for message, role in zip(message_pair, ["user", "assistant"])]
 
 
-def remove_json_codeblock(response: str):
-    """Remove any markdown json codeblock formatting if present. Useful for non schema enforceable models"""
-    return response.removeprefix("```json").removesuffix("```")
+def clean_json(response: str):
+    """Remove any markdown json codeblock and newline formatting if present. Useful for non schema enforceable models"""
+    return response.strip().replace("\n", "").removeprefix("```json").removesuffix("```")
+
+
+def clean_code_python(code: str):
+    """Remove any markdown codeblock and newline formatting if present. Useful for non schema enforceable models"""
+    return code.strip().removeprefix("```python").removesuffix("```")
+
+
+def defilter_query(query: str):
+    """Remove any query filters in query"""
+    defiltered_query = query
+    filters: List[BaseFilter] = [WordFilter(), FileFilter(), DateFilter()]
+    for filter in filters:
+        defiltered_query = filter.defilter(defiltered_query)
+    return defiltered_query
 
 
 @dataclass
@@ -375,15 +511,12 @@ def commit_conversation_trace(
     Returns the path to the repository.
     """
     # Serialize session, system message and response to yaml
-    system_message_yaml = yaml.dump(system_message, allow_unicode=True, sort_keys=False, default_flow_style=False)
-    response_yaml = yaml.dump(response, allow_unicode=True, sort_keys=False, default_flow_style=False)
+    system_message_yaml = json.dumps(system_message, ensure_ascii=False, sort_keys=False)
+    response_yaml = json.dumps(response, ensure_ascii=False, sort_keys=False)
     formatted_session = [{"role": message.role, "content": message.content} for message in session]
-    session_yaml = yaml.dump(formatted_session, allow_unicode=True, sort_keys=False, default_flow_style=False)
+    session_yaml = json.dumps(formatted_session, ensure_ascii=False, sort_keys=False)
     query = (
-        yaml.dump(session[-1].content, allow_unicode=True, sort_keys=False, default_flow_style=False)
-        .strip()
-        .removeprefix("'")
-        .removesuffix("'")
+        json.dumps(session[-1].content, ensure_ascii=False, sort_keys=False).strip().removeprefix("'").removesuffix("'")
     )  # Extract serialized query from chat session
 
     # Extract chat metadata for session

khoj/processor/embeddings.py

@@ -13,7 +13,7 @@ from tenacity import (
 )
 from torch import nn
 
-from khoj.utils.helpers import get_device, merge_dicts, timer
+from khoj.utils.helpers import fix_json_dict, get_device, merge_dicts, timer
 from khoj.utils.rawconfig import SearchResponse
 
 logger = logging.getLogger(__name__)
@@ -31,9 +31,9 @@ class EmbeddingsModel:
     ):
         default_query_encode_kwargs = {"show_progress_bar": False, "normalize_embeddings": True}
         default_docs_encode_kwargs = {"show_progress_bar": True, "normalize_embeddings": True}
-        self.query_encode_kwargs = merge_dicts(query_encode_kwargs, default_query_encode_kwargs)
-        self.docs_encode_kwargs = merge_dicts(docs_encode_kwargs, default_docs_encode_kwargs)
-        self.model_kwargs = merge_dicts(model_kwargs, {"device": get_device()})
+        self.query_encode_kwargs = merge_dicts(fix_json_dict(query_encode_kwargs), default_query_encode_kwargs)
+        self.docs_encode_kwargs = merge_dicts(fix_json_dict(docs_encode_kwargs), default_docs_encode_kwargs)
+        self.model_kwargs = merge_dicts(fix_json_dict(model_kwargs), {"device": get_device()})
         self.model_name = model_name
         self.inference_endpoint = embeddings_inference_endpoint
         self.api_key = embeddings_inference_endpoint_api_key

khoj/processor/tools/online_search.py

@@ -4,7 +4,7 @@ import logging
 import os
 import urllib.parse
 from collections import defaultdict
-from typing import Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import aiohttp
 from bs4 import BeautifulSoup
@@ -52,7 +52,9 @@ OLOSTEP_QUERY_PARAMS = {
     "expandMarkdown": "True",
     "expandHtml": "False",
 }
-MAX_WEBPAGES_TO_READ = 1
+
+DEFAULT_MAX_WEBPAGES_TO_READ = 1
+MAX_WEBPAGES_TO_INFER = 10
 
 
 async def search_online(
@@ -62,6 +64,7 @@ async def search_online(
     user: KhojUser,
     send_status_func: Optional[Callable] = None,
     custom_filters: List[str] = [],
+    max_webpages_to_read: int = DEFAULT_MAX_WEBPAGES_TO_READ,
     query_images: List[str] = None,
     agent: Agent = None,
     tracer: dict = {},
@@ -97,7 +100,7 @@ async def search_online(
     for subquery in response_dict:
         if "answerBox" in response_dict[subquery]:
             continue
-        for organic in response_dict[subquery].get("organic", [])[:MAX_WEBPAGES_TO_READ]:
+        for organic in response_dict[subquery].get("organic", [])[:max_webpages_to_read]:
             link = organic.get("link")
             if link in webpages:
                 webpages[link]["queries"].add(subquery)
@@ -155,13 +158,16 @@ async def read_webpages(
     query_images: List[str] = None,
     agent: Agent = None,
     tracer: dict = {},
+    max_webpages_to_read: int = DEFAULT_MAX_WEBPAGES_TO_READ,
 ):
     "Infer web pages to read from the query and extract relevant information from them"
     logger.info(f"Inferring web pages to read")
-    if send_status_func:
-        async for event in send_status_func(f"**Inferring web pages to read**"):
-            yield {ChatEvent.STATUS: event}
-    urls = await infer_webpage_urls(query, conversation_history, location, user, query_images)
+    urls = await infer_webpage_urls(
+        query, conversation_history, location, user, query_images, agent=agent, tracer=tracer
+    )
+
+    # Get the top 10 web pages to read
+    urls = urls[:max_webpages_to_read]
 
     logger.info(f"Reading web pages at: {urls}")
     if send_status_func:

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