Make researcher aware of no. of web, doc queries allowed per iteration

- Construct tool description dynamically based on configurable query
  count
- Inform the researcher how many webpage reads, online searches and
  document searches it can perform per iteration when it has to decide
  which next tool to use and the query to send to the tool AI.
- Pass the query counts to perform from the research AI down to the
  tool AIs
This commit is contained in:
Debanjum
2025-05-13 12:00:06 -06:00
parent 417ab42206
commit 988bde651c
6 changed files with 40 additions and 24 deletions

View File

@@ -865,8 +865,7 @@ infer_webpages_to_read = PromptTemplate.from_template(
You are Khoj, an advanced web page reading assistant. You are to construct **up to {max_webpages}, valid** webpage urls to read before answering the user's question.
- You will receive the conversation history as context.
- Add as much context from the previous questions and answers as required to construct the webpage urls.
- Use multiple web page urls if required to retrieve the relevant information.
- You have access to the the whole internet to retrieve information.
- You have access to the whole internet to retrieve information.
{personality_context}
Which webpages will you need to read to answer the user's question?
Provide web page links as a list of strings in a JSON object.

View File

@@ -64,11 +64,12 @@ async def search_online(
user: KhojUser,
send_status_func: Optional[Callable] = None,
custom_filters: List[str] = [],
max_online_searches: int = 3,
max_webpages_to_read: int = 1,
query_images: List[str] = None,
query_files: str = None,
previous_subqueries: Set = set(),
agent: Agent = None,
query_files: str = None,
tracer: dict = {},
):
query += " ".join(custom_filters)
@@ -84,9 +85,10 @@ async def search_online(
location,
user,
query_images=query_images,
query_files=query_files,
max_queries=max_online_searches,
agent=agent,
tracer=tracer,
query_files=query_files,
)
subqueries = list(new_subqueries - previous_subqueries)
response_dict: Dict[str, Dict[str, List[Dict] | Dict]] = {}

View File

@@ -1129,9 +1129,10 @@ async def chat(
user,
partial(send_event, ChatEvent.STATUS),
custom_filters,
max_online_searches=3,
query_images=uploaded_images,
agent=agent,
query_files=attached_file_context,
agent=agent,
tracer=tracer,
):
if isinstance(result, dict) and ChatEvent.STATUS in result:

View File

@@ -523,8 +523,9 @@ async def generate_online_subqueries(
location_data: LocationData,
user: KhojUser,
query_images: List[str] = None,
agent: Agent = None,
query_files: str = None,
max_queries: int = 3,
agent: Agent = None,
tracer: dict = {},
) -> Set[str]:
"""
@@ -534,7 +535,6 @@ async def generate_online_subqueries(
username = prompts.user_name.format(name=user.get_full_name()) if user.get_full_name() else ""
chat_history = construct_chat_history(conversation_history)
max_queries = 3
utc_date = datetime.now(timezone.utc).strftime("%Y-%m-%d")
personality_context = (
prompts.personality_context.format(personality=agent.personality) if agent and agent.personality else ""

View File

@@ -6,7 +6,6 @@ from enum import Enum
from typing import Callable, Dict, List, Optional, Type
import yaml
from fastapi import Request
from pydantic import BaseModel, Field
from khoj.database.adapters import AgentAdapters, EntryAdapters
@@ -14,7 +13,6 @@ from khoj.database.models import Agent, KhojUser
from khoj.processor.conversation import prompts
from khoj.processor.conversation.utils import (
InformationCollectionIteration,
construct_chat_history,
construct_iteration_history,
construct_tool_chat_history,
load_complex_json,
@@ -29,9 +27,9 @@ from khoj.routers.helpers import (
)
from khoj.utils.helpers import (
ConversationCommand,
function_calling_description_for_llm,
is_none_or_empty,
timer,
tool_description_for_research_llm,
truncate_code_context,
)
from khoj.utils.rawconfig import LocationData
@@ -79,15 +77,18 @@ async def apick_next_tool(
query: str,
conversation_history: dict,
user: KhojUser = None,
query_images: List[str] = [],
location: LocationData = None,
user_name: str = None,
agent: Agent = None,
previous_iterations: List[InformationCollectionIteration] = [],
max_iterations: int = 5,
query_images: List[str] = [],
query_files: str = None,
max_document_searches: int = 7,
max_online_searches: int = 3,
max_webpages_to_read: int = 1,
send_status_func: Optional[Callable] = None,
tracer: dict = {},
query_files: str = None,
):
"""Given a query, determine which of the available tools the agent should use in order to answer appropriately."""
@@ -96,10 +97,16 @@ async def apick_next_tool(
tool_options_str = ""
agent_tools = agent.input_tools if agent else []
user_has_entries = await EntryAdapters.auser_has_entries(user)
for tool, description in function_calling_description_for_llm.items():
for tool, description in tool_description_for_research_llm.items():
# Skip showing Notes tool as an option if user has no entries
if tool == ConversationCommand.Notes and not user_has_entries:
continue
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
@@ -210,6 +217,9 @@ async def execute_information_collection(
query_files: str = None,
cancellation_event: Optional[asyncio.Event] = None,
):
max_document_searches = 7
max_online_searches = 3
max_webpages_to_read = 1
current_iteration = 0
MAX_ITERATIONS = int(os.getenv("KHOJ_RESEARCH_ITERATIONS", 5))
previous_iterations: List[InformationCollectionIteration] = []
@@ -229,15 +239,18 @@ async def execute_information_collection(
query,
conversation_history,
user,
query_images,
location,
user_name,
agent,
previous_iterations,
MAX_ITERATIONS,
send_status_func,
tracer=tracer,
query_images=query_images,
query_files=query_files,
max_document_searches=max_document_searches,
max_online_searches=max_online_searches,
max_webpages_to_read=max_webpages_to_read,
send_status_func=send_status_func,
tracer=tracer,
):
if isinstance(result, dict) and ChatEvent.STATUS in result:
yield result[ChatEvent.STATUS]
@@ -262,7 +275,7 @@ async def execute_information_collection(
user,
construct_tool_chat_history(previous_iterations, ConversationCommand.Notes),
this_iteration.query,
7,
max_document_searches,
None,
conversation_id,
[ConversationCommand.Default],
@@ -309,6 +322,7 @@ async def execute_information_collection(
user,
send_status_func,
[],
max_online_searches=max_online_searches,
max_webpages_to_read=0,
query_images=query_images,
previous_subqueries=previous_subqueries,
@@ -334,7 +348,7 @@ async def execute_information_collection(
location,
user,
send_status_func,
max_webpages_to_read=1,
max_webpages_to_read=max_webpages_to_read,
query_images=query_images,
agent=agent,
tracer=tracer,

View File

@@ -386,10 +386,10 @@ tool_descriptions_for_llm = {
ConversationCommand.Code: e2b_tool_description if is_e2b_code_sandbox_enabled() else terrarium_tool_description,
}
function_calling_description_for_llm = {
ConversationCommand.Notes: "To search the user's personal knowledge base. Especially helpful if the question expects context from the user's notes or documents.",
ConversationCommand.Online: "To search the internet for information. Useful to get a quick, broad overview from the internet. Provide all relevant context to ensure new searches, not in previous iterations, are performed.",
ConversationCommand.Webpage: "To extract information from webpages. Useful for more detailed research from the internet. Usually used when you know the webpage links to refer to. Share the webpage links and information to extract in your query.",
tool_description_for_research_llm = {
ConversationCommand.Notes: "To search the user's personal knowledge base. Especially helpful if the question expects context from the user's notes or documents. Max {max_search_queries} search queries allowed per iteration.",
ConversationCommand.Online: "To search the internet for information. Useful to get a quick, broad overview from the internet. Provide all relevant context to ensure new searches, not in previous iterations, are performed. Max {max_search_queries} search queries allowed per iteration.",
ConversationCommand.Webpage: "To extract information from webpages. Useful for more detailed research from the internet. Usually used when you know the webpage links to refer to. Share upto {max_webpages_to_read} webpage links and what information to extract from them in your query.",
ConversationCommand.Code: e2b_tool_description if is_e2b_code_sandbox_enabled() else terrarium_tool_description,
ConversationCommand.Text: "To respond to the user once you've completed your research and have the required information.",
}