khoj 1.26.4.dev2__py3-none-any.whl → 1.26.5.dev29__py3-none-any.whl

khoj/processor/conversation/prompts.py

@@ -176,6 +176,150 @@ Improved Prompt:
 """.strip()
 )
 
+## Diagram Generation
+## --
+
+improve_diagram_description_prompt = PromptTemplate.from_template(
+    """
+you are an architect working with a novice artist using a diagramming tool.
+{personality_context}
+
+you need to convert the user's query to a description format that the novice artist can use very well. you are allowed to use primitives like
+- text
+- rectangle
+- diamond
+- ellipse
+- line
+- arrow
+- frame
+
+use these primitives to describe what sort of diagram the drawer should create. the artist must recreate the diagram every time, so include all relevant prior information in your description.
+
+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}
+
+
+""".strip()
+)
+
+excalidraw_diagram_generation_prompt = PromptTemplate.from_template(
+    """
+You are a program manager with the ability to describe diagrams to compose in professional, fine detail.
+{personality_context}
+
+You need to create a declarative description of the diagram and relevant components, using this base schema. Use the `label` property to specify the text to be rendered in the respective elements. Always use light colors for the `backgroundColor` property, like white, or light blue, green, red. "type", "x", "y", "id", are required properties for all elements.
+
+{{
+    type: string,
+    x: number,
+    y: number,
+    strokeColor: string,
+    backgroundColor: string,
+    width: number,
+    height: number,
+    id: string,
+    label: {{
+        text: string,
+    }}
+}}
+
+Valid types:
+- text
+- rectangle
+- diamond
+- ellipse
+- line
+- arrow
+
+For arrows and lines, you can use the `points` property to specify the start and end points of the arrow. You may also use the `label` property to specify the text to be rendered. You may use the `start` and `end` properties to connect the linear elements to other elements. The start and end point can either be the ID to map to an existing object, or the `type` to create a new object. Mapping to an existing object is useful if you want to connect it to multiple objects. Lines and arrows can only start and end at rectangle, text, diamond, or ellipse elements.
+
+{{
+    type: "arrow",
+    id: string,
+    x: number,
+    y: number,
+    width: number,
+    height: number,
+    strokeColor: string,
+    start: {{
+        id: string,
+        type: string,
+    }},
+    end: {{
+        id: string,
+        type: string,
+    }},
+    label: {{
+        text: string,
+    }}
+    points: [
+        [number, number],
+        [number, number],
+    ]
+}}
+
+For text, you must use the `text` property to specify the text to be rendered. You may also use `fontSize` property to specify the font size of the text. Only use the `text` element for titles, subtitles, and overviews. For labels, use the `label` property in the respective elements.
+
+{{
+    type: "text",
+    id: string,
+    x: number,
+    y: number,
+    fontSize: number,
+    text: string,
+}}
+
+For frames, use the `children` property to specify the elements that are inside the frame by their ids.
+
+{{
+    type: "frame",
+    id: string,
+    x: number,
+    y: number,
+    width: number,
+    height: number,
+    name: string,
+    children: [
+        string
+    ]
+}}
+
+Here's an example of a valid diagram:
+
+Design Description: Create a diagram describing a circular development process with 3 stages: design, implementation and feedback. The design stage is connected to the implementation stage and the implementation stage is connected to the feedback stage and the feedback stage is connected to the design stage. Each stage should be labeled with the stage name.
+
+Response:
+
+[
+    {{"type":"text","x":-150,"y":50,"width":300,"height":40,"id":"title_text","text":"Circular Development Process","fontSize":24}},
+    {{"type":"ellipse","x":-169,"y":113,"width":188,"height":202,"id":"design_ellipse", "label": {{"text": "Design"}}}},
+    {{"type":"ellipse","x":62,"y":394,"width":186,"height":188,"id":"implement_ellipse", "label": {{"text": "Implement"}}}},
+    {{"type":"ellipse","x":-348,"y":430,"width":184,"height":170,"id":"feedback_ellipse", "label": {{"text": "Feedback"}}}},
+    {{"type":"arrow","x":21,"y":273,"id":"design_to_implement_arrow","points":[[0,0],[86,105]],"start":{{"id":"design_ellipse"}}, "end":{{"id":"implement_ellipse"}}}},
+    {{"type":"arrow","x":50,"y":519,"id":"implement_to_feedback_arrow","points":[[0,0],[-198,-6]],"start":{{"id":"implement_ellipse"}}, "end":{{"id":"feedback_ellipse"}}}},
+    {{"type":"arrow","x":-228,"y":417,"id":"feedback_to_design_arrow","points":[[0,0],[85,-123]],"start":{{"id":"feedback_ellipse"}}, "end":{{"id":"design_ellipse"}}}},
+]
+
+Create a detailed diagram from the provided context and user prompt below. Return a valid JSON object:
+
+Diagram Description: {query}
+
+""".strip()
+)
+
 ## Online Search Conversation
 ## --
 online_search_conversation = PromptTemplate.from_template(

khoj/processor/conversation/utils.py

@@ -109,7 +109,7 @@ def save_to_conversation_log(
     client_application: ClientApplication = None,
     conversation_id: str = None,
     automation_id: str = None,
-    uploaded_image_url: str = None,
+    query_images: List[str] = None,
 ):
     user_message_time = user_message_time or datetime.now().strftime("%Y-%m-%d %H:%M:%S")
     updated_conversation = message_to_log(
@@ -117,7 +117,7 @@ def save_to_conversation_log(
         chat_response=chat_response,
         user_message_metadata={
             "created": user_message_time,
-            "uploadedImageData": uploaded_image_url,
+            "images": query_images,
         },
         khoj_message_metadata={
             "context": compiled_references,
@@ -145,10 +145,18 @@ Khoj: "{inferred_queries if ("text-to-image" in intent_type) else chat_response}
     )
 
 
-# Format user and system messages to chatml format
-def construct_structured_message(message, image_url, model_type, vision_enabled):
-    if image_url and vision_enabled and model_type == ChatModelOptions.ModelType.OPENAI:
-        return [{"type": "text", "text": message}, {"type": "image_url", "image_url": {"url": image_url}}]
+def construct_structured_message(message: str, images: list[str], model_type: str, vision_enabled: bool):
+    """
+    Format messages into appropriate multimedia format for supported chat model types
+    """
+    if not images or not vision_enabled:
+        return message
+
+    if model_type in [ChatModelOptions.ModelType.OPENAI, ChatModelOptions.ModelType.GOOGLE]:
+        return [
+            {"type": "text", "text": message},
+            *[{"type": "image_url", "image_url": {"url": image}} for image in images],
+        ]
     return message
 
 
@@ -160,7 +168,7 @@ def generate_chatml_messages_with_context(
     loaded_model: Optional[Llama] = None,
     max_prompt_size=None,
     tokenizer_name=None,
-    uploaded_image_url=None,
+    query_images=None,
     vision_enabled=False,
     model_type="",
 ):
@@ -181,11 +189,12 @@ def generate_chatml_messages_with_context(
         message_notes = f'\n\n Notes:\n{chat.get("context")}' if chat.get("context") else "\n"
         role = "user" if chat["by"] == "you" else "assistant"
 
-        message_content = chat["message"] + message_notes
+        if chat["by"] == "khoj" and "excalidraw" in chat["intent"].get("type"):
+            message_content = chat.get("intent").get("inferred-queries")[0] + message_notes
+        else:
+            message_content = chat["message"] + message_notes
 
-        message_content = construct_structured_message(
-            message_content, chat.get("uploadedImageData"), model_type, vision_enabled
-        )
+        message_content = construct_structured_message(message_content, chat.get("images"), model_type, vision_enabled)
 
         reconstructed_message = ChatMessage(content=message_content, role=role)
 
@@ -198,7 +207,7 @@ def generate_chatml_messages_with_context(
     if not is_none_or_empty(user_message):
         messages.append(
             ChatMessage(
-                content=construct_structured_message(user_message, uploaded_image_url, model_type, vision_enabled),
+                content=construct_structured_message(user_message, query_images, model_type, vision_enabled),
                 role="user",
             )
         )
@@ -222,7 +231,6 @@ def truncate_messages(
     tokenizer_name=None,
 ) -> list[ChatMessage]:
     """Truncate messages to fit within max prompt size supported by model"""
-
     default_tokenizer = "gpt-4o"
 
     try:
@@ -252,6 +260,7 @@ def truncate_messages(
             system_message = messages.pop(idx)
             break
 
+    # TODO: Handle truncation of multi-part message.content, i.e when message.content is a list[dict] rather than a string
     system_message_tokens = (
         len(encoder.encode(system_message.content)) if system_message and type(system_message.content) == str else 0
     )

khoj/processor/image/generate.py

@@ -26,7 +26,7 @@ async def text_to_image(
     references: List[Dict[str, Any]],
     online_results: Dict[str, Any],
     send_status_func: Optional[Callable] = None,
-    uploaded_image_url: Optional[str] = None,
+    query_images: Optional[List[str]] = None,
     agent: Agent = None,
 ):
     status_code = 200
@@ -65,7 +65,7 @@ async def text_to_image(
         note_references=references,
         online_results=online_results,
         model_type=text_to_image_config.model_type,
-        uploaded_image_url=uploaded_image_url,
+        query_images=query_images,
         user=user,
         agent=agent,
     )
@@ -87,18 +87,18 @@ async def text_to_image(
             if "content_policy_violation" in e.message:
                 logger.error(f"Image Generation blocked by OpenAI: {e}")
                 status_code = e.status_code  # type: ignore
-                message = f"Image generation blocked by OpenAI: {e.message}"  # 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
                 return
             else:
                 logger.error(f"Image Generation failed with {e}", exc_info=True)
-                message = f"Image generation failed with OpenAI error: {e.message}"  # type: ignore
+                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
                 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 with error: {e}"
+            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
             return

khoj/processor/tools/online_search.py

@@ -62,7 +62,7 @@ async def search_online(
     user: KhojUser,
     send_status_func: Optional[Callable] = None,
     custom_filters: List[str] = [],
-    uploaded_image_url: str = None,
+    query_images: List[str] = None,
     agent: Agent = None,
 ):
     query += " ".join(custom_filters)
@@ -73,7 +73,7 @@ async def search_online(
 
     # Breakdown the query into subqueries to get the correct answer
     subqueries = await generate_online_subqueries(
-        query, conversation_history, location, user, uploaded_image_url=uploaded_image_url, agent=agent
+        query, conversation_history, location, user, query_images=query_images, agent=agent
     )
     response_dict = {}
 
@@ -151,7 +151,7 @@ async def read_webpages(
     location: LocationData,
     user: KhojUser,
     send_status_func: Optional[Callable] = None,
-    uploaded_image_url: str = None,
+    query_images: List[str] = None,
     agent: Agent = None,
 ):
     "Infer web pages to read from the query and extract relevant information from them"
@@ -159,7 +159,7 @@ async def read_webpages(
     if send_status_func:
         async for event in send_status_func(f"**Inferring web pages to read**"):
             yield {ChatEvent.STATUS: event}
-    urls = await infer_webpage_urls(query, conversation_history, location, user, uploaded_image_url)
+    urls = await infer_webpage_urls(query, conversation_history, location, user, query_images)
 
     logger.info(f"Reading web pages at: {urls}")
     if send_status_func:

khoj/routers/api.py

@@ -347,7 +347,7 @@ async def extract_references_and_questions(
     conversation_commands: List[ConversationCommand] = [ConversationCommand.Default],
     location_data: LocationData = None,
     send_status_func: Optional[Callable] = None,
-    uploaded_image_url: Optional[str] = None,
+    query_images: Optional[List[str]] = None,
     agent: Agent = None,
 ):
     user = request.user.object if request.user.is_authenticated else None
@@ -438,7 +438,7 @@ async def extract_references_and_questions(
                 conversation_log=meta_log,
                 location_data=location_data,
                 user=user,
-                uploaded_image_url=uploaded_image_url,
+                query_images=query_images,
                 vision_enabled=vision_enabled,
                 personality_context=personality_context,
             )
@@ -459,12 +459,14 @@ async def extract_references_and_questions(
             chat_model = conversation_config.chat_model
             inferred_queries = extract_questions_gemini(
                 defiltered_query,
+                query_images=query_images,
                 model=chat_model,
                 api_key=api_key,
                 conversation_log=meta_log,
                 location_data=location_data,
                 max_tokens=conversation_config.max_prompt_size,
                 user=user,
+                vision_enabled=vision_enabled,
                 personality_context=personality_context,
             )
 

khoj/routers/api_agents.py

@@ -1,5 +1,7 @@
 import json
 import logging
+import random
+from datetime import datetime, timedelta, timezone
 from typing import Dict, List, Optional
 
 from asgiref.sync import sync_to_async
@@ -9,8 +11,8 @@ from fastapi.responses import Response
 from pydantic import BaseModel
 from starlette.authentication import requires
 
-from khoj.database.adapters import AgentAdapters
-from khoj.database.models import Agent, KhojUser
+from khoj.database.adapters import AgentAdapters, ConversationAdapters
+from khoj.database.models import Agent, Conversation, KhojUser
 from khoj.routers.helpers import CommonQueryParams, acheck_if_safe_prompt
 from khoj.utils.helpers import (
     ConversationCommand,
@@ -45,30 +47,49 @@ async def all_agents(
 ) -> Response:
     user: KhojUser = request.user.object if request.user.is_authenticated else None
     agents = await AgentAdapters.aget_all_accessible_agents(user)
+    default_agent = await AgentAdapters.aget_default_agent()
+    default_agent_packet = None
     agents_packet = list()
     for agent in agents:
         files = agent.fileobject_set.all()
         file_names = [file.file_name for file in files]
-        agents_packet.append(
-            {
-                "slug": agent.slug,
-                "name": agent.name,
-                "persona": agent.personality,
-                "creator": agent.creator.username if agent.creator else None,
-                "managed_by_admin": agent.managed_by_admin,
-                "color": agent.style_color,
-                "icon": agent.style_icon,
-                "privacy_level": agent.privacy_level,
-                "chat_model": agent.chat_model.chat_model,
-                "files": file_names,
-                "input_tools": agent.input_tools,
-                "output_modes": agent.output_modes,
-            }
+        agent_packet = {
+            "slug": agent.slug,
+            "name": agent.name,
+            "persona": agent.personality,
+            "creator": agent.creator.username if agent.creator else None,
+            "managed_by_admin": agent.managed_by_admin,
+            "color": agent.style_color,
+            "icon": agent.style_icon,
+            "privacy_level": agent.privacy_level,
+            "chat_model": agent.chat_model.chat_model,
+            "files": file_names,
+            "input_tools": agent.input_tools,
+            "output_modes": agent.output_modes,
+        }
+        if agent.slug == default_agent.slug:
+            default_agent_packet = agent_packet
+        else:
+            agents_packet.append(agent_packet)
+
+    # Load recent conversation sessions
+    min_date = datetime.min.replace(tzinfo=timezone.utc)
+    two_weeks_ago = datetime.today() - timedelta(weeks=2)
+    conversations = []
+    if user:
+        conversations = await sync_to_async(list[Conversation])(
+            ConversationAdapters.get_conversation_sessions(user, request.user.client_app)
+            .filter(updated_at__gte=two_weeks_ago)
+            .order_by("-updated_at")[:50]
         )
+    conversation_times = {conv.agent.slug: conv.updated_at for conv in conversations if conv.agent}
+
+    # Put default agent first, then sort by mru and finally shuffle unused randomly
+    random.shuffle(agents_packet)
+    agents_packet.sort(key=lambda x: conversation_times.get(x["slug"]) or min_date, reverse=True)
+    if default_agent_packet:
+        agents_packet.insert(0, default_agent_packet)
 
-    # Make sure that the agent named 'khoj' is first in the list. Everything else is sorted by name.
-    agents_packet.sort(key=lambda x: x["name"])
-    agents_packet.sort(key=lambda x: x["slug"] == "khoj", reverse=True)
     return Response(content=json.dumps(agents_packet), media_type="application/json", status_code=200)