Update document_generator_v3.py

document_generator_v3.py

@@ -354,10 +354,10 @@ class DocumentGenerator:
         return content.lstrip()
 
     @log_execution
-    def generate_document_outline(self, query: str, template: bool = False, image_context: str = "", max_retries: int = 3) -> Optional[Dict]:
+    def generate_document_outline(self, query: str, num_pages:int, 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, image_context=image_context)}
+            {"role": "user", "content": DOCUMENT_OUTLINE_PROMPT_USER.format(query=query, num_pages=num_pages) if not template else DOCUMENT_TEMPLATE_PROMPT_USER.format(query=query, num_pages=num_pages)}
         ]
         # Update user content to include image context if provided
         if image_context:
@@ -561,6 +561,7 @@ async def generate_document_outline_endpoint(
     vision_tools = VisionTools(ai_client)
     images = ""
     documents = ""
+    num_pages = 2
     try:
         image_context = ""
         if images:
@@ -572,6 +573,7 @@ async def generate_document_outline_endpoint(
         
         json_document = document_generator.generate_document_outline(
             query,
+            num_pages,
             template,
             image_context=image_context
         )