pvanand commited on
Commit
cdb0a43
1 Parent(s): 110356d

Update document_generator.py

Browse files
Files changed (1) hide show
  1. document_generator.py +69 -23
document_generator.py CHANGED
@@ -163,12 +163,14 @@ from typing import List, Dict, Optional, Any, Callable, Union
163
  from openai import OpenAI
164
  import logging
165
  import functools
166
- from fastapi import APIRouter, HTTPException, Request
167
  from fastapi.responses import StreamingResponse
168
  from pydantic import BaseModel
169
  from fastapi_cache.decorator import cache
170
  import psycopg2
171
  from datetime import datetime
 
 
172
 
173
 
174
  logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
@@ -266,6 +268,43 @@ class AIClient:
266
  return response.choices[0].message.content
267
 
268
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
269
  class DatabaseManager:
270
  """Manages database operations."""
271
 
@@ -408,10 +447,6 @@ class MarkdownConverter:
408
 
409
  router = APIRouter()
410
 
411
- class DocumentRequest(BaseModel):
412
- query: str
413
- template: bool = False
414
-
415
  class JsonDocumentResponse(BaseModel):
416
  json_document: Dict
417
 
@@ -419,6 +454,7 @@ class MarkdownDocumentRequest(BaseModel):
419
  json_document: Dict
420
  query: str
421
  template: bool = False
 
422
 
423
  MESSAGE_DELIMITER = b"\n---DELIMITER---\n"
424
 
@@ -487,23 +523,6 @@ async def generate_document_stream(document_generator: DocumentGenerator, docume
487
 
488
  db_manager.update_database("elevatics", query, markdown_document)
489
 
490
-
491
- @cache(expire=600*24*7)
492
- @router.post("/generate-document/json", response_model=JsonDocumentResponse)
493
- async def generate_document_outline_endpoint(request: DocumentRequest):
494
- ai_client = AIClient()
495
- document_generator = DocumentGenerator(ai_client)
496
-
497
- try:
498
- json_document = document_generator.generate_document_outline(request.query, request.template)
499
-
500
- if json_document is None:
501
- raise HTTPException(status_code=500, detail="Failed to generate a valid document outline")
502
-
503
- return JsonDocumentResponse(json_document=json_document)
504
- except Exception as e:
505
- raise HTTPException(status_code=500, detail=str(e))
506
-
507
  @router.post("/generate-document/markdown-stream")
508
  async def generate_markdown_document_stream_endpoint(request: MarkdownDocumentRequest):
509
  ai_client = AIClient()
@@ -522,6 +541,33 @@ async def generate_markdown_document_stream_endpoint(request: MarkdownDocumentRe
522
  return StreamingResponse(stream_generator(), media_type="application/octet-stream")
523
 
524
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
525
  ## OBSERVABILITY
526
  from uuid import uuid4
527
  import csv
@@ -585,7 +631,7 @@ class MarkdownDocumentResponse(BaseModel):
585
  markdown_document: str
586
 
587
  @router.post("/generate-document-test", response_model=MarkdownDocumentResponse)
588
- async def test_generate_document_endpoint(request: DocumentRequest):
589
  try:
590
  # Load JSON document from file
591
  json_path = os.path.join("output/document_generator", "ai-chatbot-prd.json")
 
163
  from openai import OpenAI
164
  import logging
165
  import functools
166
+ from fastapi import APIRouter, HTTPException, Request, UploadFile, File
167
  from fastapi.responses import StreamingResponse
168
  from pydantic import BaseModel
169
  from fastapi_cache.decorator import cache
170
  import psycopg2
171
  from datetime import datetime
172
+ import base64
173
+
174
 
175
 
176
  logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
268
  return response.choices[0].message.content
269
 
270
 
271
+ class VisionTools:
272
+ def __init__(self, ai_client):
273
+ self.ai_client = ai_client
274
+
275
+ async def extract_images_info(self, images: List[UploadFile]) -> str:
276
+ try:
277
+ image_contents = []
278
+ for image in images:
279
+ image_content = await image.read()
280
+ base64_image = base64.b64encode(image_content).decode('utf-8')
281
+ image_contents.append({
282
+ "type": "image_url",
283
+ "image_url": {
284
+ "url": f"data:image/jpeg;base64,{base64_image}"
285
+ }
286
+ })
287
+
288
+ messages = [
289
+ {
290
+ "role": "user",
291
+ "content": [
292
+ {
293
+ "type": "text",
294
+ "text": "Extract the contents of these images in detail in a structured format, focusing on any text, tables, diagrams, or visual elements that might be relevant for document generation."
295
+ },
296
+ *image_contents
297
+ ]
298
+ }
299
+ ]
300
+
301
+ image_context = self.ai_client.generate_vision_response(messages)
302
+ return image_context
303
+ except Exception as e:
304
+ print(f"Error processing images: {str(e)}")
305
+ return ""
306
+
307
+
308
  class DatabaseManager:
309
  """Manages database operations."""
310
 
 
447
 
448
  router = APIRouter()
449
 
 
 
 
 
450
  class JsonDocumentResponse(BaseModel):
451
  json_document: Dict
452
 
 
454
  json_document: Dict
455
  query: str
456
  template: bool = False
457
+ images: Optional[List[UploadFile]] = File(None)
458
 
459
  MESSAGE_DELIMITER = b"\n---DELIMITER---\n"
460
 
 
523
 
524
  db_manager.update_database("elevatics", query, markdown_document)
525
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
526
  @router.post("/generate-document/markdown-stream")
527
  async def generate_markdown_document_stream_endpoint(request: MarkdownDocumentRequest):
528
  ai_client = AIClient()
 
541
  return StreamingResponse(stream_generator(), media_type="application/octet-stream")
542
 
543
 
544
+ @cache(expire=600*24*7)
545
+ @router.post("/generate-document/json", response_model=JsonDocumentResponse)
546
+ async def generate_document_outline_endpoint(request: MarkdownDocumentRequest):
547
+ ai_client = AIClient()
548
+ document_generator = DocumentGenerator(ai_client)
549
+ vision_tools = VisionTools(ai_client)
550
+
551
+ try:
552
+ image_context = ""
553
+ if request.images:
554
+ image_context = await vision_tools.extract_images_info(request.images)
555
+
556
+ json_document = document_generator.generate_document_outline(
557
+ request.query,
558
+ request.template,
559
+ image_context=image_context
560
+ )
561
+
562
+ if json_document is None:
563
+ raise HTTPException(status_code=500, detail="Failed to generate a valid document outline")
564
+
565
+ return JsonDocumentResponse(json_document=json_document)
566
+ except Exception as e:
567
+ raise HTTPException(status_code=500, detail=str(e))
568
+
569
+
570
+
571
  ## OBSERVABILITY
572
  from uuid import uuid4
573
  import csv
 
631
  markdown_document: str
632
 
633
  @router.post("/generate-document-test", response_model=MarkdownDocumentResponse)
634
+ async def test_generate_document_endpoint(request: MarkdownDocumentRequest):
635
  try:
636
  # Load JSON document from file
637
  json_path = os.path.join("output/document_generator", "ai-chatbot-prd.json")