From 38dd02afbfa7ec6b8ae79e54d8105779533d7c0a Mon Sep 17 00:00:00 2001 From: Debanjum Date: Fri, 4 Apr 2025 09:42:13 +0530 Subject: [PATCH] Make ordering of fields expected by research planner consistent Make research planner consistently select tool before query. As the model should tune it's query for the selected tool. It got space to think about tool to use in the scratchpad already. --- src/khoj/processor/conversation/prompts.py | 2 +- src/khoj/routers/research.py | 8 +++++--- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/src/khoj/processor/conversation/prompts.py b/src/khoj/processor/conversation/prompts.py index bcaa3600..e77690f8 100644 --- a/src/khoj/processor/conversation/prompts.py +++ b/src/khoj/processor/conversation/prompts.py @@ -770,7 +770,7 @@ Which of the tool AIs listed below would you use to answer the user's question? Return the next tool AI to use and the query to ask it. Your response should always be a valid JSON object. Do not say anything else. Response format: -{{"scratchpad": "", "query": "", "tool": ""}} +{{"scratchpad": "", "tool": "", "query": ""}} """.strip() ) diff --git a/src/khoj/routers/research.py b/src/khoj/routers/research.py index b662dca9..aab0949b 100644 --- a/src/khoj/routers/research.py +++ b/src/khoj/routers/research.py @@ -41,11 +41,9 @@ logger = logging.getLogger(__name__) class PlanningResponse(BaseModel): """ Schema for the response from planning agent when deciding the next tool to pick. - The tool field is dynamically validated based on available tools. """ - scratchpad: str = Field(..., description="Reasoning about which tool to use next") - query: str = Field(..., description="Detailed query for the selected tool") + scratchpad: str = Field(..., description="Scratchpad to reason about which tool to use next") class Config: arbitrary_types_allowed = True @@ -56,6 +54,9 @@ class PlanningResponse(BaseModel): 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 @@ -68,6 +69,7 @@ class PlanningResponse(BaseModel): # Create and return a customized response model with the enum class PlanningResponseWithTool(PlanningResponse): tool: tool_enum = Field(..., description="Name of the tool to use") + query: str = Field(..., description="Detailed query for the selected tool") return PlanningResponseWithTool