khoj 1.41.1.dev97__py3-none-any.whl → 1.41.1.dev142__py3-none-any.whl

khoj/routers/api_chat.py

@@ -26,12 +26,13 @@ from khoj.database.models import Agent, KhojUser
 from khoj.processor.conversation import prompts
 from khoj.processor.conversation.prompts import help_message, no_entries_found
 from khoj.processor.conversation.utils import (
+    OperatorRun,
     ResponseWithThought,
     defilter_query,
     save_to_conversation_log,
 )
 from khoj.processor.image.generate import text_to_image
-from khoj.processor.operator.operate_browser import operate_browser
+from khoj.processor.operator import operate_environment
 from khoj.processor.speech.text_to_speech import generate_text_to_speech
 from khoj.processor.tools.online_search import (
     deduplicate_organic_results,
@@ -65,10 +66,7 @@ from khoj.routers.helpers import (
     update_telemetry_state,
     validate_chat_model,
 )
-from khoj.routers.research import (
-    InformationCollectionIteration,
-    execute_information_collection,
-)
+from khoj.routers.research import ResearchIteration, research
 from khoj.routers.storage import upload_user_image_to_bucket
 from khoj.utils import state
 from khoj.utils.helpers import (
@@ -682,11 +680,13 @@ async def chat(
     timezone = body.timezone
     raw_images = body.images
     raw_query_files = body.files
+    interrupt_flag = body.interrupt
 
     async def event_generator(q: str, images: list[str]):
         start_time = time.perf_counter()
         ttft = None
         chat_metadata: dict = {}
+        conversation = None
         user: KhojUser = request.user.object
         is_subscribed = has_required_scope(request, ["premium"])
         q = unquote(q)
@@ -720,6 +720,20 @@ async def chat(
             for file in raw_query_files:
                 query_files[file.name] = file.content
 
+        research_results: List[ResearchIteration] = []
+        online_results: Dict = dict()
+        code_results: Dict = dict()
+        operator_results: List[OperatorRun] = []
+        compiled_references: List[Any] = []
+        inferred_queries: List[Any] = []
+        attached_file_context = gather_raw_query_files(query_files)
+
+        generated_images: List[str] = []
+        generated_files: List[FileAttachment] = []
+        generated_mermaidjs_diagram: str = None
+        generated_asset_results: Dict = dict()
+        program_execution_context: List[str] = []
+
         # Create a task to monitor for disconnections
         disconnect_monitor_task = None
 
@@ -727,8 +741,34 @@ async def chat(
             try:
                 msg = await request.receive()
                 if msg["type"] == "http.disconnect":
-                    logger.debug(f"User {user} disconnected from {common.client} client.")
+                    logger.debug(f"Request cancelled. User {user} disconnected from {common.client} client.")
                     cancellation_event.set()
+                    # ensure partial chat state saved on interrupt
+                    # shield the save against task cancellation
+                    if conversation:
+                        await asyncio.shield(
+                            save_to_conversation_log(
+                                q,
+                                chat_response="",
+                                user=user,
+                                meta_log=meta_log,
+                                compiled_references=compiled_references,
+                                online_results=online_results,
+                                code_results=code_results,
+                                operator_results=operator_results,
+                                research_results=research_results,
+                                inferred_queries=inferred_queries,
+                                client_application=request.user.client_app,
+                                conversation_id=conversation_id,
+                                query_images=uploaded_images,
+                                train_of_thought=train_of_thought,
+                                raw_query_files=raw_query_files,
+                                generated_images=generated_images,
+                                raw_generated_files=generated_asset_results,
+                                generated_mermaidjs_diagram=generated_mermaidjs_diagram,
+                                tracer=tracer,
+                            )
+                        )
             except Exception as e:
                 logger.error(f"Error in disconnect monitor: {e}")
 
@@ -746,7 +786,6 @@ async def chat(
             nonlocal ttft, train_of_thought
             event_delimiter = "␃🔚␗"
             if cancellation_event.is_set():
-                logger.debug(f"User {user} disconnected from {common.client} client. Setting cancellation event.")
                 return
             try:
                 if event_type == ChatEvent.END_LLM_RESPONSE:
@@ -770,9 +809,6 @@ async def chat(
                     yield data
                 elif event_type == ChatEvent.REFERENCES or ChatEvent.METADATA or stream:
                     yield json.dumps({"type": event_type.value, "data": data}, ensure_ascii=False)
-            except asyncio.CancelledError as e:
-                if cancellation_event.is_set():
-                    logger.debug(f"Request cancelled. User {user} disconnected from {common.client} client: {e}.")
             except Exception as e:
                 if not cancellation_event.is_set():
                     logger.error(
@@ -883,21 +919,52 @@ async def chat(
         user_message_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
         meta_log = conversation.conversation_log
 
-        researched_results = ""
-        online_results: Dict = dict()
-        code_results: Dict = dict()
-        operator_results: Dict[str, str] = {}
-        generated_asset_results: Dict = dict()
-        ## Extract Document References
-        compiled_references: List[Any] = []
-        inferred_queries: List[Any] = []
-        file_filters = conversation.file_filters if conversation and conversation.file_filters else []
-        attached_file_context = gather_raw_query_files(query_files)
+        # If interrupt flag is set, wait for the previous turn to be saved before proceeding
+        if interrupt_flag:
+            max_wait_time = 20.0  # seconds
+            wait_interval = 0.3  # seconds
+            wait_start = wait_current = time.time()
+            while wait_current - wait_start < max_wait_time:
+                # Refresh conversation to check if interrupted message saved to DB
+                conversation = await ConversationAdapters.aget_conversation_by_user(
+                    user,
+                    client_application=request.user.client_app,
+                    conversation_id=conversation_id,
+                )
+                if (
+                    conversation
+                    and conversation.messages
+                    and conversation.messages[-1].by == "khoj"
+                    and not conversation.messages[-1].message
+                ):
+                    logger.info(f"Detected interrupted message save to conversation {conversation_id}.")
+                    break
+                await asyncio.sleep(wait_interval)
+                wait_current = time.time()
 
-        generated_images: List[str] = []
-        generated_files: List[FileAttachment] = []
-        generated_mermaidjs_diagram: str = None
-        program_execution_context: List[str] = []
+            if wait_current - wait_start >= max_wait_time:
+                logger.warning(
+                    f"Timeout waiting to load interrupted context from conversation {conversation_id}. Proceed without previous context."
+                )
+
+        # If interrupted message in DB
+        if (
+            conversation
+            and conversation.messages
+            and conversation.messages[-1].by == "khoj"
+            and not conversation.messages[-1].message
+        ):
+            # Populate context from interrupted message
+            last_message = conversation.messages[-1]
+            online_results = {key: val.model_dump() for key, val in last_message.onlineContext.items() or []}
+            code_results = {key: val.model_dump() for key, val in last_message.codeContext.items() or []}
+            compiled_references = [ref.model_dump() for ref in last_message.context or []]
+            research_results = [ResearchIteration(**iter_dict) for iter_dict in last_message.researchContext or []]
+            operator_results = [OperatorRun(**iter_dict) for iter_dict in last_message.operatorContext or []]
+            train_of_thought = [thought.model_dump() for thought in last_message.trainOfThought or []]
+            # Drop the interrupted message from conversation history
+            meta_log["chat"].pop()
+            logger.info(f"Loaded interrupted partial context from conversation {conversation_id}.")
 
         if conversation_commands == [ConversationCommand.Default]:
             try:
@@ -936,24 +1003,26 @@ async def chat(
                 return
 
         defiltered_query = defilter_query(q)
+        file_filters = conversation.file_filters if conversation and conversation.file_filters else []
 
         if conversation_commands == [ConversationCommand.Research]:
-            async for research_result in execute_information_collection(
+            async for research_result in research(
                 user=user,
                 query=defiltered_query,
                 conversation_id=conversation_id,
                 conversation_history=meta_log,
+                previous_iterations=list(research_results),
                 query_images=uploaded_images,
                 agent=agent,
                 send_status_func=partial(send_event, ChatEvent.STATUS),
                 user_name=user_name,
                 location=location,
-                file_filters=conversation.file_filters if conversation else [],
+                file_filters=file_filters,
                 query_files=attached_file_context,
                 tracer=tracer,
                 cancellation_event=cancellation_event,
             ):
-                if isinstance(research_result, InformationCollectionIteration):
+                if isinstance(research_result, ResearchIteration):
                     if research_result.summarizedResult:
                         if research_result.onlineContext:
                             online_results.update(research_result.onlineContext)
@@ -961,19 +1030,31 @@ async def chat(
                             code_results.update(research_result.codeContext)
                         if research_result.context:
                             compiled_references.extend(research_result.context)
-                        if research_result.operatorContext:
-                            operator_results.update(research_result.operatorContext)
-                        researched_results += research_result.summarizedResult
-
+                    if not research_results or research_results[-1] is not research_result:
+                        research_results.append(research_result)
                 else:
                     yield research_result
 
+                # Track operator results across research and operator iterations
+                # This relies on two conditions:
+                # 1. Check to append new (partial) operator results
+                #    Relies on triggering this check on every status updates.
+                #    Status updates cascade up from operator to research to chat api on every step.
+                # 2. Keep operator results in sync with each research operator step
+                #    Relies on python object references to ensure operator results
+                #    are implicitly kept in sync after the initial append
+                if (
+                    research_results
+                    and research_results[-1].operatorContext
+                    and (not operator_results or operator_results[-1] is not research_results[-1].operatorContext)
+                ):
+                    operator_results.append(research_results[-1].operatorContext)
+
             # researched_results = await extract_relevant_info(q, researched_results, agent)
             if state.verbose > 1:
-                logger.debug(f"Researched Results: {researched_results}")
+                logger.debug(f'Researched Results: {"".join(r.summarizedResult for r in research_results)}')
 
         used_slash_summarize = conversation_commands == [ConversationCommand.Summarize]
-        file_filters = conversation.file_filters if conversation else []
         # Skip trying to summarize if
         if (
             # summarization intent was inferred
@@ -1221,11 +1302,12 @@ async def chat(
                 )
         if ConversationCommand.Operator in conversation_commands:
             try:
-                async for result in operate_browser(
+                async for result in operate_environment(
                     defiltered_query,
                     user,
                     meta_log,
                     location,
+                    list(operator_results)[-1] if operator_results else None,
                     query_images=uploaded_images,
                     query_files=attached_file_context,
                     send_status_func=partial(send_event, ChatEvent.STATUS),
@@ -1235,16 +1317,17 @@ async def chat(
                 ):
                     if isinstance(result, dict) and ChatEvent.STATUS in result:
                         yield result[ChatEvent.STATUS]
-                    else:
-                        operator_results = {result["query"]: result["result"]}
+                    elif isinstance(result, OperatorRun):
+                        if not operator_results or operator_results[-1] is not result:
+                            operator_results.append(result)
                         # Add webpages visited while operating browser to references
-                        if result.get("webpages"):
+                        if result.webpages:
                             if not online_results.get(defiltered_query):
-                                online_results[defiltered_query] = {"webpages": result["webpages"]}
+                                online_results[defiltered_query] = {"webpages": result.webpages}
                             elif not online_results[defiltered_query].get("webpages"):
-                                online_results[defiltered_query]["webpages"] = result["webpages"]
+                                online_results[defiltered_query]["webpages"] = result.webpages
                             else:
-                                online_results[defiltered_query]["webpages"] += result["webpages"]
+                                online_results[defiltered_query]["webpages"] += result.webpages
             except ValueError as e:
                 program_execution_context.append(f"Browser operation error: {e}")
                 logger.warning(f"Failed to operate browser with {e}", exc_info=True)
@@ -1262,7 +1345,6 @@ async def chat(
                 "context": compiled_references,
                 "onlineContext": unique_online_results,
                 "codeContext": code_results,
-                "operatorContext": operator_results,
             },
         ):
             yield result
@@ -1362,7 +1444,7 @@ async def chat(
 
         # Check if the user has disconnected
         if cancellation_event.is_set():
-            logger.debug(f"User {user} disconnected from {common.client} client. Stopping LLM response.")
+            logger.debug(f"Stopping LLM response to user {user} on {common.client} client.")
             # Cancel the disconnect monitor task if it is still running
             await cancel_disconnect_monitor()
             return
@@ -1379,13 +1461,13 @@ async def chat(
             online_results,
             code_results,
             operator_results,
+            research_results,
             inferred_queries,
             conversation_commands,
             user,
             request.user.client_app,
             location,
             user_name,
-            researched_results,
             uploaded_images,
             train_of_thought,
             attached_file_context,

khoj/routers/api_model.py

@@ -72,7 +72,7 @@ async def update_chat_model(
     if chat_model is None:
         return Response(status_code=404, content=json.dumps({"status": "error", "message": "Chat model not found"}))
     if not subscribed and chat_model.price_tier != PriceTier.FREE:
-        raise Response(
+        return Response(
             status_code=403,
             content=json.dumps({"status": "error", "message": "Subscribe to switch to this chat model"}),
         )
@@ -108,7 +108,7 @@ async def update_voice_model(
     if voice_model is None:
         return Response(status_code=404, content=json.dumps({"status": "error", "message": "Voice model not found"}))
     if not subscribed and voice_model.price_tier != PriceTier.FREE:
-        raise Response(
+        return Response(
             status_code=403,
             content=json.dumps({"status": "error", "message": "Subscribe to switch to this voice model"}),
         )
@@ -143,7 +143,7 @@ async def update_paint_model(
     if image_model is None:
         return Response(status_code=404, content=json.dumps({"status": "error", "message": "Image model not found"}))
     if not subscribed and image_model.price_tier != PriceTier.FREE:
-        raise Response(
+        return Response(
             status_code=403,
             content=json.dumps({"status": "error", "message": "Subscribe to switch to this image model"}),
         )

khoj/routers/helpers.py

@@ -94,6 +94,8 @@ from khoj.processor.conversation.openai.gpt import (
 )
 from khoj.processor.conversation.utils import (
     ChatEvent,
+    OperatorRun,
+    ResearchIteration,
     ResponseWithThought,
     clean_json,
     clean_mermaidjs,
@@ -384,7 +386,7 @@ async def aget_data_sources_and_output_format(
         if len(agent_outputs) == 0 or output.value in agent_outputs:
             output_options_str += f'- "{output.value}": "{description}"\n'
 
-    chat_history = construct_chat_history(conversation_history)
+    chat_history = construct_chat_history(conversation_history, n=6)
 
     if query_images:
         query = f"[placeholder for {len(query_images)} user attached images]\n{query}"
@@ -1173,12 +1175,7 @@ async def send_message_to_model_wrapper(
     if vision_available and query_images:
         logger.info(f"Using {chat_model.name} model to understand {len(query_images)} images.")
 
-    subscribed = await ais_user_subscribed(user) if user else False
-    max_tokens = (
-        chat_model.subscribed_max_prompt_size
-        if subscribed and chat_model.subscribed_max_prompt_size
-        else chat_model.max_prompt_size
-    )
+    max_tokens = await ConversationAdapters.aget_max_context_size(chat_model, user)
     chat_model_name = chat_model.name
     tokenizer = chat_model.tokenizer
     model_type = chat_model.model_type
@@ -1270,12 +1267,7 @@ def send_message_to_model_wrapper_sync(
     if chat_model is None:
         raise HTTPException(status_code=500, detail="Contact the server administrator to set a default chat model.")
 
-    subscribed = is_user_subscribed(user) if user else False
-    max_tokens = (
-        chat_model.subscribed_max_prompt_size
-        if subscribed and chat_model.subscribed_max_prompt_size
-        else chat_model.max_prompt_size
-    )
+    max_tokens = ConversationAdapters.get_max_context_size(chat_model, user)
     chat_model_name = chat_model.name
     model_type = chat_model.model_type
     vision_available = chat_model.vision_enabled
@@ -1354,14 +1346,14 @@ async def agenerate_chat_response(
     compiled_references: List[Dict] = [],
     online_results: Dict[str, Dict] = {},
     code_results: Dict[str, Dict] = {},
-    operator_results: Dict[str, str] = {},
+    operator_results: List[OperatorRun] = [],
+    research_results: List[ResearchIteration] = [],
     inferred_queries: List[str] = [],
     conversation_commands: List[ConversationCommand] = [ConversationCommand.Default],
     user: KhojUser = None,
     client_application: ClientApplication = None,
     location_data: LocationData = None,
     user_name: Optional[str] = None,
-    meta_research: str = "",
     query_images: Optional[List[str]] = None,
     train_of_thought: List[Any] = [],
     query_files: str = None,
@@ -1391,6 +1383,7 @@ async def agenerate_chat_response(
             online_results=online_results,
             code_results=code_results,
             operator_results=operator_results,
+            research_results=research_results,
             inferred_queries=inferred_queries,
             client_application=client_application,
             conversation_id=str(conversation.id),
@@ -1405,12 +1398,14 @@ async def agenerate_chat_response(
 
         query_to_run = q
         deepthought = False
-        if meta_research:
-            query_to_run = f"<query>{q}</query>\n<collected_research>\n{meta_research}\n</collected_research>"
+        if research_results:
+            compiled_research = "".join([r.summarizedResult for r in research_results if r.summarizedResult])
+            if compiled_research:
+                query_to_run = f"<query>{q}</query>\n<collected_research>\n{compiled_research}\n</collected_research>"
             compiled_references = []
             online_results = {}
             code_results = {}
-            operator_results = {}
+            operator_results = []
             deepthought = True
 
         chat_model = await ConversationAdapters.aget_valid_chat_model(user, conversation, is_subscribed)

khoj/routers/research.py

@@ -1,6 +1,7 @@
 import asyncio
 import logging
 import os
+from copy import deepcopy
 from datetime import datetime
 from enum import Enum
 from typing import Callable, Dict, List, Optional, Type
@@ -12,12 +13,13 @@ from khoj.database.adapters import AgentAdapters, EntryAdapters
 from khoj.database.models import Agent, KhojUser
 from khoj.processor.conversation import prompts
 from khoj.processor.conversation.utils import (
-    InformationCollectionIteration,
+    OperatorRun,
+    ResearchIteration,
     construct_iteration_history,
     construct_tool_chat_history,
     load_complex_json,
 )
-from khoj.processor.operator.operate_browser import operate_browser
+from khoj.processor.operator import operate_environment
 from khoj.processor.tools.online_search import read_webpages, search_online
 from khoj.processor.tools.run_code import run_code
 from khoj.routers.api import extract_references_and_questions
@@ -82,7 +84,7 @@ async def apick_next_tool(
     location: LocationData = None,
     user_name: str = None,
     agent: Agent = None,
-    previous_iterations: List[InformationCollectionIteration] = [],
+    previous_iterations: List[ResearchIteration] = [],
     max_iterations: int = 5,
     query_images: List[str] = [],
     query_files: str = None,
@@ -94,6 +96,24 @@ async def apick_next_tool(
 ):
     """Given a query, determine which of the available tools the agent should use in order to answer appropriately."""
 
+    # Continue with previous iteration if a multi-step tool use is in progress
+    if (
+        previous_iterations
+        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,
+            codeContext=previous_iteration.codeContext,
+            operatorContext=previous_iteration.operatorContext,
+            warning=previous_iteration.warning,
+        )
+        return
+
     # Construct tool options for the agent to choose from
     tool_options = dict()
     tool_options_str = ""
@@ -141,7 +161,7 @@ async def apick_next_tool(
         query = f"[placeholder for user attached images]\n{query}"
 
     # Construct chat history with user and iteration history with researcher agent for context
-    previous_iterations_history = construct_iteration_history(query, previous_iterations, prompts.previous_iteration)
+    previous_iterations_history = construct_iteration_history(previous_iterations, prompts.previous_iteration, query)
     iteration_chat_log = {"chat": conversation_history.get("chat", []) + previous_iterations_history}
 
     # Plan function execution for the next tool
@@ -164,7 +184,7 @@ async def apick_next_tool(
             )
     except Exception as e:
         logger.error(f"Failed to infer information sources to refer: {e}", exc_info=True)
-        yield InformationCollectionIteration(
+        yield ResearchIteration(
             tool=None,
             query=None,
             warning="Failed to infer information sources to refer. Skipping iteration. Try again.",
@@ -193,25 +213,26 @@ async def apick_next_tool(
             async for event in send_status_func(f"{scratchpad}"):
                 yield {ChatEvent.STATUS: event}
 
-        yield InformationCollectionIteration(
+        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 InformationCollectionIteration(
+        yield ResearchIteration(
             tool=None,
             query=None,
             warning=f"Invalid response for determining relevant tools: {response}. Skipping iteration. Fix error: {e}",
         )
 
 
-async def execute_information_collection(
+async def research(
     user: KhojUser,
     query: str,
     conversation_id: str,
     conversation_history: dict,
+    previous_iterations: List[ResearchIteration],
     query_images: List[str],
     agent: Agent = None,
     send_status_func: Optional[Callable] = None,
@@ -227,23 +248,32 @@ async def execute_information_collection(
     max_webpages_to_read = 1
     current_iteration = 0
     MAX_ITERATIONS = int(os.getenv("KHOJ_RESEARCH_ITERATIONS", 5))
-    previous_iterations: List[InformationCollectionIteration] = []
+
+    # Incorporate previous partial research into current research chat history
+    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, prompts.previous_iteration)
+        research_conversation_history["chat"] = (
+            research_conversation_history.get("chat", []) + previous_iterations_history
+        )
+
     while current_iteration < MAX_ITERATIONS:
         # Check for cancellation at the start of each iteration
         if cancellation_event and cancellation_event.is_set():
-            logger.debug(f"User {user} disconnected client. Research cancelled.")
+            logger.debug(f"Research cancelled. User {user} disconnected client.")
             break
 
         online_results: Dict = dict()
         code_results: Dict = dict()
         document_results: List[Dict[str, str]] = []
-        operator_results: Dict[str, str] = {}
+        operator_results: OperatorRun = None
         summarize_files: str = ""
-        this_iteration = InformationCollectionIteration(tool=None, query=query)
+        this_iteration = ResearchIteration(tool=None, query=query)
 
         async for result in apick_next_tool(
             query,
-            conversation_history,
+            research_conversation_history,
             user,
             location,
             user_name,
@@ -260,8 +290,9 @@ async def execute_information_collection(
         ):
             if isinstance(result, dict) and ChatEvent.STATUS in result:
                 yield result[ChatEvent.STATUS]
-            elif isinstance(result, InformationCollectionIteration):
+            elif isinstance(result, ResearchIteration):
                 this_iteration = result
+                yield this_iteration
 
         # Skip running iteration if warning present in iteration
         if this_iteration.warning:
@@ -406,12 +437,13 @@ async def execute_information_collection(
 
         elif this_iteration.tool == ConversationCommand.Operator:
             try:
-                async for result in operate_browser(
+                async for result in operate_environment(
                     this_iteration.query,
                     user,
                     construct_tool_chat_history(previous_iterations, ConversationCommand.Operator),
                     location,
-                    send_status_func,
+                    previous_iterations[-1].operatorContext if previous_iterations else None,
+                    send_status_func=send_status_func,
                     query_images=query_images,
                     agent=agent,
                     query_files=query_files,
@@ -420,17 +452,17 @@ async def execute_information_collection(
                 ):
                     if isinstance(result, dict) and ChatEvent.STATUS in result:
                         yield result[ChatEvent.STATUS]
-                    else:
-                        operator_results = {result["query"]: result["result"]}
+                    elif isinstance(result, OperatorRun):
+                        operator_results = result
                         this_iteration.operatorContext = operator_results
                         # Add webpages visited while operating browser to references
-                        if result.get("webpages"):
+                        if result.webpages:
                             if not online_results.get(this_iteration.query):
-                                online_results[this_iteration.query] = {"webpages": result["webpages"]}
+                                online_results[this_iteration.query] = {"webpages": result.webpages}
                             elif not online_results[this_iteration.query].get("webpages"):
-                                online_results[this_iteration.query]["webpages"] = result["webpages"]
+                                online_results[this_iteration.query]["webpages"] = result.webpages
                             else:
-                                online_results[this_iteration.query]["webpages"] += result["webpages"]
+                                online_results[this_iteration.query]["webpages"] += result.webpages
                             this_iteration.onlineContext = online_results
             except Exception as e:
                 this_iteration.warning = f"Error operating browser: {e}"
@@ -478,7 +510,9 @@ async def execute_information_collection(
             if code_results:
                 results_data += f"\n<code_results>\n{yaml.dump(truncate_code_context(code_results), allow_unicode=True, sort_keys=False, default_flow_style=False)}\n</code_results>"
             if operator_results:
-                results_data += f"\n<browser_operator_results>\n{next(iter(operator_results.values()))}\n</browser_operator_results>"
+                results_data += (
+                    f"\n<browser_operator_results>\n{operator_results.response}\n</browser_operator_results>"
+                )
             if summarize_files:
                 results_data += f"\n<summarized_files>\n{yaml.dump(summarize_files, allow_unicode=True, sort_keys=False, default_flow_style=False)}\n</summarized_files>"
             if this_iteration.warning:

khoj/utils/constants.py

@@ -18,8 +18,8 @@ default_offline_chat_models = [
     "bartowski/Qwen2.5-14B-Instruct-GGUF",
 ]
 default_openai_chat_models = ["gpt-4o-mini", "gpt-4.1"]
-default_gemini_chat_models = ["gemini-2.0-flash", "gemini-2.5-flash-preview-04-17", "gemini-2.5-pro-preview-03-25"]
-default_anthropic_chat_models = ["claude-3-7-sonnet-latest", "claude-3-5-haiku-latest"]
+default_gemini_chat_models = ["gemini-2.0-flash", "gemini-2.5-flash-preview-05-20", "gemini-2.5-pro-preview-05-06"]
+default_anthropic_chat_models = ["claude-sonnet-4-0", "claude-3-5-haiku-latest"]
 
 empty_config = {
     "search-type": {
@@ -63,10 +63,10 @@ model_to_cost: Dict[str, Dict[str, float]] = {
     "claude-3-7-sonnet-20250219": {"input": 3.0, "output": 15.0, "cache_read": 0.3, "cache_write": 3.75},
     "claude-3-7-sonnet@20250219": {"input": 3.0, "output": 15.0, "cache_read": 0.3, "cache_write": 3.75},
     "claude-3-7-sonnet-latest": {"input": 3.0, "output": 15.0, "cache_read": 0.3, "cache_write": 3.75},
-    "claude-sonnet-4": {"input": 3.0, "output": 15.0, "cache_read": 0.3, "cache_write": 3.75},
+    "claude-sonnet-4-0": {"input": 3.0, "output": 15.0, "cache_read": 0.3, "cache_write": 3.75},
     "claude-sonnet-4-20250514": {"input": 3.0, "output": 15.0, "cache_read": 0.3, "cache_write": 3.75},
     "claude-sonnet-4@20250514": {"input": 3.0, "output": 15.0, "cache_read": 0.3, "cache_write": 3.75},
-    "claude-opus-4": {"input": 15.0, "output": 75.0, "cache_read": 1.50, "cache_write": 18.75},
+    "claude-opus-4-0": {"input": 15.0, "output": 75.0, "cache_read": 1.50, "cache_write": 18.75},
     "claude-opus-4-20250514": {"input": 15.0, "output": 75.0, "cache_read": 1.50, "cache_write": 18.75},
     "claude-opus-4@20250514": {"input": 15.0, "output": 75.0, "cache_read": 1.50, "cache_write": 18.75},
     # Grok pricing: https://docs.x.ai/docs/models