ocr/core/wrappers.py

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')


class ErrorOcrResponse(BaseModel):
    message: str


class OcrResponseWrapper(BaseModel, Generic[T]):
    data: Optional[T] = None
    successful: bool = True
    error: Optional[ErrorOcrResponse] = None

    def response(self, status_code: int):
        return JSONResponse(
            status_code=status_code,
            content={
                "data": self.data,
                "successful": self.successful,
                "error": self.error.dict() if self.error else None
            }
        )


def exception_wrapper(http_error: int, error_message: str):
    def decorator(func):
        @wraps(func)
        async def wrapper(*args, **kwargs):
            try:
                return await func(*args, **kwargs)
            except Exception as e:
                raise HTTPException(status_code=http_error, detail=error_message) from e

        return wrapper

    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)
        async def wrapper(*args, **kwargs) -> str:
            try:
                result = await func(*args, **kwargs)
                return result
            except Exception as e:
                pass

        return wrapper

    return decorator