idreamers/api/calculation/openai_request.py

from typing import Literal

from idreamers.api.calculation.prompts import CalculationPrompts
from idreamers.api.calculation.schemas import PilotRequest
from idreamers.core.wrappers import openai_wrapper


@openai_wrapper(is_json=True, return_='score')
async def calculate_scores(data: PilotRequest, type_: Literal["business", "competitive", "supplier", "critical"]):
    if type_ == "competitive":
        prompt = CalculationPrompts.competitive
    elif type_ == "supplier":
        prompt = CalculationPrompts.supplier
    elif type_ == "business":
        prompt = CalculationPrompts.business
    else:
        prompt = CalculationPrompts.critical

    if data.context:
        data.context = f'\n\n**Context**:\n```\n{data.context}\n```'

    messages = [
        {
            "role": "system",
            "content": prompt
            .replace("{supplier}", data.supplier)
            .replace("{category}", data.category)
            .replace("{buying_company}", data.buyingCompany)
            .replace("{context}", data.context)
        }
    ]
    return messages