khoj 1.30.11.dev13__py3-none-any.whl → 1.30.11.dev56__py3-none-any.whl

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

@@ -136,7 +136,7 @@ def validate_conversation_config(user: KhojUser):
     if default_config is None:
         raise HTTPException(status_code=500, detail="Contact the server administrator to add a chat model.")
 
-    if default_config.model_type == "openai" and not default_config.openai_config:
+    if default_config.model_type == "openai" and not default_config.ai_model_api:
         raise HTTPException(status_code=500, detail="Contact the server administrator to add a chat model.")
 
 
@@ -163,7 +163,7 @@ async def is_ready_to_chat(user: KhojUser):
                 ChatModelOptions.ModelType.GOOGLE,
             ]
         )
-        and user_conversation_config.openai_config
+        and user_conversation_config.ai_model_api
     ):
         return True
 
@@ -990,7 +990,7 @@ async def send_message_to_model_wrapper(
         )
 
     elif model_type == ChatModelOptions.ModelType.OPENAI:
-        openai_chat_config = conversation_config.openai_config
+        openai_chat_config = conversation_config.ai_model_api
         api_key = openai_chat_config.api_key
         api_base_url = openai_chat_config.api_base_url
         truncated_messages = generate_chatml_messages_with_context(
@@ -1015,7 +1015,7 @@ async def send_message_to_model_wrapper(
             tracer=tracer,
         )
     elif model_type == ChatModelOptions.ModelType.ANTHROPIC:
-        api_key = conversation_config.openai_config.api_key
+        api_key = conversation_config.ai_model_api.api_key
         truncated_messages = generate_chatml_messages_with_context(
             user_message=query,
             context_message=context,
@@ -1037,7 +1037,7 @@ async def send_message_to_model_wrapper(
             tracer=tracer,
         )
     elif model_type == ChatModelOptions.ModelType.GOOGLE:
-        api_key = conversation_config.openai_config.api_key
+        api_key = conversation_config.ai_model_api.api_key
         truncated_messages = generate_chatml_messages_with_context(
             user_message=query,
             context_message=context,
@@ -1102,7 +1102,7 @@ def send_message_to_model_wrapper_sync(
         )
 
     elif conversation_config.model_type == ChatModelOptions.ModelType.OPENAI:
-        api_key = conversation_config.openai_config.api_key
+        api_key = conversation_config.ai_model_api.api_key
         truncated_messages = generate_chatml_messages_with_context(
             user_message=message,
             system_message=system_message,
@@ -1124,7 +1124,7 @@ def send_message_to_model_wrapper_sync(
         return openai_response
 
     elif conversation_config.model_type == ChatModelOptions.ModelType.ANTHROPIC:
-        api_key = conversation_config.openai_config.api_key
+        api_key = conversation_config.ai_model_api.api_key
         truncated_messages = generate_chatml_messages_with_context(
             user_message=message,
             system_message=system_message,
@@ -1144,7 +1144,7 @@ def send_message_to_model_wrapper_sync(
         )
 
     elif conversation_config.model_type == ChatModelOptions.ModelType.GOOGLE:
-        api_key = conversation_config.openai_config.api_key
+        api_key = conversation_config.ai_model_api.api_key
         truncated_messages = generate_chatml_messages_with_context(
             user_message=message,
             system_message=system_message,
@@ -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,11 +1250,12 @@ def generate_chat_response(
                 user_name=user_name,
                 agent=agent,
                 query_files=query_files,
+                generated_files=raw_generated_files,
                 tracer=tracer,
             )
 
         elif conversation_config.model_type == ChatModelOptions.ModelType.OPENAI:
-            openai_chat_config = conversation_config.openai_config
+            openai_chat_config = conversation_config.ai_model_api
             api_key = openai_chat_config.api_key
             chat_model = conversation_config.chat_model
             chat_response = converse(
@@ -1269,11 +1277,15 @@ 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.ANTHROPIC:
-            api_key = conversation_config.openai_config.api_key
+            api_key = conversation_config.ai_model_api.api_key
             chat_response = converse_anthropic(
                 compiled_references,
                 query_to_run,
@@ -1292,10 +1304,14 @@ 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:
-            api_key = conversation_config.openai_config.api_key
+            api_key = conversation_config.ai_model_api.api_key
             chat_response = converse_gemini(
                 compiled_references,
                 query_to_run,
@@ -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/utils/initialization.py

@@ -6,9 +6,9 @@ import openai
 
 from khoj.database.adapters import ConversationAdapters
 from khoj.database.models import (
+    AiModelApi,
     ChatModelOptions,
     KhojUser,
-    OpenAIProcessorConversationConfig,
     SpeechToTextModelOptions,
     TextToImageModelConfig,
 )
@@ -98,7 +98,7 @@ def initialization(interactive: bool = True):
             TextToImageModelConfig.objects.create(
                 model_name=openai_text_to_image_model,
                 model_type=TextToImageModelConfig.ModelType.OPENAI,
-                openai_config=openai_provider,
+                ai_model_api=openai_provider,
             )
 
         # Set up Google's Gemini online chat models
@@ -177,7 +177,7 @@ def initialization(interactive: bool = True):
         vision_enabled: bool = False,
         is_offline: bool = False,
         provider_name: str = None,
-    ) -> Tuple[bool, OpenAIProcessorConversationConfig]:
+    ) -> Tuple[bool, AiModelApi]:
         supported_vision_models = (
             default_openai_chat_models + default_anthropic_chat_models + default_gemini_chat_models
         )
@@ -192,16 +192,14 @@ def initialization(interactive: bool = True):
 
         logger.info(f"️💬 Setting up your {provider_name} chat configuration")
 
-        chat_provider = None
+        ai_model_api = None
         if not is_offline:
             if interactive:
                 user_api_key = input(f"Enter your {provider_name} API key (default: {default_api_key}): ")
                 api_key = user_api_key if user_api_key != "" else default_api_key
             else:
                 api_key = default_api_key
-            chat_provider = OpenAIProcessorConversationConfig.objects.create(
-                api_key=api_key, name=provider_name, api_base_url=api_base_url
-            )
+            ai_model_api = AiModelApi.objects.create(api_key=api_key, name=provider_name, api_base_url=api_base_url)
 
         if interactive:
             chat_model_names = input(
@@ -223,19 +221,19 @@ def initialization(interactive: bool = True):
                 "max_prompt_size": default_max_tokens,
                 "vision_enabled": vision_enabled,
                 "tokenizer": default_tokenizer,
-                "openai_config": chat_provider,
+                "ai_model_api": ai_model_api,
             }
 
             ChatModelOptions.objects.create(**chat_model_options)
 
         logger.info(f"🗣️ {provider_name} chat model configuration complete")
-        return True, chat_provider
+        return True, ai_model_api
 
     def _update_chat_model_options():
         """Update available chat models for OpenAI-compatible APIs"""
         try:
             # Get OpenAI configs with custom base URLs
-            custom_configs = OpenAIProcessorConversationConfig.objects.exclude(api_base_url__isnull=True)
+            custom_configs = AiModelApi.objects.exclude(api_base_url__isnull=True)
 
             for config in custom_configs:
                 try:
@@ -247,7 +245,7 @@ def initialization(interactive: bool = True):
 
                     # Get existing chat model options for this config
                     existing_models = ChatModelOptions.objects.filter(
-                        openai_config=config, model_type=ChatModelOptions.ModelType.OPENAI
+                        ai_model_api=config, model_type=ChatModelOptions.ModelType.OPENAI
                     )
 
                     # Add new models
@@ -259,7 +257,7 @@ def initialization(interactive: bool = True):
                                 max_prompt_size=model_to_prompt_size.get(model),
                                 vision_enabled=model in default_openai_chat_models,
                                 tokenizer=model_to_tokenizer.get(model),
-                                openai_config=config,
+                                ai_model_api=config,
                             )
 
                     # Remove models that are no longer available

{khoj-1.30.11.dev13.dist-info → khoj-1.30.11.dev56.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: khoj
-Version: 1.30.11.dev13
+Version: 1.30.11.dev56
 Summary: Your Second Brain
 Project-URL: Homepage, https://khoj.dev
 Project-URL: Documentation, https://docs.khoj.dev
@@ -31,6 +31,7 @@ Requires-Dist: dateparser>=1.1.1
 Requires-Dist: defusedxml==0.7.1
 Requires-Dist: django-apscheduler==0.6.2
 Requires-Dist: django-phonenumber-field==7.3.0
+Requires-Dist: django-unfold==0.42.0
 Requires-Dist: django==5.0.9
 Requires-Dist: docx2txt==0.8
 Requires-Dist: einops==0.8.0