khoj 1.26.2__py3-none-any.whl → 1.26.5.dev34__py3-none-any.whl

khoj/processor/conversation/google/gemini_chat.py

@@ -6,14 +6,17 @@ from typing import Dict, Optional
 
 from langchain.schema import ChatMessage
 
-from khoj.database.models import Agent, KhojUser
+from khoj.database.models import Agent, ChatModelOptions, KhojUser
 from khoj.processor.conversation import prompts
 from khoj.processor.conversation.google.utils import (
     format_messages_for_gemini,
     gemini_chat_completion_with_backoff,
     gemini_completion_with_backoff,
 )
-from khoj.processor.conversation.utils import generate_chatml_messages_with_context
+from khoj.processor.conversation.utils import (
+    construct_structured_message,
+    generate_chatml_messages_with_context,
+)
 from khoj.utils.helpers import ConversationCommand, is_none_or_empty
 from khoj.utils.rawconfig import LocationData
 
@@ -29,6 +32,8 @@ def extract_questions_gemini(
     max_tokens=None,
     location_data: LocationData = None,
     user: KhojUser = None,
+    query_images: Optional[list[str]] = None,
+    vision_enabled: bool = False,
     personality_context: Optional[str] = None,
 ):
     """
@@ -70,17 +75,17 @@ def extract_questions_gemini(
         text=text,
     )
 
-    messages = [ChatMessage(content=prompt, role="user")]
+    prompt = construct_structured_message(
+        message=prompt,
+        images=query_images,
+        model_type=ChatModelOptions.ModelType.GOOGLE,
+        vision_enabled=vision_enabled,
+    )
 
-    model_kwargs = {"response_mime_type": "application/json"}
+    messages = [ChatMessage(content=prompt, role="user"), ChatMessage(content=system_prompt, role="system")]
 
-    response = gemini_completion_with_backoff(
-        messages=messages,
-        system_prompt=system_prompt,
-        model_name=model,
-        temperature=temperature,
-        api_key=api_key,
-        model_kwargs=model_kwargs,
+    response = gemini_send_message_to_model(
+        messages, api_key, model, response_type="json_object", temperature=temperature
     )
 
     # Extract, Clean Message from Gemini's Response
@@ -102,7 +107,7 @@ def extract_questions_gemini(
     return questions
 
 
-def gemini_send_message_to_model(messages, api_key, model, response_type="text"):
+def gemini_send_message_to_model(messages, api_key, model, response_type="text", temperature=0, model_kwargs=None):
     """
     Send message to model
     """
@@ -114,7 +119,12 @@ def gemini_send_message_to_model(messages, api_key, model, response_type="text")
 
     # Get Response from Gemini
     return gemini_completion_with_backoff(
-        messages=messages, system_prompt=system_prompt, model_name=model, api_key=api_key, model_kwargs=model_kwargs
+        messages=messages,
+        system_prompt=system_prompt,
+        model_name=model,
+        api_key=api_key,
+        temperature=temperature,
+        model_kwargs=model_kwargs,
     )
 
 
@@ -133,6 +143,8 @@ def converse_gemini(
     location_data: LocationData = None,
     user_name: str = None,
     agent: Agent = None,
+    query_images: Optional[list[str]] = None,
+    vision_available: bool = False,
 ):
     """
     Converse with user using Google's Gemini
@@ -187,6 +199,9 @@ def converse_gemini(
         model_name=model,
         max_prompt_size=max_prompt_size,
         tokenizer_name=tokenizer_name,
+        query_images=query_images,
+        vision_enabled=vision_available,
+        model_type=ChatModelOptions.ModelType.GOOGLE,
     )
 
     messages, system_prompt = format_messages_for_gemini(messages, system_prompt)

khoj/processor/conversation/google/utils.py

@@ -1,8 +1,11 @@
 import logging
 import random
+from io import BytesIO
 from threading import Thread
 
 import google.generativeai as genai
+import PIL.Image
+import requests
 from google.generativeai.types.answer_types import FinishReason
 from google.generativeai.types.generation_types import StopCandidateException
 from google.generativeai.types.safety_types import (
@@ -53,14 +56,14 @@ def gemini_completion_with_backoff(
         },
     )
 
-    formatted_messages = [{"role": message.role, "parts": [message.content]} for message in messages]
+    formatted_messages = [{"role": message.role, "parts": message.content} for message in messages]
 
     # Start chat session. All messages up to the last are considered to be part of the chat history
     chat_session = model.start_chat(history=formatted_messages[0:-1])
 
     try:
         # Generate the response. The last message is considered to be the current prompt
-        aggregated_response = chat_session.send_message(formatted_messages[-1]["parts"][0])
+        aggregated_response = chat_session.send_message(formatted_messages[-1]["parts"])
         return aggregated_response.text
     except StopCandidateException as e:
         response_message, _ = handle_gemini_response(e.args)
@@ -117,11 +120,11 @@ def gemini_llm_thread(g, messages, system_prompt, model_name, temperature, api_k
             },
         )
 
-        formatted_messages = [{"role": message.role, "parts": [message.content]} for message in messages]
+        formatted_messages = [{"role": message.role, "parts": message.content} for message in messages]
         # all messages up to the last are considered to be part of the chat history
         chat_session = model.start_chat(history=formatted_messages[0:-1])
         # the last message is considered to be the current prompt
-        for chunk in chat_session.send_message(formatted_messages[-1]["parts"][0], stream=True):
+        for chunk in chat_session.send_message(formatted_messages[-1]["parts"], stream=True):
             message, stopped = handle_gemini_response(chunk.candidates, chunk.prompt_feedback)
             message = message or chunk.text
             g.send(message)
@@ -191,14 +194,6 @@ def generate_safety_response(safety_ratings):
 
 
 def format_messages_for_gemini(messages: list[ChatMessage], system_prompt: str = None) -> tuple[list[str], str]:
-    if len(messages) == 1:
-        messages[0].role = "user"
-        return messages, system_prompt
-
-    for message in messages:
-        if message.role == "assistant":
-            message.role = "model"
-
     # Extract system message
     system_prompt = system_prompt or ""
     for message in messages.copy():
@@ -207,4 +202,31 @@ def format_messages_for_gemini(messages: list[ChatMessage], system_prompt: str =
             messages.remove(message)
     system_prompt = None if is_none_or_empty(system_prompt) else system_prompt
 
+    for message in messages:
+        # Convert message content to string list from chatml dictionary list
+        if isinstance(message.content, list):
+            # Convert image_urls to PIL.Image and place them at beginning of list (better for Gemini)
+            message.content = [
+                get_image_from_url(item["image_url"]["url"]) if item["type"] == "image_url" else item["text"]
+                for item in sorted(message.content, key=lambda x: 0 if x["type"] == "image_url" else 1)
+            ]
+        elif isinstance(message.content, str):
+            message.content = [message.content]
+
+        if message.role == "assistant":
+            message.role = "model"
+
+    if len(messages) == 1:
+        messages[0].role = "user"
+
     return messages, system_prompt
+
+
+def get_image_from_url(image_url: str) -> PIL.Image:
+    try:
+        response = requests.get(image_url)
+        response.raise_for_status()  # Check if the request was successful
+        return PIL.Image.open(BytesIO(response.content))
+    except requests.exceptions.RequestException as e:
+        logger.error(f"Failed to get image from URL {image_url}: {e}")
+        return None

khoj/processor/conversation/openai/gpt.py

@@ -30,7 +30,7 @@ def extract_questions(
     api_base_url=None,
     location_data: LocationData = None,
     user: KhojUser = None,
-    uploaded_image_url: Optional[str] = None,
+    query_images: Optional[list[str]] = None,
     vision_enabled: bool = False,
     personality_context: Optional[str] = None,
 ):
@@ -74,7 +74,7 @@ def extract_questions(
 
     prompt = construct_structured_message(
         message=prompt,
-        image_url=uploaded_image_url,
+        images=query_images,
         model_type=ChatModelOptions.ModelType.OPENAI,
         vision_enabled=vision_enabled,
     )
@@ -135,7 +135,7 @@ def converse(
     location_data: LocationData = None,
     user_name: str = None,
     agent: Agent = None,
-    image_url: Optional[str] = None,
+    query_images: Optional[list[str]] = None,
     vision_available: bool = False,
 ):
     """
@@ -191,7 +191,7 @@ def converse(
         model_name=model,
         max_prompt_size=max_prompt_size,
         tokenizer_name=tokenizer_name,
-        uploaded_image_url=image_url,
+        query_images=query_images,
         vision_enabled=vision_available,
         model_type=ChatModelOptions.ModelType.OPENAI,
     )

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

@@ -21,6 +21,7 @@ from starlette.authentication import has_required_scope, requires
 from khoj.configure import initialize_content
 from khoj.database import adapters
 from khoj.database.adapters import (
+    AgentAdapters,
     AutomationAdapters,
     ConversationAdapters,
     EntryAdapters,
@@ -114,10 +115,16 @@ async def execute_search(
     dedupe: Optional[bool] = True,
     agent: Optional[Agent] = None,
 ):
-    start_time = time.time()
-
     # Run validation checks
     results: List[SearchResponse] = []
+
+    start_time = time.time()
+
+    # Ensure the agent, if present, is accessible by the user
+    if user and agent and not await AgentAdapters.ais_agent_accessible(agent, user):
+        logger.error(f"Agent {agent.slug} is not accessible by user {user}")
+        return results
+
     if q is None or q == "":
         logger.warning(f"No query param (q) passed in API call to initiate search")
         return results
@@ -340,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
@@ -431,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,
             )
@@ -452,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,
             )