khoj 1.42.1.dev10__py3-none-any.whl → 1.42.2.dev16__py3-none-any.whl

khoj/processor/image/generate.py

@@ -10,7 +10,12 @@ from google import genai
 from google.genai import types as gtypes
 
 from khoj.database.adapters import ConversationAdapters
-from khoj.database.models import Agent, KhojUser, TextToImageModelConfig
+from khoj.database.models import (
+    Agent,
+    ChatMessageModel,
+    KhojUser,
+    TextToImageModelConfig,
+)
 from khoj.routers.helpers import ChatEvent, generate_better_image_prompt
 from khoj.routers.storage import upload_generated_image_to_bucket
 from khoj.utils import state
@@ -23,7 +28,7 @@ logger = logging.getLogger(__name__)
 async def text_to_image(
     message: str,
     user: KhojUser,
-    conversation_log: dict,
+    chat_history: List[ChatMessageModel],
     location_data: LocationData,
     references: List[Dict[str, Any]],
     online_results: Dict[str, Any],
@@ -46,14 +51,14 @@ async def text_to_image(
         return
 
     text2image_model = text_to_image_config.model_name
-    chat_history = ""
-    for chat in conversation_log.get("chat", [])[-4:]:
-        if chat["by"] == "khoj" and chat["intent"].get("type") in ["remember", "reminder"]:
-            chat_history += f"Q: {chat['intent']['query']}\n"
-            chat_history += f"A: {chat['message']}\n"
-        elif chat["by"] == "khoj" and chat.get("images"):
-            chat_history += f"Q: {chat['intent']['query']}\n"
-            chat_history += f"A: Improved Prompt: {chat['intent']['inferred-queries'][0]}\n"
+    chat_history_str = ""
+    for chat in chat_history[-4:]:
+        if chat.by == "khoj" and chat.intent and chat.intent.type in ["remember", "reminder"]:
+            chat_history_str += f"Q: {chat.intent.query or ''}\n"
+            chat_history_str += f"A: {chat.message}\n"
+        elif chat.by == "khoj" and chat.images:
+            chat_history_str += f"Q: {chat.intent.query}\n"
+            chat_history_str += f"A: Improved Prompt: {chat.intent.inferred_queries[0]}\n"
 
     if send_status_func:
         async for event in send_status_func("**Enhancing the Painting Prompt**"):
@@ -63,7 +68,7 @@ async def text_to_image(
     # Use the user's message, chat history, and other context
     image_prompt = await generate_better_image_prompt(
         message,
-        chat_history,
+        chat_history_str,
         location_data=location_data,
         note_references=references,
         online_results=online_results,

khoj/processor/operator/__init__.py

@@ -5,10 +5,9 @@ import os
 from typing import Callable, List, Optional
 
 from khoj.database.adapters import AgentAdapters, ConversationAdapters
-from khoj.database.models import Agent, ChatModel, KhojUser
+from khoj.database.models import Agent, ChatMessageModel, ChatModel, KhojUser
 from khoj.processor.conversation.utils import (
     OperatorRun,
-    construct_chat_history,
     construct_chat_history_for_operator,
 )
 from khoj.processor.operator.operator_actions import *
@@ -34,7 +33,7 @@ logger = logging.getLogger(__name__)
 async def operate_environment(
     query: str,
     user: KhojUser,
-    conversation_log: dict,
+    conversation_log: List[ChatMessageModel],
     location_data: LocationData,
     previous_trajectory: Optional[OperatorRun] = None,
     environment_type: EnvironmentType = EnvironmentType.COMPUTER,

khoj/processor/operator/operator_agent_binary.py

@@ -4,7 +4,7 @@ from datetime import datetime
 from textwrap import dedent
 from typing import List, Optional
 
-from khoj.database.models import ChatModel
+from khoj.database.models import ChatMessageModel, ChatModel
 from khoj.processor.conversation.utils import (
     AgentMessage,
     OperatorRun,
@@ -119,13 +119,13 @@ class BinaryOperatorAgent(OperatorAgent):
             query_screenshot = self._get_message_images(current_message)
 
         # Construct input for visual reasoner history
-        visual_reasoner_history = {"chat": self._format_message_for_api(self.messages)}
+        visual_reasoner_history = self._format_message_for_api(self.messages)
         try:
             natural_language_action = await send_message_to_model_wrapper(
                 query=query_text,
                 query_images=query_screenshot,
                 system_message=reasoning_system_prompt,
-                conversation_log=visual_reasoner_history,
+                chat_history=visual_reasoner_history,
                 agent_chat_model=self.reasoning_model,
                 tracer=self.tracer,
             )
@@ -238,11 +238,11 @@ class BinaryOperatorAgent(OperatorAgent):
 
     async def summarize(self, env_state: EnvState, summarize_prompt: str = None) -> str:
         summarize_prompt = summarize_prompt or self.summarize_prompt
-        conversation_history = {"chat": self._format_message_for_api(self.messages)}
+        conversation_history = self._format_message_for_api(self.messages)
         try:
             summary = await send_message_to_model_wrapper(
                 query=summarize_prompt,
-                conversation_log=conversation_history,
+                chat_history=conversation_history,
                 agent_chat_model=self.reasoning_model,
                 tracer=self.tracer,
             )
@@ -296,14 +296,14 @@ class BinaryOperatorAgent(OperatorAgent):
             images = [item["image_url"]["url"] for item in message.content if item["type"] == "image_url"]
         return images
 
-    def _format_message_for_api(self, messages: list[AgentMessage]) -> List[dict]:
+    def _format_message_for_api(self, messages: list[AgentMessage]) -> List[ChatMessageModel]:
         """Format operator agent messages into the Khoj conversation history format."""
         formatted_messages = [
-            {
-                "message": self._get_message_text(message),
-                "images": self._get_message_images(message),
-                "by": "you" if message.role in ["user", "environment"] else message.role,
-            }
+            ChatMessageModel(
+                message=self._get_message_text(message),
+                images=self._get_message_images(message),
+                by="you" if message.role in ["user", "environment"] else message.role,
+            )
             for message in messages
         ]
         return formatted_messages

khoj/processor/tools/online_search.py

@@ -10,7 +10,13 @@ from bs4 import BeautifulSoup
 from markdownify import markdownify
 
 from khoj.database.adapters import ConversationAdapters
-from khoj.database.models import Agent, KhojUser, ServerChatSettings, WebScraper
+from khoj.database.models import (
+    Agent,
+    ChatMessageModel,
+    KhojUser,
+    ServerChatSettings,
+    WebScraper,
+)
 from khoj.processor.conversation import prompts
 from khoj.routers.helpers import (
     ChatEvent,
@@ -59,7 +65,7 @@ OLOSTEP_QUERY_PARAMS = {
 
 async def search_online(
     query: str,
-    conversation_history: dict,
+    conversation_history: List[ChatMessageModel],
     location: LocationData,
     user: KhojUser,
     send_status_func: Optional[Callable] = None,
@@ -361,7 +367,7 @@ async def search_with_serper(query: str, location: LocationData) -> Tuple[str, D
 
 async def read_webpages(
     query: str,
-    conversation_history: dict,
+    conversation_history: List[ChatMessageModel],
     location: LocationData,
     user: KhojUser,
     send_status_func: Optional[Callable] = None,

khoj/processor/tools/run_code.py

@@ -20,7 +20,7 @@ from tenacity import (
 )
 
 from khoj.database.adapters import FileObjectAdapters
-from khoj.database.models import Agent, FileObject, KhojUser
+from khoj.database.models import Agent, ChatMessageModel, FileObject, KhojUser
 from khoj.processor.conversation import prompts
 from khoj.processor.conversation.utils import (
     ChatEvent,
@@ -50,7 +50,7 @@ class GeneratedCode(NamedTuple):
 
 async def run_code(
     query: str,
-    conversation_history: dict,
+    conversation_history: List[ChatMessageModel],
     context: str,
     location_data: LocationData,
     user: KhojUser,
@@ -116,7 +116,7 @@ async def run_code(
 
 async def generate_python_code(
     q: str,
-    conversation_history: dict,
+    chat_history: List[ChatMessageModel],
     context: str,
     location_data: LocationData,
     user: KhojUser,
@@ -127,7 +127,7 @@ async def generate_python_code(
 ) -> GeneratedCode:
     location = f"{location_data}" if location_data else "Unknown"
     username = prompts.user_name.format(name=user.get_full_name()) if user.get_full_name() else ""
-    chat_history = construct_chat_history(conversation_history)
+    chat_history_str = construct_chat_history(chat_history)
 
     utc_date = datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%d")
     personality_context = (
@@ -143,7 +143,7 @@ async def generate_python_code(
 
     code_generation_prompt = prompts.python_code_generation_prompt.format(
         query=q,
-        chat_history=chat_history,
+        chat_history=chat_history_str,
         context=context,
         has_network_access=network_access_context,
         current_date=utc_date,