khoj 1.42.9.dev26__py3-none-any.whl → 1.42.10.dev2__py3-none-any.whl

khoj/routers/research.py

@@ -3,9 +3,11 @@ import logging
 import os
 from copy import deepcopy
 from datetime import datetime
-from typing import Callable, Dict, List, Optional
+from enum import Enum
+from typing import Callable, Dict, List, Optional, Type
 
 import yaml
+from pydantic import BaseModel, Field
 
 from khoj.database.adapters import AgentAdapters, EntryAdapters
 from khoj.database.models import Agent, ChatMessageModel, KhojUser
@@ -13,31 +15,25 @@ from khoj.processor.conversation import prompts
 from khoj.processor.conversation.utils import (
     OperatorRun,
     ResearchIteration,
-    ToolCall,
     construct_iteration_history,
     construct_tool_chat_history,
     load_complex_json,
 )
 from khoj.processor.operator import operate_environment
-from khoj.processor.tools.online_search import read_webpages_content, search_online
+from khoj.processor.tools.online_search import read_webpages, search_online
 from khoj.processor.tools.run_code import run_code
 from khoj.routers.helpers import (
     ChatEvent,
     generate_summary_from_files,
-    grep_files,
-    list_files,
     search_documents,
     send_message_to_model_wrapper,
-    view_file_content,
 )
 from khoj.utils.helpers import (
     ConversationCommand,
-    ToolDefinition,
-    dict_to_tuple,
     is_none_or_empty,
     is_operator_enabled,
     timer,
-    tools_for_research_llm,
+    tool_description_for_research_llm,
     truncate_code_context,
 )
 from khoj.utils.rawconfig import LocationData
@@ -45,6 +41,47 @@ from khoj.utils.rawconfig import LocationData
 logger = logging.getLogger(__name__)
 
 
+class PlanningResponse(BaseModel):
+    """
+    Schema for the response from planning agent when deciding the next tool to pick.
+    """
+
+    scratchpad: str = Field(..., description="Scratchpad to reason about which tool to use next")
+
+    class Config:
+        arbitrary_types_allowed = True
+
+    @classmethod
+    def create_model_with_enum(cls: Type["PlanningResponse"], tool_options: dict) -> Type["PlanningResponse"]:
+        """
+        Factory method that creates a customized PlanningResponse model
+        with a properly typed tool field based on available tools.
+
+        The tool field is dynamically generated based on available tools.
+        The query field should be filled by the model after the tool field for a more logical reasoning flow.
+
+        Args:
+            tool_options: Dictionary mapping tool names to values
+
+        Returns:
+            A customized PlanningResponse class
+        """
+        # Create dynamic enum from tool options
+        tool_enum = Enum("ToolEnum", tool_options)  # type: ignore
+
+        # Create and return a customized response model with the enum
+        class PlanningResponseWithTool(PlanningResponse):
+            """
+            Use the scratchpad to reason about which tool to use next and the query to send to the tool.
+            Pick tool from provided options and your query to send to the tool.
+            """
+
+            tool: tool_enum = Field(..., description="Name of the tool to use")
+            query: str = Field(..., description="Detailed query for the selected tool")
+
+        return PlanningResponseWithTool
+
+
 async def apick_next_tool(
     query: str,
     conversation_history: List[ChatMessageModel],
@@ -67,13 +104,12 @@ async def apick_next_tool(
     # Continue with previous iteration if a multi-step tool use is in progress
     if (
         previous_iterations
-        and previous_iterations[-1].query
-        and isinstance(previous_iterations[-1].query, ToolCall)
-        and previous_iterations[-1].query.name == ConversationCommand.Operator
+        and previous_iterations[-1].tool == ConversationCommand.Operator
         and not previous_iterations[-1].summarizedResult
     ):
         previous_iteration = previous_iterations[-1]
         yield ResearchIteration(
+            tool=previous_iteration.tool,
             query=query,
             context=previous_iteration.context,
             onlineContext=previous_iteration.onlineContext,
@@ -84,40 +120,30 @@ async def apick_next_tool(
         return
 
     # Construct tool options for the agent to choose from
-    tools = []
+    tool_options = dict()
     tool_options_str = ""
     agent_tools = agent.input_tools if agent else []
     user_has_entries = await EntryAdapters.auser_has_entries(user)
-    for tool, tool_data in tools_for_research_llm.items():
+    for tool, description in tool_description_for_research_llm.items():
         # Skip showing operator tool as an option if not enabled
         if tool == ConversationCommand.Operator and not is_operator_enabled():
             continue
-        # Skip showing document related tools if user has no documents
-        if (
-            tool == ConversationCommand.SemanticSearchFiles
-            or tool == ConversationCommand.RegexSearchFiles
-            or tool == ConversationCommand.ViewFile
-            or tool == ConversationCommand.ListFiles
-        ) and not user_has_entries:
-            continue
-        if tool == ConversationCommand.SemanticSearchFiles:
-            description = tool_data.description.format(max_search_queries=max_document_searches)
-        elif tool == ConversationCommand.Webpage:
-            description = tool_data.description.format(max_webpages_to_read=max_webpages_to_read)
-        elif tool == ConversationCommand.Online:
-            description = tool_data.description.format(max_search_queries=max_online_searches)
-        else:
-            description = tool_data.description
+        # Skip showing Notes tool as an option if user has no entries
+        if tool == ConversationCommand.Notes:
+            if not user_has_entries:
+                continue
+            description = description.format(max_search_queries=max_document_searches)
+        if tool == ConversationCommand.Webpage:
+            description = description.format(max_webpages_to_read=max_webpages_to_read)
+        if tool == ConversationCommand.Online:
+            description = description.format(max_search_queries=max_online_searches)
         # Add tool if agent does not have any tools defined or the tool is supported by the agent.
         if len(agent_tools) == 0 or tool.value in agent_tools:
+            tool_options[tool.name] = tool.value
             tool_options_str += f'- "{tool.value}": "{description}"\n'
-            tools.append(
-                ToolDefinition(
-                    name=tool.value,
-                    description=description,
-                    schema=tool_data.schema,
-                )
-            )
+
+    # Create planning reponse model with dynamically populated tool enum class
+    planning_response_model = PlanningResponse.create_model_with_enum(tool_options)
 
     today = datetime.today()
     location_data = f"{location}" if location else "Unknown"
@@ -136,17 +162,24 @@ async def apick_next_tool(
         max_iterations=max_iterations,
     )
 
+    if query_images:
+        query = f"[placeholder for user attached images]\n{query}"
+
     # Construct chat history with user and iteration history with researcher agent for context
-    iteration_chat_history = construct_iteration_history(previous_iterations, query, query_images, query_files)
+    iteration_chat_history = construct_iteration_history(previous_iterations, prompts.previous_iteration, query)
     chat_and_research_history = conversation_history + iteration_chat_history
 
+    # Plan function execution for the next tool
+    query = prompts.plan_function_execution_next_tool.format(query=query) if iteration_chat_history else query
+
     try:
         with timer("Chat actor: Infer information sources to refer", logger):
             response = await send_message_to_model_wrapper(
-                query="",
+                query=query,
                 system_message=function_planning_prompt,
                 chat_history=chat_and_research_history,
-                tools=tools,
+                response_type="json_object",
+                response_schema=planning_response_model,
                 deepthought=True,
                 user=user,
                 query_images=query_images,
@@ -157,38 +190,48 @@ async def apick_next_tool(
     except Exception as e:
         logger.error(f"Failed to infer information sources to refer: {e}", exc_info=True)
         yield ResearchIteration(
+            tool=None,
             query=None,
             warning="Failed to infer information sources to refer. Skipping iteration. Try again.",
         )
         return
 
     try:
-        # Try parse the response as function call response to infer next tool to use.
-        # TODO: Handle multiple tool calls.
-        response_text = response.text
-        parsed_response = [ToolCall(**item) for item in load_complex_json(response_text)][0]
-    except Exception as e:
-        # Otherwise assume the model has decided to end the research run and respond to the user.
-        parsed_response = ToolCall(name=ConversationCommand.Text, args={"response": response_text}, id=None)
-
-    # If we have a valid response, extract the tool and query.
-    warning = None
-    logger.info(f"Response for determining relevant tools: {parsed_response.name}({parsed_response.args})")
-
-    # Detect selection of previously used query, tool combination.
-    previous_tool_query_combinations = {
-        (i.query.name, dict_to_tuple(i.query.args))
-        for i in previous_iterations
-        if i.warning is None and isinstance(i.query, ToolCall)
-    }
-    if (parsed_response.name, dict_to_tuple(parsed_response.args)) in previous_tool_query_combinations:
-        warning = f"Repeated tool, query combination detected. Skipping iteration. Try something different."
-    # Only send client status updates if we'll execute this iteration and model has thoughts to share.
-    elif send_status_func and not is_none_or_empty(response.thought):
-        async for event in send_status_func(response.thought):
-            yield {ChatEvent.STATUS: event}
+        response = load_complex_json(response)
+        if not isinstance(response, dict):
+            raise ValueError(f"Expected dict response, got {type(response).__name__}: {response}")
+        selected_tool = response.get("tool", None)
+        generated_query = response.get("query", None)
+        scratchpad = response.get("scratchpad", None)
+        warning = None
+        logger.info(f"Response for determining relevant tools: {response}")
+
+        # Detect selection of previously used query, tool combination.
+        previous_tool_query_combinations = {(i.tool, i.query) for i in previous_iterations if i.warning is None}
+        if (selected_tool, generated_query) in previous_tool_query_combinations:
+            warning = f"Repeated tool, query combination detected. Skipping iteration. Try something different."
+        # Only send client status updates if we'll execute this iteration
+        elif send_status_func:
+            determined_tool_message = "**Determined Tool**: "
+            determined_tool_message += (
+                f"{selected_tool}({generated_query})." if selected_tool != ConversationCommand.Text else "respond."
+            )
+            determined_tool_message += f"\nReason: {scratchpad}" if scratchpad else ""
+            async for event in send_status_func(f"{scratchpad}"):
+                yield {ChatEvent.STATUS: event}
 
-    yield ResearchIteration(query=parsed_response, warning=warning, raw_response=response.raw_content)
+        yield ResearchIteration(
+            tool=selected_tool,
+            query=generated_query,
+            warning=warning,
+        )
+    except Exception as e:
+        logger.error(f"Invalid response for determining relevant tools: {response}. {e}", exc_info=True)
+        yield ResearchIteration(
+            tool=None,
+            query=None,
+            warning=f"Invalid response for determining relevant tools: {response}. Skipping iteration. Fix error: {e}",
+        )
 
 
 async def research(
@@ -214,10 +257,10 @@ async def research(
     MAX_ITERATIONS = int(os.getenv("KHOJ_RESEARCH_ITERATIONS", 5))
 
     # Incorporate previous partial research into current research chat history
-    research_conversation_history = [chat for chat in deepcopy(conversation_history) if chat.message]
+    research_conversation_history = deepcopy(conversation_history)
     if current_iteration := len(previous_iterations) > 0:
         logger.info(f"Continuing research with the previous {len(previous_iterations)} iteration results.")
-        previous_iterations_history = construct_iteration_history(previous_iterations)
+        previous_iterations_history = construct_iteration_history(previous_iterations, prompts.previous_iteration)
         research_conversation_history += previous_iterations_history
 
     while current_iteration < MAX_ITERATIONS:
@@ -230,7 +273,7 @@ async def research(
         code_results: Dict = dict()
         document_results: List[Dict[str, str]] = []
         operator_results: OperatorRun = None
-        this_iteration = ResearchIteration(query=query)
+        this_iteration = ResearchIteration(tool=None, query=query)
 
         async for result in apick_next_tool(
             query,
@@ -260,30 +303,26 @@ async def research(
             logger.warning(f"Research mode: {this_iteration.warning}.")
 
         # Terminate research if selected text tool or query, tool not set for next iteration
-        elif (
-            not this_iteration.query
-            or isinstance(this_iteration.query, str)
-            or this_iteration.query.name == ConversationCommand.Text
-        ):
+        elif not this_iteration.query or not this_iteration.tool or this_iteration.tool == ConversationCommand.Text:
             current_iteration = MAX_ITERATIONS
 
-        elif this_iteration.query.name == ConversationCommand.SemanticSearchFiles:
+        elif this_iteration.tool == ConversationCommand.Notes:
             this_iteration.context = []
             document_results = []
             previous_inferred_queries = {
                 c["query"] for iteration in previous_iterations if iteration.context for c in iteration.context
             }
             async for result in search_documents(
-                **this_iteration.query.args,
-                n=max_document_searches,
-                d=None,
-                user=user,
-                chat_history=construct_tool_chat_history(previous_iterations, ConversationCommand.SemanticSearchFiles),
-                conversation_id=conversation_id,
-                conversation_commands=[ConversationCommand.Default],
-                location_data=location,
-                send_status_func=send_status_func,
-                query_images=query_images,
+                this_iteration.query,
+                max_document_searches,
+                None,
+                user,
+                construct_tool_chat_history(previous_iterations, ConversationCommand.Notes),
+                conversation_id,
+                [ConversationCommand.Default],
+                location,
+                send_status_func,
+                query_images,
                 previous_inferred_queries=previous_inferred_queries,
                 agent=agent,
                 tracer=tracer,
@@ -311,7 +350,7 @@ async def research(
             else:
                 this_iteration.warning = "No matching document references found"
 
-        elif this_iteration.query.name == ConversationCommand.SearchWeb:
+        elif this_iteration.tool == ConversationCommand.Online:
             previous_subqueries = {
                 subquery
                 for iteration in previous_iterations
@@ -320,12 +359,12 @@ async def research(
             }
             try:
                 async for result in search_online(
-                    **this_iteration.query.args,
-                    conversation_history=construct_tool_chat_history(previous_iterations, ConversationCommand.Online),
-                    location=location,
-                    user=user,
-                    send_status_func=send_status_func,
-                    custom_filters=[],
+                    this_iteration.query,
+                    construct_tool_chat_history(previous_iterations, ConversationCommand.Online),
+                    location,
+                    user,
+                    send_status_func,
+                    [],
                     max_online_searches=max_online_searches,
                     max_webpages_to_read=0,
                     query_images=query_images,
@@ -344,15 +383,19 @@ async def research(
                 this_iteration.warning = f"Error searching online: {e}"
                 logger.error(this_iteration.warning, exc_info=True)
 
-        elif this_iteration.query.name == ConversationCommand.ReadWebpage:
+        elif this_iteration.tool == ConversationCommand.Webpage:
             try:
-                async for result in read_webpages_content(
-                    **this_iteration.query.args,
-                    user=user,
-                    send_status_func=send_status_func,
-                    # max_webpages_to_read=max_webpages_to_read,
+                async for result in read_webpages(
+                    this_iteration.query,
+                    construct_tool_chat_history(previous_iterations, ConversationCommand.Webpage),
+                    location,
+                    user,
+                    send_status_func,
+                    max_webpages_to_read=max_webpages_to_read,
+                    query_images=query_images,
                     agent=agent,
                     tracer=tracer,
+                    query_files=query_files,
                 ):
                     if isinstance(result, dict) and ChatEvent.STATUS in result:
                         yield result[ChatEvent.STATUS]
@@ -373,15 +416,15 @@ async def research(
                 this_iteration.warning = f"Error reading webpages: {e}"
                 logger.error(this_iteration.warning, exc_info=True)
 
-        elif this_iteration.query.name == ConversationCommand.RunCode:
+        elif this_iteration.tool == ConversationCommand.Code:
             try:
                 async for result in run_code(
-                    **this_iteration.query.args,
-                    conversation_history=construct_tool_chat_history(previous_iterations, ConversationCommand.Code),
-                    context="",
-                    location_data=location,
-                    user=user,
-                    send_status_func=send_status_func,
+                    this_iteration.query,
+                    construct_tool_chat_history(previous_iterations, ConversationCommand.Code),
+                    "",
+                    location,
+                    user,
+                    send_status_func,
                     query_images=query_images,
                     agent=agent,
                     query_files=query_files,
@@ -398,14 +441,14 @@ async def research(
                 this_iteration.warning = f"Error running code: {e}"
                 logger.warning(this_iteration.warning, exc_info=True)
 
-        elif this_iteration.query.name == ConversationCommand.OperateComputer:
+        elif this_iteration.tool == ConversationCommand.Operator:
             try:
                 async for result in operate_environment(
-                    **this_iteration.query.args,
-                    user=user,
-                    conversation_log=construct_tool_chat_history(previous_iterations, ConversationCommand.Operator),
-                    location_data=location,
-                    previous_trajectory=previous_iterations[-1].operatorContext if previous_iterations else None,
+                    this_iteration.query,
+                    user,
+                    construct_tool_chat_history(previous_iterations, ConversationCommand.Operator),
+                    location,
+                    previous_iterations[-1].operatorContext if previous_iterations else None,
                     send_status_func=send_status_func,
                     query_images=query_images,
                     agent=agent,
@@ -431,63 +474,6 @@ async def research(
                 this_iteration.warning = f"Error operating browser: {e}"
                 logger.error(this_iteration.warning, exc_info=True)
 
-        elif this_iteration.query.name == ConversationCommand.ViewFile:
-            try:
-                async for result in view_file_content(
-                    **this_iteration.query.args,
-                    user=user,
-                ):
-                    if isinstance(result, dict) and ChatEvent.STATUS in result:
-                        yield result[ChatEvent.STATUS]
-                    else:
-                        if this_iteration.context is None:
-                            this_iteration.context = []
-                        document_results: List[Dict[str, str]] = result  # type: ignore
-                        this_iteration.context += document_results
-                async for result in send_status_func(f"**Viewed file**: {this_iteration.query.args['path']}"):
-                    yield result
-            except Exception as e:
-                this_iteration.warning = f"Error viewing file: {e}"
-                logger.error(this_iteration.warning, exc_info=True)
-
-        elif this_iteration.query.name == ConversationCommand.ListFiles:
-            try:
-                async for result in list_files(
-                    **this_iteration.query.args,
-                    user=user,
-                ):
-                    if isinstance(result, dict) and ChatEvent.STATUS in result:
-                        yield result[ChatEvent.STATUS]
-                    else:
-                        if this_iteration.context is None:
-                            this_iteration.context = []
-                        document_results: List[Dict[str, str]] = [result]  # type: ignore
-                        this_iteration.context += document_results
-                async for result in send_status_func(result["query"]):
-                    yield result
-            except Exception as e:
-                this_iteration.warning = f"Error listing files: {e}"
-                logger.error(this_iteration.warning, exc_info=True)
-
-        elif this_iteration.query.name == ConversationCommand.RegexSearchFiles:
-            try:
-                async for result in grep_files(
-                    **this_iteration.query.args,
-                    user=user,
-                ):
-                    if isinstance(result, dict) and ChatEvent.STATUS in result:
-                        yield result[ChatEvent.STATUS]
-                    else:
-                        if this_iteration.context is None:
-                            this_iteration.context = []
-                        document_results: List[Dict[str, str]] = [result]  # type: ignore
-                        this_iteration.context += document_results
-                async for result in send_status_func(result["query"]):
-                    yield result
-            except Exception as e:
-                this_iteration.warning = f"Error searching with regex: {e}"
-                logger.error(this_iteration.warning, exc_info=True)
-
         else:
             # No valid tools. This is our exit condition.
             current_iteration = MAX_ITERATIONS
@@ -495,7 +481,7 @@ async def research(
         current_iteration += 1
 
         if document_results or online_results or code_results or operator_results or this_iteration.warning:
-            results_data = f"\n<iteration_{current_iteration}_results>"
+            results_data = f"\n<iteration>{current_iteration}\n<tool>{this_iteration.tool}</tool>\n<query>{this_iteration.query}</query>\n<results>"
             if document_results:
                 results_data += f"\n<document_references>\n{yaml.dump(document_results, allow_unicode=True, sort_keys=False, default_flow_style=False)}\n</document_references>"
             if online_results:
@@ -508,7 +494,7 @@ async def research(
                 )
             if this_iteration.warning:
                 results_data += f"\n<warning>\n{this_iteration.warning}\n</warning>"
-            results_data += f"\n</results>\n</iteration_{current_iteration}_results>"
+            results_data += "\n</results>\n</iteration>"
 
             # intermediate_result = await extract_relevant_info(this_iteration.query, results_data, agent)
             this_iteration.summarizedResult = results_data

khoj/search_type/text_search.py

@@ -157,7 +157,6 @@ def collate_results(hits, dedupe=True):
                     "additional": {
                         "source": hit.file_source,
                         "file": hit.file_path,
-                        "uri": hit.url,
                         "compiled": hit.compiled,
                         "heading": hit.heading,
                     },
@@ -181,7 +180,6 @@ def deduplicated_search_responses(hits: List[SearchResponse]):
                     "additional": {
                         "source": hit.additional["source"],
                         "file": hit.additional["file"],
-                        "uri": hit.additional["uri"],
                         "query": hit.additional["query"],
                         "compiled": hit.additional["compiled"],
                         "heading": hit.additional["heading"],