khoj 2.0.0b14.dev51__py3-none-any.whl → 2.0.0b15.dev22__py3-none-any.whl

khoj/processor/conversation/google/gemini_chat.py

@@ -1,22 +1,16 @@
 import logging
-from datetime import datetime
-from typing import AsyncGenerator, Dict, List, Optional
+from typing import AsyncGenerator, List, Optional
+
+from langchain_core.messages.chat import ChatMessage
 
-from khoj.database.models import Agent, ChatMessageModel, ChatModel
-from khoj.processor.conversation import prompts
 from khoj.processor.conversation.google.utils import (
     gemini_chat_completion_with_backoff,
     gemini_completion_with_backoff,
 )
 from khoj.processor.conversation.utils import (
-    OperatorRun,
     ResponseWithThought,
-    generate_chatml_messages_with_context,
     messages_to_print,
 )
-from khoj.utils.helpers import is_none_or_empty, truncate_code_context
-from khoj.utils.rawconfig import FileAttachment, LocationData
-from khoj.utils.yaml import yaml_dump
 
 logger = logging.getLogger(__name__)
 
@@ -61,95 +55,18 @@ def gemini_send_message_to_model(
 
 async def converse_gemini(
     # Query
-    user_query: str,
-    # Context
-    references: list[dict],
-    online_results: Optional[Dict[str, Dict]] = None,
-    code_results: Optional[Dict[str, Dict]] = None,
-    operator_results: Optional[List[OperatorRun]] = None,
-    query_images: Optional[list[str]] = None,
-    query_files: str = None,
-    generated_files: List[FileAttachment] = None,
-    generated_asset_results: Dict[str, Dict] = {},
-    program_execution_context: List[str] = None,
-    location_data: LocationData = None,
-    user_name: str = None,
-    chat_history: List[ChatMessageModel] = [],
+    messages: List[ChatMessage],
     # Model
     model: Optional[str] = "gemini-2.5-flash",
     api_key: Optional[str] = None,
     api_base_url: Optional[str] = None,
     temperature: float = 1.0,
-    max_prompt_size=None,
-    tokenizer_name=None,
-    agent: Agent = None,
-    vision_available: bool = False,
     deepthought: Optional[bool] = False,
     tracer={},
 ) -> AsyncGenerator[ResponseWithThought, None]:
     """
     Converse with user using Google's Gemini
     """
-    # Initialize Variables
-    current_date = datetime.now()
-
-    if agent and agent.personality:
-        system_prompt = prompts.custom_personality.format(
-            name=agent.name,
-            bio=agent.personality,
-            current_date=current_date.strftime("%Y-%m-%d"),
-            day_of_week=current_date.strftime("%A"),
-        )
-    else:
-        system_prompt = prompts.personality.format(
-            current_date=current_date.strftime("%Y-%m-%d"),
-            day_of_week=current_date.strftime("%A"),
-        )
-
-    system_prompt += f"{system_prompt}\n\n{prompts.gemini_verbose_language_personality}"
-    if location_data:
-        location_prompt = prompts.user_location.format(location=f"{location_data}")
-        system_prompt = f"{system_prompt}\n{location_prompt}"
-
-    if user_name:
-        user_name_prompt = prompts.user_name.format(name=user_name)
-        system_prompt = f"{system_prompt}\n{user_name_prompt}"
-
-    context_message = ""
-    if not is_none_or_empty(references):
-        context_message = f"{prompts.notes_conversation.format(query=user_query, references=yaml_dump(references))}\n\n"
-    if not is_none_or_empty(online_results):
-        context_message += f"{prompts.online_search_conversation.format(online_results=yaml_dump(online_results))}\n\n"
-    if not is_none_or_empty(code_results):
-        context_message += (
-            f"{prompts.code_executed_context.format(code_results=truncate_code_context(code_results))}\n\n"
-        )
-    if not is_none_or_empty(operator_results):
-        operator_content = [
-            {"query": oc.query, "response": oc.response, "webpages": oc.webpages} for oc in operator_results
-        ]
-        context_message += (
-            f"{prompts.operator_execution_context.format(operator_results=yaml_dump(operator_content))}\n\n"
-        )
-    context_message = context_message.strip()
-
-    # Setup Prompt with Primer or Conversation History
-    messages = generate_chatml_messages_with_context(
-        user_query,
-        context_message=context_message,
-        chat_history=chat_history,
-        model_name=model,
-        max_prompt_size=max_prompt_size,
-        tokenizer_name=tokenizer_name,
-        query_images=query_images,
-        vision_enabled=vision_available,
-        model_type=ChatModel.ModelType.GOOGLE,
-        query_files=query_files,
-        generated_files=generated_files,
-        generated_asset_results=generated_asset_results,
-        program_execution_context=program_execution_context,
-    )
-
     logger.debug(f"Conversation Context for Gemini: {messages_to_print(messages)}")
 
     # Get Response from Google AI
@@ -159,7 +76,6 @@ async def converse_gemini(
         temperature=temperature,
         api_key=api_key,
         api_base_url=api_base_url,
-        system_prompt=system_prompt,
         deepthought=deepthought,
         tracer=tracer,
     ):

khoj/processor/conversation/google/utils.py

@@ -308,7 +308,7 @@ async def gemini_chat_completion_with_backoff(
     temperature: float,
     api_key: str,
     api_base_url: str,
-    system_prompt: str,
+    system_prompt: str = "",
     model_kwargs=None,
     deepthought=False,
     tracer: dict = {},
@@ -472,9 +472,12 @@ def format_messages_for_gemini(
     for message in messages.copy():
         if message.role == "system":
             if isinstance(message.content, list):
-                system_prompt += "\n".join([part["text"] for part in message.content if part["type"] == "text"])
+                system_prompt += "\n\n" + "\n".join(
+                    [part["text"] for part in message.content if part["type"] == "text"]
+                )
             else:
-                system_prompt += message.content
+                system_prompt += "\n\n" + message.content
+            system_prompt = system_prompt.strip()
             messages.remove(message)
     system_prompt = None if is_none_or_empty(system_prompt) else system_prompt
 

khoj/processor/conversation/openai/gpt.py

@@ -1,29 +1,25 @@
 import logging
-from datetime import datetime
 from typing import Any, AsyncGenerator, Dict, List, Optional
 
-from khoj.database.models import Agent, ChatMessageModel, ChatModel
-from khoj.processor.conversation import prompts
+from langchain_core.messages.chat import ChatMessage
+
 from khoj.processor.conversation.openai.utils import (
     chat_completion_with_backoff,
     clean_response_schema,
     completion_with_backoff,
     get_structured_output_support,
-    is_openai_api,
+    is_cerebras_api,
     responses_chat_completion_with_backoff,
     responses_completion_with_backoff,
+    supports_responses_api,
     to_openai_tools,
 )
 from khoj.processor.conversation.utils import (
-    OperatorRun,
     ResponseWithThought,
     StructuredOutputSupport,
-    generate_chatml_messages_with_context,
     messages_to_print,
 )
-from khoj.utils.helpers import ToolDefinition, is_none_or_empty, truncate_code_context
-from khoj.utils.rawconfig import FileAttachment, LocationData
-from khoj.utils.yaml import yaml_dump
+from khoj.utils.helpers import ToolDefinition
 
 logger = logging.getLogger(__name__)
 
@@ -45,16 +41,19 @@ def send_message_to_model(
 
     model_kwargs: Dict[str, Any] = {}
     json_support = get_structured_output_support(model, api_base_url)
+    strict = not is_cerebras_api(api_base_url)
     if tools and json_support == StructuredOutputSupport.TOOL:
-        model_kwargs["tools"] = to_openai_tools(tools, use_responses_api=is_openai_api(api_base_url))
+        model_kwargs["tools"] = to_openai_tools(
+            tools, use_responses_api=supports_responses_api(model, api_base_url), strict=strict
+        )
     elif response_schema and json_support >= StructuredOutputSupport.SCHEMA:
         # Drop unsupported fields from schema passed to OpenAI APi
         cleaned_response_schema = clean_response_schema(response_schema)
-        if is_openai_api(api_base_url):
+        if supports_responses_api(model, api_base_url):
             model_kwargs["text"] = {
                 "format": {
                     "type": "json_schema",
-                    "strict": True,
+                    "strict": strict,
                     "name": response_schema.__name__,
                     "schema": cleaned_response_schema,
                 }
@@ -65,14 +64,14 @@ def send_message_to_model(
                 "json_schema": {
                     "schema": cleaned_response_schema,
                     "name": response_schema.__name__,
-                    "strict": True,
+                    "strict": strict,
                 },
             }
     elif response_type == "json_object" and json_support == StructuredOutputSupport.OBJECT:
         model_kwargs["response_format"] = {"type": response_type}
 
     # Get Response from GPT
-    if is_openai_api(api_base_url):
+    if supports_responses_api(model, api_base_url):
         return responses_completion_with_backoff(
             messages=messages,
             model_name=model,
@@ -96,98 +95,22 @@ def send_message_to_model(
 
 async def converse_openai(
     # Query
-    user_query: str,
-    # Context
-    references: list[dict],
-    online_results: Optional[Dict[str, Dict]] = None,
-    code_results: Optional[Dict[str, Dict]] = None,
-    operator_results: Optional[List[OperatorRun]] = None,
-    query_images: Optional[list[str]] = None,
-    query_files: str = None,
-    generated_files: List[FileAttachment] = None,
-    generated_asset_results: Dict[str, Dict] = {},
-    program_execution_context: List[str] = None,
-    location_data: LocationData = None,
-    chat_history: list[ChatMessageModel] = [],
+    messages: List[ChatMessage],
+    # Model
     model: str = "gpt-4.1-mini",
     api_key: Optional[str] = None,
     api_base_url: Optional[str] = None,
     temperature: float = 0.6,
-    max_prompt_size=None,
-    tokenizer_name=None,
-    user_name: str = None,
-    agent: Agent = None,
-    vision_available: bool = False,
     deepthought: Optional[bool] = False,
     tracer: dict = {},
 ) -> AsyncGenerator[ResponseWithThought, None]:
     """
     Converse with user using OpenAI's ChatGPT
     """
-    # Initialize Variables
-    current_date = datetime.now()
-
-    if agent and agent.personality:
-        system_prompt = prompts.custom_personality.format(
-            name=agent.name,
-            bio=agent.personality,
-            current_date=current_date.strftime("%Y-%m-%d"),
-            day_of_week=current_date.strftime("%A"),
-        )
-    else:
-        system_prompt = prompts.personality.format(
-            current_date=current_date.strftime("%Y-%m-%d"),
-            day_of_week=current_date.strftime("%A"),
-        )
-
-    if location_data:
-        location_prompt = prompts.user_location.format(location=f"{location_data}")
-        system_prompt = f"{system_prompt}\n{location_prompt}"
-
-    if user_name:
-        user_name_prompt = prompts.user_name.format(name=user_name)
-        system_prompt = f"{system_prompt}\n{user_name_prompt}"
-
-    context_message = ""
-    if not is_none_or_empty(references):
-        context_message = f"{prompts.notes_conversation.format(references=yaml_dump(references))}\n\n"
-    if not is_none_or_empty(online_results):
-        context_message += f"{prompts.online_search_conversation.format(online_results=yaml_dump(online_results))}\n\n"
-    if not is_none_or_empty(code_results):
-        context_message += (
-            f"{prompts.code_executed_context.format(code_results=truncate_code_context(code_results))}\n\n"
-        )
-    if not is_none_or_empty(operator_results):
-        operator_content = [
-            {"query": oc.query, "response": oc.response, "webpages": oc.webpages} for oc in operator_results
-        ]
-        context_message += (
-            f"{prompts.operator_execution_context.format(operator_results=yaml_dump(operator_content))}\n\n"
-        )
-
-    context_message = context_message.strip()
-
-    # Setup Prompt with Primer or Conversation History
-    messages = generate_chatml_messages_with_context(
-        user_query,
-        system_prompt,
-        chat_history,
-        context_message=context_message,
-        model_name=model,
-        max_prompt_size=max_prompt_size,
-        tokenizer_name=tokenizer_name,
-        query_images=query_images,
-        vision_enabled=vision_available,
-        model_type=ChatModel.ModelType.OPENAI,
-        query_files=query_files,
-        generated_files=generated_files,
-        generated_asset_results=generated_asset_results,
-        program_execution_context=program_execution_context,
-    )
     logger.debug(f"Conversation Context for GPT: {messages_to_print(messages)}")
 
     # Get Response from GPT
-    if is_openai_api(api_base_url):
+    if supports_responses_api(model, api_base_url):
         async for chunk in responses_chat_completion_with_backoff(
             messages=messages,
             model_name=model,

khoj/processor/conversation/openai/utils.py

@@ -111,14 +111,16 @@ def completion_with_backoff(
     model_kwargs["temperature"] = temperature
     model_kwargs["top_p"] = model_kwargs.get("top_p", 0.95)
 
-    formatted_messages = format_message_for_api(messages, api_base_url)
+    formatted_messages = format_message_for_api(messages, model_name, api_base_url)
 
     # Tune reasoning models arguments
     if is_openai_reasoning_model(model_name, api_base_url):
         model_kwargs["temperature"] = 1
         reasoning_effort = "medium" if deepthought else "low"
         model_kwargs["reasoning_effort"] = reasoning_effort
+        # Remove unsupported params for reasoning models
         model_kwargs.pop("top_p", None)
+        model_kwargs.pop("stop", None)
     elif is_twitter_reasoning_model(model_name, api_base_url):
         model_kwargs.pop("temperature", None)
         reasoning_effort = "high" if deepthought else "low"
@@ -294,7 +296,7 @@ async def chat_completion_with_backoff(
 
     model_kwargs["top_p"] = model_kwargs.get("top_p", 0.95)
 
-    formatted_messages = format_message_for_api(messages, api_base_url)
+    formatted_messages = format_message_for_api(messages, model_name, api_base_url)
 
     # Configure thinking for openai reasoning models
     if is_openai_reasoning_model(model_name, api_base_url):
@@ -304,19 +306,6 @@ async def chat_completion_with_backoff(
         # Remove unsupported params for reasoning models
         model_kwargs.pop("top_p", None)
         model_kwargs.pop("stop", None)
-
-        # Get the first system message and add the string `Formatting re-enabled` to it.
-        # See https://platform.openai.com/docs/guides/reasoning-best-practices
-        if len(formatted_messages) > 0:
-            system_messages = [
-                (i, message) for i, message in enumerate(formatted_messages) if message["role"] == "system"
-            ]
-            if len(system_messages) > 0:
-                first_system_message_index, first_system_message = system_messages[0]
-                first_system_message_content = first_system_message["content"]
-                formatted_messages[first_system_message_index]["content"] = (
-                    f"{first_system_message_content}\nFormatting re-enabled"
-                )
     elif is_twitter_reasoning_model(model_name, api_base_url):
         reasoning_effort = "high" if deepthought else "low"
         # Grok-4 models do not support reasoning_effort parameter
@@ -325,7 +314,7 @@ async def chat_completion_with_backoff(
     elif (
         model_name.startswith("deepseek-chat")
         or model_name.startswith("deepseek-reasoner")
-        or "deepseek-r1" in model_name
+        or "deepseek-r1" in model_name.lower()
     ):
         # Official Deepseek models and some inference APIs like vLLM return structured thinking output.
         # Others like DeepInfra return it in response stream.
@@ -459,7 +448,7 @@ def responses_completion_with_backoff(
         client = get_openai_client(openai_api_key, api_base_url)
         openai_clients[client_key] = client
 
-    formatted_messages = format_message_for_api(messages, api_base_url)
+    formatted_messages = format_message_for_api(messages, model_name, api_base_url)
     # Move the first system message to Responses API instructions
     instructions: Optional[str] = None
     if formatted_messages and formatted_messages[0].get("role") == "system":
@@ -472,8 +461,10 @@ def responses_completion_with_backoff(
     if is_openai_reasoning_model(model_name, api_base_url):
         temperature = 1
         reasoning_effort = "medium" if deepthought else "low"
-        model_kwargs["reasoning"] = {"effort": reasoning_effort, "summary": "auto"}
-        model_kwargs["include"] = ["reasoning.encrypted_content"]
+        model_kwargs["reasoning"] = {"effort": reasoning_effort}
+        if is_openai_api(api_base_url):
+            model_kwargs["reasoning"]["summary"] = "auto"
+            model_kwargs["include"] = ["reasoning.encrypted_content"]
         # Remove unsupported params for reasoning models
         model_kwargs.pop("top_p", None)
         model_kwargs.pop("stop", None)
@@ -570,7 +561,7 @@ async def responses_chat_completion_with_backoff(
         client = get_openai_async_client(openai_api_key, api_base_url)
         openai_async_clients[client_key] = client
 
-    formatted_messages = format_message_for_api(messages, api_base_url)
+    formatted_messages = format_message_for_api(messages, model_name, api_base_url)
     # Move the first system message to Responses API instructions
     instructions: Optional[str] = None
     if formatted_messages and formatted_messages[0].get("role") == "system":
@@ -583,7 +574,10 @@ async def responses_chat_completion_with_backoff(
     if is_openai_reasoning_model(model_name, api_base_url):
         temperature = 1
         reasoning_effort = "medium" if deepthought else "low"
-        model_kwargs["reasoning"] = {"effort": reasoning_effort, "summary": "auto"}
+        model_kwargs["reasoning"] = {"effort": reasoning_effort}
+        if is_openai_api(api_base_url):
+            model_kwargs["reasoning"]["summary"] = "auto"
+            model_kwargs["include"] = ["reasoning.encrypted_content"]
         # Remove unsupported params for reasoning models
         model_kwargs.pop("top_p", None)
         model_kwargs.pop("stop", None)
@@ -716,7 +710,7 @@ def get_structured_output_support(model_name: str, api_base_url: str = None) ->
     return StructuredOutputSupport.TOOL
 
 
-def format_message_for_api(raw_messages: List[ChatMessage], api_base_url: str) -> List[dict]:
+def format_message_for_api(raw_messages: List[ChatMessage], model_name: str, api_base_url: str) -> List[dict]:
     """
     Format messages to send to chat model served over OpenAI (compatible) API.
     """
@@ -726,7 +720,7 @@ def format_message_for_api(raw_messages: List[ChatMessage], api_base_url: str) -
         # Handle tool call and tool result message types
         message_type = message.additional_kwargs.get("message_type")
         if message_type == "tool_call":
-            if is_openai_api(api_base_url):
+            if supports_responses_api(model_name, api_base_url):
                 for part in message.content:
                     if "status" in part:
                         part.pop("status")  # Drop unsupported tool call status field
@@ -770,7 +764,7 @@ def format_message_for_api(raw_messages: List[ChatMessage], api_base_url: str) -
                 if not tool_call_id:
                     logger.warning(f"Dropping tool result without valid tool_call_id: {part.get('name')}")
                     continue
-                if is_openai_api(api_base_url):
+                if supports_responses_api(model_name, api_base_url):
                     formatted_messages.append(
                         {
                             "type": "function_call_output",
@@ -788,7 +782,7 @@ def format_message_for_api(raw_messages: List[ChatMessage], api_base_url: str) -
                         }
                     )
             continue
-        if isinstance(message.content, list) and not is_openai_api(api_base_url):
+        if isinstance(message.content, list) and not supports_responses_api(model_name, api_base_url):
             assistant_texts = []
             has_images = False
             for idx, part in enumerate(message.content):
@@ -801,7 +795,7 @@ def format_message_for_api(raw_messages: List[ChatMessage], api_base_url: str) -
                 if (
                     part.get("type") == "text"
                     and message.role == "assistant"
-                    and api_base_url.startswith("https://api.deepinfra.com/v1")
+                    and (api_base_url.startswith("https://api.deepinfra.com/v1") or is_cerebras_api(api_base_url))
                 ):
                     assistant_texts += [part["text"]]
                     message.content.pop(idx)
@@ -844,6 +838,13 @@ def is_openai_api(api_base_url: str = None) -> bool:
     return api_base_url is None or api_base_url.startswith("https://api.openai.com/v1")
 
 
+def supports_responses_api(model_name: str, api_base_url: str = None) -> bool:
+    """
+    Check if the model, ai api supports the OpenAI Responses API
+    """
+    return is_openai_api(api_base_url)
+
+
 def is_openai_reasoning_model(model_name: str, api_base_url: str = None) -> bool:
     """
     Check if the model is an OpenAI reasoning model
@@ -851,7 +852,7 @@ def is_openai_reasoning_model(model_name: str, api_base_url: str = None) -> bool
     return (
         is_openai_api(api_base_url)
         and (model_name.lower().startswith("o") or model_name.lower().startswith("gpt-5"))
-        or model_name.lower().startswith("gpt-oss")
+        or "gpt-oss" in model_name.lower()
     )
 
 
@@ -875,6 +876,13 @@ def is_twitter_reasoning_model(model_name: str, api_base_url: str = None) -> boo
     )
 
 
+def is_cerebras_api(api_base_url: str = None) -> bool:
+    """
+    Check if the model is served over the Cerebras API
+    """
+    return api_base_url is not None and api_base_url.startswith("https://api.cerebras.ai/v1")
+
+
 def is_groq_api(api_base_url: str = None) -> bool:
     """
     Check if the model is served over the Groq API
@@ -1211,7 +1219,7 @@ def add_qwen_no_think_tag(formatted_messages: List[dict]) -> None:
                         break
 
 
-def to_openai_tools(tools: List[ToolDefinition], use_responses_api: bool) -> List[Dict] | None:
+def to_openai_tools(tools: List[ToolDefinition], use_responses_api: bool, strict: bool) -> List[Dict] | None:
     "Transform tool definitions from standard format to OpenAI format."
     if use_responses_api:
         openai_tools = [
@@ -1220,7 +1228,7 @@ def to_openai_tools(tools: List[ToolDefinition], use_responses_api: bool) -> Lis
                 "name": tool.name,
                 "description": tool.description,
                 "parameters": clean_response_schema(tool.schema),
-                "strict": True,
+                "strict": strict,
             }
             for tool in tools
         ]
@@ -1232,7 +1240,7 @@ def to_openai_tools(tools: List[ToolDefinition], use_responses_api: bool) -> Lis
                     "name": tool.name,
                     "description": tool.description,
                     "parameters": clean_response_schema(tool.schema),
-                    "strict": True,
+                    "strict": strict,
                 },
             }
             for tool in tools

khoj/processor/conversation/prompts.py

@@ -18,12 +18,11 @@ Today is {day_of_week}, {current_date} in UTC.
 
 # Style
 - Your responses should be helpful, conversational and tuned to the user's communication style.
-- Make sure to use the specific LaTeX math mode delimiters for your response. LaTex math mode specific delimiters as following
-    - inline math mode : \\( and \\)
-    - display math mode: insert linebreak after opening $$, \\[ and before closing $$, \\]
 - Provide inline citations to documents and websites referenced. Add them inline in markdown format to directly support your claim.
   For example: "The weather today is sunny [1](https://weather.com)."
-- Mention generated assets like images by reference, e.g ![chart](/visualization/image.png). Do not manually output raw, b64 encoded bytes in your response.
+- KaTeX is used to render LaTeX expressions. Make sure you only use the KaTeX math mode delimiters specified below:
+  - inline math mode : \\( and \\)
+  - display math mode: insert linebreak after opening $$, \\[ and before closing $$, \\]
 - Do not respond with raw programs or scripts in your final response unless you know the user is a programmer or has explicitly requested code.
 """.strip()
 )
@@ -41,12 +40,11 @@ Today is {day_of_week}, {current_date} in UTC.
 - Users can share files and other information with you using the Khoj Web, Desktop, Obsidian or Emacs app. They can also drag and drop their files into the chat window.
 
 # Style
-- Make sure to use the specific LaTeX math mode delimiters for your response. LaTex math mode specific delimiters as following
-    - inline math mode : `\\(` and `\\)`
-    - display math mode: insert linebreak after opening `$$`, `\\[` and before closing `$$`, `\\]`
 - Provide inline citations to documents and websites referenced. Add them inline in markdown format to directly support your claim.
   For example: "The weather today is sunny [1](https://weather.com)."
-- Mention generated assets like images by reference, e.g ![chart](/visualization/image.png). Do not manually output raw, b64 encoded bytes in your response.
+- KaTeX is used to render LaTeX expressions. Make sure you only use the KaTeX math mode delimiters specified below:
+  - inline math mode : \\( and \\)
+  - display math mode: insert linebreak after opening $$, \\[ and before closing $$, \\]
 
 # Instructions:\n{bio}
 """.strip()
@@ -115,45 +113,38 @@ User's Notes:
 ## Image Generation
 ## --
 
-image_generation_improve_prompt_base = """
+enhance_image_system_message = PromptTemplate.from_template(
+    """
 You are a talented media artist with the ability to describe images to compose in professional, fine detail.
+Your image description will be transformed into an image by an AI model on your team.
 {personality_context}
-Generate a vivid description of the image to be rendered using the provided context and user prompt below:
-
-Today's Date: {current_date}
-User's Location: {location}
-
-User's Notes:
-{references}
-
-Online References:
-{online_results}
 
-Conversation Log:
-{chat_history}
-
-User Prompt: "{query}"
-
-Now generate an professional description of the image to generate in vivid, fine detail.
-- Use today's date, user's location, user's notes and online references to weave in any context that will improve the image generation.
-- Retain any important information and follow any instructions in the conversation log or user prompt.
+# Instructions
+- Retain important information and follow instructions by the user when composing the image description.
+- Weave in the context provided below if it will enhance the image.
+- Specify desired elements, lighting, mood, and composition in the description.
+- Decide the shape best suited to render the image. It can be one of square, portrait or landscape.
 - Add specific, fine position details. Mention painting style, camera parameters to compose the image.
-- Ensure your improved prompt is in prose format."""
+- Transform any negations in user instructions into positive alternatives.
+  Instead of saying what should NOT be in the image, describe what SHOULD be there instead.
+  Examples:
+  - "no sun" → "overcast cloudy sky"
+  - "don't include people" → "empty landscape" or "solitary scene"
+- Ensure your image description is in prose format (e.g no lists, links).
+- If any text is to be rendered in the image put it within double quotes in your image description.
 
-image_generation_improve_prompt_dalle = PromptTemplate.from_template(
-    f"""
-{image_generation_improve_prompt_base}
+# Context
 
-Improved Prompt:
-""".strip()
-)
+## User Location: {location}
 
-image_generation_improve_prompt_sd = PromptTemplate.from_template(
-    f"""
-{image_generation_improve_prompt_base}
-- If any text is to be rendered in the image put it within double quotes in your improved prompt.
+## User Documents
+{references}
+
+## Online References
+{online_results}
 
-Improved Prompt:
+Now generate a vivid description of the image and image shape to be rendered.
+Your response should be a JSON object with 'description' and 'shape' fields specified.
 """.strip()
 )