Improve the image diagramming prompts and response parsing

src/interface/web/app/components/excalidraw/excalidrawWrapper.tsx

@@ -58,7 +58,7 @@ export default function ExcalidrawWrapper(props: ExcalidrawWrapperProps) {
 
         for (const element of props.data) {
             if (isValidExcalidrawElement(element as ExcalidrawElementSkeleton)) {
-                basicValidSkeletons.push(element as ExcalidrawElementSkeleton);
+                basicValidSkeletons.push(element);
             }
         }
 
@@ -68,11 +68,22 @@ export default function ExcalidrawWrapper(props: ExcalidrawWrapperProps) {
                 continue;
             }
             if (element.type === "arrow") {
-                const start = basicValidSkeletons.find((child) => child.id === element.start?.id);
+                if (element.start) {
-                const end = basicValidSkeletons.find((child) => child.id === element.end?.id);
+                    const start = basicValidSkeletons.find(
-                if (start && end) {
+                        (child) => child.id === element.start?.id,
-                    validSkeletons.push(element);
+                    );
+                    if (!start) {
+                        continue;
+                    }
                 }
+                if (element.end) {
+                    const end = basicValidSkeletons.find((child) => child.id === element.end?.id);
+                    if (!end) {
+                        continue;
+                    }
+                }
+
+                validSkeletons.push(element);
             } else {
                 validSkeletons.push(element);
             }

src/khoj/processor/conversation/prompts.py

@@ -183,20 +183,23 @@ Improved Prompt:
 
 improve_diagram_description_prompt = PromptTemplate.from_template(
     """
-you are an architect working with a novice artist using a diagramming tool.
+you are an architect working with a novice digital artist using a diagramming software.
 {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
 
 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.
+- include the full, exact description. the artist does not have much experience, so be precise.
+- describe the layout.
+- you can only use straight lines.
+- use simple, concise language.
+- keep it simple and easy to understand. the artist is easily distracted.
 
 Today's Date: {current_date}
 User's Location: {location}
@@ -218,19 +221,23 @@ Query: {query}
 
 excalidraw_diagram_generation_prompt = PromptTemplate.from_template(
     """
-You are a program manager with the ability to describe diagrams to compose in professional, fine detail.
+You are a program manager with the ability to describe diagrams to compose in professional, fine detail. You LOVE getting into the details and making tedious labels, lines, and shapes look beautiful. You make everything look perfect.
 {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.
+You need to create a declarative description of the diagram and relevant components, using this base schema.
+- `label`: 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
+- **ALWAYS Required properties for ALL elements**: `type`, `x`, `y`, `id`.
+- Be very generous with spacing and composition. Use ample space between elements.
 
 {{
     type: string,
     x: number,
     y: number,
-    strokeColor: string,
-    backgroundColor: string,
     width: number,
     height: number,
+    strokeColor: string,
+    backgroundColor: string,
     id: string,
     label: {{
         text: string,
@@ -240,28 +247,30 @@ You need to create a declarative description of the diagram and relevant compone
 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.
+For arrows and lines,
+- `points`: specify the start and end points of the arrow
+- **ALWAYS Required properties for ALL elements**: `type`, `x`, `y`, `id`.
+- `start` and `end` properties: 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` and `text` 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, or ellipse elements. Even if you're using the `start` and `end` properties, you still need to specify the `x` and `y` properties for the start and end points.
 
 {{
     type: "arrow",
     id: string,
     x: number,
     y: number,
-    width: number,
-    height: number,
     strokeColor: string,
     start: {{
         id: string,
         type: string,
+        text: string,
     }},
     end: {{
         id: string,
         type: string,
+        text: string,
     }},
     label: {{
         text: string,
@@ -272,7 +281,11 @@ For arrows and lines, you can use the `points` property to specify the start and
     ]
 }}
 
-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.
+For text,
+- `text`: specify the text to be rendered
+- **ALWAYS Required properties for ALL elements**: `type`, `x`, `y`, `id`.
+- `fontSize`: optional property to specify the font size of the text
+- Use this element only for titles, subtitles, and overviews. For labels, use the `label` property in the respective elements.
 
 {{
     type: "text",
@@ -287,19 +300,25 @@ 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:
+Example Response:
-
+```json
-[
+{{
-    {{"type":"text","x":-150,"y":50,"width":300,"height":40,"id":"title_text","text":"Circular Development Process","fontSize":24}},
+    "scratchpad": "The diagram represents a circular development process with 3 stages: design, implementation and feedback. Each stage is connected to the next stage using an arrow, forming a circular process.",
-    {{"type":"ellipse","x":-169,"y":113,"width":188,"height":202,"id":"design_ellipse", "label": {{"text": "Design"}}}},
+    "elements": [
-    {{"type":"ellipse","x":62,"y":394,"width":186,"height":188,"id":"implement_ellipse", "label": {{"text": "Implement"}}}},
+    {{"type":"text","x":-150,"y":50,"id":"title_text","text":"Circular Development Process","fontSize":24}},
-    {{"type":"ellipse","x":-348,"y":430,"width":184,"height":170,"id":"feedback_ellipse", "label": {{"text": "Feedback"}}}},
+    {{"type":"ellipse","x":-169,"y":113,"id":"design_ellipse", "label": {{"text": "Design"}}}},
+    {{"type":"ellipse","x":62,"y":394,"id":"implement_ellipse", "label": {{"text": "Implement"}}}},
+    {{"type":"ellipse","x":-348,"y":430,"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:
+Think about spacing and composition. Use ample space between elements. Double the amount of space you think you need. Create a detailed diagram from the provided context and user prompt below.
+
+Return a valid JSON object, where the drawing is in `elements` and your thought process is in `scratchpad`. If you can't make the whole diagram in one response, you can split it into multiple responses. If you need to simplify for brevity, simply do so in the `scratchpad` field. DO NOT add additional info in the `elements` field.
 
 Diagram Description: {query}
 

src/khoj/routers/helpers.py

@@ -753,7 +753,11 @@ async def generate_excalidraw_diagram(
         yield None, None
         return
 
-    yield better_diagram_description_prompt, excalidraw_diagram_description
+    scratchpad = excalidraw_diagram_description.get("scratchpad")
+
+    inferred_queries = f"Instruction: {better_diagram_description_prompt}\n\nScratchpad: {scratchpad}"
+
+    yield inferred_queries, excalidraw_diagram_description.get("elements")
 
 
 async def generate_better_diagram_description(
@@ -838,10 +842,18 @@ async def generate_excalidraw_diagram_from_description(
         )
         raw_response = clean_json(raw_response)
         try:
+            # Expect response to have `elements` and `scratchpad` keys
             response: Dict[str, str] = json.loads(raw_response)
+            if (
+                not response
+                or not isinstance(response, Dict)
+                or not response.get("elements")
+                or not response.get("scratchpad")
+            ):
+                raise AssertionError(f"Invalid response for generating Excalidraw diagram: {response}")
         except Exception:
             raise AssertionError(f"Invalid response for generating Excalidraw diagram: {raw_response}")
-        if not response or not isinstance(response, List) or not isinstance(response[0], Dict):
+        if not response or not isinstance(response["elements"], List) or not isinstance(response["elements"][0], Dict):
             # TODO Some additional validation here that it's a valid Excalidraw diagram
             raise AssertionError(f"Invalid response for improving diagram description: {response}")