Update document_generator.py

document_generator.py

@@ -161,14 +161,12 @@ from typing import List, Dict, Optional, Any, Callable, Union
 from openai import OpenAI
 import logging
 import functools
-from fastapi import APIRouter, HTTPException, Request, UploadFile, File
+from fastapi import APIRouter, HTTPException, Request
 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,43 +264,6 @@ 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."""
 
@@ -352,14 +313,11 @@ 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, template: bool = False, 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) if not template else DOCUMENT_TEMPLATE_PROMPT_USER.format(query=query)}
         ]
-        # 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")
@@ -448,13 +406,12 @@ class MarkdownConverter:
 
 router = APIRouter()
 
-class JsonDocumentResponse(BaseModel):
-    json_document: Dict
-
-class JsonDocumentRequest(BaseModel):
+class DocumentRequest(BaseModel):
     query: str
     template: bool = False
-    images: Optional[List[UploadFile]] = File(None)
+
+class JsonDocumentResponse(BaseModel):
+    json_document: Dict
 
 class MarkdownDocumentRequest(BaseModel):
     json_document: Dict
@@ -529,6 +486,22 @@ 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()
@@ -547,34 +520,6 @@ 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")
-async def generate_document_outline_endpoint(request: jsonDocumentRequest):
-):
-    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
@@ -631,12 +576,14 @@ async def get_last_observations(limit: int = 10, format: str = "json"):
         raise HTTPException(status_code=500, detail=f"Failed to retrieve observations: {str(e)}")
 
 
+
+
 ###########################################
 class MarkdownDocumentResponse(BaseModel):
     markdown_document: str
 
 @router.post("/generate-document-test", response_model=MarkdownDocumentResponse)
-async def test_generate_document_endpoint(request: MarkdownDocumentRequest):
+async def test_generate_document_endpoint(request: DocumentRequest):
     try:
         # Load JSON document from file
         json_path = os.path.join("output/document_generator", "ai-chatbot-prd.json")