Update document_generator.py

document_generator.py

@@ -453,12 +453,16 @@ router = APIRouter()
 class JsonDocumentResponse(BaseModel):
     json_document: Dict
 
-class MarkdownDocumentRequest(BaseModel):
-    json_document: Dict
+
+class JsonDocumentRequest(BaseModel):
     query: str
     template: bool = False
     images: Optional[List[UploadFile]] = File(None)
 
+class MarkdownDocumentRequest(BaseModel):
+    json_document: Dict
+    query: str
+
 MESSAGE_DELIMITER = b"\n---DELIMITER---\n"
 
 def yield_message(message):
@@ -547,7 +551,7 @@ async def generate_markdown_document_stream_endpoint(request: MarkdownDocumentRe
 @cache(expire=600*24*7)
 @router.post("/generate-document/json", response_model=JsonDocumentResponse)
 async def generate_document_outline_endpoint(
-    request: MarkdownDocumentRequest = Depends(),
+    request: JsonDocumentRequest = Depends(),
     images: List[UploadFile] = File(None)
 ):
     ai_client = AIClient()