Update document_generator.py

document_generator.py

@@ -354,11 +354,14 @@ class DocumentGenerator:
         return content.lstrip()
 
     @log_execution
-    def generate_document_outline(self, query: str, template: bool = False, max_retries: int = 3) -> Optional[Dict]:
+    def generate_document_outline(self, query: str, template: bool = False, image_context: str = "", max_retries: int = 3) -> Optional[Dict]:
         messages = [
             {"role": "system", "content": DOCUMENT_OUTLINE_PROMPT_SYSTEM if not template else DOCUMENT_TEMPLATE_OUTLINE_PROMPT_SYSTEM},
-            {"role": "user", "content": DOCUMENT_OUTLINE_PROMPT_USER.format(query=query) if not template else DOCUMENT_TEMPLATE_PROMPT_USER.format(query=query)}
+            {"role": "user", "content": DOCUMENT_OUTLINE_PROMPT_USER.format(query=query) if not template else DOCUMENT_TEMPLATE_PROMPT_USER.format(query=query, image_context=image_context)}
         ]
+        # Update user content to include image context if provided
+        if image_context:
+            messages[1]["content"] += f"<attached_images>\n\n{image_context}\n\n</attached_images>"
         
         for attempt in range(max_retries):
             outline_response = self.ai_client.generate_response(messages, model="openai/gpt-4o")