khoj 1.33.3.dev35__py3-none-any.whl → 1.33.3.dev51__py3-none-any.whl

khoj/processor/conversation/prompts.py

@@ -194,7 +194,7 @@ Limit your response to 3 sentences max. Be succinct, clear, and informative.
 ## Diagram Generation
 ## --
 
-improve_diagram_description_prompt = PromptTemplate.from_template(
+improve_excalidraw_diagram_description_prompt = PromptTemplate.from_template(
     """
 You are an architect working with a novice digital artist using a diagramming software.
 {personality_context}
@@ -338,6 +338,123 @@ Diagram Description: {query}
 """.strip()
 )
 
+improve_mermaid_js_diagram_description_prompt = PromptTemplate.from_template(
+    """
+You are a senior architect working with an illustrator using a diagramming software.
+{personality_context}
+
+Given a particular request, you need to translate it to to a detailed description that the illustrator can use to create a diagram.
+
+You can use the following diagram types in your instructions:
+- Flowchart
+- Sequence Diagram
+- Gantt Chart (only for time-based queries after 0 AD)
+- State Diagram
+- Pie Chart
+
+Use these primitives to describe what sort of diagram the drawer should create in natural language, not special syntax. We must recreate the diagram every time, so include all relevant prior information in your description.
+
+- Describe the layout, components, and connections.
+- Use simple, concise language.
+
+Today's Date: {current_date}
+User's Location: {location}
+
+User's Notes:
+{references}
+
+Online References:
+{online_results}
+
+Conversation Log:
+{chat_history}
+
+Query: {query}
+
+Enhanced Description:
+""".strip()
+)
+
+mermaid_js_diagram_generation_prompt = PromptTemplate.from_template(
+    """
+You are a designer with the ability to describe diagrams to compose in professional, fine detail. You dive into the details and make labels, connections, and shapes to represent complex systems.
+{personality_context}
+
+----Goals----
+You need to create a declarative description of the diagram and relevant components, using the Mermaid.js syntax.
+
+You can choose from the following diagram types:
+- Flowchart
+- Sequence Diagram
+- State Diagram
+- Gantt Chart
+- Pie Chart
+
+----Examples----
+---
+title: Node
+---
+
+flowchart LR
+    id["This is the start"] --> id2["This is the end"]
+
+sequenceDiagram
+    Alice->>John: Hello John, how are you?
+    John-->>Alice: Great!
+    Alice-)John: See you later!
+
+stateDiagram-v2
+    [*] --> Still
+    Still --> [*]
+
+    Still --> Moving
+    Moving --> Still
+    Moving --> Crash
+    Crash --> [*]
+
+gantt
+    title A Gantt Diagram
+    dateFormat YYYY-MM-DD
+    section Section
+        A task          :a1, 2014-01-01, 30d
+        Another task    :after a1, 20d
+    section Another
+        Task in Another :2014-01-12, 12d
+        another task    :24d
+
+pie title Pets adopted by volunteers
+    "Dogs" : 10
+    "Cats" : 30
+    "Rats" : 60
+
+flowchart TB
+    subgraph "Group 1"
+        a1["Start Node"] --> a2["End Node"]
+    end
+    subgraph "Group 2"
+        b1["Process 1"] --> b2["Process 2"]
+    end
+    subgraph "Group 3"
+        c1["Input"] --> c2["Output"]
+    end
+    a["Group 1"] --> b["Group 2"]
+    c["Group 3"] --> d["Group 2"]
+
+----Process----
+Create your diagram with great composition and intuitiveness from the provided context and user prompt below.
+- You may use subgraphs to group elements together. Each subgraph must have a title.
+- **You must wrap ALL entity and node labels in double quotes**, example: "My Node Label"
+- **All nodes MUST use the id["label"] format**. For example: node1["My Node Label"]
+- Custom style are not permitted. Default styles only.
+- JUST provide the diagram, no additional text or context. Say nothing else in your response except the diagram.
+- Keep diagrams simple - maximum 15 nodes
+- Every node inside a subgraph MUST use square bracket notation: id["label"]
+
+output: {query}
+
+""".strip()
+)
+
 failed_diagram_generation = PromptTemplate.from_template(
     """
 You attempted to programmatically generate a diagram but failed due to a system issue. You are normally able to generate diagrams, but you encountered a system issue this time.

khoj/processor/conversation/utils.py

@@ -266,7 +266,7 @@ def save_to_conversation_log(
     raw_query_files: List[FileAttachment] = [],
     generated_images: List[str] = [],
     raw_generated_files: List[FileAttachment] = [],
-    generated_excalidraw_diagram: str = None,
+    generated_mermaidjs_diagram: str = None,
     train_of_thought: List[Any] = [],
     tracer: Dict[str, Any] = {},
 ):
@@ -290,8 +290,8 @@ def save_to_conversation_log(
         "queryFiles": [file.model_dump(mode="json") for file in raw_generated_files],
     }
 
-    if generated_excalidraw_diagram:
-        khoj_message_metadata["excalidrawDiagram"] = generated_excalidraw_diagram
+    if generated_mermaidjs_diagram:
+        khoj_message_metadata["mermaidjsDiagram"] = generated_mermaidjs_diagram
 
     updated_conversation = message_to_log(
         user_message=q,
@@ -441,7 +441,7 @@ def generate_chatml_messages_with_context(
                 "query": chat.get("intent", {}).get("inferred-queries", [user_message])[0],
             }
 
-        if not is_none_or_empty(chat.get("excalidrawDiagram")) and role == "assistant":
+        if not is_none_or_empty(chat.get("mermaidjsDiagram")) and role == "assistant":
             generated_assets["diagram"] = {
                 "query": chat.get("intent", {}).get("inferred-queries", [user_message])[0],
             }
@@ -593,6 +593,11 @@ def clean_json(response: str):
     return response.strip().replace("\n", "").removeprefix("```json").removesuffix("```")
 
 
+def clean_mermaidjs(response: str):
+    """Remove any markdown mermaidjs codeblock and newline formatting if present. Useful for non schema enforceable models"""
+    return response.strip().removeprefix("```mermaid").removesuffix("```")
+
+
 def clean_code_python(code: str):
     """Remove any markdown codeblock and newline formatting if present. Useful for non schema enforceable models"""
     return code.strip().removeprefix("```python").removesuffix("```")

khoj/processor/image/generate.py

@@ -111,7 +111,7 @@ async def text_to_image(
         image_url = upload_image(webp_image_bytes, user.uuid)
 
     if not image_url:
-        image = base64.b64encode(webp_image_bytes).decode("utf-8")
+        image = f"data:image/webp;base64,{base64.b64encode(webp_image_bytes).decode('utf-8')}"
 
     yield image_url or image, status_code, image_prompt
 
@@ -119,25 +119,27 @@ async def text_to_image(
 def generate_image_with_openai(
     improved_image_prompt: str, text_to_image_config: TextToImageModelConfig, text2image_model: str
 ):
-    "Generate image using OpenAI API"
+    "Generate image using OpenAI (compatible) API"
 
-    # Get the API key from the user's configuration
+    # Get the API config from the user's configuration
+    api_key = None
     if text_to_image_config.api_key:
         api_key = text_to_image_config.api_key
+        openai_client = openai.OpenAI(api_key=api_key)
     elif text_to_image_config.ai_model_api:
         api_key = text_to_image_config.ai_model_api.api_key
+        api_base_url = text_to_image_config.ai_model_api.api_base_url
+        openai_client = openai.OpenAI(api_key=api_key, base_url=api_base_url)
     elif state.openai_client:
-        api_key = state.openai_client.api_key
-    auth_header = {"Authorization": f"Bearer {api_key}"} if api_key else {}
+        openai_client = state.openai_client
 
     # Generate image using OpenAI API
     OPENAI_IMAGE_GEN_STYLE = "vivid"
-    response = state.openai_client.images.generate(
+    response = openai_client.images.generate(
         prompt=improved_image_prompt,
         model=text2image_model,
         style=OPENAI_IMAGE_GEN_STYLE,
         response_format="b64_json",
-        extra_headers=auth_header,
     )
 
     # Extract the base64 image from the response

khoj/routers/api.py

@@ -9,6 +9,7 @@ import uuid
 from typing import Any, Callable, List, Optional, Set, Union
 
 import cron_descriptor
+import openai
 import pytz
 from apscheduler.job import Job
 from apscheduler.triggers.cron import CronTrigger
@@ -264,12 +265,21 @@ async def transcribe(
         if not speech_to_text_config:
             # If the user has not configured a speech to text model, return an unsupported on server error
             status_code = 501
-        elif state.openai_client and speech_to_text_config.model_type == SpeechToTextModelOptions.ModelType.OPENAI:
-            speech2text_model = speech_to_text_config.model_name
-            user_message = await transcribe_audio(audio_file, speech2text_model, client=state.openai_client)
         elif speech_to_text_config.model_type == SpeechToTextModelOptions.ModelType.OFFLINE:
             speech2text_model = speech_to_text_config.model_name
             user_message = await transcribe_audio_offline(audio_filename, speech2text_model)
+        elif speech_to_text_config.model_type == SpeechToTextModelOptions.ModelType.OPENAI:
+            speech2text_model = speech_to_text_config.model_name
+            if speech_to_text_config.ai_model_api:
+                api_key = speech_to_text_config.ai_model_api.api_key
+                api_base_url = speech_to_text_config.ai_model_api.api_base_url
+                openai_client = openai.OpenAI(api_key=api_key, base_url=api_base_url)
+            elif state.openai_client:
+                openai_client = state.openai_client
+            if openai_client:
+                user_message = await transcribe_audio(audio_file, speech2text_model, client=openai_client)
+            else:
+                status_code = 501
     finally:
         # Close and Delete the temporary audio file
         audio_file.close()

khoj/routers/api_chat.py

@@ -51,7 +51,7 @@ from khoj.routers.helpers import (
     construct_automation_created_message,
     create_automation,
     gather_raw_query_files,
-    generate_excalidraw_diagram,
+    generate_mermaidjs_diagram,
     generate_summary_from_files,
     get_conversation_command,
     is_query_empty,
@@ -781,7 +781,7 @@ async def chat(
 
         generated_images: List[str] = []
         generated_files: List[FileAttachment] = []
-        generated_excalidraw_diagram: str = None
+        generated_mermaidjs_diagram: str = None
         program_execution_context: List[str] = []
 
         if conversation_commands == [ConversationCommand.Default]:
@@ -1161,7 +1161,7 @@ async def chat(
             inferred_queries = []
             diagram_description = ""
 
-            async for result in generate_excalidraw_diagram(
+            async for result in generate_mermaidjs_diagram(
                 q=defiltered_query,
                 conversation_history=meta_log,
                 location_data=location,
@@ -1177,12 +1177,12 @@ async def chat(
                 if isinstance(result, dict) and ChatEvent.STATUS in result:
                     yield result[ChatEvent.STATUS]
                 else:
-                    better_diagram_description_prompt, excalidraw_diagram_description = result
-                    if better_diagram_description_prompt and excalidraw_diagram_description:
+                    better_diagram_description_prompt, mermaidjs_diagram_description = result
+                    if better_diagram_description_prompt and mermaidjs_diagram_description:
                         inferred_queries.append(better_diagram_description_prompt)
-                        diagram_description = excalidraw_diagram_description
+                        diagram_description = mermaidjs_diagram_description
 
-                        generated_excalidraw_diagram = diagram_description
+                        generated_mermaidjs_diagram = diagram_description
 
                         generated_asset_results["diagrams"] = {
                             "query": better_diagram_description_prompt,
@@ -1191,7 +1191,7 @@ async def chat(
                         async for result in send_event(
                             ChatEvent.GENERATED_ASSETS,
                             {
-                                "excalidrawDiagram": excalidraw_diagram_description,
+                                "mermaidjsDiagram": mermaidjs_diagram_description,
                             },
                         ):
                             yield result
@@ -1231,7 +1231,7 @@ async def chat(
             raw_query_files,
             generated_images,
             generated_files,
-            generated_excalidraw_diagram,
+            generated_mermaidjs_diagram,
             program_execution_context,
             generated_asset_results,
             tracer,

khoj/routers/helpers.py

@@ -97,6 +97,7 @@ from khoj.processor.conversation.utils import (
     ChatEvent,
     ThreadedGenerator,
     clean_json,
+    clean_mermaidjs,
     construct_chat_history,
     generate_chatml_messages_with_context,
     save_to_conversation_log,
@@ -823,7 +824,7 @@ async def generate_better_diagram_description(
             elif online_results[result].get("webpages"):
                 simplified_online_results[result] = online_results[result]["webpages"]
 
-    improve_diagram_description_prompt = prompts.improve_diagram_description_prompt.format(
+    improve_diagram_description_prompt = prompts.improve_excalidraw_diagram_description_prompt.format(
         query=q,
         chat_history=chat_history,
         location=location,
@@ -887,6 +888,133 @@ async def generate_excalidraw_diagram_from_description(
     return response
 
 
+async def generate_mermaidjs_diagram(
+    q: str,
+    conversation_history: Dict[str, Any],
+    location_data: LocationData,
+    note_references: List[Dict[str, Any]],
+    online_results: Optional[dict] = None,
+    query_images: List[str] = None,
+    user: KhojUser = None,
+    agent: Agent = None,
+    send_status_func: Optional[Callable] = None,
+    query_files: str = None,
+    tracer: dict = {},
+):
+    if send_status_func:
+        async for event in send_status_func("**Enhancing the Diagramming Prompt**"):
+            yield {ChatEvent.STATUS: event}
+
+    better_diagram_description_prompt = await generate_better_mermaidjs_diagram_description(
+        q=q,
+        conversation_history=conversation_history,
+        location_data=location_data,
+        note_references=note_references,
+        online_results=online_results,
+        query_images=query_images,
+        user=user,
+        agent=agent,
+        query_files=query_files,
+        tracer=tracer,
+    )
+
+    if send_status_func:
+        async for event in send_status_func(f"**Diagram to Create:**:\n{better_diagram_description_prompt}"):
+            yield {ChatEvent.STATUS: event}
+
+    mermaidjs_diagram_description = await generate_mermaidjs_diagram_from_description(
+        q=better_diagram_description_prompt,
+        user=user,
+        agent=agent,
+        tracer=tracer,
+    )
+
+    inferred_queries = f"Instruction: {better_diagram_description_prompt}"
+
+    yield inferred_queries, mermaidjs_diagram_description
+
+
+async def generate_better_mermaidjs_diagram_description(
+    q: str,
+    conversation_history: Dict[str, Any],
+    location_data: LocationData,
+    note_references: List[Dict[str, Any]],
+    online_results: Optional[dict] = None,
+    query_images: List[str] = None,
+    user: KhojUser = None,
+    agent: Agent = None,
+    query_files: str = None,
+    tracer: dict = {},
+) -> str:
+    """
+    Generate a diagram description from the given query and context
+    """
+
+    today_date = datetime.now(tz=timezone.utc).strftime("%Y-%m-%d, %A")
+    personality_context = (
+        prompts.personality_context.format(personality=agent.personality) if agent and agent.personality else ""
+    )
+
+    location = f"{location_data}" if location_data else "Unknown"
+
+    user_references = "\n\n".join([f"# {item['compiled']}" for item in note_references])
+
+    chat_history = construct_chat_history(conversation_history)
+
+    simplified_online_results = {}
+
+    if online_results:
+        for result in online_results:
+            if online_results[result].get("answerBox"):
+                simplified_online_results[result] = online_results[result]["answerBox"]
+            elif online_results[result].get("webpages"):
+                simplified_online_results[result] = online_results[result]["webpages"]
+
+    improve_diagram_description_prompt = prompts.improve_mermaid_js_diagram_description_prompt.format(
+        query=q,
+        chat_history=chat_history,
+        location=location,
+        current_date=today_date,
+        references=user_references,
+        online_results=simplified_online_results,
+        personality_context=personality_context,
+    )
+
+    with timer("Chat actor: Generate better Mermaid.js diagram description", logger):
+        response = await send_message_to_model_wrapper(
+            improve_diagram_description_prompt,
+            query_images=query_images,
+            user=user,
+            query_files=query_files,
+            tracer=tracer,
+        )
+        response = response.strip()
+        if response.startswith(('"', "'")) and response.endswith(('"', "'")):
+            response = response[1:-1]
+
+    return response
+
+
+async def generate_mermaidjs_diagram_from_description(
+    q: str,
+    user: KhojUser = None,
+    agent: Agent = None,
+    tracer: dict = {},
+) -> str:
+    personality_context = (
+        prompts.personality_context.format(personality=agent.personality) if agent and agent.personality else ""
+    )
+
+    mermaidjs_diagram_generation = prompts.mermaid_js_diagram_generation_prompt.format(
+        personality_context=personality_context,
+        query=q,
+    )
+
+    with timer("Chat actor: Generate Mermaid.js diagram", logger):
+        raw_response = await send_message_to_model_wrapper(query=mermaidjs_diagram_generation, user=user, tracer=tracer)
+        return clean_mermaidjs(raw_response.strip())
+
+
 async def generate_better_image_prompt(
     q: str,
     conversation_history: str,
@@ -1224,7 +1352,7 @@ def generate_chat_response(
     raw_query_files: List[FileAttachment] = None,
     generated_images: List[str] = None,
     raw_generated_files: List[FileAttachment] = [],
-    generated_excalidraw_diagram: str = None,
+    generated_mermaidjs_diagram: str = None,
     program_execution_context: List[str] = [],
     generated_asset_results: Dict[str, Dict] = {},
     tracer: dict = {},
@@ -1252,7 +1380,7 @@ def generate_chat_response(
             raw_query_files=raw_query_files,
             generated_images=generated_images,
             raw_generated_files=raw_generated_files,
-            generated_excalidraw_diagram=generated_excalidraw_diagram,
+            generated_mermaidjs_diagram=generated_mermaidjs_diagram,
             tracer=tracer,
         )
 
@@ -1967,7 +2095,7 @@ class MessageProcessor:
         self.raw_response = ""
         self.generated_images = []
         self.generated_files = []
-        self.generated_excalidraw_diagram = []
+        self.generated_mermaidjs_diagram = []
 
     def convert_message_chunk_to_json(self, raw_chunk: str) -> Dict[str, Any]:
         if raw_chunk.startswith("{") and raw_chunk.endswith("}"):
@@ -2014,8 +2142,8 @@ class MessageProcessor:
                         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]
+                    elif key == "mermaidjsDiagram":
+                        self.generated_mermaidjs_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:
@@ -2052,7 +2180,7 @@ async def read_chat_stream(response_iterator: AsyncGenerator[str, None]) -> Dict
         "usage": processor.usage,
         "images": processor.generated_images,
         "files": processor.generated_files,
-        "excalidrawDiagram": processor.generated_excalidraw_diagram,
+        "mermaidjsDiagram": processor.generated_mermaidjs_diagram,
     }
 
 

khoj/utils/initialization.py

@@ -43,14 +43,14 @@ def initialization(interactive: bool = True):
             "🗣️ Configure chat models available to your server. You can always update these at /server/admin using your admin account"
         )
 
-        openai_api_base = os.getenv("OPENAI_API_BASE")
-        provider = "Ollama" if openai_api_base and openai_api_base.endswith(":11434/v1/") else "OpenAI"
-        openai_api_key = os.getenv("OPENAI_API_KEY", "placeholder" if openai_api_base else None)
+        openai_base_url = os.getenv("OPENAI_BASE_URL")
+        provider = "Ollama" if openai_base_url and openai_base_url.endswith(":11434/v1/") else "OpenAI"
+        openai_api_key = os.getenv("OPENAI_API_KEY", "placeholder" if openai_base_url else None)
         default_chat_models = default_openai_chat_models
-        if openai_api_base:
+        if openai_base_url:
             # Get available chat models from OpenAI compatible API
             try:
-                openai_client = openai.OpenAI(api_key=openai_api_key, base_url=openai_api_base)
+                openai_client = openai.OpenAI(api_key=openai_api_key, base_url=openai_base_url)
                 default_chat_models = [model.id for model in openai_client.models.list()]
                 # Put the available default OpenAI models at the top
                 valid_default_models = [model for model in default_openai_chat_models if model in default_chat_models]
@@ -66,7 +66,7 @@ def initialization(interactive: bool = True):
             ChatModel.ModelType.OPENAI,
             default_chat_models,
             default_api_key=openai_api_key,
-            api_base_url=openai_api_base,
+            api_base_url=openai_base_url,
             vision_enabled=True,
             is_offline=False,
             interactive=interactive,

{khoj-1.33.3.dev35.dist-info → khoj-1.33.3.dev51.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: khoj
-Version: 1.33.3.dev35
+Version: 1.33.3.dev51
 Summary: Your Second Brain
 Project-URL: Homepage, https://khoj.dev
 Project-URL: Documentation, https://docs.khoj.dev