PyGetGPT/model/util.py

import importlib
from configUtil import ConfigUtil


class InputData:
    def __init__(self, message):
        self.message = message

class OutputData:
    def __init__(self, message,code,tokenUsed):
        self.message = message
        self.code = code
        self.tokenUsed = tokenUsed
        


def getModels() -> list:
    model_config = ConfigUtil("data/models.ini")
    out = []
    for model in model_config.getSectionList():
        if model_config.getBool(model, "enabled") == False:
            continue
        out.append(model)
    return out


def getModelsInfo() -> list:
    model_config = ConfigUtil("data/models.ini")
    out = []
    for model in model_config.getSectionList():
        if model_config.getBool(model, "enabled") == False:
            continue
        util = {
            "name":model_config.get(model, "name"),
            "id":model,
        }
        out.append(util)
    return out


def getModelAPIKey(model: str) -> str:
    model_config = ConfigUtil("data/models.ini")
    try:
        return model_config.get(model, "key")
    except:
        return "Model API key not found"

def requestModel(model: str, input_data: InputData) -> OutputData:
    if model not in getModels(): 
        ret = OutputData("Model not found",404,0)
    else:
        module = importlib.import_module(f"model.{model}")
        ret = module.predict(input_data)
    resq = {}
    for key in ret.__dict__: resq[key] = ret.__dict__[key]
    return resq