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

khoj/routers/helpers.py

@@ -43,6 +43,7 @@ from khoj.database.adapters import (
     AutomationAdapters,
     ConversationAdapters,
     EntryAdapters,
+    FileObjectAdapters,
     ais_user_subscribed,
     create_khoj_token,
     get_khoj_tokens,
@@ -87,9 +88,11 @@ from khoj.processor.conversation.offline.chat_model import (
 )
 from khoj.processor.conversation.openai.gpt import converse, send_message_to_model
 from khoj.processor.conversation.utils import (
+    ChatEvent,
     ThreadedGenerator,
+    clean_json,
+    construct_chat_history,
     generate_chatml_messages_with_context,
-    remove_json_codeblock,
     save_to_conversation_log,
 )
 from khoj.processor.speech.text_to_speech import is_eleven_labs_enabled
@@ -137,7 +140,7 @@ def validate_conversation_config(user: KhojUser):
 async def is_ready_to_chat(user: KhojUser):
     user_conversation_config = await ConversationAdapters.aget_user_conversation_config(user)
     if user_conversation_config == None:
-        user_conversation_config = await ConversationAdapters.aget_default_conversation_config()
+        user_conversation_config = await ConversationAdapters.aget_default_conversation_config(user)
 
     if user_conversation_config and user_conversation_config.model_type == ChatModelOptions.ModelType.OFFLINE:
         chat_model = user_conversation_config.chat_model
@@ -210,21 +213,6 @@ def get_next_url(request: Request) -> str:
     return urljoin(str(request.base_url).rstrip("/"), next_path)
 
 
-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"
-            chat_history += f"{agent_name}: {chat['message']}\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 get_conversation_command(query: str, any_references: bool = False) -> ConversationCommand:
     if query.startswith("/notes"):
         return ConversationCommand.Notes
@@ -244,6 +232,10 @@ def get_conversation_command(query: str, any_references: bool = False) -> Conver
         return ConversationCommand.Summarize
     elif query.startswith("/diagram"):
         return ConversationCommand.Diagram
+    elif query.startswith("/code"):
+        return ConversationCommand.Code
+    elif query.startswith("/research"):
+        return ConversationCommand.Research
     # If no relevant notes found for the given query
     elif not any_references:
         return ConversationCommand.General
@@ -301,6 +293,7 @@ async def aget_relevant_information_sources(
     user: KhojUser,
     query_images: List[str] = None,
     agent: Agent = None,
+    tracer: dict = {},
 ):
     """
     Given a query, determine which of the available tools the agent should use in order to answer appropriately.
@@ -337,11 +330,11 @@ async def aget_relevant_information_sources(
             relevant_tools_prompt,
             response_type="json_object",
             user=user,
+            tracer=tracer,
         )
 
     try:
-        response = response.strip()
-        response = remove_json_codeblock(response)
+        response = clean_json(response)
         response = json.loads(response)
         response = [q.strip() for q in response["source"] if q.strip()]
         if not isinstance(response, list) or not response or len(response) == 0:
@@ -378,6 +371,7 @@ async def aget_relevant_output_modes(
     user: KhojUser = None,
     query_images: List[str] = None,
     agent: Agent = None,
+    tracer: dict = {},
 ):
     """
     Given a query, determine which of the available tools the agent should use in order to answer appropriately.
@@ -413,11 +407,12 @@ async def aget_relevant_output_modes(
     )
 
     with timer("Chat actor: Infer output mode for chat response", logger):
-        response = await send_message_to_model_wrapper(relevant_mode_prompt, response_type="json_object", user=user)
+        response = await send_message_to_model_wrapper(
+            relevant_mode_prompt, response_type="json_object", user=user, tracer=tracer
+        )
 
     try:
-        response = response.strip()
-        response = remove_json_codeblock(response)
+        response = clean_json(response)
         response = json.loads(response)
 
         if is_none_or_empty(response):
@@ -444,6 +439,7 @@ async def infer_webpage_urls(
     user: KhojUser,
     query_images: List[str] = None,
     agent: Agent = None,
+    tracer: dict = {},
 ) -> List[str]:
     """
     Infer webpage links from the given query
@@ -468,12 +464,16 @@ async def infer_webpage_urls(
 
     with timer("Chat actor: Infer webpage urls to read", logger):
         response = await send_message_to_model_wrapper(
-            online_queries_prompt, query_images=query_images, response_type="json_object", user=user
+            online_queries_prompt,
+            query_images=query_images,
+            response_type="json_object",
+            user=user,
+            tracer=tracer,
         )
 
     # Validate that the response is a non-empty, JSON-serializable list of URLs
     try:
-        response = response.strip()
+        response = clean_json(response)
         urls = json.loads(response)
         valid_unique_urls = {str(url).strip() for url in urls["links"] if is_valid_url(url)}
         if is_none_or_empty(valid_unique_urls):
@@ -490,6 +490,7 @@ async def generate_online_subqueries(
     user: KhojUser,
     query_images: List[str] = None,
     agent: Agent = None,
+    tracer: dict = {},
 ) -> List[str]:
     """
     Generate subqueries from the given query
@@ -514,13 +515,16 @@ async def generate_online_subqueries(
 
     with timer("Chat actor: Generate online search subqueries", logger):
         response = await send_message_to_model_wrapper(
-            online_queries_prompt, query_images=query_images, response_type="json_object", user=user
+            online_queries_prompt,
+            query_images=query_images,
+            response_type="json_object",
+            user=user,
+            tracer=tracer,
         )
 
     # Validate that the response is a non-empty, JSON-serializable list
     try:
-        response = response.strip()
-        response = remove_json_codeblock(response)
+        response = clean_json(response)
         response = json.loads(response)
         response = [q.strip() for q in response["queries"] if q.strip()]
         if not isinstance(response, list) or not response or len(response) == 0:
@@ -533,7 +537,7 @@ async def generate_online_subqueries(
 
 
 async def schedule_query(
-    q: str, conversation_history: dict, user: KhojUser, query_images: List[str] = None
+    q: str, conversation_history: dict, user: KhojUser, query_images: List[str] = None, tracer: dict = {}
 ) -> Tuple[str, ...]:
     """
     Schedule the date, time to run the query. Assume the server timezone is UTC.
@@ -546,7 +550,7 @@ async def schedule_query(
     )
 
     raw_response = await send_message_to_model_wrapper(
-        crontime_prompt, query_images=query_images, response_type="json_object", user=user
+        crontime_prompt, query_images=query_images, response_type="json_object", user=user, tracer=tracer
     )
 
     # Validate that the response is a non-empty, JSON-serializable list
@@ -561,7 +565,7 @@ async def schedule_query(
 
 
 async def extract_relevant_info(
-    qs: set[str], corpus: str, user: KhojUser = None, agent: Agent = None
+    qs: set[str], corpus: str, user: KhojUser = None, agent: Agent = None, tracer: dict = {}
 ) -> Union[str, None]:
     """
     Extract relevant information for a given query from the target corpus
@@ -584,6 +588,7 @@ async def extract_relevant_info(
         extract_relevant_information,
         prompts.system_prompt_extract_relevant_information,
         user=user,
+        tracer=tracer,
     )
     return response.strip()
 
@@ -595,6 +600,7 @@ async def extract_relevant_summary(
     query_images: List[str] = None,
     user: KhojUser = None,
     agent: Agent = None,
+    tracer: dict = {},
 ) -> Union[str, None]:
     """
     Extract relevant information for a given query from the target corpus
@@ -622,10 +628,58 @@ async def extract_relevant_summary(
             prompts.system_prompt_extract_relevant_summary,
             user=user,
             query_images=query_images,
+            tracer=tracer,
         )
     return response.strip()
 
 
+async def generate_summary_from_files(
+    q: str,
+    user: KhojUser,
+    file_filters: List[str],
+    meta_log: dict,
+    query_images: List[str] = None,
+    agent: Agent = None,
+    send_status_func: Optional[Callable] = None,
+    tracer: dict = {},
+):
+    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.pop(), 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."
+            yield response_log
+            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_status_func(f"**Constructing Summary Using:** {file_object[0].file_name}"):
+            yield {ChatEvent.STATUS: result}
+
+        response = await extract_relevant_summary(
+            q,
+            contextual_data,
+            conversation_history=meta_log,
+            query_images=query_images,
+            user=user,
+            agent=agent,
+            tracer=tracer,
+        )
+
+        yield str(response)
+    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)
+        yield result
+
+
 async def generate_excalidraw_diagram(
     q: str,
     conversation_history: Dict[str, Any],
@@ -636,6 +690,7 @@ async def generate_excalidraw_diagram(
     user: KhojUser = None,
     agent: Agent = None,
     send_status_func: Optional[Callable] = None,
+    tracer: dict = {},
 ):
     if send_status_func:
         async for event in send_status_func("**Enhancing the Diagramming Prompt**"):
@@ -650,6 +705,7 @@ async def generate_excalidraw_diagram(
         query_images=query_images,
         user=user,
         agent=agent,
+        tracer=tracer,
     )
 
     if send_status_func:
@@ -660,6 +716,7 @@ async def generate_excalidraw_diagram(
         q=better_diagram_description_prompt,
         user=user,
         agent=agent,
+        tracer=tracer,
     )
 
     yield better_diagram_description_prompt, excalidraw_diagram_description
@@ -674,6 +731,7 @@ async def generate_better_diagram_description(
     query_images: List[str] = None,
     user: KhojUser = None,
     agent: Agent = None,
+    tracer: dict = {},
 ) -> str:
     """
     Generate a diagram description from the given query and context
@@ -711,7 +769,7 @@ async def generate_better_diagram_description(
 
     with timer("Chat actor: Generate better diagram description", logger):
         response = await send_message_to_model_wrapper(
-            improve_diagram_description_prompt, query_images=query_images, user=user
+            improve_diagram_description_prompt, query_images=query_images, user=user, tracer=tracer
         )
         response = response.strip()
         if response.startswith(('"', "'")) and response.endswith(('"', "'")):
@@ -724,6 +782,7 @@ async def generate_excalidraw_diagram_from_description(
     q: str,
     user: KhojUser = None,
     agent: Agent = None,
+    tracer: dict = {},
 ) -> str:
     personality_context = (
         prompts.personality_context.format(personality=agent.personality) if agent and agent.personality else ""
@@ -735,9 +794,10 @@ async def generate_excalidraw_diagram_from_description(
     )
 
     with timer("Chat actor: Generate excalidraw diagram", logger):
-        raw_response = await send_message_to_model_wrapper(message=excalidraw_diagram_generation, user=user)
-        raw_response = raw_response.strip()
-        raw_response = remove_json_codeblock(raw_response)
+        raw_response = await send_message_to_model_wrapper(
+            query=excalidraw_diagram_generation, user=user, tracer=tracer
+        )
+        raw_response = clean_json(raw_response)
         response: Dict[str, str] = json.loads(raw_response)
         if not response or not isinstance(response, List) or not isinstance(response[0], Dict):
             # TODO Some additional validation here that it's a valid Excalidraw diagram
@@ -756,6 +816,7 @@ async def generate_better_image_prompt(
     query_images: Optional[List[str]] = None,
     user: KhojUser = None,
     agent: Agent = None,
+    tracer: dict = {},
 ) -> str:
     """
     Generate a better image prompt from the given query
@@ -802,7 +863,9 @@ async def generate_better_image_prompt(
         )
 
     with timer("Chat actor: Generate contextual image prompt", logger):
-        response = await send_message_to_model_wrapper(image_prompt, query_images=query_images, user=user)
+        response = await send_message_to_model_wrapper(
+            image_prompt, query_images=query_images, user=user, tracer=tracer
+        )
         response = response.strip()
         if response.startswith(('"', "'")) and response.endswith(('"', "'")):
             response = response[1:-1]
@@ -811,11 +874,13 @@ async def generate_better_image_prompt(
 
 
 async def send_message_to_model_wrapper(
-    message: str,
+    query: str,
     system_message: str = "",
     response_type: str = "text",
     user: KhojUser = None,
     query_images: List[str] = None,
+    context: str = "",
+    tracer: dict = {},
 ):
     conversation_config: ChatModelOptions = await ConversationAdapters.aget_default_conversation_config(user)
     vision_available = conversation_config.vision_enabled
@@ -845,7 +910,8 @@ async def send_message_to_model_wrapper(
 
         loaded_model = state.offline_chat_processor_config.loaded_model
         truncated_messages = generate_chatml_messages_with_context(
-            user_message=message,
+            user_message=query,
+            context_message=context,
             system_message=system_message,
             model_name=chat_model,
             loaded_model=loaded_model,
@@ -862,6 +928,7 @@ async def send_message_to_model_wrapper(
             max_prompt_size=max_tokens,
             streaming=False,
             response_type=response_type,
+            tracer=tracer,
         )
 
     elif model_type == ChatModelOptions.ModelType.OPENAI:
@@ -869,7 +936,8 @@ async def send_message_to_model_wrapper(
         api_key = openai_chat_config.api_key
         api_base_url = openai_chat_config.api_base_url
         truncated_messages = generate_chatml_messages_with_context(
-            user_message=message,
+            user_message=query,
+            context_message=context,
             system_message=system_message,
             model_name=chat_model,
             max_prompt_size=max_tokens,
@@ -885,11 +953,13 @@ async def send_message_to_model_wrapper(
             model=chat_model,
             response_type=response_type,
             api_base_url=api_base_url,
+            tracer=tracer,
         )
     elif model_type == ChatModelOptions.ModelType.ANTHROPIC:
         api_key = conversation_config.openai_config.api_key
         truncated_messages = generate_chatml_messages_with_context(
-            user_message=message,
+            user_message=query,
+            context_message=context,
             system_message=system_message,
             model_name=chat_model,
             max_prompt_size=max_tokens,
@@ -903,11 +973,14 @@ async def send_message_to_model_wrapper(
             messages=truncated_messages,
             api_key=api_key,
             model=chat_model,
+            response_type=response_type,
+            tracer=tracer,
         )
     elif model_type == ChatModelOptions.ModelType.GOOGLE:
         api_key = conversation_config.openai_config.api_key
         truncated_messages = generate_chatml_messages_with_context(
-            user_message=message,
+            user_message=query,
+            context_message=context,
             system_message=system_message,
             model_name=chat_model,
             max_prompt_size=max_tokens,
@@ -918,7 +991,7 @@ async def send_message_to_model_wrapper(
         )
 
         return gemini_send_message_to_model(
-            messages=truncated_messages, api_key=api_key, model=chat_model, response_type=response_type
+            messages=truncated_messages, api_key=api_key, model=chat_model, response_type=response_type, tracer=tracer
         )
     else:
         raise HTTPException(status_code=500, detail="Invalid conversation config")
@@ -929,6 +1002,7 @@ def send_message_to_model_wrapper_sync(
     system_message: str = "",
     response_type: str = "text",
     user: KhojUser = None,
+    tracer: dict = {},
 ):
     conversation_config: ChatModelOptions = ConversationAdapters.get_default_conversation_config(user)
 
@@ -961,6 +1035,7 @@ def send_message_to_model_wrapper_sync(
             max_prompt_size=max_tokens,
             streaming=False,
             response_type=response_type,
+            tracer=tracer,
         )
 
     elif conversation_config.model_type == ChatModelOptions.ModelType.OPENAI:
@@ -975,7 +1050,11 @@ def send_message_to_model_wrapper_sync(
         )
 
         openai_response = send_message_to_model(
-            messages=truncated_messages, api_key=api_key, model=chat_model, response_type=response_type
+            messages=truncated_messages,
+            api_key=api_key,
+            model=chat_model,
+            response_type=response_type,
+            tracer=tracer,
         )
 
         return openai_response
@@ -995,6 +1074,8 @@ def send_message_to_model_wrapper_sync(
             messages=truncated_messages,
             api_key=api_key,
             model=chat_model,
+            response_type=response_type,
+            tracer=tracer,
         )
 
     elif conversation_config.model_type == ChatModelOptions.ModelType.GOOGLE:
@@ -1013,6 +1094,7 @@ def send_message_to_model_wrapper_sync(
             api_key=api_key,
             model=chat_model,
             response_type=response_type,
+            tracer=tracer,
         )
     else:
         raise HTTPException(status_code=500, detail="Invalid conversation config")
@@ -1024,6 +1106,7 @@ def generate_chat_response(
     conversation: Conversation,
     compiled_references: List[Dict] = [],
     online_results: Dict[str, Dict] = {},
+    code_results: Dict[str, Dict] = {},
     inferred_queries: List[str] = [],
     conversation_commands: List[ConversationCommand] = [ConversationCommand.Default],
     user: KhojUser = None,
@@ -1031,7 +1114,10 @@ def generate_chat_response(
     conversation_id: str = None,
     location_data: LocationData = None,
     user_name: Optional[str] = None,
+    meta_research: str = "",
     query_images: Optional[List[str]] = None,
+    tracer: dict = {},
+    train_of_thought: List[Any] = [],
 ) -> Tuple[Union[ThreadedGenerator, Iterator[str]], Dict[str, str]]:
     # Initialize Variables
     chat_response = None
@@ -1039,6 +1125,9 @@ def generate_chat_response(
 
     metadata = {}
     agent = AgentAdapters.get_conversation_agent_by_id(conversation.agent.id) if conversation.agent else None
+    query_to_run = q
+    if meta_research:
+        query_to_run = f"AI Research: {meta_research} {q}"
     try:
         partial_completion = partial(
             save_to_conversation_log,
@@ -1047,10 +1136,13 @@ def generate_chat_response(
             meta_log=meta_log,
             compiled_references=compiled_references,
             online_results=online_results,
+            code_results=code_results,
             inferred_queries=inferred_queries,
             client_application=client_application,
             conversation_id=conversation_id,
             query_images=query_images,
+            tracer=tracer,
+            train_of_thought=train_of_thought,
         )
 
         conversation_config = ConversationAdapters.get_valid_conversation_config(user, conversation)
@@ -1064,9 +1156,9 @@ def generate_chat_response(
         if conversation_config.model_type == "offline":
             loaded_model = state.offline_chat_processor_config.loaded_model
             chat_response = converse_offline(
+                user_query=query_to_run,
                 references=compiled_references,
                 online_results=online_results,
-                user_query=q,
                 loaded_model=loaded_model,
                 conversation_log=meta_log,
                 completion_func=partial_completion,
@@ -1077,6 +1169,7 @@ def generate_chat_response(
                 location_data=location_data,
                 user_name=user_name,
                 agent=agent,
+                tracer=tracer,
             )
 
         elif conversation_config.model_type == ChatModelOptions.ModelType.OPENAI:
@@ -1085,9 +1178,10 @@ def generate_chat_response(
             chat_model = conversation_config.chat_model
             chat_response = converse(
                 compiled_references,
-                q,
+                query_to_run,
                 query_images=query_images,
                 online_results=online_results,
+                code_results=code_results,
                 conversation_log=meta_log,
                 model=chat_model,
                 api_key=api_key,
@@ -1100,15 +1194,17 @@ def generate_chat_response(
                 user_name=user_name,
                 agent=agent,
                 vision_available=vision_available,
+                tracer=tracer,
             )
 
         elif conversation_config.model_type == ChatModelOptions.ModelType.ANTHROPIC:
             api_key = conversation_config.openai_config.api_key
             chat_response = converse_anthropic(
                 compiled_references,
-                q,
+                query_to_run,
                 query_images=query_images,
                 online_results=online_results,
+                code_results=code_results,
                 conversation_log=meta_log,
                 model=conversation_config.chat_model,
                 api_key=api_key,
@@ -1120,15 +1216,16 @@ def generate_chat_response(
                 user_name=user_name,
                 agent=agent,
                 vision_available=vision_available,
+                tracer=tracer,
             )
         elif conversation_config.model_type == ChatModelOptions.ModelType.GOOGLE:
             api_key = conversation_config.openai_config.api_key
             chat_response = converse_gemini(
                 compiled_references,
-                q,
-                query_images=query_images,
-                online_results=online_results,
-                conversation_log=meta_log,
+                query_to_run,
+                online_results,
+                code_results,
+                meta_log,
                 model=conversation_config.chat_model,
                 api_key=api_key,
                 completion_func=partial_completion,
@@ -1139,6 +1236,7 @@ def generate_chat_response(
                 user_name=user_name,
                 agent=agent,
                 vision_available=vision_available,
+                tracer=tracer,
             )
 
         metadata.update({"chat_model": conversation_config.chat_model})
@@ -1495,9 +1593,15 @@ def scheduled_chat(
 
 
 async def create_automation(
-    q: str, timezone: str, user: KhojUser, calling_url: URL, meta_log: dict = {}, conversation_id: str = None
+    q: str,
+    timezone: str,
+    user: KhojUser,
+    calling_url: URL,
+    meta_log: dict = {},
+    conversation_id: str = None,
+    tracer: dict = {},
 ):
-    crontime, query_to_run, subject = await schedule_query(q, meta_log, user)
+    crontime, query_to_run, subject = await schedule_query(q, meta_log, user, tracer=tracer)
     job = await schedule_automation(query_to_run, subject, crontime, timezone, q, user, calling_url, conversation_id)
     return job, crontime, query_to_run, subject
 
@@ -1575,14 +1679,6 @@ Manage your automations [here](/automations).
     """.strip()
 
 
-class ChatEvent(Enum):
-    START_LLM_RESPONSE = "start_llm_response"
-    END_LLM_RESPONSE = "end_llm_response"
-    MESSAGE = "message"
-    REFERENCES = "references"
-    STATUS = "status"
-
-
 class MessageProcessor:
     def __init__(self):
         self.references = {}