Update document_generator.py

document_generator.py

@@ -248,6 +248,12 @@ class MarkdownDocumentRequest(BaseModel):
     json_document: Dict
     query: str
 
+MESSAGE_DELIMITER = b"\n---DELIMITER---\n"
+
+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):
     document_generator.document_outline = document_outline
     db_manager = DatabaseManager()
@@ -271,14 +277,14 @@ async def generate_document_stream(document_generator: DocumentGenerator, docume
         logging.info(f"Generating content for section: {section_title}")
         content = document_generator.generate_content(section_title, content_instruction, section_number)
         section["Content"] = content
-        yield json.dumps({
+        yield yield_message({
             "type": "document_section",
             "content": {
                 "section_number": section_number,
                 "section_title": section_title,
                 "content": content
             }
-        }, ensure_ascii=False) + "\n"
+        })
 
         for subsection in section.get("Subsections", []):
             subsection_title = subsection.get("Title", "")
@@ -287,27 +293,28 @@ async def generate_document_stream(document_generator: DocumentGenerator, docume
             logging.info(f"Generating content for subsection: {subsection_title}")
             content = document_generator.generate_content(subsection_title, subsection_content_instruction, subsection_number)
             subsection["Content"] = content
-            yield json.dumps({
+            yield yield_message({
                 "type": "document_section",
                 "content": {
                     "section_number": subsection_number,
                     "section_title": subsection_title,
                     "content": content
                 }
-            }, ensure_ascii=False) + "\n"
+            })
 
     markdown_document = MarkdownConverter.convert_to_markdown(document_generator.document_outline["Document"])
-    yield json.dumps({
+    
+    yield yield_message({
         "type": "complete_document",
-        "content": markdown_document,
-    }, ensure_ascii=False) + "\n"
+        "content": {
+            "markdown": markdown_document,
+            "json": document_generator.document_outline
+                    },
+            });
 
-    yield json.dumps({
-        "type": "complete_document",
-        "json": document_generator.document_outline
-    }, ensure_ascii=False) + "\n"
     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):
@@ -334,12 +341,12 @@ async def generate_markdown_document_stream_endpoint(request: MarkdownDocumentRe
             async for chunk in generate_document_stream(document_generator, request.json_document, request.query):
                 yield chunk
         except Exception as e:
-            yield json.dumps({
+            yield yield_message({
                 "type": "error",
                 "content": str(e)
-            }) + "\n"
+            })
 
-    return StreamingResponse(stream_generator(), media_type="text/plain")
+    return StreamingResponse(stream_generator(), media_type="application/octet-stream")
 
 ###########################################
 class MarkdownDocumentResponse(BaseModel):