pvanand commited on
Commit
6ef6156
1 Parent(s): bf1f3eb

Update document_generator.py

Browse files
Files changed (1) hide show
  1. document_generator.py +25 -28
document_generator.py CHANGED
@@ -40,7 +40,6 @@ FORMAT YOUR OUTPUT AS MARKDOWN ENCLOSED IN <response></response> tags
40
  DOCUMENT_SECTION_PROMPT_USER = """<prompt>Output the content for the section "{section_or_subsection_title}" formatted as markdown. Follow this instruction: {content_instruction}</prompt>"""
41
 
42
  # File: app.py
43
-
44
  import os
45
  import json
46
  import re
@@ -53,7 +52,6 @@ import functools
53
  from fastapi import APIRouter, HTTPException
54
  from pydantic import BaseModel
55
  from fastapi_cache.decorator import cache
56
- #from prompts import *
57
 
58
  logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
59
  logger = logging.getLogger(__name__)
@@ -77,6 +75,7 @@ class AIClient:
77
  base_url="https://openrouter.ai/api/v1",
78
  api_key="sk-or-v1-"+os.environ['OPENROUTER_API_KEY']
79
  )
 
80
  @log_execution
81
  def generate_response(
82
  self,
@@ -161,12 +160,9 @@ class DocumentGenerator:
161
  return content
162
 
163
  @log_execution
164
- def generate_document(self, query: str) -> Dict:
165
- self.generate_document_outline(query)
166
 
167
- if self.document_outline is None:
168
- raise ValueError("Failed to generate a valid document outline")
169
-
170
  overall_objective = query
171
  document_layout = json.dumps(self.document_outline, indent=2)
172
 
@@ -196,7 +192,6 @@ class DocumentGenerator:
196
 
197
  return self.document_outline
198
 
199
-
200
  class MarkdownConverter:
201
  @staticmethod
202
  def slugify(text: str) -> str:
@@ -250,51 +245,56 @@ class MarkdownConverter:
250
 
251
  markdown += "</div>"
252
  return markdown
253
-
254
 
255
  router = APIRouter()
256
 
257
- class DocumentRequest(BaseModel):
258
  query: str
259
 
260
- class DocumentResponse(BaseModel):
261
- json_document: Dict
262
- markdown_document: str
263
-
264
  class JsonDocumentResponse(BaseModel):
265
  json_document: Dict
266
 
 
 
 
 
267
  class MarkdownDocumentResponse(BaseModel):
268
  markdown_document: str
269
 
270
  @cache(expire=600*24*7)
271
  @router.post("/generate-document/json", response_model=JsonDocumentResponse)
272
- async def generate_json_document_endpoint(request: DocumentRequest):
273
  ai_client = AIClient()
274
  document_generator = DocumentGenerator(ai_client)
 
275
  try:
276
- # Generate the JSON document
277
- json_document = document_generator.generate_document(request.query)
 
 
 
 
278
  return JsonDocumentResponse(json_document=json_document)
279
  except Exception as e:
280
  raise HTTPException(status_code=500, detail=str(e))
281
 
282
- @cache(expire=600*24*7)
283
  @router.post("/generate-document/markdown", response_model=MarkdownDocumentResponse)
284
- async def generate_markdown_document_endpoint(request: DocumentRequest):
285
  ai_client = AIClient()
286
  document_generator = DocumentGenerator(ai_client)
 
287
  try:
288
- # Generate the JSON document
289
- json_document = document_generator.generate_document(request.query)
 
290
  # Convert to Markdown
291
- markdown_document = MarkdownConverter.convert_to_markdown(json_document["Document"])
 
292
  return MarkdownDocumentResponse(markdown_document=markdown_document)
293
  except Exception as e:
294
  raise HTTPException(status_code=500, detail=str(e))
295
 
296
-
297
- @router.post("/generate-document-test", response_model=DocumentResponse)
298
  async def test_generate_document_endpoint(request: DocumentRequest):
299
  try:
300
  # Load JSON document from file
@@ -307,10 +307,7 @@ async def test_generate_document_endpoint(request: DocumentRequest):
307
  with open(md_path, "r") as md_file:
308
  markdown_document = md_file.read()
309
 
310
- return DocumentResponse(
311
- json_document=json_document,
312
- markdown_document=markdown_document
313
- )
314
  except FileNotFoundError:
315
  raise HTTPException(status_code=404, detail="Test files not found")
316
  except json.JSONDecodeError:
 
40
  DOCUMENT_SECTION_PROMPT_USER = """<prompt>Output the content for the section "{section_or_subsection_title}" formatted as markdown. Follow this instruction: {content_instruction}</prompt>"""
41
 
42
  # File: app.py
 
43
  import os
44
  import json
45
  import re
 
52
  from fastapi import APIRouter, HTTPException
53
  from pydantic import BaseModel
54
  from fastapi_cache.decorator import cache
 
55
 
56
  logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
57
  logger = logging.getLogger(__name__)
 
75
  base_url="https://openrouter.ai/api/v1",
76
  api_key="sk-or-v1-"+os.environ['OPENROUTER_API_KEY']
77
  )
78
+
79
  @log_execution
80
  def generate_response(
81
  self,
 
160
  return content
161
 
162
  @log_execution
163
+ def generate_full_document(self, document_outline: Dict, query: str) -> Dict:
164
+ self.document_outline = document_outline
165
 
 
 
 
166
  overall_objective = query
167
  document_layout = json.dumps(self.document_outline, indent=2)
168
 
 
192
 
193
  return self.document_outline
194
 
 
195
  class MarkdownConverter:
196
  @staticmethod
197
  def slugify(text: str) -> str:
 
245
 
246
  markdown += "</div>"
247
  return markdown
 
248
 
249
  router = APIRouter()
250
 
251
+ class DocumentRequest(BaseModel):
252
  query: str
253
 
 
 
 
 
254
  class JsonDocumentResponse(BaseModel):
255
  json_document: Dict
256
 
257
+ class MarkdownDocumentRequest(BaseModel):
258
+ json_document: Dict
259
+ query: str
260
+
261
  class MarkdownDocumentResponse(BaseModel):
262
  markdown_document: str
263
 
264
  @cache(expire=600*24*7)
265
  @router.post("/generate-document/json", response_model=JsonDocumentResponse)
266
+ async def generate_document_outline_endpoint(request: DocumentRequest):
267
  ai_client = AIClient()
268
  document_generator = DocumentGenerator(ai_client)
269
+
270
  try:
271
+ # Generate the document outline
272
+ json_document = document_generator.generate_document_outline(request.query)
273
+
274
+ if json_document is None:
275
+ raise HTTPException(status_code=500, detail="Failed to generate a valid document outline")
276
+
277
  return JsonDocumentResponse(json_document=json_document)
278
  except Exception as e:
279
  raise HTTPException(status_code=500, detail=str(e))
280
 
 
281
  @router.post("/generate-document/markdown", response_model=MarkdownDocumentResponse)
282
+ async def generate_markdown_document_endpoint(request: MarkdownDocumentRequest):
283
  ai_client = AIClient()
284
  document_generator = DocumentGenerator(ai_client)
285
+
286
  try:
287
+ # Generate the full document content
288
+ full_document = document_generator.generate_full_document(request.json_document, request.query)
289
+
290
  # Convert to Markdown
291
+ markdown_document = MarkdownConverter.convert_to_markdown(full_document["Document"])
292
+
293
  return MarkdownDocumentResponse(markdown_document=markdown_document)
294
  except Exception as e:
295
  raise HTTPException(status_code=500, detail=str(e))
296
 
297
+ @router.post("/generate-document-test", response_model=MarkdownDocumentResponse)
 
298
  async def test_generate_document_endpoint(request: DocumentRequest):
299
  try:
300
  # Load JSON document from file
 
307
  with open(md_path, "r") as md_file:
308
  markdown_document = md_file.read()
309
 
310
+ return MarkdownDocumentResponse(markdown_document=markdown_document)
 
 
 
311
  except FileNotFoundError:
312
  raise HTTPException(status_code=404, detail="Test files not found")
313
  except json.JSONDecodeError: