khoj 1.30.11.dev15__py3-none-any.whl → 1.30.11.dev46__py3-none-any.whl

khoj/processor/conversation/utils.py

@@ -154,7 +154,7 @@ def construct_chat_history(conversation_history: dict, n: int = 4, agent_name="A
                 chat_history += f'{agent_name}: {{"queries": {chat["intent"].get("inferred-queries")}}}\n'
 
             chat_history += f"{agent_name}: {chat['message']}\n\n"
-        elif chat["by"] == "khoj" and ("text-to-image" in chat["intent"].get("type")):
+        elif chat["by"] == "khoj" and chat.get("images"):
             chat_history += f"User: {chat['intent']['query']}\n"
             chat_history += f"{agent_name}: [generated image redacted for space]\n"
         elif chat["by"] == "khoj" and ("excalidraw" in chat["intent"].get("type")):
@@ -213,6 +213,7 @@ class ChatEvent(Enum):
     END_LLM_RESPONSE = "end_llm_response"
     MESSAGE = "message"
     REFERENCES = "references"
+    GENERATED_ASSETS = "generated_assets"
     STATUS = "status"
     METADATA = "metadata"
     USAGE = "usage"
@@ -225,7 +226,6 @@ def message_to_log(
     user_message_metadata={},
     khoj_message_metadata={},
     conversation_log=[],
-    train_of_thought=[],
 ):
     """Create json logs from messages, metadata for conversation log"""
     default_khoj_message_metadata = {
@@ -234,6 +234,10 @@ def message_to_log(
     }
     khoj_response_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
 
+    # Filter out any fields that are set to None
+    user_message_metadata = {k: v for k, v in user_message_metadata.items() if v is not None}
+    khoj_message_metadata = {k: v for k, v in khoj_message_metadata.items() if v is not None}
+
     # Create json log from Human's message
     human_log = merge_dicts({"message": user_message, "by": "you"}, user_message_metadata)
 
@@ -261,31 +265,41 @@ def save_to_conversation_log(
     automation_id: str = None,
     query_images: List[str] = None,
     raw_query_files: List[FileAttachment] = [],
+    generated_images: List[str] = [],
+    raw_generated_files: List[FileAttachment] = [],
+    generated_excalidraw_diagram: str = None,
     train_of_thought: List[Any] = [],
     tracer: Dict[str, Any] = {},
 ):
     user_message_time = user_message_time or datetime.now().strftime("%Y-%m-%d %H:%M:%S")
     turn_id = tracer.get("mid") or str(uuid.uuid4())
+
+    user_message_metadata = {"created": user_message_time, "images": query_images, "turnId": turn_id}
+
+    if raw_query_files and len(raw_query_files) > 0:
+        user_message_metadata["queryFiles"] = [file.model_dump(mode="json") for file in raw_query_files]
+
+    khoj_message_metadata = {
+        "context": compiled_references,
+        "intent": {"inferred-queries": inferred_queries, "type": intent_type},
+        "onlineContext": online_results,
+        "codeContext": code_results,
+        "automationId": automation_id,
+        "trainOfThought": train_of_thought,
+        "turnId": turn_id,
+        "images": generated_images,
+        "queryFiles": [file.model_dump(mode="json") for file in raw_generated_files],
+    }
+
+    if generated_excalidraw_diagram:
+        khoj_message_metadata["excalidrawDiagram"] = generated_excalidraw_diagram
+
     updated_conversation = message_to_log(
         user_message=q,
         chat_response=chat_response,
-        user_message_metadata={
-            "created": user_message_time,
-            "images": query_images,
-            "turnId": turn_id,
-            "queryFiles": [file.model_dump(mode="json") for file in raw_query_files],
-        },
-        khoj_message_metadata={
-            "context": compiled_references,
-            "intent": {"inferred-queries": inferred_queries, "type": intent_type},
-            "onlineContext": online_results,
-            "codeContext": code_results,
-            "automationId": automation_id,
-            "trainOfThought": train_of_thought,
-            "turnId": turn_id,
-        },
+        user_message_metadata=user_message_metadata,
+        khoj_message_metadata=khoj_message_metadata,
         conversation_log=meta_log.get("chat", []),
-        train_of_thought=train_of_thought,
     )
     ConversationAdapters.save_conversation(
         user,
@@ -303,13 +317,13 @@ def save_to_conversation_log(
 Saved Conversation Turn
 You ({user.username}): "{q}"
 
-Khoj: "{inferred_queries if ("text-to-image" in intent_type) else chat_response}"
+Khoj: "{chat_response}"
 """.strip()
     )
 
 
 def construct_structured_message(
-    message: str, images: list[str], model_type: str, vision_enabled: bool, attached_file_context: str
+    message: str, 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
@@ -327,7 +341,8 @@ def construct_structured_message(
             constructed_messages.append({"type": "text", "text": attached_file_context})
         if vision_enabled and images:
             for image in images:
-                constructed_messages.append({"type": "image_url", "image_url": {"url": image}})
+                if image.startswith("https://"):
+                    constructed_messages.append({"type": "image_url", "image_url": {"url": image}})
         return constructed_messages
 
     if not is_none_or_empty(attached_file_context):
@@ -365,6 +380,10 @@ def generate_chatml_messages_with_context(
     model_type="",
     context_message="",
     query_files: str = None,
+    generated_images: Optional[list[str]] = None,
+    generated_files: List[FileAttachment] = None,
+    generated_excalidraw_diagram: str = None,
+    program_execution_context: List[str] = [],
 ):
     """Generate chat messages with appropriate context from previous conversation to send to the chat model"""
     # Set max prompt size from user config or based on pre-configured for model and machine specs
@@ -384,6 +403,7 @@ def generate_chatml_messages_with_context(
         message_attached_files = ""
 
         chat_message = chat.get("message")
+        role = "user" if chat["by"] == "you" else "assistant"
 
         if chat["by"] == "khoj" and "excalidraw" in chat["intent"].get("type", ""):
             chat_message = chat["intent"].get("inferred-queries")[0]
@@ -404,7 +424,7 @@ def generate_chatml_messages_with_context(
                 query_files_dict[file["name"]] = file["content"]
 
             message_attached_files = gather_raw_query_files(query_files_dict)
-            chatml_messages.append(ChatMessage(content=message_attached_files, role="user"))
+            chatml_messages.append(ChatMessage(content=message_attached_files, role=role))
 
         if not is_none_or_empty(chat.get("onlineContext")):
             message_context += f"{prompts.online_search_conversation.format(online_results=chat.get('onlineContext'))}"
@@ -413,10 +433,20 @@ def generate_chatml_messages_with_context(
             reconstructed_context_message = ChatMessage(content=message_context, role="user")
             chatml_messages.insert(0, reconstructed_context_message)
 
-        role = "user" if chat["by"] == "you" else "assistant"
-        message_content = construct_structured_message(
-            chat_message, chat.get("images"), model_type, vision_enabled, attached_file_context=query_files
-        )
+        if chat.get("images"):
+            if role == "assistant":
+                # Issue: the assistant role cannot accept an image as a message content, so send it in a separate user message.
+                file_attachment_message = construct_structured_message(
+                    message=prompts.generated_image_attachment.format(),
+                    images=chat.get("images"),
+                    model_type=model_type,
+                    vision_enabled=vision_enabled,
+                )
+                chatml_messages.append(ChatMessage(content=file_attachment_message, role="user"))
+            else:
+                message_content = construct_structured_message(
+                    chat_message, chat.get("images"), model_type, vision_enabled
+                )
 
         reconstructed_message = ChatMessage(content=message_content, role=role)
         chatml_messages.insert(0, reconstructed_message)
@@ -425,6 +455,7 @@ def generate_chatml_messages_with_context(
             break
 
     messages = []
+
     if not is_none_or_empty(user_message):
         messages.append(
             ChatMessage(
@@ -437,6 +468,31 @@ def generate_chatml_messages_with_context(
     if not is_none_or_empty(context_message):
         messages.append(ChatMessage(content=context_message, role="user"))
 
+    if generated_images:
+        messages.append(
+            ChatMessage(
+                content=construct_structured_message(
+                    prompts.generated_image_attachment.format(), generated_images, model_type, vision_enabled
+                ),
+                role="user",
+            )
+        )
+
+    if generated_files:
+        message_attached_files = gather_raw_query_files({file.name: file.content for file in generated_files})
+        messages.append(ChatMessage(content=message_attached_files, role="assistant"))
+
+    if generated_excalidraw_diagram:
+        messages.append(ChatMessage(content=prompts.generated_diagram_attachment.format(), role="assistant"))
+
+    if program_execution_context:
+        messages.append(
+            ChatMessage(
+                content=prompts.additional_program_context.format(context="\n".join(program_execution_context)),
+                role="assistant",
+            )
+        )
+
     if len(chatml_messages) > 0:
         messages += chatml_messages
 

khoj/processor/image/generate.py

@@ -12,7 +12,7 @@ from khoj.database.models import Agent, KhojUser, TextToImageModelConfig
 from khoj.routers.helpers import ChatEvent, generate_better_image_prompt
 from khoj.routers.storage import upload_image
 from khoj.utils import state
-from khoj.utils.helpers import ImageIntentType, convert_image_to_webp, timer
+from khoj.utils.helpers import convert_image_to_webp, timer
 from khoj.utils.rawconfig import LocationData
 
 logger = logging.getLogger(__name__)
@@ -34,14 +34,13 @@ async def text_to_image(
     status_code = 200
     image = None
     image_url = None
-    intent_type = ImageIntentType.TEXT_TO_IMAGE_V3
 
     text_to_image_config = await ConversationAdapters.aget_user_text_to_image_model(user)
     if not text_to_image_config:
         # If the user has not configured a text to image model, return an unsupported on server error
         status_code = 501
         message = "Failed to generate image. Setup image generation on the server."
-        yield image_url or image, status_code, message, intent_type.value
+        yield image_url or image, status_code, message
         return
 
     text2image_model = text_to_image_config.model_name
@@ -50,8 +49,8 @@ async def text_to_image(
         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 "text-to-image" in chat["intent"].get("type"):
-            chat_history += f"Q: Prompt: {chat['intent']['query']}\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"
 
     if send_status_func:
@@ -92,31 +91,29 @@ async def text_to_image(
                 logger.error(f"Image Generation blocked by OpenAI: {e}")
                 status_code = e.status_code  # type: ignore
                 message = f"Image generation blocked by OpenAI due to policy violation"  # type: ignore
-                yield image_url or image, status_code, message, intent_type.value
+                yield image_url or image, status_code, message
                 return
             else:
                 logger.error(f"Image Generation failed with {e}", exc_info=True)
                 message = f"Image generation failed using OpenAI"  # type: ignore
                 status_code = e.status_code  # type: ignore
-                yield image_url or image, status_code, message, intent_type.value
+                yield image_url or image, status_code, message
                 return
         except requests.RequestException as e:
             logger.error(f"Image Generation failed with {e}", exc_info=True)
             message = f"Image generation using {text2image_model} via {text_to_image_config.model_type} failed due to a network error."
             status_code = 502
-            yield image_url or image, status_code, message, intent_type.value
+            yield image_url or image, status_code, message
             return
 
     # Decide how to store the generated image
     with timer("Upload image to S3", logger):
         image_url = upload_image(webp_image_bytes, user.uuid)
-    if image_url:
-        intent_type = ImageIntentType.TEXT_TO_IMAGE2
-    else:
-        intent_type = ImageIntentType.TEXT_TO_IMAGE_V3
+
+    if not image_url:
         image = base64.b64encode(webp_image_bytes).decode("utf-8")
 
-    yield image_url or image, status_code, image_prompt, intent_type.value
+    yield image_url or image, status_code, image_prompt
 
 
 def generate_image_with_openai(

khoj/routers/api_chat.py

@@ -77,6 +77,7 @@ from khoj.utils.helpers import (
 )
 from khoj.utils.rawconfig import (
     ChatRequestBody,
+    FileAttachment,
     FileFilterRequest,
     FilesFilterRequest,
     LocationData,
@@ -770,6 +771,11 @@ async def chat(
         file_filters = conversation.file_filters if conversation and conversation.file_filters else []
         attached_file_context = gather_raw_query_files(query_files)
 
+        generated_images: List[str] = []
+        generated_files: List[FileAttachment] = []
+        generated_excalidraw_diagram: str = None
+        program_execution_context: List[str] = []
+
         if conversation_commands == [ConversationCommand.Default] or is_automated_task:
             chosen_io = await aget_data_sources_and_output_format(
                 q,
@@ -875,21 +881,17 @@ async def chat(
                             async for result in send_llm_response(response, tracer.get("usage")):
                                 yield result
 
-            await sync_to_async(save_to_conversation_log)(
-                q,
-                response_log,
-                user,
-                meta_log,
-                user_message_time,
-                intent_type="summarize",
-                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,
-                tracer=tracer,
+            summarized_document = FileAttachment(
+                name="Summarized Document",
+                content=response_log,
+                type="text/plain",
+                size=len(response_log.encode("utf-8")),
             )
-            return
+
+            async for result in send_event(ChatEvent.GENERATED_ASSETS, {"files": [summarized_document.model_dump()]}):
+                yield result
+
+            generated_files.append(summarized_document)
 
         custom_filters = []
         if conversation_commands == [ConversationCommand.Help]:
@@ -1078,6 +1080,7 @@ async def chat(
                 async for result in send_event(ChatEvent.STATUS, f"**Ran code snippets**: {len(code_results)}"):
                     yield result
             except ValueError as e:
+                program_execution_context.append(f"Failed to run code")
                 logger.warning(
                     f"Failed to use code tool: {e}. Attempting to respond without code results",
                     exc_info=True,
@@ -1115,51 +1118,28 @@ async def chat(
                 if isinstance(result, dict) and ChatEvent.STATUS in result:
                     yield result[ChatEvent.STATUS]
                 else:
-                    generated_image, status_code, improved_image_prompt, intent_type = result
+                    generated_image, status_code, improved_image_prompt = result
 
+            inferred_queries.append(improved_image_prompt)
             if generated_image is None or status_code != 200:
-                content_obj = {
-                    "content-type": "application/json",
-                    "intentType": intent_type,
-                    "detail": improved_image_prompt,
-                    "image": None,
-                }
-                async for result in send_llm_response(json.dumps(content_obj), tracer.get("usage")):
+                program_execution_context.append(f"Failed to generate image with {improved_image_prompt}")
+                async for result in send_event(ChatEvent.STATUS, f"Failed to generate image"):
                     yield result
-                return
+            else:
+                generated_images.append(generated_image)
 
-            await sync_to_async(save_to_conversation_log)(
-                q,
-                generated_image,
-                user,
-                meta_log,
-                user_message_time,
-                intent_type=intent_type,
-                inferred_queries=[improved_image_prompt],
-                client_application=request.user.client_app,
-                conversation_id=conversation_id,
-                compiled_references=compiled_references,
-                online_results=online_results,
-                code_results=code_results,
-                query_images=uploaded_images,
-                train_of_thought=train_of_thought,
-                raw_query_files=raw_query_files,
-                tracer=tracer,
-            )
-            content_obj = {
-                "intentType": intent_type,
-                "inferredQueries": [improved_image_prompt],
-                "image": generated_image,
-            }
-            async for result in send_llm_response(json.dumps(content_obj), tracer.get("usage")):
-                yield result
-            return
+                async for result in send_event(
+                    ChatEvent.GENERATED_ASSETS,
+                    {
+                        "images": [generated_image],
+                    },
+                ):
+                    yield result
 
         if ConversationCommand.Diagram in conversation_commands:
             async for result in send_event(ChatEvent.STATUS, f"Creating diagram"):
                 yield result
 
-            intent_type = "excalidraw"
             inferred_queries = []
             diagram_description = ""
 
@@ -1183,62 +1163,29 @@ async def chat(
                     if better_diagram_description_prompt and excalidraw_diagram_description:
                         inferred_queries.append(better_diagram_description_prompt)
                         diagram_description = excalidraw_diagram_description
+
+                        generated_excalidraw_diagram = diagram_description
+
+                        async for result in send_event(
+                            ChatEvent.GENERATED_ASSETS,
+                            {
+                                "excalidrawDiagram": excalidraw_diagram_description,
+                            },
+                        ):
+                            yield result
                     else:
                         error_message = "Failed to generate diagram. Please try again later."
-                        async for result in send_llm_response(error_message, tracer.get("usage")):
-                            yield result
-
-                        await sync_to_async(save_to_conversation_log)(
-                            q,
-                            error_message,
-                            user,
-                            meta_log,
-                            user_message_time,
-                            inferred_queries=[better_diagram_description_prompt],
-                            client_application=request.user.client_app,
-                            conversation_id=conversation_id,
-                            compiled_references=compiled_references,
-                            online_results=online_results,
-                            code_results=code_results,
-                            query_images=uploaded_images,
-                            train_of_thought=train_of_thought,
-                            raw_query_files=raw_query_files,
-                            tracer=tracer,
+                        program_execution_context.append(
+                            f"AI attempted to programmatically generate a diagram but failed due to a program issue. Generally, it is able to do so, but encountered a system issue this time. AI can suggest text description or rendering of the diagram or user can try again with a simpler prompt."
                         )
-                        return
-
-            content_obj = {
-                "intentType": intent_type,
-                "inferredQueries": inferred_queries,
-                "image": diagram_description,
-            }
 
-            await sync_to_async(save_to_conversation_log)(
-                q,
-                excalidraw_diagram_description,
-                user,
-                meta_log,
-                user_message_time,
-                intent_type="excalidraw",
-                inferred_queries=[better_diagram_description_prompt],
-                client_application=request.user.client_app,
-                conversation_id=conversation_id,
-                compiled_references=compiled_references,
-                online_results=online_results,
-                code_results=code_results,
-                query_images=uploaded_images,
-                train_of_thought=train_of_thought,
-                raw_query_files=raw_query_files,
-                tracer=tracer,
-            )
-
-            async for result in send_llm_response(json.dumps(content_obj), tracer.get("usage")):
-                yield result
-            return
+                        async for result in send_event(ChatEvent.STATUS, error_message):
+                            yield result
 
         ## Generate Text Output
         async for result in send_event(ChatEvent.STATUS, f"**Generating a well-informed response**"):
             yield result
+
         llm_response, chat_metadata = await agenerate_chat_response(
             defiltered_query,
             meta_log,
@@ -1258,6 +1205,10 @@ async def chat(
             train_of_thought,
             attached_file_context,
             raw_query_files,
+            generated_images,
+            generated_files,
+            generated_excalidraw_diagram,
+            program_execution_context,
             tracer,
         )
 

khoj/routers/helpers.py

@@ -1185,6 +1185,10 @@ def generate_chat_response(
     train_of_thought: List[Any] = [],
     query_files: str = None,
     raw_query_files: List[FileAttachment] = None,
+    generated_images: List[str] = None,
+    raw_generated_files: List[FileAttachment] = [],
+    generated_excalidraw_diagram: str = None,
+    program_execution_context: List[str] = [],
     tracer: dict = {},
 ) -> Tuple[Union[ThreadedGenerator, Iterator[str]], Dict[str, str]]:
     # Initialize Variables
@@ -1208,6 +1212,9 @@ def generate_chat_response(
             query_images=query_images,
             train_of_thought=train_of_thought,
             raw_query_files=raw_query_files,
+            generated_images=generated_images,
+            raw_generated_files=raw_generated_files,
+            generated_excalidraw_diagram=generated_excalidraw_diagram,
             tracer=tracer,
         )
 
@@ -1243,6 +1250,7 @@ def generate_chat_response(
                 user_name=user_name,
                 agent=agent,
                 query_files=query_files,
+                generated_files=raw_generated_files,
                 tracer=tracer,
             )
 
@@ -1269,6 +1277,10 @@ def generate_chat_response(
                 agent=agent,
                 vision_available=vision_available,
                 query_files=query_files,
+                generated_files=raw_generated_files,
+                generated_images=generated_images,
+                generated_excalidraw_diagram=generated_excalidraw_diagram,
+                program_execution_context=program_execution_context,
                 tracer=tracer,
             )
 
@@ -1292,6 +1304,10 @@ def generate_chat_response(
                 agent=agent,
                 vision_available=vision_available,
                 query_files=query_files,
+                generated_files=raw_generated_files,
+                generated_images=generated_images,
+                generated_excalidraw_diagram=generated_excalidraw_diagram,
+                program_execution_context=program_execution_context,
                 tracer=tracer,
             )
         elif conversation_config.model_type == ChatModelOptions.ModelType.GOOGLE:
@@ -1314,6 +1330,10 @@ def generate_chat_response(
                 query_images=query_images,
                 vision_available=vision_available,
                 query_files=query_files,
+                generated_files=raw_generated_files,
+                generated_images=generated_images,
+                generated_excalidraw_diagram=generated_excalidraw_diagram,
+                program_execution_context=program_execution_context,
                 tracer=tracer,
             )
 
@@ -1785,6 +1805,9 @@ class MessageProcessor:
         self.references = {}
         self.usage = {}
         self.raw_response = ""
+        self.generated_images = []
+        self.generated_files = []
+        self.generated_excalidraw_diagram = []
 
     def convert_message_chunk_to_json(self, raw_chunk: str) -> Dict[str, Any]:
         if raw_chunk.startswith("{") and raw_chunk.endswith("}"):
@@ -1823,6 +1846,16 @@ class MessageProcessor:
                     self.raw_response += chunk_data
             else:
                 self.raw_response += chunk_data
+        elif chunk_type == ChatEvent.GENERATED_ASSETS:
+            chunk_data = chunk["data"]
+            if isinstance(chunk_data, dict):
+                for key in chunk_data:
+                    if key == "images":
+                        self.generated_images = chunk_data[key]
+                    elif key == "files":
+                        self.generated_files = chunk_data[key]
+                    elif key == "excalidrawDiagram":
+                        self.generated_excalidraw_diagram = chunk_data[key]
 
     def handle_json_response(self, json_data: Dict[str, str]) -> str | Dict[str, str]:
         if "image" in json_data or "details" in json_data:
@@ -1853,7 +1886,14 @@ async def read_chat_stream(response_iterator: AsyncGenerator[str, None]) -> Dict
     if buffer:
         processor.process_message_chunk(buffer)
 
-    return {"response": processor.raw_response, "references": processor.references, "usage": processor.usage}
+    return {
+        "response": processor.raw_response,
+        "references": processor.references,
+        "usage": processor.usage,
+        "images": processor.generated_images,
+        "files": processor.generated_files,
+        "excalidrawDiagram": processor.generated_excalidraw_diagram,
+    }
 
 
 def get_user_config(user: KhojUser, request: Request, is_detailed: bool = False):

{khoj-1.30.11.dev15.dist-info → khoj-1.30.11.dev46.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: khoj
-Version: 1.30.11.dev15
+Version: 1.30.11.dev46
 Summary: Your Second Brain
 Project-URL: Homepage, https://khoj.dev
 Project-URL: Documentation, https://docs.khoj.dev