Revert "remove ai"

This reverts commit 19897cd83faa46db9cf6c4047d3af2f5cedcd2f7.

ocr/api/message/openai_request.py

@@ -0,0 +1,17 @@
+from ocr.api.message.prompts import OCRPrompts
+from ocr.core.wrappers import openai_wrapper
+
+
+@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/prompts.py

@@ -0,0 +1,44 @@
+class OCRPrompts:
+    generate_general_answer = """## Task
+
+You must analyze the attached medical document and generate a comprehensive report in **Markdown2** format. Ensure that every detail provided in the document is included, and do not omit or modify any information. Your output must strictly follow the required format.
+
+## Report Structure
+
+The report should be structured as follows, with each section containing only relevant information from the document:
+
+```markdown
+## Patient Information
+
+- Name: [Patient Name]
+- Age: [Patient Age]
+- Date of Scan: [Date]
+- Indication: [Reason for the CT scan]
+
+## Findings
+
+**Primary findings**:
+[Describe significant abnormalities or findings relevant to the indication]
+
+** Secondary findings**:
+[List incidental findings, e.g., "Mild hepatic steatosis noted."]
+**No abnormalities**:
+[Mention organs or systems without abnormalities, e.g., "No evidence of lymphadenopathy or pleural effusion."]
+
+## Impression
+
+[Summarize the findings concisely, e.g., "Findings suggest a primary lung tumor. Biopsy recommended for further evaluation."]
+
+## Recommendations
+
+[Include next steps or further tests, e.g., "PET scan and consultation with oncology recommended."]
+```
+
+[INST]
+
+## Instructions
+
+- **Do not invent or infer any information.** Only use data provided in the document.
+- Ensure that the format is followed strictly, and the output is complete without any deviations.
+
+[/INST]"""

ocr/api/message/views.py

@@ -1,6 +1,7 @@
 from fastapi import File, UploadFile
 
 from ocr.api.message import ocr_router
+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
@@ -12,6 +13,9 @@ async def get_all_chat_messages(
 ) -> OcrResponseWrapper[OcrResponse]:
     try:
         contents = await file.read()
-        return OcrResponseWrapper(data=OcrResponse(text=clean_response("## Coming soon")))
+        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()

ocr/core/wrappers.py

@@ -1,10 +1,14 @@
+import json
 from functools import wraps
 from typing import Generic, Optional, TypeVar
 
+import pydash
 from fastapi import HTTPException
 from pydantic import BaseModel
 from starlette.responses import JSONResponse
 
+from ocr.core.config import settings
+
 T = TypeVar('T')
 
 
@@ -42,6 +46,32 @@ def exception_wrapper(http_error: int, error_message: str):
     return decorator
 
 
+def openai_wrapper(
+        temperature: int | float = 0, model: str = "gpt-4o-mini", is_json: bool = False, return_: str = None
+):
+    def decorator(func):
+        @wraps(func)
+        async def wrapper(*args, **kwargs) -> str:
+            messages = await func(*args, **kwargs)
+            completion = await settings.OPENAI_CLIENT.chat.completions.create(
+                messages=messages,
+                temperature=temperature,
+                n=1,
+                model=model,
+                response_format={"type": "json_object"} if is_json else {"type": "text"}
+            )
+            response = completion.choices[0].message.content
+            if is_json:
+                response = json.loads(response)
+                if return_:
+                    return pydash.get(response, return_)
+            return response
+
+        return wrapper
+
+    return decorator
+
+
 def background_task():
     def decorator(func):
         @wraps(func)