Update document_generator.py

document_generator.py

@@ -40,7 +40,7 @@ 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 os 
 import json
 import re
 import time
@@ -52,16 +52,17 @@ import functools
 from fastapi import APIRouter, HTTPException
 from pydantic import BaseModel
 from fastapi_cache.decorator import cache
+from starlette.responses import StreamingResponse
 
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
 def log_execution(func: Callable) -> Callable:
     @functools.wraps(func)
-    def wrapper(*args: Any, **kwargs: Any) -> Any:
+    async def wrapper(*args: Any, **kwargs: Any) -> Any:
         logger.info(f"Executing {func.__name__}")
         try:
-            result = func(*args, **kwargs)
+            result = await func(*args, **kwargs)
             logger.info(f"{func.__name__} completed successfully")
             return result
         except Exception as e:
@@ -77,7 +78,7 @@ class AIClient:
         )
     
     @log_execution
-    def generate_response(
+    async def generate_response(
         self,
         messages: List[Dict[str, str]],
         model: str = "openai/gpt-4o-mini",
@@ -85,12 +86,14 @@ class AIClient:
     ) -> Optional[str]:
         if not messages:
             return None
-        response = self.client.chat.completions.create(
+        loop = asyncio.get_event_loop()
+        response = await loop.run_in_executor(None, functools.partial(
+            self.client.chat.completions.create,
             model=model,
             messages=messages,
             max_tokens=max_tokens,
             stream=False
-        )
+        ))
         return response.choices[0].message.content
 
 class DocumentGenerator:
@@ -120,14 +123,14 @@ class DocumentGenerator:
         return content.lstrip()
 
     @log_execution
-    def generate_document_outline(self, query: str, max_retries: int = 3) -> Optional[Dict]:
+    async def generate_document_outline(self, query: str, max_retries: int = 3) -> Optional[Dict]:
         messages = [
             {"role": "system", "content": DOCUMENT_OUTLINE_PROMPT_SYSTEM},
             {"role": "user", "content": DOCUMENT_OUTLINE_PROMPT_USER.format(query=query)}
         ]
         
         for attempt in range(max_retries):
-            outline_response = self.ai_client.generate_response(messages, model="openai/gpt-4o")
+            outline_response = await self.ai_client.generate_response(messages, model="openai/gpt-4o")
             outline_json_text = self.extract_between_tags(outline_response, "output")
             
             try:
@@ -142,7 +145,7 @@ class DocumentGenerator:
                     return None
 
     @log_execution
-    def generate_content(self, title: str, content_instruction: str, section_number: str) -> str:
+    async def generate_content(self, title: str, content_instruction: str, section_number: str) -> str:
         self.content_messages.append({
             "role": "user",
             "content": DOCUMENT_SECTION_PROMPT_USER.format(
@@ -150,7 +153,7 @@ class DocumentGenerator:
                 content_instruction=content_instruction
             )
         })
-        section_response = self.ai_client.generate_response(self.content_messages)
+        section_response = await self.ai_client.generate_response(self.content_messages)
         content = self.extract_between_tags(section_response, "response")
         content = self.remove_duplicate_title(content, title, section_number)
         self.content_messages.append({
@@ -160,7 +163,7 @@ class DocumentGenerator:
         return content
 
     @log_execution
-    def generate_full_document(self, document_outline: Dict, query: str) -> Dict:
+    async def generate_full_document(self, document_outline: Dict, query: str):
         self.document_outline = document_outline
         
         overall_objective = query
@@ -181,16 +184,21 @@ class DocumentGenerator:
             section_number = section.get("SectionNumber", "")
             content_instruction = section.get("Content", "")
             logger.info(f"Generating content for section: {section_title}")
-            section["Content"] = self.generate_content(section_title, content_instruction, section_number)
+            section["Content"] = await self.generate_content(section_title, content_instruction, section_number)
+            yield json.dumps({"type": "document_section", "content": section}) + "\n"
 
             for subsection in section.get("Subsections", []):
                 subsection_title = subsection.get("Title", "")
                 subsection_number = subsection.get("SectionNumber", "")
                 subsection_content_instruction = subsection.get("Content", "")
                 logger.info(f"Generating content for subsection: {subsection_title}")
-                subsection["Content"] = self.generate_content(subsection_title, subsection_content_instruction, subsection_number)
+                subsection["Content"] = await self.generate_content(subsection_title, subsection_content_instruction, subsection_number)
+                yield json.dumps({"type": "document_subsection", "content": subsection}) + "\n"
 
-        return self.document_outline
+        # Generate the complete markdown document
+        full_document = self.document_outline
+        markdown_document = MarkdownConverter.convert_to_markdown(full_document["Document"])
+        yield json.dumps({"type": "complete_document", "content": markdown_document}) + "\n"
 
 class MarkdownConverter:
     @staticmethod
@@ -258,9 +266,6 @@ 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_document_outline_endpoint(request: DocumentRequest):
@@ -269,7 +274,7 @@ async def generate_document_outline_endpoint(request: DocumentRequest):
     
     try:
         # Generate the document outline
-        json_document = document_generator.generate_document_outline(request.query)
+        json_document = await 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")
@@ -278,21 +283,20 @@ async def generate_document_outline_endpoint(request: DocumentRequest):
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
-@router.post("/generate-document/markdown", response_model=MarkdownDocumentResponse)
+@router.post("/generate-document/markdown")
 async def generate_markdown_document_endpoint(request: MarkdownDocumentRequest):
     ai_client = AIClient()
     document_generator = DocumentGenerator(ai_client)
     
-    try:
-        # 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(full_document["Document"])
-        
-        return MarkdownDocumentResponse(markdown_document=markdown_document)
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+    async def event_stream():
+        try:
+            # Generate the full document content and stream it
+            async for section in document_generator.generate_full_document(request.json_document, request.query):
+                yield section
+        except Exception as e:
+            yield json.dumps({"type": "error", "message": str(e)}) + "\n"
+
+    return StreamingResponse(event_stream(), media_type="application/json")
 
 @router.post("/generate-document-test", response_model=MarkdownDocumentResponse)
 async def test_generate_document_endpoint(request: DocumentRequest):