wehub-resource-sync
164 lines
7.3 KiB
Python
164 lines
7.3 KiB
Python
from sqlalchemy import Column, Integer, String, and_, distinct
|
|
from superagi.lib.logger import logger
|
|
from superagi.models.base_model import DBBaseModel
|
|
from superagi.models.organisation import Organisation
|
|
from superagi.models.project import Project
|
|
from superagi.models.models import Models
|
|
from superagi.llms.openai import OpenAi
|
|
from superagi.helper.encyption_helper import encrypt_data, decrypt_data
|
|
from fastapi import HTTPException
|
|
import logging
|
|
|
|
class ModelsConfig(DBBaseModel):
|
|
"""
|
|
Represents a Model Config record in the database.
|
|
|
|
Attributes:
|
|
id (Integer): The unique identifier of the event.
|
|
provider (String): The name of the model provider.
|
|
api_key (String): The api_key for individual model providers for every Organisation
|
|
org_id (Integer): The ID of the organisation.
|
|
"""
|
|
|
|
__tablename__ = 'models_config'
|
|
|
|
id = Column(Integer, primary_key=True)
|
|
provider = Column(String, nullable=False)
|
|
api_key = Column(String, nullable=False)
|
|
org_id = Column(Integer, nullable=False)
|
|
|
|
def __repr__(self):
|
|
"""
|
|
Returns a string representation of the ModelsConfig instance.
|
|
"""
|
|
return f"ModelsConfig(id={self.id}, provider={self.provider}, " \
|
|
f"org_id={self.org_id})"
|
|
|
|
@classmethod
|
|
def fetch_value_by_agent_id(cls, session, agent_id: int, model: str):
|
|
"""
|
|
Fetches the configuration of an agent.
|
|
|
|
Args:
|
|
session: The database session object.
|
|
agent_id (int): The ID of the agent.
|
|
model (str): The model of the configuration.
|
|
|
|
Returns:
|
|
dict: Parsed configuration.
|
|
|
|
"""
|
|
from superagi.models.agent import Agent
|
|
agent = session.query(Agent).filter(Agent.id == agent_id).first()
|
|
if not agent:
|
|
raise HTTPException(status_code=404, detail="Agent not found")
|
|
|
|
project = session.query(Project).filter(Project.id == agent.project_id).first()
|
|
if not project:
|
|
raise HTTPException(status_code=404, detail="Project not found")
|
|
|
|
organisation = session.query(Organisation).filter(Organisation.id == project.organisation_id).first()
|
|
if not organisation:
|
|
raise HTTPException(status_code=404, detail="Organisation not found")
|
|
|
|
model_provider = session.query(Models).filter(Models.org_id == organisation.id, Models.model_name == model).first()
|
|
if not model_provider:
|
|
raise HTTPException(status_code=404, detail="Model provider not found")
|
|
|
|
config = session.query(ModelsConfig.provider, ModelsConfig.api_key).filter(ModelsConfig.org_id == organisation.id, ModelsConfig.id == model_provider.model_provider_id).first()
|
|
|
|
if not config:
|
|
return None
|
|
|
|
if config.provider == 'Local LLM':
|
|
return {"provider": config.provider, "api_key": config.api_key} if config else None
|
|
|
|
return {"provider": config.provider, "api_key": decrypt_data(config.api_key)} if config else None
|
|
|
|
@classmethod
|
|
def store_api_key(cls, session, organisation_id, model_provider, model_api_key):
|
|
existing_entry = session.query(ModelsConfig).filter(and_(ModelsConfig.org_id == organisation_id,
|
|
ModelsConfig.provider == model_provider)).first()
|
|
if existing_entry:
|
|
existing_entry.api_key = encrypt_data(model_api_key)
|
|
session.commit()
|
|
session.flush()
|
|
if model_provider == 'OpenAI':
|
|
cls.storeGptModels(session, organisation_id, existing_entry.id, model_api_key)
|
|
result = {'message': 'The API key was successfully updated'}
|
|
else:
|
|
new_entry = ModelsConfig(org_id=organisation_id, provider=model_provider,
|
|
api_key=encrypt_data(model_api_key))
|
|
session.add(new_entry)
|
|
session.commit()
|
|
session.flush()
|
|
if model_provider == 'OpenAI':
|
|
cls.storeGptModels(session, organisation_id, new_entry.id, model_api_key)
|
|
result = {'message': 'The API key was successfully stored', 'model_provider_id': new_entry.id}
|
|
|
|
return result
|
|
|
|
@classmethod
|
|
def storeGptModels(cls, session, organisation_id, model_provider_id, model_api_key):
|
|
default_models = {"gpt-3.5-turbo": 4032, "gpt-4": 8092, "gpt-3.5-turbo-16k": 16184}
|
|
models = OpenAi(api_key=model_api_key).get_models()
|
|
installed_models = [model[0] for model in session.query(Models.model_name).filter(Models.org_id == organisation_id).all()]
|
|
for model in models:
|
|
if model not in installed_models and model in default_models:
|
|
result = Models.store_model_details(session, organisation_id, model, model, '',
|
|
model_provider_id, default_models[model], 'Custom', '', 0)
|
|
|
|
@classmethod
|
|
def fetch_api_keys(cls, session, organisation_id):
|
|
api_key_info = session.query(ModelsConfig.provider, ModelsConfig.api_key).filter(
|
|
ModelsConfig.org_id == organisation_id).all()
|
|
|
|
if not api_key_info:
|
|
logging.error("No API key found for the provided model provider")
|
|
return []
|
|
|
|
api_keys = [{"provider": provider, "api_key": decrypt_data(api_key)} for provider, api_key in api_key_info if api_key != 'EMPTY']
|
|
|
|
return api_keys
|
|
|
|
@classmethod
|
|
def fetch_api_key(cls, session, organisation_id, model_provider):
|
|
api_key_data = session.query(ModelsConfig.id, ModelsConfig.provider, ModelsConfig.api_key).filter(
|
|
and_(ModelsConfig.org_id == organisation_id, ModelsConfig.provider == model_provider)).first()
|
|
|
|
logger.info(api_key_data)
|
|
if api_key_data is None:
|
|
return []
|
|
elif api_key_data.provider == 'Local LLM':
|
|
api_key = [{'id': api_key_data.id, 'provider': api_key_data.provider,
|
|
'api_key': api_key_data.api_key}]
|
|
return api_key
|
|
else:
|
|
api_key = [{'id': api_key_data.id, 'provider': api_key_data.provider,
|
|
'api_key': decrypt_data(api_key_data.api_key)}]
|
|
return api_key
|
|
|
|
@classmethod
|
|
def fetch_model_by_id(cls, session, organisation_id, model_provider_id):
|
|
model = session.query(ModelsConfig.provider).filter(ModelsConfig.id == model_provider_id,
|
|
ModelsConfig.org_id == organisation_id).first()
|
|
if model is None:
|
|
return {"error": "Model not found"}
|
|
else:
|
|
return {"provider": model.provider}
|
|
|
|
@classmethod
|
|
def fetch_model_by_id_marketplace(cls, session, model_provider_id):
|
|
model = session.query(ModelsConfig.provider).filter(ModelsConfig.id == model_provider_id).first()
|
|
if model is None:
|
|
return {"error": "Model not found"}
|
|
else:
|
|
return {"provider": model.provider}
|
|
|
|
@classmethod
|
|
def add_llm_config(cls, session, organisation_id):
|
|
existing_models_config = session.query(ModelsConfig).filter(ModelsConfig.org_id == organisation_id, ModelsConfig.provider == 'Local LLM').first()
|
|
if existing_models_config is None:
|
|
models_config = ModelsConfig(org_id=organisation_id, provider='Local LLM', api_key="EMPTY")
|
|
session.add(models_config)
|
|
session.commit() |