Update document_generator_v2.py

document_generator_v2.py

@@ -464,7 +464,7 @@ class MarkdownDocumentRequest(BaseModel):
     json_document: Dict
     query: str
     template: bool = False
-    conversation_id: str = None
+    conversation_id: str = ""
 
 MESSAGE_DELIMITER = b"\n---DELIMITER---\n"
 
@@ -472,12 +472,14 @@ def yield_message(message):
     message_json = json.dumps(message, ensure_ascii=False).encode('utf-8')
     return message_json + MESSAGE_DELIMITER
 
-async def generate_document_stream(document_generator: DocumentGenerator, document_outline: Dict, query: str, template: bool = False):
+async def generate_document_stream(document_generator: DocumentGenerator, document_outline: Dict, query: str, template: bool = False, conversation_id: str = ""):
     document_generator.document_outline = document_outline
     db_manager = DatabaseManager()
     overall_objective = query
     document_layout = json.dumps(document_generator.document_outline, indent=2)
-
+    cache_key = f"image_context_{conversation_id}"
+    image_context = await FastAPICache.get_backend().get(cache_key)
+    
     SECTION_PROMPT_SYSTEM = DOCUMENT_SECTION_PROMPT_SYSTEM if not template else DOCUMENT_TEMPLATE_SECTION_PROMPT_SYSTEM
     document_generator.content_messages = [
         {
@@ -488,6 +490,8 @@ async def generate_document_stream(document_generator: DocumentGenerator, docume
             )
         }
     ]
+    if image_context:
+        document_generator.content_messages[0]["content"] += f"<attached_images>\n\n{image_context}\n\n</attached_images>"
 
     for section in document_generator.document_outline["Document"].get("Sections", []):
         section_title = section.get("Title", "")
@@ -540,7 +544,7 @@ async def generate_markdown_document_stream_endpoint(request: MarkdownDocumentRe
     
     async def stream_generator():
         try:
-            async for chunk in generate_document_stream(document_generator, request.json_document, request.query, request.template):
+            async for chunk in generate_document_stream(document_generator, request.json_document, request.query, request.template, request.conversation_id):
                 yield chunk
         except Exception as e:
             yield yield_message({
@@ -568,6 +572,10 @@ async def generate_document_outline_endpoint(
         image_context = ""
         if images:
             image_context = await vision_tools.extract_images_info(images)
+            
+            # Store the image_context in the cache
+            cache_key = f"image_context_{conversation_id}"
+            await FastAPICache.get_backend().set(cache_key, image_context, expire=3600)  # Cache for 1 hour
         
         json_document = document_generator.generate_document_outline(
             query,
@@ -583,7 +591,6 @@ async def generate_document_outline_endpoint(
         raise HTTPException(status_code=500, detail=str(e))
 
 
-
 ## OBSERVABILITY
 from uuid import uuid4
 import csv