idreamers/core/wrappers.py

import json
from functools import wraps
from typing import Generic, Optional, TypeVar

import pydash
from pydantic import BaseModel
from starlette.responses import JSONResponse

from idreamers.core.config import settings

T = TypeVar('T')


class ErrorPilotResponse(BaseModel):
    message: str


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

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


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