fix

ocr/api/message/openai_request.py

@@ -1,28 +1,17 @@
-import asyncio
-import io
-
-from starlette.datastructures import UploadFile
-
 from ocr.api.message.prompts import OCRPrompts
-from ocr.api.message.utils import clean_assistant_response
-from ocr.core.config import settings
+from ocr.core.wrappers import openai_wrapper
 
 
-async def analyze_uploaded_document(file: UploadFile):
-    contents = await file.read()
-    openai_file = io.BytesIO(contents)
-    openai_file.name = file.filename
-    thread, openai_file = await asyncio.gather(
-        settings.OPENAI_CLIENT.beta.threads.create(),
-        settings.OPENAI_CLIENT.files.create(purpose='assistants', file=openai_file)
-    )
-    await settings.OPENAI_CLIENT.beta.threads.messages.create(
-        attachments=[{"file_id": openai_file.id, "tools": [{"type": "file_search"}]}],
-        thread_id=thread.id,
-        role="user",
-        content='Generate a report on the attached document'
-    )
-    run = await settings.OPENAI_CLIENT.beta.threads.runs.create_and_poll(
-        assistant_id=settings.ASSISTANT_ID, thread_id=thread.id, instructions=OCRPrompts.generate_general_answer
-    )
-    return await clean_assistant_response(thread.id, run.id)
+@openai_wrapper(model='gpt-4o-mini')
+async def generate_report(request_content: list[dict]):
+    messages = [
+        {
+            "role": "system",
+            "content": OCRPrompts.generate_general_answer
+        },
+        {
+            "role": "user",
+            "content": request_content
+        }
+    ]
+    return messages

ocr/api/message/utils.py

@@ -1,18 +1,40 @@
+import base64
+import io
 import re
 
-from ocr.core.config import settings
+from pdf2image import convert_from_bytes
 
 
-async def clean_assistant_response(thread_id: str, run_id: str):
-    result = ''
-    async for message in settings.OPENAI_CLIENT.beta.threads.messages.list(thread_id=thread_id, run_id=run_id):
-        message_content = message.content[0].text
-        annotations = message_content.annotations
-        for annotation in annotations:
-            message_content.value = message_content.value.replace(annotation.text, f"")
-        result = message_content.value
+def divide_images(contents: bytes) -> list[bytes]:
+    images = convert_from_bytes(contents, dpi=250)
+    image_bytes_list = []
+    for image in images:
+        img_byte_array = io.BytesIO()
+        image.save(img_byte_array, format='PNG')
+        img_byte_array.seek(0)
+        image_bytes_list.append(img_byte_array.read())
+    return image_bytes_list
+
+
+def prepare_request_content(images: list[bytes]):
+    content = [
+        {"type": "text", "text": "Generate a report on the attached document"},
+        *[
+            {
+                "type": "image_url",
+                "image_url": {
+                    "url": f"data:image/jpeg;base64,{base64.b64encode(image).decode('utf-8')}",
+                },
+            }
+            for image in images
+        ]
+    ]
+    return content
+
+
+def clean_response(text: str) -> str:
     try:
-        result = re.search(r'```markdown\s*(.*?)\s*```', result, re.DOTALL).group(1)
+        text = re.search(r'```markdown\s*(.*?)\s*```', text, re.DOTALL).group(1)
     except Exception as e:
         pass
-    return result
+    return text

ocr/api/message/views.py

@@ -1,8 +1,9 @@
 from fastapi import File, UploadFile
 
 from ocr.api.message import ocr_router
-from ocr.api.message.openai_request import analyze_uploaded_document
+from ocr.api.message.openai_request import generate_report
 from ocr.api.message.schemas import OcrResponse
+from ocr.api.message.utils import divide_images, prepare_request_content, clean_response
 from ocr.core.wrappers import OcrResponseWrapper
 
 
@@ -10,5 +11,11 @@ from ocr.core.wrappers import OcrResponseWrapper
 async def get_all_chat_messages(
         file: UploadFile = File(...)
 ) -> OcrResponseWrapper[OcrResponse]:
-    response = await analyze_uploaded_document(file)
-    return OcrResponseWrapper(data=OcrResponse(text=response))
+    try:
+        contents = await file.read()
+        images = divide_images(contents)
+        request_content = prepare_request_content(images)
+        response = await generate_report(request_content)
+        return OcrResponseWrapper(data=OcrResponse(text=clean_response(response)))
+    finally:
+        await file.close()