Update document_generator.py

document_generator.py

@@ -39,6 +39,49 @@ 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>"""
 
+##########################################
+
+DOCUMENT_TEMPLATE_OUTLINE_PROMPT_SYSTEM = """You are a document template generator. Provide the outline of the document requested in <prompt></prompt> in JSON format.
+Include sections and subsections if required. Use the "Content" field to provide a specific prompt or instruction for generating template (with empty input fields) that particular section or subsection. Specify that input fields should be empty in each prompt.
+make sure the Sections follow a logical flow and each prompt's content does not overlap with other sections.
+OUTPUT IN FOLLOWING JSON FORMAT enclosed in <output> tags
+<output>
+{
+    "Document": {
+      "Title": "Document Title",
+      "Author": "Author Name",
+      "Date": "YYYY-MM-DD",
+      "Version": "1.0",
+
+      "Sections": [
+        {
+          "SectionNumber": "1",
+          "Title": "Section Title",
+          "Content": "Specific prompt or instruction for generating empty template for this section",
+          "Subsections": [
+            {
+              "SectionNumber": "1.1",
+              "Title": "Subsection Title",
+              "Content": "Specific prompt or instruction for generating empty template for this subsection"
+            }
+          ]
+        }
+      ]
+    }
+  }
+</output>"""
+
+DOCUMENT_TEMPLATE_PROMPT_USER = """<prompt>{query}</prompt>"""
+
+DOCUMENT_TEMPLATE_SECTION_PROMPT_SYSTEM = """You are a document template generator, You need to output only the content requested in the section in the prompt.
+FORMAT YOUR OUTPUT AS MARKDOWN ENCLOSED IN <response></response> tags
+<overall_objective>{overall_objective}</overall_objective>
+<document_layout>{document_layout}</document_layout>"""
+
+DOCUMENT_TEMPLATE_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
@@ -146,10 +189,10 @@ class DocumentGenerator:
         return content.lstrip()
 
     @log_execution
-    def generate_document_outline(self, query: 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},
-            {"role": "user", "content": DOCUMENT_OUTLINE_PROMPT_USER.format(query=query)}
+            {"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)}
         ]
         
         for attempt in range(max_retries):
@@ -168,10 +211,11 @@ class DocumentGenerator:
                     return None
 
     @log_execution
-    def generate_content(self, title: str, content_instruction: str, section_number: str) -> str:
+    def generate_content(self, title: str, content_instruction: str, section_number: str, template: bool = False) -> str:
+        SECTION_PROMPT_USER = DOCUMENT_SECTION_PROMPT_USER if not template else DOCUMENT_TEMPLATE_SECTION_PROMPT_USER
         self.content_messages.append({
             "role": "user",
-            "content": DOCUMENT_SECTION_PROMPT_USER.format(
+            "content": SECTION_PROMPT_USER.format(
                 section_or_subsection_title=title,
                 content_instruction=content_instruction
             )
@@ -240,6 +284,7 @@ router = APIRouter()
 
 class DocumentRequest(BaseModel):
     query: str
+    template: bool = False
 
 class JsonDocumentResponse(BaseModel):
     json_document: Dict
@@ -247,6 +292,7 @@ class JsonDocumentResponse(BaseModel):
 class MarkdownDocumentRequest(BaseModel):
     json_document: Dict
     query: str
+    template: bool = False
 
 MESSAGE_DELIMITER = b"\n---DELIMITER---\n"
 
@@ -254,16 +300,17 @@ 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):
+async def generate_document_stream(document_generator: DocumentGenerator, document_outline: Dict, query: str, template: bool = False):
     document_generator.document_outline = document_outline
     db_manager = DatabaseManager()
     overall_objective = query
     document_layout = json.dumps(document_generator.document_outline, indent=2)
 
+    SECTION_PROMPT_SYSTEM = DOCUMENT_SECTION_PROMPT_SYSTEM if not template else DOCUMENT_TEMPLATE_SECTION_PROMPT_SYSTEM
     document_generator.content_messages = [
         {
             "role": "system",
-            "content": DOCUMENT_SECTION_PROMPT_SYSTEM.format(
+            "content": SECTION_PROMPT_SYSTEM.format(
                 overall_objective=overall_objective,
                 document_layout=document_layout
             )
@@ -275,7 +322,7 @@ async def generate_document_stream(document_generator: DocumentGenerator, docume
         section_number = section.get("SectionNumber", "")
         content_instruction = section.get("Content", "")
         logging.info(f"Generating content for section: {section_title}")
-        content = document_generator.generate_content(section_title, content_instruction, section_number)
+        content = document_generator.generate_content(section_title, content_instruction, section_number, template)
         section["Content"] = content
         yield yield_message({
             "type": "document_section",
@@ -291,7 +338,7 @@ async def generate_document_stream(document_generator: DocumentGenerator, docume
             subsection_number = subsection.get("SectionNumber", "")
             subsection_content_instruction = subsection.get("Content", "")
             logging.info(f"Generating content for subsection: {subsection_title}")
-            content = document_generator.generate_content(subsection_title, subsection_content_instruction, subsection_number)
+            content = document_generator.generate_content(subsection_title, subsection_content_instruction, subsection_number, template)
             subsection["Content"] = content
             yield yield_message({
                 "type": "document_section",
@@ -322,7 +369,7 @@ async def generate_document_outline_endpoint(request: DocumentRequest):
     document_generator = DocumentGenerator(ai_client)
     
     try:
-        json_document = document_generator.generate_document_outline(request.query)
+        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")
@@ -338,7 +385,7 @@ async def generate_markdown_document_stream_endpoint(request: MarkdownDocumentRe
     
     async def stream_generator():
         try:
-            async for chunk in generate_document_stream(document_generator, request.json_document, request.query):
+            async for chunk in generate_document_stream(document_generator, request.json_document, request.query, request.template):
                 yield chunk
         except Exception as e:
             yield yield_message({