Update document_generator.py

document_generator.py

@@ -40,7 +40,6 @@ FORMAT YOUR OUTPUT AS MARKDOWN ENCLOSED IN <response></response> tags
 DOCUMENT_SECTION_PROMPT_USER = """<prompt>Output the content for the section "{section_or_subsection_title}" formatted as markdown. Follow this instruction: {content_instruction}</prompt>"""
 
 # File: app.py
-
 import os
 import json
 import re
@@ -53,7 +52,6 @@ import functools
 from fastapi import APIRouter, HTTPException
 from pydantic import BaseModel
 from fastapi_cache.decorator import cache
-#from prompts import *
 
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
@@ -77,6 +75,7 @@ class AIClient:
             base_url="https://openrouter.ai/api/v1",
             api_key="sk-or-v1-"+os.environ['OPENROUTER_API_KEY']
         )
+    
     @log_execution
     def generate_response(
         self,
@@ -161,12 +160,9 @@ class DocumentGenerator:
         return content
 
     @log_execution
-    def generate_document(self, query: str) -> Dict:
-        self.generate_document_outline(query)
+    def generate_full_document(self, document_outline: Dict, query: str) -> Dict:
+        self.document_outline = document_outline
         
-        if self.document_outline is None:
-            raise ValueError("Failed to generate a valid document outline")
-
         overall_objective = query
         document_layout = json.dumps(self.document_outline, indent=2)
 
@@ -196,7 +192,6 @@ class DocumentGenerator:
 
         return self.document_outline
 
-
 class MarkdownConverter:
     @staticmethod
     def slugify(text: str) -> str:
@@ -250,51 +245,56 @@ class MarkdownConverter:
         
         markdown += "</div>"
         return markdown
-        
 
 router = APIRouter()
 
-class DocumentRequest(BaseModel): 
+class DocumentRequest(BaseModel):
     query: str
 
-class DocumentResponse(BaseModel): 
-    json_document: Dict 
-    markdown_document: str 
-
 class JsonDocumentResponse(BaseModel):
     json_document: Dict
 
+class MarkdownDocumentRequest(BaseModel):
+    json_document: Dict
+    query: str
+
 class MarkdownDocumentResponse(BaseModel):
     markdown_document: str
 
 @cache(expire=600*24*7)
 @router.post("/generate-document/json", response_model=JsonDocumentResponse)
-async def generate_json_document_endpoint(request: DocumentRequest):
+async def generate_document_outline_endpoint(request: DocumentRequest):
     ai_client = AIClient()
     document_generator = DocumentGenerator(ai_client)
+    
     try:
-        # Generate the JSON document
-        json_document = document_generator.generate_document(request.query)
+        # Generate the document outline
+        json_document = document_generator.generate_document_outline(request.query)
+        
+        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))
 
-@cache(expire=600*24*7)
 @router.post("/generate-document/markdown", response_model=MarkdownDocumentResponse)
-async def generate_markdown_document_endpoint(request: DocumentRequest):
+async def generate_markdown_document_endpoint(request: MarkdownDocumentRequest):
     ai_client = AIClient()
     document_generator = DocumentGenerator(ai_client)
+    
     try:
-        # Generate the JSON document
-        json_document = document_generator.generate_document(request.query)
+        # Generate the full document content
+        full_document = document_generator.generate_full_document(request.json_document, request.query)
+        
         # Convert to Markdown
-        markdown_document = MarkdownConverter.convert_to_markdown(json_document["Document"])
+        markdown_document = MarkdownConverter.convert_to_markdown(full_document["Document"])
+        
         return MarkdownDocumentResponse(markdown_document=markdown_document)
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
-
-@router.post("/generate-document-test", response_model=DocumentResponse)
+@router.post("/generate-document-test", response_model=MarkdownDocumentResponse)
 async def test_generate_document_endpoint(request: DocumentRequest):
     try:
         # Load JSON document from file
@@ -307,10 +307,7 @@ async def test_generate_document_endpoint(request: DocumentRequest):
         with open(md_path, "r") as md_file:
             markdown_document = md_file.read()
 
-        return DocumentResponse(
-            json_document=json_document,
-            markdown_document=markdown_document
-        )
+        return MarkdownDocumentResponse(markdown_document=markdown_document)
     except FileNotFoundError:
         raise HTTPException(status_code=404, detail="Test files not found")
     except json.JSONDecodeError: