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

khoj/processor/conversation/utils.py

@@ -10,7 +10,7 @@ from dataclasses import dataclass
 from datetime import datetime
 from enum import Enum
 from io import BytesIO
-from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union
+from typing import Any, Callable, Dict, List, Literal, Optional, Union
 
 import PIL.Image
 import pyjson5
@@ -137,83 +137,60 @@ class OperatorRun:
         }
 
 
-class ToolCall:
-    def __init__(self, name: str, args: dict, id: str):
-        self.name = name
-        self.args = args
-        self.id = id
-
-
 class ResearchIteration:
     def __init__(
         self,
-        query: ToolCall | dict | str,
+        tool: str,
+        query: str,
         context: list = None,
         onlineContext: dict = None,
         codeContext: dict = None,
         operatorContext: dict | OperatorRun = None,
         summarizedResult: str = None,
         warning: str = None,
-        raw_response: list = None,
     ):
-        self.query = ToolCall(**query) if isinstance(query, dict) else query
+        self.tool = tool
+        self.query = query
         self.context = context
         self.onlineContext = onlineContext
         self.codeContext = codeContext
         self.operatorContext = OperatorRun(**operatorContext) if isinstance(operatorContext, dict) else operatorContext
         self.summarizedResult = summarizedResult
         self.warning = warning
-        self.raw_response = raw_response
 
     def to_dict(self) -> dict:
         data = vars(self).copy()
-        data["query"] = self.query.__dict__ if isinstance(self.query, ToolCall) else self.query
         data["operatorContext"] = self.operatorContext.to_dict() if self.operatorContext else None
         return data
 
 
 def construct_iteration_history(
     previous_iterations: List[ResearchIteration],
+    previous_iteration_prompt: str,
     query: str = None,
-    query_images: List[str] = None,
-    query_files: str = None,
 ) -> list[ChatMessageModel]:
     iteration_history: list[ChatMessageModel] = []
-    query_message_content = construct_structured_message(query, query_images, attached_file_context=query_files)
-    if query_message_content:
-        iteration_history.append(ChatMessageModel(by="you", message=query_message_content))
+    previous_iteration_messages: list[dict] = []
+    for idx, iteration in enumerate(previous_iterations):
+        iteration_data = previous_iteration_prompt.format(
+            tool=iteration.tool,
+            query=iteration.query,
+            result=iteration.summarizedResult,
+            index=idx + 1,
+        )
 
-    for iteration in previous_iterations:
-        if not iteration.query or isinstance(iteration.query, str):
-            iteration_history.append(
-                ChatMessageModel(
-                    by="you",
-                    message=iteration.summarizedResult
-                    or iteration.warning
-                    or "Please specify what you want to do next.",
-                )
-            )
-            continue
-        iteration_history += [
+        previous_iteration_messages.append({"type": "text", "text": iteration_data})
+
+    if previous_iteration_messages:
+        if query:
+            iteration_history.append(ChatMessageModel(by="you", message=query))
+        iteration_history.append(
             ChatMessageModel(
                 by="khoj",
-                message=iteration.raw_response or [iteration.query.__dict__],
-                intent=Intent(type="tool_call", query=query),
-            ),
-            ChatMessageModel(
-                by="you",
-                intent=Intent(type="tool_result"),
-                message=[
-                    {
-                        "type": "tool_result",
-                        "id": iteration.query.id,
-                        "name": iteration.query.name,
-                        "content": iteration.summarizedResult,
-                    }
-                ],
-            ),
-        ]
-
+                intent=Intent(type="remember", query=query),
+                message=previous_iteration_messages,
+            )
+        )
     return iteration_history
 
 
@@ -325,44 +302,33 @@ def construct_tool_chat_history(
         ConversationCommand.Notes: (
             lambda iteration: [c["query"] for c in iteration.context] if iteration.context else []
         ),
-        ConversationCommand.SearchWeb: (
+        ConversationCommand.Online: (
             lambda iteration: list(iteration.onlineContext.keys()) if iteration.onlineContext else []
         ),
-        ConversationCommand.ReadWebpage: (
+        ConversationCommand.Webpage: (
             lambda iteration: list(iteration.onlineContext.keys()) if iteration.onlineContext else []
         ),
-        ConversationCommand.RunCode: (
+        ConversationCommand.Code: (
             lambda iteration: list(iteration.codeContext.keys()) if iteration.codeContext else []
         ),
     }
     for iteration in previous_iterations:
-        if not iteration.query or isinstance(iteration.query, str):
-            chat_history.append(
-                ChatMessageModel(
-                    by="you",
-                    message=iteration.summarizedResult
-                    or iteration.warning
-                    or "Please specify what you want to do next.",
-                )
-            )
-            continue
-
         # If a tool is provided use the inferred query extractor for that tool if available
         # If no tool is provided, use inferred query extractor for the tool used in the iteration
         # Fallback to base extractor if the tool does not have an inferred query extractor
         inferred_query_extractor = extract_inferred_query_map.get(
-            tool or ConversationCommand(iteration.query.name), base_extractor
+            tool or ConversationCommand(iteration.tool), base_extractor
         )
         chat_history += [
             ChatMessageModel(
                 by="you",
-                message=yaml.dump(iteration.query.args, default_flow_style=False),
+                message=iteration.query,
             ),
             ChatMessageModel(
                 by="khoj",
                 intent=Intent(
                     type="remember",
-                    query=yaml.dump(iteration.query.args, default_flow_style=False),
+                    query=iteration.query,
                     inferred_queries=inferred_query_extractor(iteration),
                     memory_type="notes",
                 ),
@@ -515,32 +481,28 @@ Khoj: "{chat_response}"
 
 def construct_structured_message(
     message: list[dict] | str,
-    images: list[str] = None,
-    model_type: str = None,
-    vision_enabled: bool = True,
+    images: list[str],
+    model_type: str,
+    vision_enabled: bool,
     attached_file_context: str = None,
 ):
     """
-    Format messages into appropriate multimedia format for supported chat model types.
-
-    Assume vision is enabled and chat model provider supports messages in chatml format, unless specified otherwise.
+    Format messages into appropriate multimedia format for supported chat model types
     """
-    if not model_type or model_type in [
+    if model_type in [
         ChatModel.ModelType.OPENAI,
         ChatModel.ModelType.GOOGLE,
         ChatModel.ModelType.ANTHROPIC,
     ]:
-        constructed_messages: List[dict[str, Any]] = []
-        if not is_none_or_empty(message):
-            constructed_messages += [{"type": "text", "text": message}] if isinstance(message, str) else message
-        # Drop image message passed by caller if chat model does not have vision enabled
-        if not vision_enabled:
-            constructed_messages = [m for m in constructed_messages if m.get("type") != "image_url"]
+        constructed_messages: List[dict[str, Any]] = (
+            [{"type": "text", "text": message}] if isinstance(message, str) else message
+        )
+
         if not is_none_or_empty(attached_file_context):
-            constructed_messages += [{"type": "text", "text": attached_file_context}]
+            constructed_messages.append({"type": "text", "text": attached_file_context})
         if vision_enabled and images:
             for image in images:
-                constructed_messages += [{"type": "image_url", "image_url": {"url": image}}]
+                constructed_messages.append({"type": "image_url", "image_url": {"url": image}})
         return constructed_messages
 
     message = message if isinstance(message, str) else "\n\n".join(m["text"] for m in message)
@@ -646,7 +608,7 @@ def generate_chatml_messages_with_context(
 
         if not is_none_or_empty(chat.context):
             references = "\n\n".join(
-                {f"# URI: {item.uri}\n## {item.compiled}\n" for item in chat.context or [] if isinstance(item, dict)}
+                {f"# File: {item.file}\n## {item.compiled}\n" for item in chat.context or [] if isinstance(item, dict)}
             )
             message_context += [{"type": "text", "text": f"{prompts.notes_conversation.format(references=references)}"}]
 
@@ -676,11 +638,7 @@ def generate_chatml_messages_with_context(
             chat_message, chat.images if role == "user" else [], model_type, vision_enabled
         )
 
-        reconstructed_message = ChatMessage(
-            content=message_content,
-            role=role,
-            additional_kwargs={"message_type": chat.intent.type if chat.intent else None},
-        )
+        reconstructed_message = ChatMessage(content=message_content, role=role)
         chatml_messages.insert(0, reconstructed_message)
 
         if len(chatml_messages) >= 3 * lookback_turns:
@@ -779,21 +737,10 @@ def count_tokens(
         message_content_parts: list[str] = []
         # Collate message content into single string to ease token counting
         for part in message_content:
-            if isinstance(part, dict) and part.get("type") == "image_url":
-                image_count += 1
-            elif isinstance(part, dict) and part.get("type") == "text":
+            if isinstance(part, dict) and part.get("type") == "text":
                 message_content_parts.append(part["text"])
-            elif isinstance(part, dict) and hasattr(part, "model_dump"):
-                message_content_parts.append(json.dumps(part.model_dump()))
-            elif isinstance(part, dict) and hasattr(part, "__dict__"):
-                message_content_parts.append(json.dumps(part.__dict__))
-            elif isinstance(part, dict):
-                # If part is a dict but not a recognized type, convert to JSON string
-                try:
-                    message_content_parts.append(json.dumps(part))
-                except (TypeError, ValueError) as e:
-                    logger.warning(f"Failed to serialize part {part} to JSON: {e}. Skipping.")
-                    image_count += 1  # Treat as an image/binary if serialization fails
+            elif isinstance(part, dict) and part.get("type") == "image_url":
+                image_count += 1
             elif isinstance(part, str):
                 message_content_parts.append(part)
             else:
@@ -806,15 +753,6 @@ def count_tokens(
         return len(encoder.encode(json.dumps(message_content)))
 
 
-def count_total_tokens(messages: list[ChatMessage], encoder, system_message: Optional[ChatMessage]) -> Tuple[int, int]:
-    """Count total tokens in messages including system message"""
-    system_message_tokens = count_tokens(system_message.content, encoder) if system_message else 0
-    message_tokens = sum([count_tokens(message.content, encoder) for message in messages])
-    # Reserves 4 tokens to demarcate each message (e.g <|im_start|>user, <|im_end|>, <|endoftext|> etc.)
-    total_tokens = message_tokens + system_message_tokens + 4 * len(messages)
-    return total_tokens, system_message_tokens
-
-
 def truncate_messages(
     messages: list[ChatMessage],
     max_prompt_size: int,
@@ -833,30 +771,23 @@ def truncate_messages(
             break
 
     # Drop older messages until under max supported prompt size by model
-    total_tokens, system_message_tokens = count_total_tokens(messages, encoder, system_message)
+    # Reserves 4 tokens to demarcate each message (e.g <|im_start|>user, <|im_end|>, <|endoftext|> etc.)
+    system_message_tokens = count_tokens(system_message.content, encoder) if system_message else 0
+    tokens = sum([count_tokens(message.content, encoder) for message in messages])
+    total_tokens = tokens + system_message_tokens + 4 * len(messages)
 
     while total_tokens > max_prompt_size and (len(messages) > 1 or len(messages[0].content) > 1):
-        # If the last message has more than one content part, pop the oldest content part.
-        # For tool calls, the whole message should dropped, assistant's tool call content being truncated annoys AI APIs.
-        if len(messages[-1].content) > 1 and messages[-1].additional_kwargs.get("message_type") != "tool_call":
+        if len(messages[-1].content) > 1:
             # The oldest content part is earlier in content list. So pop from the front.
             messages[-1].content.pop(0)
-        # Otherwise, pop the last message if it has only one content part or is a tool call.
         else:
             # The oldest message is the last one. So pop from the back.
-            dropped_message = messages.pop()
-            # Drop tool result pair of tool call, if tool call message has been removed
-            if (
-                dropped_message.additional_kwargs.get("message_type") == "tool_call"
-                and messages
-                and messages[-1].additional_kwargs.get("message_type") == "tool_result"
-            ):
-                messages.pop()
-
-        total_tokens, _ = count_total_tokens(messages, encoder, system_message)
+            messages.pop()
+        tokens = sum([count_tokens(message.content, encoder) for message in messages])
+        total_tokens = tokens + system_message_tokens + 4 * len(messages)
 
     # Truncate current message if still over max supported prompt size by model
-    total_tokens, _ = count_total_tokens(messages, encoder, system_message)
+    total_tokens = tokens + system_message_tokens + 4 * len(messages)
     if total_tokens > max_prompt_size:
         # At this point, a single message with a single content part of type dict should remain
         assert (
@@ -1218,15 +1149,13 @@ def messages_to_print(messages: list[ChatMessage], max_length: int = 70) -> str:
     return "\n".join([f"{json.dumps(safe_serialize(message.content))[:max_length]}..." for message in messages])
 
 
-class StructuredOutputSupport(int, Enum):
+class JsonSupport(int, Enum):
     NONE = 0
     OBJECT = 1
     SCHEMA = 2
-    TOOL = 3
 
 
 class ResponseWithThought:
-    def __init__(self, text: str = None, thought: str = None, raw_content: list = None):
-        self.text = text
+    def __init__(self, response: str = None, thought: str = None):
+        self.response = response
         self.thought = thought
-        self.raw_content = raw_content

khoj/processor/operator/grounding_agent.py

@@ -73,7 +73,7 @@ class GroundingAgent:
         grounding_user_prompt = self.get_instruction(instruction, self.environment_type)
         screenshots = [f"data:image/webp;base64,{current_state.screenshot}"]
         grounding_messages_content = construct_structured_message(
-            grounding_user_prompt, screenshots, self.model.model_type, vision_enabled=True
+            grounding_user_prompt, screenshots, self.model.name, vision_enabled=True
         )
         return [{"role": "user", "content": grounding_messages_content}]
 

khoj/processor/operator/operator_agent_binary.py

@@ -121,7 +121,7 @@ class BinaryOperatorAgent(OperatorAgent):
         # Construct input for visual reasoner history
         visual_reasoner_history = self._format_message_for_api(self.messages)
         try:
-            raw_response = await send_message_to_model_wrapper(
+            natural_language_action = await send_message_to_model_wrapper(
                 query=query_text,
                 query_images=query_screenshot,
                 system_message=reasoning_system_prompt,
@@ -129,7 +129,6 @@ class BinaryOperatorAgent(OperatorAgent):
                 agent_chat_model=self.reasoning_model,
                 tracer=self.tracer,
             )
-            natural_language_action = raw_response.text
 
             if not isinstance(natural_language_action, str) or not natural_language_action.strip():
                 raise ValueError(f"Natural language action is empty or not a string. Got {natural_language_action}")
@@ -256,10 +255,10 @@ class BinaryOperatorAgent(OperatorAgent):
 
         # Append summary messages to history
         trigger_summary = AgentMessage(role="user", content=summarize_prompt)
-        summary_message = AgentMessage(role="assistant", content=summary.text)
+        summary_message = AgentMessage(role="assistant", content=summary)
         self.messages.extend([trigger_summary, summary_message])
 
-        return summary.text
+        return summary
 
     def _compile_response(self, response_content: str | List) -> str:
         """Compile response content into a string, handling OpenAI message structures."""

khoj/processor/tools/online_search.py

@@ -390,25 +390,7 @@ async def read_webpages(
         query_files=query_files,
         tracer=tracer,
     )
-    async for result in read_webpages_content(
-        query,
-        urls,
-        user,
-        send_status_func=send_status_func,
-        agent=agent,
-        tracer=tracer,
-    ):
-        yield result
-
 
-async def read_webpages_content(
-    query: str,
-    urls: List[str],
-    user: KhojUser,
-    send_status_func: Optional[Callable] = None,
-    agent: Agent = None,
-    tracer: dict = {},
-):
     logger.info(f"Reading web pages at: {urls}")
     if send_status_func:
         webpage_links_str = "\n- " + "\n- ".join(list(urls))

khoj/processor/tools/run_code.py

@@ -161,7 +161,7 @@ async def generate_python_code(
     )
 
     # Extract python code wrapped in markdown code blocks from the response
-    code_blocks = re.findall(r"```(?:python)?\n(.*?)```", response.text, re.DOTALL)
+    code_blocks = re.findall(r"```(?:python)?\n(.*?)```", response, re.DOTALL)
 
     if not code_blocks:
         raise ValueError("No Python code blocks found in response")

khoj/routers/api_chat.py

@@ -1390,7 +1390,7 @@ async def chat(
                 continue
             if cancellation_event.is_set():
                 break
-            message = item.text
+            message = item.response
             full_response += message if message else ""
             if item.thought:
                 async for result in send_event(ChatEvent.THOUGHT, item.thought):

khoj/routers/api_content.py

@@ -101,9 +101,9 @@ async def put_content(
     host: Optional[str] = Header(None),
     indexed_data_limiter: ApiIndexedDataLimiter = Depends(
         ApiIndexedDataLimiter(
-            incoming_entries_size_limit=10,
-            subscribed_incoming_entries_size_limit=75,
-            total_entries_size_limit=10,
+            incoming_entries_size_limit=50,
+            subscribed_incoming_entries_size_limit=100,
+            total_entries_size_limit=50,
             subscribed_total_entries_size_limit=500,
         )
     ),
@@ -123,9 +123,9 @@ async def patch_content(
     host: Optional[str] = Header(None),
     indexed_data_limiter: ApiIndexedDataLimiter = Depends(
         ApiIndexedDataLimiter(
-            incoming_entries_size_limit=10,
-            subscribed_incoming_entries_size_limit=75,
-            total_entries_size_limit=10,
+            incoming_entries_size_limit=50,
+            subscribed_incoming_entries_size_limit=100,
+            total_entries_size_limit=50,
             subscribed_total_entries_size_limit=500,
         )
     ),