Update document_generator.py

document_generator.py

@@ -163,12 +163,14 @@ from typing import List, Dict, Optional, Any, Callable, Union
 from openai import OpenAI
 import logging
 import functools
-from fastapi import APIRouter, HTTPException, Request
+from fastapi import APIRouter, HTTPException, Request, UploadFile, File
 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
 from fastapi_cache.decorator import cache
 import psycopg2
 from datetime import datetime
+import base64
+
 
 
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
@@ -266,6 +268,43 @@ class AIClient:
         return response.choices[0].message.content
 
 
+class VisionTools:
+    def __init__(self, ai_client):
+        self.ai_client = ai_client
+
+    async def extract_images_info(self, images: List[UploadFile]) -> str:
+        try:
+            image_contents = []
+            for image in images:
+                image_content = await image.read()
+                base64_image = base64.b64encode(image_content).decode('utf-8')
+                image_contents.append({
+                    "type": "image_url",
+                    "image_url": {
+                        "url": f"data:image/jpeg;base64,{base64_image}"
+                    }
+                })
+            
+            messages = [
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "text",
+                            "text": "Extract the contents of these images in detail in a structured format, focusing on any text, tables, diagrams, or visual elements that might be relevant for document generation."
+                        },
+                        *image_contents
+                    ]
+                }
+            ]
+            
+            image_context = self.ai_client.generate_vision_response(messages)
+            return image_context
+        except Exception as e:
+            print(f"Error processing images: {str(e)}")
+            return ""
+
+
 class DatabaseManager:
     """Manages database operations."""
 
@@ -408,10 +447,6 @@ class MarkdownConverter:
 
 router = APIRouter()
 
-class DocumentRequest(BaseModel):
-    query: str
-    template: bool = False
-
 class JsonDocumentResponse(BaseModel):
     json_document: Dict
 
@@ -419,6 +454,7 @@ class MarkdownDocumentRequest(BaseModel):
     json_document: Dict
     query: str
     template: bool = False
+    images: Optional[List[UploadFile]] = File(None)
 
 MESSAGE_DELIMITER = b"\n---DELIMITER---\n"
 
@@ -487,23 +523,6 @@ async def generate_document_stream(document_generator: DocumentGenerator, docume
 
     db_manager.update_database("elevatics", query, markdown_document)
 
-
-@cache(expire=600*24*7)
-@router.post("/generate-document/json", response_model=JsonDocumentResponse)
-async def generate_document_outline_endpoint(request: DocumentRequest):
-    ai_client = AIClient()
-    document_generator = DocumentGenerator(ai_client)
-    
-    try:
-        json_document = document_generator.generate_document_outline(request.query, request.template)
-        
-        if json_document is None:
-            raise HTTPException(status_code=500, detail="Failed to generate a valid document outline")
-        
-        return JsonDocumentResponse(json_document=json_document)
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
-
 @router.post("/generate-document/markdown-stream")
 async def generate_markdown_document_stream_endpoint(request: MarkdownDocumentRequest):
     ai_client = AIClient()
@@ -522,6 +541,33 @@ async def generate_markdown_document_stream_endpoint(request: MarkdownDocumentRe
     return StreamingResponse(stream_generator(), media_type="application/octet-stream")
 
 
+@cache(expire=600*24*7)
+@router.post("/generate-document/json", response_model=JsonDocumentResponse)
+async def generate_document_outline_endpoint(request: MarkdownDocumentRequest):
+    ai_client = AIClient()
+    document_generator = DocumentGenerator(ai_client)
+    vision_tools = VisionTools(ai_client)
+    
+    try:
+        image_context = ""
+        if request.images:
+            image_context = await vision_tools.extract_images_info(request.images)
+        
+        json_document = document_generator.generate_document_outline(
+            request.query,
+            request.template,
+            image_context=image_context
+        )
+        
+        if json_document is None:
+            raise HTTPException(status_code=500, detail="Failed to generate a valid document outline")
+        
+        return JsonDocumentResponse(json_document=json_document)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+
 ## OBSERVABILITY
 from uuid import uuid4
 import csv
@@ -585,7 +631,7 @@ class MarkdownDocumentResponse(BaseModel):
     markdown_document: str
 
 @router.post("/generate-document-test", response_model=MarkdownDocumentResponse)
-async def test_generate_document_endpoint(request: DocumentRequest):
+async def test_generate_document_endpoint(request: MarkdownDocumentRequest):
     try:
         # Load JSON document from file
         json_path = os.path.join("output/document_generator", "ai-chatbot-prd.json")