wehub-resource-sync
238 lines
9.7 KiB
Python
238 lines
9.7 KiB
Python
import yaml
|
|
from sqlalchemy import Column, Integer, String, and_
|
|
from sqlalchemy.sql import func
|
|
from typing import List, Dict, Union
|
|
from superagi.models.base_model import DBBaseModel
|
|
from superagi.controllers.types.models_types import ModelsTypes
|
|
from superagi.helper.encyption_helper import decrypt_data
|
|
import requests, logging
|
|
from superagi.lib.logger import logger
|
|
|
|
marketplace_url = "https://app.superagi.com/api"
|
|
# marketplace_url = "http://localhost:8001"
|
|
|
|
|
|
class Models(DBBaseModel):
|
|
"""
|
|
Represents a Model record in the database
|
|
|
|
Attributes:
|
|
id (Integer): The unique identifier of the event.
|
|
model_name (String): The name of the model.
|
|
description (String): The description for the model.
|
|
end_point (String): The end_point for the model.3001
|
|
model_provider_id (Integer): The unique id of the model_provider from the models_config table.
|
|
token_limit (Integer): The maximum number of tokens for a model.
|
|
type (Strng): The place it is added from.
|
|
version (String): The version of the replicate model.
|
|
org_id (Integer): The ID of the organisation.
|
|
model_features (String): The Features of the Model.
|
|
"""
|
|
|
|
__tablename__ = 'models'
|
|
|
|
id = Column(Integer, primary_key=True)
|
|
model_name = Column(String, nullable=False)
|
|
description = Column(String, nullable=True)
|
|
end_point = Column(String, nullable=False)
|
|
model_provider_id = Column(Integer, nullable=False)
|
|
token_limit = Column(Integer, nullable=False)
|
|
type = Column(String, nullable=False)
|
|
version = Column(String, nullable=False)
|
|
org_id = Column(Integer, nullable=False)
|
|
model_features = Column(String, nullable=False)
|
|
context_length = Column(Integer, nullable=True)
|
|
|
|
def __repr__(self):
|
|
"""
|
|
Returns a string representation of the Models instance.
|
|
"""
|
|
return f"Models(id={self.id}, model_name={self.model_name}, " \
|
|
f"end_point={self.end_point}, model_provider_id={self.model_provider_id}, " \
|
|
f"token_limit={self.token_limit}, " \
|
|
f"type={self.type}, " \
|
|
f"version={self.version}, " \
|
|
f"org_id={self.org_id}, " \
|
|
f"model_features={self.model_features})"
|
|
|
|
@classmethod
|
|
def fetch_marketplace_list(cls, page):
|
|
headers = {'Content-Type': 'application/json'}
|
|
response = requests.get(
|
|
marketplace_url + f"/models_controller/marketplace/list/{str(page)}",
|
|
headers=headers, timeout=10)
|
|
if response.status_code == 200:
|
|
return response.json()
|
|
else:
|
|
return []
|
|
|
|
@classmethod
|
|
def get_model_install_details(cls, session, marketplace_models, organisation_id, type=ModelsTypes.CUSTOM.value):
|
|
from superagi.models.models_config import ModelsConfig
|
|
installed_models = session.query(Models).filter(Models.org_id == organisation_id).all()
|
|
model_counts_dict = dict(
|
|
session.query(Models.model_name, func.count(Models.org_id)).group_by(Models.model_name).all()
|
|
)
|
|
installed_models_dict = {model.model_name: True for model in installed_models}
|
|
|
|
for model in marketplace_models:
|
|
try:
|
|
if type == ModelsTypes.MARKETPLACE.value:
|
|
model["is_installed"] = False
|
|
else:
|
|
model["is_installed"] = installed_models_dict.get(model["model_name"], False)
|
|
model["installs"] = model_counts_dict.get(model["model_name"], 0)
|
|
except TypeError as e:
|
|
logging.error("Error Occurred: %s", e)
|
|
|
|
return marketplace_models
|
|
|
|
@classmethod
|
|
def fetch_model_tokens(cls, session, organisation_id) -> Dict[str, int]:
|
|
try:
|
|
models = session.query(
|
|
Models.model_name, Models.token_limit
|
|
).filter(
|
|
Models.org_id == organisation_id
|
|
).all()
|
|
|
|
if models:
|
|
return dict(models)
|
|
else:
|
|
return {"error": "No models found for the given organisation ID."}
|
|
|
|
except Exception as e:
|
|
logging.error(f"Unexpected Error Occured: {e}")
|
|
return {"error": "Unexpected Error Occured"}
|
|
|
|
@classmethod
|
|
def store_model_details(cls, session, organisation_id, model_name, description, end_point, model_provider_id, token_limit, type, version, context_length):
|
|
from superagi.models.models_config import ModelsConfig
|
|
if not model_name:
|
|
return {"error": "Model Name is empty or undefined"}
|
|
if not description:
|
|
return {"error": "Description is empty or undefined"}
|
|
if not model_provider_id:
|
|
return {"error": "Model Provider Id is null or undefined or 0"}
|
|
if not token_limit:
|
|
return {"error": "Token Limit is null or undefined or 0"}
|
|
|
|
# Check if model_name already exists in the database
|
|
existing_model = session.query(Models).filter(Models.model_name == model_name, Models.org_id == organisation_id).first()
|
|
if existing_model:
|
|
return {"error": "Model Name already exists"}
|
|
|
|
# Get the provider of the model
|
|
if type == 'Marketplace':
|
|
model = ModelsConfig.fetch_model_by_id_marketplace(session, model_provider_id)
|
|
else:
|
|
model = ModelsConfig.fetch_model_by_id(session, organisation_id, model_provider_id)
|
|
|
|
if "error" in model:
|
|
return model # Return error message if model not found
|
|
|
|
# Check the 'provider' from ModelsConfig table
|
|
if not end_point and model["provider"] not in ['OpenAI', 'Google Palm', 'Replicate','Local LLM']:
|
|
return {"error": "End Point is empty or undefined"}
|
|
|
|
if context_length is None:
|
|
context_length = 0
|
|
|
|
try:
|
|
model = Models(
|
|
model_name=model_name,
|
|
description=description,
|
|
end_point=end_point,
|
|
token_limit=token_limit,
|
|
model_provider_id=model_provider_id,
|
|
type=type,
|
|
version=version,
|
|
org_id=organisation_id,
|
|
model_features='',
|
|
context_length=context_length
|
|
)
|
|
session.add(model)
|
|
session.commit()
|
|
session.flush()
|
|
|
|
except Exception as e:
|
|
logging.error(f"Unexpected Error Occured: {e}")
|
|
return {"error": "Unexpected Error Occured"}
|
|
|
|
return {"success": "Model Details stored successfully", "model_id": model.id}
|
|
|
|
@classmethod
|
|
def api_key_from_configurations(cls, session, organisation_id):
|
|
try:
|
|
from superagi.models.models_config import ModelsConfig
|
|
from superagi.models.configuration import Configuration
|
|
|
|
model_provider = session.query(ModelsConfig).filter(ModelsConfig.provider == "OpenAI",
|
|
ModelsConfig.org_id == organisation_id).first()
|
|
if model_provider is None:
|
|
configurations = session.query(Configuration).filter(Configuration.key == 'model_api_key',
|
|
Configuration.organisation_id == organisation_id).first()
|
|
|
|
if configurations is None:
|
|
return {"error": "API Key is Missing"}
|
|
else:
|
|
model_api_key = decrypt_data(configurations.value)
|
|
model_details = ModelsConfig.store_api_key(session, organisation_id, "OpenAI", model_api_key)
|
|
except Exception as e:
|
|
logging.error(f"Exception has been raised while checking API Key:: {e}")
|
|
|
|
|
|
@classmethod
|
|
def fetch_models(cls, session, organisation_id) -> Union[Dict[str, str], List[Dict[str, Union[str, int]]]]:
|
|
try:
|
|
from superagi.models.models_config import ModelsConfig
|
|
cls.api_key_from_configurations(session, organisation_id)
|
|
|
|
models = session.query(Models.id, Models.model_name, Models.description, ModelsConfig.provider).join(
|
|
ModelsConfig, Models.model_provider_id == ModelsConfig.id).filter(
|
|
Models.org_id == organisation_id).all()
|
|
|
|
result = []
|
|
for model in models:
|
|
result.append({
|
|
"id": model[0],
|
|
"name": model[1],
|
|
"description": model[2],
|
|
"model_provider": model[3]
|
|
})
|
|
|
|
except Exception as e:
|
|
logging.error(f"Unexpected Error Occurred: {e}")
|
|
return {"error": "Unexpected Error Occurred"}
|
|
|
|
return result
|
|
|
|
@classmethod
|
|
def fetch_model_details(cls, session, organisation_id, model_id: int) -> Dict[str, Union[str, int]]:
|
|
try:
|
|
from superagi.models.models_config import ModelsConfig
|
|
model = session.query(
|
|
Models.id, Models.model_name, Models.description, Models.end_point, Models.token_limit, Models.type,
|
|
ModelsConfig.provider,
|
|
).join(
|
|
ModelsConfig, Models.model_provider_id == ModelsConfig.id
|
|
).filter(
|
|
and_(Models.org_id == organisation_id, Models.id == model_id)
|
|
).first()
|
|
|
|
if model:
|
|
return {
|
|
"id": model[0],
|
|
"name": model[1],
|
|
"description": model[2],
|
|
"end_point": model[3],
|
|
"token_limit": model[4],
|
|
"type": model[5],
|
|
"model_provider": model[6]
|
|
}
|
|
else:
|
|
return {"error": "Model with the given ID doesn't exist."}
|
|
|
|
except Exception as e:
|
|
logging.error(f"Unexpected Error Occured: {e}")
|
|
return {"error": "Unexpected Error Occured"} |