jyapayne/dify
Watch
1
0
Fork
You've already forked dify
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

chore(api/services): apply ruff reformatting (#7599)

Browse source 
Co-authored-by: -LAN- <laipz8200@outlook.com>
This commit is contained in:
Bowen Liang 2024-08-26 13:43:57 +08:00 • committed by GitHub
commit 17fd773a30
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
49 changed files with 2640 additions and 2665 deletions
Download patch file Download diff file Expand all files Collapse all files

1
api/pyproject.toml
View file

@ -74,7 +74,6 @@ exclude = [
"controllers/**/*.py",
"models/**/*.py",
"migrations/**/*",
"services/**/*.py",
]
[tool.pytest_env]

2
api/services/__init__.py
View file

@ -1,3 +1,3 @@
from . import errors
__all__ = ['errors']
__all__ = ["errors"]

288
api/services/account_service.py
View file

@ -39,12 +39,7 @@ from tasks.mail_reset_password_task import send_reset_password_mail_task
class AccountService:
reset_password_rate_limiter = RateLimiter(
prefix="reset_password_rate_limit",
max_attempts=5,
time_window=60 * 60
)
reset_password_rate_limiter = RateLimiter(prefix="reset_password_rate_limit", max_attempts=5, time_window=60 * 60)
@staticmethod
def load_user(user_id: str) -> None | Account:
@ -55,12 +50,15 @@ class AccountService:
if account.status in [AccountStatus.BANNED.value, AccountStatus.CLOSED.value]:
raise Unauthorized("Account is banned or closed.")
current_tenant: TenantAccountJoin = TenantAccountJoin.query.filter_by(account_id=account.id, current=True).first()
current_tenant: TenantAccountJoin = TenantAccountJoin.query.filter_by(
account_id=account.id, current=True
).first()
if current_tenant:
account.current_tenant_id = current_tenant.tenant_id
else:
available_ta = TenantAccountJoin.query.filter_by(account_id=account.id) \
.order_by(TenantAccountJoin.id.asc()).first()
available_ta = (
TenantAccountJoin.query.filter_by(account_id=account.id).order_by(TenantAccountJoin.id.asc()).first()
)
if not available_ta:
return None
@ -74,14 +72,13 @@ class AccountService:
return account
@staticmethod
def get_account_jwt_token(account, *, exp: timedelta = timedelta(days=30)):
payload = {
"user_id": account.id,
"exp": datetime.now(timezone.utc).replace(tzinfo=None) + exp,
"iss": dify_config.EDITION,
"sub": 'Console API Passport',
"sub": "Console API Passport",
}
token = PassportService().issue(payload)
@ -93,10 +90,10 @@ class AccountService:
account = Account.query.filter_by(email=email).first()
if not account:
raise AccountLoginError('Invalid email or password.')
raise AccountLoginError("Invalid email or password.")
if account.status == AccountStatus.BANNED.value or account.status == AccountStatus.CLOSED.value:
raise AccountLoginError('Account is banned or closed.')
raise AccountLoginError("Account is banned or closed.")
if account.status == AccountStatus.PENDING.value:
account.status = AccountStatus.ACTIVE.value
@ -104,7 +101,7 @@ class AccountService:
db.session.commit()
if account.password is None or not compare_password(password, account.password, account.password_salt):
raise AccountLoginError('Invalid email or password.')
raise AccountLoginError("Invalid email or password.")
return account
@staticmethod
@ -129,11 +126,9 @@ class AccountService:
return account
@staticmethod
def create_account(email: str,
name: str,
interface_language: str,
password: Optional[str] = None,
interface_theme: str = 'light') -> Account:
def create_account(
email: str, name: str, interface_language: str, password: Optional[str] = None, interface_theme: str = "light"
) -> Account:
"""create account"""
account = Account()
account.email = email
@ -155,7 +150,7 @@ class AccountService:
account.interface_theme = interface_theme
# Set timezone based on language
account.timezone = language_timezone_mapping.get(interface_language, 'UTC')
account.timezone = language_timezone_mapping.get(interface_language, "UTC")
db.session.add(account)
db.session.commit()
@ -166,8 +161,9 @@ class AccountService:
"""Link account integrate"""
try:
# Query whether there is an existing binding record for the same provider
account_integrate: Optional[AccountIntegrate] = AccountIntegrate.query.filter_by(account_id=account.id,
provider=provider).first()
account_integrate: Optional[AccountIntegrate] = AccountIntegrate.query.filter_by(
account_id=account.id, provider=provider
).first()
if account_integrate:
# If it exists, update the record
@ -176,15 +172,16 @@ class AccountService:
account_integrate.updated_at = datetime.now(timezone.utc).replace(tzinfo=None)
else:
# If it does not exist, create a new record
account_integrate = AccountIntegrate(account_id=account.id, provider=provider, open_id=open_id,
encrypted_token="")
account_integrate = AccountIntegrate(
account_id=account.id, provider=provider, open_id=open_id, encrypted_token=""
)
db.session.add(account_integrate)
db.session.commit()
logging.info(f'Account {account.id} linked {provider} account {open_id}.')
logging.info(f"Account {account.id} linked {provider} account {open_id}.")
except Exception as e:
logging.exception(f'Failed to link {provider} account {open_id} to Account {account.id}')
raise LinkAccountIntegrateError('Failed to link account.') from e
logging.exception(f"Failed to link {provider} account {open_id} to Account {account.id}")
raise LinkAccountIntegrateError("Failed to link account.") from e
@staticmethod
def close_account(account: Account) -> None:
@ -218,7 +215,7 @@ class AccountService:
AccountService.update_last_login(account, ip_address=ip_address)
exp = timedelta(days=30)
token = AccountService.get_account_jwt_token(account, exp=exp)
redis_client.set(_get_login_cache_key(account_id=account.id, token=token), '1', ex=int(exp.total_seconds()))
redis_client.set(_get_login_cache_key(account_id=account.id, token=token), "1", ex=int(exp.total_seconds()))
return token
@staticmethod
@ -236,22 +233,18 @@ class AccountService:
if cls.reset_password_rate_limiter.is_rate_limited(account.email):
raise RateLimitExceededError(f"Rate limit exceeded for email: {account.email}. Please try again later.")
token = TokenManager.generate_token(account, 'reset_password')
send_reset_password_mail_task.delay(
language=account.interface_language,
to=account.email,
token=token
)
token = TokenManager.generate_token(account, "reset_password")
send_reset_password_mail_task.delay(language=account.interface_language, to=account.email, token=token)
cls.reset_password_rate_limiter.increment_rate_limit(account.email)
return token
@classmethod
def revoke_reset_password_token(cls, token: str):
TokenManager.revoke_token(token, 'reset_password')
TokenManager.revoke_token(token, "reset_password")
@classmethod
def get_reset_password_data(cls, token: str) -> Optional[dict[str, Any]]:
return TokenManager.get_token_data(token, 'reset_password')
return TokenManager.get_token_data(token, "reset_password")
def _get_login_cache_key(*, account_id: str, token: str):
@ -259,7 +252,6 @@ def _get_login_cache_key(*, account_id: str, token: str):
class TenantService:
@staticmethod
def create_tenant(name: str) -> Tenant:
"""Create tenant"""
@ -275,31 +267,28 @@ class TenantService:
@staticmethod
def create_owner_tenant_if_not_exist(account: Account):
"""Create owner tenant if not exist"""
available_ta = TenantAccountJoin.query.filter_by(account_id=account.id) \
.order_by(TenantAccountJoin.id.asc()).first()
available_ta = (
TenantAccountJoin.query.filter_by(account_id=account.id).order_by(TenantAccountJoin.id.asc()).first()
)
if available_ta:
return
tenant = TenantService.create_tenant(f"{account.name}'s Workspace")
TenantService.create_tenant_member(tenant, account, role='owner')
TenantService.create_tenant_member(tenant, account, role="owner")
account.current_tenant = tenant
db.session.commit()
tenant_was_created.send(tenant)
@staticmethod
def create_tenant_member(tenant: Tenant, account: Account, role: str = 'normal') -> TenantAccountJoin:
def create_tenant_member(tenant: Tenant, account: Account, role: str = "normal") -> TenantAccountJoin:
"""Create tenant member"""
if role == TenantAccountJoinRole.OWNER.value:
if TenantService.has_roles(tenant, [TenantAccountJoinRole.OWNER]):
logging.error(f'Tenant {tenant.id} has already an owner.')
raise Exception('Tenant already has an owner.')
logging.error(f"Tenant {tenant.id} has already an owner.")
raise Exception("Tenant already has an owner.")
ta = TenantAccountJoin(
tenant_id=tenant.id,
account_id=account.id,
role=role
)
ta = TenantAccountJoin(tenant_id=tenant.id, account_id=account.id, role=role)
db.session.add(ta)
db.session.commit()
return ta
@ -307,9 +296,12 @@ class TenantService:
@staticmethod
def get_join_tenants(account: Account) -> list[Tenant]:
"""Get account join tenants"""
return db.session.query(Tenant).join(
TenantAccountJoin, Tenant.id == TenantAccountJoin.tenant_id
).filter(TenantAccountJoin.account_id == account.id, Tenant.status == TenantStatus.NORMAL).all()
return (
db.session.query(Tenant)
.join(TenantAccountJoin, Tenant.id == TenantAccountJoin.tenant_id)
.filter(TenantAccountJoin.account_id == account.id, Tenant.status == TenantStatus.NORMAL)
.all()
)
@staticmethod
def get_current_tenant_by_account(account: Account):
@ -333,16 +325,23 @@ class TenantService:
if tenant_id is None:
raise ValueError("Tenant ID must be provided.")
tenant_account_join = db.session.query(TenantAccountJoin).join(Tenant, TenantAccountJoin.tenant_id == Tenant.id).filter(
TenantAccountJoin.account_id == account.id,
TenantAccountJoin.tenant_id == tenant_id,
Tenant.status == TenantStatus.NORMAL,
).first()
tenant_account_join = (
db.session.query(TenantAccountJoin)
.join(Tenant, TenantAccountJoin.tenant_id == Tenant.id)
.filter(
TenantAccountJoin.account_id == account.id,
TenantAccountJoin.tenant_id == tenant_id,
Tenant.status == TenantStatus.NORMAL,
)
.first()
)
if not tenant_account_join:
raise AccountNotLinkTenantError("Tenant not found or account is not a member of the tenant.")
else:
TenantAccountJoin.query.filter(TenantAccountJoin.account_id == account.id, TenantAccountJoin.tenant_id != tenant_id).update({'current': False})
TenantAccountJoin.query.filter(
TenantAccountJoin.account_id == account.id, TenantAccountJoin.tenant_id != tenant_id
).update({"current": False})
tenant_account_join.current = True
# Set the current tenant for the account
account.current_tenant_id = tenant_account_join.tenant_id
@ -354,9 +353,7 @@ class TenantService:
query = (
db.session.query(Account, TenantAccountJoin.role)
.select_from(Account)
.join(
TenantAccountJoin, Account.id == TenantAccountJoin.account_id
)
.join(TenantAccountJoin, Account.id == TenantAccountJoin.account_id)
.filter(TenantAccountJoin.tenant_id == tenant.id)
)
@ -375,11 +372,9 @@ class TenantService:
query = (
db.session.query(Account, TenantAccountJoin.role)
.select_from(Account)
.join(
TenantAccountJoin, Account.id == TenantAccountJoin.account_id
)
.join(TenantAccountJoin, Account.id == TenantAccountJoin.account_id)
.filter(TenantAccountJoin.tenant_id == tenant.id)
.filter(TenantAccountJoin.role == 'dataset_operator')
.filter(TenantAccountJoin.role == "dataset_operator")
)
# Initialize an empty list to store the updated accounts
@ -395,20 +390,25 @@ class TenantService:
def has_roles(tenant: Tenant, roles: list[TenantAccountJoinRole]) -> bool:
"""Check if user has any of the given roles for a tenant"""
if not all(isinstance(role, TenantAccountJoinRole) for role in roles):
raise ValueError('all roles must be TenantAccountJoinRole')
raise ValueError("all roles must be TenantAccountJoinRole")
return db.session.query(TenantAccountJoin).filter(
TenantAccountJoin.tenant_id == tenant.id,
TenantAccountJoin.role.in_([role.value for role in roles])
).first() is not None
return (
db.session.query(TenantAccountJoin)
.filter(
TenantAccountJoin.tenant_id == tenant.id, TenantAccountJoin.role.in_([role.value for role in roles])
)
.first()
is not None
)
@staticmethod
def get_user_role(account: Account, tenant: Tenant) -> Optional[TenantAccountJoinRole]:
"""Get the role of the current account for a given tenant"""
join = db.session.query(TenantAccountJoin).filter(
TenantAccountJoin.tenant_id == tenant.id,
TenantAccountJoin.account_id == account.id
).first()
join = (
db.session.query(TenantAccountJoin)
.filter(TenantAccountJoin.tenant_id == tenant.id, TenantAccountJoin.account_id == account.id)
.first()
)
return join.role if join else None
@staticmethod
@ -420,29 +420,26 @@ class TenantService:
def check_member_permission(tenant: Tenant, operator: Account, member: Account, action: str) -> None:
"""Check member permission"""
perms = {
'add': [TenantAccountRole.OWNER, TenantAccountRole.ADMIN],
'remove': [TenantAccountRole.OWNER],
'update': [TenantAccountRole.OWNER]
"add": [TenantAccountRole.OWNER, TenantAccountRole.ADMIN],
"remove": [TenantAccountRole.OWNER],
"update": [TenantAccountRole.OWNER],
}
if action not in ['add', 'remove', 'update']:
if action not in ["add", "remove", "update"]:
raise InvalidActionError("Invalid action.")
if member:
if operator.id == member.id:
raise CannotOperateSelfError("Cannot operate self.")
ta_operator = TenantAccountJoin.query.filter_by(
tenant_id=tenant.id,
account_id=operator.id
).first()
ta_operator = TenantAccountJoin.query.filter_by(tenant_id=tenant.id, account_id=operator.id).first()
if not ta_operator or ta_operator.role not in perms[action]:
raise NoPermissionError(f'No permission to {action} member.')
raise NoPermissionError(f"No permission to {action} member.")
@staticmethod
def remove_member_from_tenant(tenant: Tenant, account: Account, operator: Account) -> None:
"""Remove member from tenant"""
if operator.id == account.id and TenantService.check_member_permission(tenant, operator, account, 'remove'):
if operator.id == account.id and TenantService.check_member_permission(tenant, operator, account, "remove"):
raise CannotOperateSelfError("Cannot operate self.")
ta = TenantAccountJoin.query.filter_by(tenant_id=tenant.id, account_id=account.id).first()
@ -455,23 +452,17 @@ class TenantService:
@staticmethod
def update_member_role(tenant: Tenant, member: Account, new_role: str, operator: Account) -> None:
"""Update member role"""
TenantService.check_member_permission(tenant, operator, member, 'update')
TenantService.check_member_permission(tenant, operator, member, "update")
target_member_join = TenantAccountJoin.query.filter_by(
tenant_id=tenant.id,
account_id=member.id
).first()
target_member_join = TenantAccountJoin.query.filter_by(tenant_id=tenant.id, account_id=member.id).first()
if target_member_join.role == new_role:
raise RoleAlreadyAssignedError("The provided role is already assigned to the member.")
if new_role == 'owner':
if new_role == "owner":
# Find the current owner and change their role to 'admin'
current_owner_join = TenantAccountJoin.query.filter_by(
tenant_id=tenant.id,
role='owner'
).first()
current_owner_join.role = 'admin'
current_owner_join = TenantAccountJoin.query.filter_by(tenant_id=tenant.id, role="owner").first()
current_owner_join.role = "admin"
# Update the role of the target member
target_member_join.role = new_role
@ -480,8 +471,8 @@ class TenantService:
@staticmethod
def dissolve_tenant(tenant: Tenant, operator: Account) -> None:
"""Dissolve tenant"""
if not TenantService.check_member_permission(tenant, operator, operator, 'remove'):
raise NoPermissionError('No permission to dissolve tenant.')
if not TenantService.check_member_permission(tenant, operator, operator, "remove"):
raise NoPermissionError("No permission to dissolve tenant.")
db.session.query(TenantAccountJoin).filter_by(tenant_id=tenant.id).delete()
db.session.delete(tenant)
db.session.commit()
@ -494,10 +485,9 @@ class TenantService:
class RegisterService:
@classmethod
def _get_invitation_token_key(cls, token: str) -> str:
return f'member_invite:token:{token}'
return f"member_invite:token:{token}"
@classmethod
def setup(cls, email: str, name: str, password: str, ip_address: str) -> None:
@ -523,9 +513,7 @@ class RegisterService:
TenantService.create_owner_tenant_if_not_exist(account)
dify_setup = DifySetup(
version=dify_config.CURRENT_VERSION
)
dify_setup = DifySetup(version=dify_config.CURRENT_VERSION)
db.session.add(dify_setup)
db.session.commit()
except Exception as e:
@ -535,34 +523,35 @@ class RegisterService:
db.session.query(Tenant).delete()
db.session.commit()
logging.exception(f'Setup failed: {e}')
raise ValueError(f'Setup failed: {e}')
logging.exception(f"Setup failed: {e}")
raise ValueError(f"Setup failed: {e}")
@classmethod
def register(cls, email, name,
password: Optional[str] = None,
open_id: Optional[str] = None,
provider: Optional[str] = None,
language: Optional[str] = None,
status: Optional[AccountStatus] = None) -> Account:
def register(
cls,
email,
name,
password: Optional[str] = None,
open_id: Optional[str] = None,
provider: Optional[str] = None,
language: Optional[str] = None,
status: Optional[AccountStatus] = None,
) -> Account:
db.session.begin_nested()
"""Register account"""
try:
account = AccountService.create_account(
email=email,
name=name,
interface_language=language if language else languages[0],
password=password
email=email, name=name, interface_language=language if language else languages[0], password=password
)
account.status = AccountStatus.ACTIVE.value if not status else status.value
account.initialized_at = datetime.now(timezone.utc).replace(tzinfo=None)
if open_id is not None or provider is not None:
AccountService.link_account_integrate(provider, open_id, account)
if dify_config.EDITION != 'SELF_HOSTED':
if dify_config.EDITION != "SELF_HOSTED":
tenant = TenantService.create_tenant(f"{account.name}'s Workspace")
TenantService.create_tenant_member(tenant, account, role='owner')
TenantService.create_tenant_member(tenant, account, role="owner")
account.current_tenant = tenant
tenant_was_created.send(tenant)
@ -570,30 +559,29 @@ class RegisterService:
db.session.commit()
except Exception as e:
db.session.rollback()
logging.error(f'Register failed: {e}')
raise AccountRegisterError(f'Registration failed: {e}') from e
logging.error(f"Register failed: {e}")
raise AccountRegisterError(f"Registration failed: {e}") from e
return account
@classmethod
def invite_new_member(cls, tenant: Tenant, email: str, language: str, role: str = 'normal', inviter: Account = None) -> str:
def invite_new_member(
cls, tenant: Tenant, email: str, language: str, role: str = "normal", inviter: Account = None
) -> str:
"""Invite new member"""
account = Account.query.filter_by(email=email).first()
if not account:
TenantService.check_member_permission(tenant, inviter, None, 'add')
name = email.split('@')[0]
TenantService.check_member_permission(tenant, inviter, None, "add")
name = email.split("@")[0]
account = cls.register(email=email, name=name, language=language, status=AccountStatus.PENDING)
# Create new tenant member for invited tenant
TenantService.create_tenant_member(tenant, account, role)
TenantService.switch_tenant(account, tenant.id)
else:
TenantService.check_member_permission(tenant, inviter, account, 'add')
ta = TenantAccountJoin.query.filter_by(
tenant_id=tenant.id,
account_id=account.id
).first()
TenantService.check_member_permission(tenant, inviter, account, "add")
ta = TenantAccountJoin.query.filter_by(tenant_id=tenant.id, account_id=account.id).first()
if not ta:
TenantService.create_tenant_member(tenant, account, role)
@ -609,7 +597,7 @@ class RegisterService:
language=account.interface_language,
to=email,
token=token,
inviter_name=inviter.name if inviter else 'Dify',
inviter_name=inviter.name if inviter else "Dify",
workspace_name=tenant.name,
)
@ -619,23 +607,19 @@ class RegisterService:
def generate_invite_token(cls, tenant: Tenant, account: Account) -> str:
token = str(uuid.uuid4())
invitation_data = {
'account_id': account.id,
'email': account.email,
'workspace_id': tenant.id,
"account_id": account.id,
"email": account.email,
"workspace_id": tenant.id,
}
expiryHours = dify_config.INVITE_EXPIRY_HOURS
redis_client.setex(
cls._get_invitation_token_key(token),
expiryHours * 60 * 60,
json.dumps(invitation_data)
)
redis_client.setex(cls._get_invitation_token_key(token), expiryHours * 60 * 60, json.dumps(invitation_data))
return token
@classmethod
def revoke_token(cls, workspace_id: str, email: str, token: str):
if workspace_id and email:
email_hash = sha256(email.encode()).hexdigest()
cache_key = 'member_invite_token:{}, {}:{}'.format(workspace_id, email_hash, token)
cache_key = "member_invite_token:{}, {}:{}".format(workspace_id, email_hash, token)
redis_client.delete(cache_key)
else:
redis_client.delete(cls._get_invitation_token_key(token))
@ -646,17 +630,21 @@ class RegisterService:
if not invitation_data:
return None
tenant = db.session.query(Tenant).filter(
Tenant.id == invitation_data['workspace_id'],
Tenant.status == 'normal'
).first()
tenant = (
db.session.query(Tenant)
.filter(Tenant.id == invitation_data["workspace_id"], Tenant.status == "normal")
.first()
)
if not tenant:
return None
tenant_account = db.session.query(Account, TenantAccountJoin.role).join(
TenantAccountJoin, Account.id == TenantAccountJoin.account_id
).filter(Account.email == invitation_data['email'], TenantAccountJoin.tenant_id == tenant.id).first()
tenant_account = (
db.session.query(Account, TenantAccountJoin.role)
.join(TenantAccountJoin, Account.id == TenantAccountJoin.account_id)
.filter(Account.email == invitation_data["email"], TenantAccountJoin.tenant_id == tenant.id)
.first()
)
if not tenant_account:
return None
@ -665,29 +653,29 @@ class RegisterService:
if not account:
return None
if invitation_data['account_id'] != str(account.id):
if invitation_data["account_id"] != str(account.id):
return None
return {
'account': account,
'data': invitation_data,
'tenant': tenant,
"account": account,
"data": invitation_data,
"tenant": tenant,
}
@classmethod
def _get_invitation_by_token(cls, token: str, workspace_id: str, email: str) -> Optional[dict[str, str]]:
if workspace_id is not None and email is not None:
email_hash = sha256(email.encode()).hexdigest()
cache_key = f'member_invite_token:{workspace_id}, {email_hash}:{token}'
cache_key = f"member_invite_token:{workspace_id}, {email_hash}:{token}"
account_id = redis_client.get(cache_key)
if not account_id:
return None
return {
'account_id': account_id.decode('utf-8'),
'email': email,
'workspace_id': workspace_id,
"account_id": account_id.decode("utf-8"),
"email": email,
"workspace_id": workspace_id,
}
else:
data = redis_client.get(cls._get_invitation_token_key(token))

64
api/services/advanced_prompt_template_service.py
View file

@ -1,4 +1,3 @@
import copy
from core.prompt.prompt_templates.advanced_prompt_templates import (
@ -17,59 +16,78 @@ from models.model import AppMode
class AdvancedPromptTemplateService:
@classmethod
def get_prompt(cls, args: dict) -> dict:
app_mode = args['app_mode']
model_mode = args['model_mode']
model_name = args['model_name']
has_context = args['has_context']
app_mode = args["app_mode"]
model_mode = args["model_mode"]
model_name = args["model_name"]
has_context = args["has_context"]
if 'baichuan' in model_name.lower():
if "baichuan" in model_name.lower():
return cls.get_baichuan_prompt(app_mode, model_mode, has_context)
else:
return cls.get_common_prompt(app_mode, model_mode, has_context)
@classmethod
def get_common_prompt(cls, app_mode: str, model_mode:str, has_context: str) -> dict:
def get_common_prompt(cls, app_mode: str, model_mode: str, has_context: str) -> dict:
context_prompt = copy.deepcopy(CONTEXT)
if app_mode == AppMode.CHAT.value:
if model_mode == "completion":
return cls.get_completion_prompt(copy.deepcopy(CHAT_APP_COMPLETION_PROMPT_CONFIG), has_context, context_prompt)
return cls.get_completion_prompt(
copy.deepcopy(CHAT_APP_COMPLETION_PROMPT_CONFIG), has_context, context_prompt
)
elif model_mode == "chat":
return cls.get_chat_prompt(copy.deepcopy(CHAT_APP_CHAT_PROMPT_CONFIG), has_context, context_prompt)
elif app_mode == AppMode.COMPLETION.value:
if model_mode == "completion":
return cls.get_completion_prompt(copy.deepcopy(COMPLETION_APP_COMPLETION_PROMPT_CONFIG), has_context, context_prompt)
return cls.get_completion_prompt(
copy.deepcopy(COMPLETION_APP_COMPLETION_PROMPT_CONFIG), has_context, context_prompt
)
elif model_mode == "chat":
return cls.get_chat_prompt(copy.deepcopy(COMPLETION_APP_CHAT_PROMPT_CONFIG), has_context, context_prompt)
return cls.get_chat_prompt(
copy.deepcopy(COMPLETION_APP_CHAT_PROMPT_CONFIG), has_context, context_prompt
)
@classmethod
def get_completion_prompt(cls, prompt_template: dict, has_context: str, context: str) -> dict:
if has_context == 'true':
prompt_template['completion_prompt_config']['prompt']['text'] = context + prompt_template['completion_prompt_config']['prompt']['text']
if has_context == "true":
prompt_template["completion_prompt_config"]["prompt"]["text"] = (
context + prompt_template["completion_prompt_config"]["prompt"]["text"]
)
return prompt_template
@classmethod
def get_chat_prompt(cls, prompt_template: dict, has_context: str, context: str) -> dict:
if has_context == 'true':
prompt_template['chat_prompt_config']['prompt'][0]['text'] = context + prompt_template['chat_prompt_config']['prompt'][0]['text']
if has_context == "true":
prompt_template["chat_prompt_config"]["prompt"][0]["text"] = (
context + prompt_template["chat_prompt_config"]["prompt"][0]["text"]
)
return prompt_template
@classmethod
def get_baichuan_prompt(cls, app_mode: str, model_mode:str, has_context: str) -> dict:
def get_baichuan_prompt(cls, app_mode: str, model_mode: str, has_context: str) -> dict:
baichuan_context_prompt = copy.deepcopy(BAICHUAN_CONTEXT)
if app_mode == AppMode.CHAT.value:
if model_mode == "completion":
return cls.get_completion_prompt(copy.deepcopy(BAICHUAN_CHAT_APP_COMPLETION_PROMPT_CONFIG), has_context, baichuan_context_prompt)
return cls.get_completion_prompt(
copy.deepcopy(BAICHUAN_CHAT_APP_COMPLETION_PROMPT_CONFIG), has_context, baichuan_context_prompt
)
elif model_mode == "chat":
return cls.get_chat_prompt(copy.deepcopy(BAICHUAN_CHAT_APP_CHAT_PROMPT_CONFIG), has_context, baichuan_context_prompt)
return cls.get_chat_prompt(
copy.deepcopy(BAICHUAN_CHAT_APP_CHAT_PROMPT_CONFIG), has_context, baichuan_context_prompt
)
elif app_mode == AppMode.COMPLETION.value:
if model_mode == "completion":
return cls.get_completion_prompt(copy.deepcopy(BAICHUAN_COMPLETION_APP_COMPLETION_PROMPT_CONFIG), has_context, baichuan_context_prompt)
return cls.get_completion_prompt(
copy.deepcopy(BAICHUAN_COMPLETION_APP_COMPLETION_PROMPT_CONFIG),
has_context,
baichuan_context_prompt,
)
elif model_mode == "chat":
return cls.get_chat_prompt(copy.deepcopy(BAICHUAN_COMPLETION_APP_CHAT_PROMPT_CONFIG), has_context, baichuan_context_prompt)
return cls.get_chat_prompt(
copy.deepcopy(BAICHUAN_COMPLETION_APP_CHAT_PROMPT_CONFIG), has_context, baichuan_context_prompt
)

126
api/services/agent_service.py
View file

@ -10,59 +10,65 @@ from models.model import App, Conversation, EndUser, Message, MessageAgentThough
class AgentService:
@classmethod
def get_agent_logs(cls, app_model: App,
conversation_id: str,
message_id: str) -> dict:
def get_agent_logs(cls, app_model: App, conversation_id: str, message_id: str) -> dict:
"""
Service to get agent logs
"""
conversation: Conversation = db.session.query(Conversation).filter(
Conversation.id == conversation_id,
Conversation.app_id == app_model.id,
).first()
conversation: Conversation = (
db.session.query(Conversation)
.filter(
Conversation.id == conversation_id,
Conversation.app_id == app_model.id,
)
.first()
)
if not conversation:
raise ValueError(f"Conversation not found: {conversation_id}")
message: Message = db.session.query(Message).filter(
Message.id == message_id,
Message.conversation_id == conversation_id,
).first()
message: Message = (
db.session.query(Message)
.filter(
Message.id == message_id,
Message.conversation_id == conversation_id,
)
.first()
)
if not message:
raise ValueError(f"Message not found: {message_id}")
agent_thoughts: list[MessageAgentThought] = message.agent_thoughts
if conversation.from_end_user_id:
# only select name field
executor = db.session.query(EndUser, EndUser.name).filter(
EndUser.id == conversation.from_end_user_id
).first()
executor = (
db.session.query(EndUser, EndUser.name).filter(EndUser.id == conversation.from_end_user_id).first()
)
else:
executor = db.session.query(Account, Account.name).filter(
Account.id == conversation.from_account_id
).first()
executor = (
db.session.query(Account, Account.name).filter(Account.id == conversation.from_account_id).first()
)
if executor:
executor = executor.name
else:
executor = 'Unknown'
executor = "Unknown"
timezone = pytz.timezone(current_user.timezone)
result = {
'meta': {
'status': 'success',
'executor': executor,
'start_time': message.created_at.astimezone(timezone).isoformat(),
'elapsed_time': message.provider_response_latency,
'total_tokens': message.answer_tokens + message.message_tokens,
'agent_mode': app_model.app_model_config.agent_mode_dict.get('strategy', 'react'),
'iterations': len(agent_thoughts),
"meta": {
"status": "success",
"executor": executor,
"start_time": message.created_at.astimezone(timezone).isoformat(),
"elapsed_time": message.provider_response_latency,
"total_tokens": message.answer_tokens + message.message_tokens,
"agent_mode": app_model.app_model_config.agent_mode_dict.get("strategy", "react"),
"iterations": len(agent_thoughts),
},
'iterations': [],
'files': message.files,
"iterations": [],
"files": message.files,
}
agent_config = AgentConfigManager.convert(app_model.app_model_config.to_dict())
@ -86,12 +92,12 @@ class AgentService:
tool_input = tool_inputs.get(tool_name, {})
tool_output = tool_outputs.get(tool_name, {})
tool_meta_data = tool_meta.get(tool_name, {})
tool_config = tool_meta_data.get('tool_config', {})
if tool_config.get('tool_provider_type', '') != 'dataset-retrieval':
tool_config = tool_meta_data.get("tool_config", {})
if tool_config.get("tool_provider_type", "") != "dataset-retrieval":
tool_icon = ToolManager.get_tool_icon(
tenant_id=app_model.tenant_id,
provider_type=tool_config.get('tool_provider_type', ''),
provider_id=tool_config.get('tool_provider', ''),
provider_type=tool_config.get("tool_provider_type", ""),
provider_id=tool_config.get("tool_provider", ""),
)
if not tool_icon:
tool_entity = find_agent_tool(tool_name)
@ -102,30 +108,34 @@ class AgentService:
provider_id=tool_entity.provider_id,
)
else:
tool_icon = ''
tool_icon = ""
tool_calls.append({
'status': 'success' if not tool_meta_data.get('error') else 'error',
'error': tool_meta_data.get('error'),
'time_cost': tool_meta_data.get('time_cost', 0),
'tool_name': tool_name,
'tool_label': tool_label,
'tool_input': tool_input,
'tool_output': tool_output,
'tool_parameters': tool_meta_data.get('tool_parameters', {}),
'tool_icon': tool_icon,
})
tool_calls.append(
{
"status": "success" if not tool_meta_data.get("error") else "error",
"error": tool_meta_data.get("error"),
"time_cost": tool_meta_data.get("time_cost", 0),
"tool_name": tool_name,
"tool_label": tool_label,
"tool_input": tool_input,
"tool_output": tool_output,
"tool_parameters": tool_meta_data.get("tool_parameters", {}),
"tool_icon": tool_icon,
}
)
result['iterations'].append({
'tokens': agent_thought.tokens,
'tool_calls': tool_calls,
'tool_raw': {
'inputs': agent_thought.tool_input,
'outputs': agent_thought.observation,
},
'thought': agent_thought.thought,
'created_at': agent_thought.created_at.isoformat(),
'files': agent_thought.files,
})
result["iterations"].append(
{
"tokens": agent_thought.tokens,
"tool_calls": tool_calls,
"tool_raw": {
"inputs": agent_thought.tool_input,
"outputs": agent_thought.observation,
},
"thought": agent_thought.thought,
"created_at": agent_thought.created_at.isoformat(),
"files": agent_thought.files,
}
)
return result
return result

372
api/services/annotation_service.py
View file

@ -23,21 +23,18 @@ class AppAnnotationService:
@classmethod
def up_insert_app_annotation_from_message(cls, args: dict, app_id: str) -> MessageAnnotation:
# get app info
app = db.session.query(App).filter(
App.id == app_id,
App.tenant_id == current_user.current_tenant_id,
App.status == 'normal'
).first()
app = (
db.session.query(App)
.filter(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal")
.first()
)
if not app:
raise NotFound("App not found")
if args.get('message_id'):
message_id = str(args['message_id'])
if args.get("message_id"):
message_id = str(args["message_id"])
# get message info
message = db.session.query(Message).filter(
Message.id == message_id,
Message.app_id == app.id
).first()
message = db.session.query(Message).filter(Message.id == message_id, Message.app_id == app.id).first()
if not message:
raise NotFound("Message Not Exists.")
@ -45,159 +42,166 @@ class AppAnnotationService:
annotation = message.annotation
# save the message annotation
if annotation:
annotation.content = args['answer']
annotation.question = args['question']
annotation.content = args["answer"]
annotation.question = args["question"]
else:
annotation = MessageAnnotation(
app_id=app.id,
conversation_id=message.conversation_id,
message_id=message.id,
content=args['answer'],
question=args['question'],
account_id=current_user.id
content=args["answer"],
question=args["question"],
account_id=current_user.id,
)
else:
annotation = MessageAnnotation(
app_id=app.id,
content=args['answer'],
question=args['question'],
account_id=current_user.id
app_id=app.id, content=args["answer"], question=args["question"], account_id=current_user.id
)
db.session.add(annotation)
db.session.commit()
# if annotation reply is enabled , add annotation to index
annotation_setting = db.session.query(AppAnnotationSetting).filter(
AppAnnotationSetting.app_id == app_id).first()
annotation_setting = (
db.session.query(AppAnnotationSetting).filter(AppAnnotationSetting.app_id == app_id).first()
)
if annotation_setting:
add_annotation_to_index_task.delay(annotation.id, args['question'], current_user.current_tenant_id,
app_id, annotation_setting.collection_binding_id)
add_annotation_to_index_task.delay(
annotation.id,
args["question"],
current_user.current_tenant_id,
app_id,
annotation_setting.collection_binding_id,
)
return annotation
@classmethod
def enable_app_annotation(cls, args: dict, app_id: str) -> dict:
enable_app_annotation_key = 'enable_app_annotation_{}'.format(str(app_id))
enable_app_annotation_key = "enable_app_annotation_{}".format(str(app_id))
cache_result = redis_client.get(enable_app_annotation_key)
if cache_result is not None:
return {
'job_id': cache_result,
'job_status': 'processing'
}
return {"job_id": cache_result, "job_status": "processing"}
# async job
job_id = str(uuid.uuid4())
enable_app_annotation_job_key = 'enable_app_annotation_job_{}'.format(str(job_id))
enable_app_annotation_job_key = "enable_app_annotation_job_{}".format(str(job_id))
# send batch add segments task
redis_client.setnx(enable_app_annotation_job_key, 'waiting')
enable_annotation_reply_task.delay(str(job_id), app_id, current_user.id, current_user.current_tenant_id,
args['score_threshold'],
args['embedding_provider_name'], args['embedding_model_name'])
return {
'job_id': job_id,
'job_status': 'waiting'
}
redis_client.setnx(enable_app_annotation_job_key, "waiting")
enable_annotation_reply_task.delay(
str(job_id),
app_id,
current_user.id,
current_user.current_tenant_id,
args["score_threshold"],
args["embedding_provider_name"],
args["embedding_model_name"],
)
return {"job_id": job_id, "job_status": "waiting"}
@classmethod
def disable_app_annotation(cls, app_id: str) -> dict:
disable_app_annotation_key = 'disable_app_annotation_{}'.format(str(app_id))
disable_app_annotation_key = "disable_app_annotation_{}".format(str(app_id))
cache_result = redis_client.get(disable_app_annotation_key)
if cache_result is not None:
return {
'job_id': cache_result,
'job_status': 'processing'
}
return {"job_id": cache_result, "job_status": "processing"}
# async job
job_id = str(uuid.uuid4())
disable_app_annotation_job_key = 'disable_app_annotation_job_{}'.format(str(job_id))
disable_app_annotation_job_key = "disable_app_annotation_job_{}".format(str(job_id))
# send batch add segments task
redis_client.setnx(disable_app_annotation_job_key, 'waiting')
redis_client.setnx(disable_app_annotation_job_key, "waiting")
disable_annotation_reply_task.delay(str(job_id), app_id, current_user.current_tenant_id)
return {
'job_id': job_id,
'job_status': 'waiting'
}
return {"job_id": job_id, "job_status": "waiting"}
@classmethod
def get_annotation_list_by_app_id(cls, app_id: str, page: int, limit: int, keyword: str):
# get app info
app = db.session.query(App).filter(
App.id == app_id,
App.tenant_id == current_user.current_tenant_id,
App.status == 'normal'
).first()
app = (
db.session.query(App)
.filter(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal")
.first()
)
if not app:
raise NotFound("App not found")
if keyword:
annotations = (db.session.query(MessageAnnotation)
.filter(MessageAnnotation.app_id == app_id)
.filter(
or_(
MessageAnnotation.question.ilike('%{}%'.format(keyword)),
MessageAnnotation.content.ilike('%{}%'.format(keyword))
annotations = (
db.session.query(MessageAnnotation)
.filter(MessageAnnotation.app_id == app_id)
.filter(
or_(
MessageAnnotation.question.ilike("%{}%".format(keyword)),
MessageAnnotation.content.ilike("%{}%".format(keyword)),
)
)
.order_by(MessageAnnotation.created_at.desc())
.paginate(page=page, per_page=limit, max_per_page=100, error_out=False)
)
.order_by(MessageAnnotation.created_at.desc())
.paginate(page=page, per_page=limit, max_per_page=100, error_out=False))
else:
annotations = (db.session.query(MessageAnnotation)
.filter(MessageAnnotation.app_id == app_id)
.order_by(MessageAnnotation.created_at.desc())
.paginate(page=page, per_page=limit, max_per_page=100, error_out=False))
annotations = (
db.session.query(MessageAnnotation)
.filter(MessageAnnotation.app_id == app_id)
.order_by(MessageAnnotation.created_at.desc())
.paginate(page=page, per_page=limit, max_per_page=100, error_out=False)
)
return annotations.items, annotations.total
@classmethod
def export_annotation_list_by_app_id(cls, app_id: str):
# get app info
app = db.session.query(App).filter(
App.id == app_id,
App.tenant_id == current_user.current_tenant_id,
App.status == 'normal'
).first()
app = (
db.session.query(App)
.filter(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal")
.first()
)
if not app:
raise NotFound("App not found")
annotations = (db.session.query(MessageAnnotation)
.filter(MessageAnnotation.app_id == app_id)
.order_by(MessageAnnotation.created_at.desc()).all())
annotations = (
db.session.query(MessageAnnotation)
.filter(MessageAnnotation.app_id == app_id)
.order_by(MessageAnnotation.created_at.desc())
.all()
)
return annotations
@classmethod
def insert_app_annotation_directly(cls, args: dict, app_id: str) -> MessageAnnotation:
# get app info
app = db.session.query(App).filter(
App.id == app_id,
App.tenant_id == current_user.current_tenant_id,
App.status == 'normal'
).first()
app = (
db.session.query(App)
.filter(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal")
.first()
)
if not app:
raise NotFound("App not found")
annotation = MessageAnnotation(
app_id=app.id,
content=args['answer'],
question=args['question'],
account_id=current_user.id
app_id=app.id, content=args["answer"], question=args["question"], account_id=current_user.id
)
db.session.add(annotation)
db.session.commit()
# if annotation reply is enabled , add annotation to index
annotation_setting = db.session.query(AppAnnotationSetting).filter(
AppAnnotationSetting.app_id == app_id).first()
annotation_setting = (
db.session.query(AppAnnotationSetting).filter(AppAnnotationSetting.app_id == app_id).first()
)
if annotation_setting:
add_annotation_to_index_task.delay(annotation.id, args['question'], current_user.current_tenant_id,
app_id, annotation_setting.collection_binding_id)
add_annotation_to_index_task.delay(
annotation.id,
args["question"],
current_user.current_tenant_id,
app_id,
annotation_setting.collection_binding_id,
)
return annotation
@classmethod
def update_app_annotation_directly(cls, args: dict, app_id: str, annotation_id: str):
# get app info
app = db.session.query(App).filter(
App.id == app_id,
App.tenant_id == current_user.current_tenant_id,
App.status == 'normal'
).first()
app = (
db.session.query(App)
.filter(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal")
.first()
)
if not app:
raise NotFound("App not found")
@ -207,30 +211,34 @@ class AppAnnotationService:
if not annotation:
raise NotFound("Annotation not found")
annotation.content = args['answer']
annotation.question = args['question']
annotation.content = args["answer"]
annotation.question = args["question"]
db.session.commit()
# if annotation reply is enabled , add annotation to index
app_annotation_setting = db.session.query(AppAnnotationSetting).filter(
AppAnnotationSetting.app_id == app_id
).first()
app_annotation_setting = (
db.session.query(AppAnnotationSetting).filter(AppAnnotationSetting.app_id == app_id).first()
)
if app_annotation_setting:
update_annotation_to_index_task.delay(annotation.id, annotation.question,
current_user.current_tenant_id,
app_id, app_annotation_setting.collection_binding_id)
update_annotation_to_index_task.delay(
annotation.id,
annotation.question,
current_user.current_tenant_id,
app_id,
app_annotation_setting.collection_binding_id,
)
return annotation
@classmethod
def delete_app_annotation(cls, app_id: str, annotation_id: str):
# get app info
app = db.session.query(App).filter(
App.id == app_id,
App.tenant_id == current_user.current_tenant_id,
App.status == 'normal'
).first()
app = (
db.session.query(App)
.filter(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal")
.first()
)
if not app:
raise NotFound("App not found")
@ -242,33 +250,34 @@ class AppAnnotationService:
db.session.delete(annotation)
annotation_hit_histories = (db.session.query(AppAnnotationHitHistory)
.filter(AppAnnotationHitHistory.annotation_id == annotation_id)
.all()
)
annotation_hit_histories = (
db.session.query(AppAnnotationHitHistory)
.filter(AppAnnotationHitHistory.annotation_id == annotation_id)
.all()
)
if annotation_hit_histories:
for annotation_hit_history in annotation_hit_histories:
db.session.delete(annotation_hit_history)
db.session.commit()
# if annotation reply is enabled , delete annotation index
app_annotation_setting = db.session.query(AppAnnotationSetting).filter(
AppAnnotationSetting.app_id == app_id
).first()
app_annotation_setting = (
db.session.query(AppAnnotationSetting).filter(AppAnnotationSetting.app_id == app_id).first()
)
if app_annotation_setting:
delete_annotation_index_task.delay(annotation.id, app_id,
current_user.current_tenant_id,
app_annotation_setting.collection_binding_id)
delete_annotation_index_task.delay(
annotation.id, app_id, current_user.current_tenant_id, app_annotation_setting.collection_binding_id
)
@classmethod
def batch_import_app_annotations(cls, app_id, file: FileStorage) -> dict:
# get app info
app = db.session.query(App).filter(
App.id == app_id,
App.tenant_id == current_user.current_tenant_id,
App.status == 'normal'
).first()
app = (
db.session.query(App)
.filter(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal")
.first()
)
if not app:
raise NotFound("App not found")
@ -278,10 +287,7 @@ class AppAnnotationService:
df = pd.read_csv(file)
result = []
for index, row in df.iterrows():
content = {
'question': row[0],
'answer': row[1]
}
content = {"question": row[0], "answer": row[1]}
result.append(content)
if len(result) == 0:
raise ValueError("The CSV file is empty.")
@ -293,28 +299,24 @@ class AppAnnotationService:
raise ValueError("The number of annotations exceeds the limit of your subscription.")
# async job
job_id = str(uuid.uuid4())
indexing_cache_key = 'app_annotation_batch_import_{}'.format(str(job_id))
indexing_cache_key = "app_annotation_batch_import_{}".format(str(job_id))
# send batch add segments task
redis_client.setnx(indexing_cache_key, 'waiting')
batch_import_annotations_task.delay(str(job_id), result, app_id,
current_user.current_tenant_id, current_user.id)
redis_client.setnx(indexing_cache_key, "waiting")
batch_import_annotations_task.delay(
str(job_id), result, app_id, current_user.current_tenant_id, current_user.id
)
except Exception as e:
return {
'error_msg': str(e)
}
return {
'job_id': job_id,
'job_status': 'waiting'
}
return {"error_msg": str(e)}
return {"job_id": job_id, "job_status": "waiting"}
@classmethod
def get_annotation_hit_histories(cls, app_id: str, annotation_id: str, page, limit):
# get app info
app = db.session.query(App).filter(
App.id == app_id,
App.tenant_id == current_user.current_tenant_id,
App.status == 'normal'
).first()
app = (
db.session.query(App)
.filter(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal")
.first()
)
if not app:
raise NotFound("App not found")
@ -324,12 +326,15 @@ class AppAnnotationService:
if not annotation:
raise NotFound("Annotation not found")
annotation_hit_histories = (db.session.query(AppAnnotationHitHistory)
.filter(AppAnnotationHitHistory.app_id == app_id,
AppAnnotationHitHistory.annotation_id == annotation_id,
)
.order_by(AppAnnotationHitHistory.created_at.desc())
.paginate(page=page, per_page=limit, max_per_page=100, error_out=False))
annotation_hit_histories = (
db.session.query(AppAnnotationHitHistory)
.filter(
AppAnnotationHitHistory.app_id == app_id,
AppAnnotationHitHistory.annotation_id == annotation_id,
)
.order_by(AppAnnotationHitHistory.created_at.desc())
.paginate(page=page, per_page=limit, max_per_page=100, error_out=False)
)
return annotation_hit_histories.items, annotation_hit_histories.total
@classmethod
@ -341,15 +346,21 @@ class AppAnnotationService:
return annotation
@classmethod
def add_annotation_history(cls, annotation_id: str, app_id: str, annotation_question: str,
annotation_content: str, query: str, user_id: str,
message_id: str, from_source: str, score: float):
def add_annotation_history(
cls,
annotation_id: str,
app_id: str,
annotation_question: str,
annotation_content: str,
query: str,
user_id: str,
message_id: str,
from_source: str,
score: float,
):
# add hit count to annotation
db.session.query(MessageAnnotation).filter(
MessageAnnotation.id == annotation_id
).update(
{MessageAnnotation.hit_count: MessageAnnotation.hit_count + 1},
synchronize_session=False
db.session.query(MessageAnnotation).filter(MessageAnnotation.id == annotation_id).update(
{MessageAnnotation.hit_count: MessageAnnotation.hit_count + 1}, synchronize_session=False
)
annotation_hit_history = AppAnnotationHitHistory(
@ -361,7 +372,7 @@ class AppAnnotationService:
score=score,
message_id=message_id,
annotation_question=annotation_question,
annotation_content=annotation_content
annotation_content=annotation_content,
)
db.session.add(annotation_hit_history)
db.session.commit()
@ -369,17 +380,18 @@ class AppAnnotationService:
@classmethod
def get_app_annotation_setting_by_app_id(cls, app_id: str):
# get app info
app = db.session.query(App).filter(
App.id == app_id,
App.tenant_id == current_user.current_tenant_id,
App.status == 'normal'
).first()
app = (
db.session.query(App)
.filter(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal")
.first()
)
if not app:
raise NotFound("App not found")
annotation_setting = db.session.query(AppAnnotationSetting).filter(
AppAnnotationSetting.app_id == app_id).first()
annotation_setting = (
db.session.query(AppAnnotationSetting).filter(AppAnnotationSetting.app_id == app_id).first()
)
if annotation_setting:
collection_binding_detail = annotation_setting.collection_binding_detail
return {
@ -388,32 +400,34 @@ class AppAnnotationService:
"score_threshold": annotation_setting.score_threshold,
"embedding_model": {
"embedding_provider_name": collection_binding_detail.provider_name,
"embedding_model_name": collection_binding_detail.model_name
}
"embedding_model_name": collection_binding_detail.model_name,
},
}
return {
"enabled": False
}
return {"enabled": False}
@classmethod
def update_app_annotation_setting(cls, app_id: str, annotation_setting_id: str, args: dict):
# get app info
app = db.session.query(App).filter(
App.id == app_id,
App.tenant_id == current_user.current_tenant_id,
App.status == 'normal'
).first()
app = (
db.session.query(App)
.filter(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal")
.first()
)
if not app:
raise NotFound("App not found")
annotation_setting = db.session.query(AppAnnotationSetting).filter(
AppAnnotationSetting.app_id == app_id,
AppAnnotationSetting.id == annotation_setting_id,
).first()
annotation_setting = (
db.session.query(AppAnnotationSetting)
.filter(
AppAnnotationSetting.app_id == app_id,
AppAnnotationSetting.id == annotation_setting_id,
)
.first()
)
if not annotation_setting:
raise NotFound("App annotation not found")
annotation_setting.score_threshold = args['score_threshold']
annotation_setting.score_threshold = args["score_threshold"]
annotation_setting.updated_user_id = current_user.id
annotation_setting.updated_at = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)
db.session.add(annotation_setting)
@ -427,6 +441,6 @@ class AppAnnotationService:
"score_threshold": annotation_setting.score_threshold,
"embedding_model": {
"embedding_provider_name": collection_binding_detail.provider_name,
"embedding_model_name": collection_binding_detail.model_name
}
"embedding_model_name": collection_binding_detail.model_name,
},
}

39
api/services/api_based_extension_service.py
View file

@ -5,13 +5,14 @@ from models.api_based_extension import APIBasedExtension, APIBasedExtensionPoint
class APIBasedExtensionService:
@staticmethod
def get_all_by_tenant_id(tenant_id: str) -> list[APIBasedExtension]:
extension_list = db.session.query(APIBasedExtension) \
.filter_by(tenant_id=tenant_id) \
.order_by(APIBasedExtension.created_at.desc()) \
.all()
extension_list = (
db.session.query(APIBasedExtension)
.filter_by(tenant_id=tenant_id)
.order_by(APIBasedExtension.created_at.desc())
.all()
)
for extension in extension_list:
extension.api_key = decrypt_token(extension.tenant_id, extension.api_key)
@ -35,10 +36,12 @@ class APIBasedExtensionService:
@staticmethod
def get_with_tenant_id(tenant_id: str, api_based_extension_id: str) -> APIBasedExtension:
extension = db.session.query(APIBasedExtension) \
.filter_by(tenant_id=tenant_id) \
.filter_by(id=api_based_extension_id) \
extension = (
db.session.query(APIBasedExtension)
.filter_by(tenant_id=tenant_id)
.filter_by(id=api_based_extension_id)
.first()
)
if not extension:
raise ValueError("API based extension is not found")
@ -55,20 +58,24 @@ class APIBasedExtensionService:
if not extension_data.id:
# case one: check new data, name must be unique
is_name_existed = db.session.query(APIBasedExtension) \
.filter_by(tenant_id=extension_data.tenant_id) \
.filter_by(name=extension_data.name) \
is_name_existed = (
db.session.query(APIBasedExtension)
.filter_by(tenant_id=extension_data.tenant_id)
.filter_by(name=extension_data.name)
.first()
)
if is_name_existed:
raise ValueError("name must be unique, it is already existed")
else:
# case two: check existing data, name must be unique
is_name_existed = db.session.query(APIBasedExtension) \
.filter_by(tenant_id=extension_data.tenant_id) \
.filter_by(name=extension_data.name) \
.filter(APIBasedExtension.id != extension_data.id) \
is_name_existed = (
db.session.query(APIBasedExtension)
.filter_by(tenant_id=extension_data.tenant_id)
.filter_by(name=extension_data.name)
.filter(APIBasedExtension.id != extension_data.id)
.first()
)
if is_name_existed:
raise ValueError("name must be unique, it is already existed")
@ -92,7 +99,7 @@ class APIBasedExtensionService:
try:
client = APIBasedExtensionRequestor(extension_data.api_endpoint, extension_data.api_key)
resp = client.request(point=APIBasedExtensionPoint.PING, params={})
if resp.get('result') != 'pong':
if resp.get("result") != "pong":
raise ValueError(resp)
except Exception as e:
raise ValueError("connection error: {}".format(e))

185
api/services/app_dsl_service.py
View file

@ -75,43 +75,44 @@ class AppDslService:
# check or repair dsl version
import_data = cls._check_or_fix_dsl(import_data)
app_data = import_data.get('app')
app_data = import_data.get("app")
if not app_data:
raise ValueError("Missing app in data argument")
# get app basic info
name = args.get("name") if args.get("name") else app_data.get('name')
description = args.get("description") if args.get("description") else app_data.get('description', '')
icon_type = args.get("icon_type") if args.get("icon_type") else app_data.get('icon_type')
icon = args.get("icon") if args.get("icon") else app_data.get('icon')
icon_background = args.get("icon_background") if args.get("icon_background") \
else app_data.get('icon_background')
name = args.get("name") if args.get("name") else app_data.get("name")
description = args.get("description") if args.get("description") else app_data.get("description", "")
icon_type = args.get("icon_type") if args.get("icon_type") else app_data.get("icon_type")
icon = args.get("icon") if args.get("icon") else app_data.get("icon")
icon_background = (
args.get("icon_background") if args.get("icon_background") else app_data.get("icon_background")
)
# import dsl and create app
app_mode = AppMode.value_of(app_data.get('mode'))
app_mode = AppMode.value_of(app_data.get("mode"))
if app_mode in [AppMode.ADVANCED_CHAT, AppMode.WORKFLOW]:
app = cls._import_and_create_new_workflow_based_app(
tenant_id=tenant_id,
app_mode=app_mode,
workflow_data=import_data.get('workflow'),
workflow_data=import_data.get("workflow"),
account=account,
name=name,
description=description,
icon_type=icon_type,
icon=icon,
icon_background=icon_background
icon_background=icon_background,
)
elif app_mode in [AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.COMPLETION]:
app = cls._import_and_create_new_model_config_based_app(
tenant_id=tenant_id,
app_mode=app_mode,
model_config_data=import_data.get('model_config'),
model_config_data=import_data.get("model_config"),
account=account,
name=name,
description=description,
icon_type=icon_type,
icon=icon,
icon_background=icon_background
icon_background=icon_background,
)
else:
raise ValueError("Invalid app mode")
@ -134,27 +135,26 @@ class AppDslService:
# check or repair dsl version
import_data = cls._check_or_fix_dsl(import_data)
app_data = import_data.get('app')
app_data = import_data.get("app")
if not app_data:
raise ValueError("Missing app in data argument")
# import dsl and overwrite app
app_mode = AppMode.value_of(app_data.get('mode'))
app_mode = AppMode.value_of(app_data.get("mode"))
if app_mode not in [AppMode.ADVANCED_CHAT, AppMode.WORKFLOW]:
raise ValueError("Only support import workflow in advanced-chat or workflow app.")
if app_data.get('mode') != app_model.mode:
raise ValueError(
f"App mode {app_data.get('mode')} is not matched with current app mode {app_mode.value}")
if app_data.get("mode") != app_model.mode:
raise ValueError(f"App mode {app_data.get('mode')} is not matched with current app mode {app_mode.value}")
return cls._import_and_overwrite_workflow_based_app(
app_model=app_model,
workflow_data=import_data.get('workflow'),
workflow_data=import_data.get("workflow"),
account=account,
)
@classmethod
def export_dsl(cls, app_model: App, include_secret:bool = False) -> str:
def export_dsl(cls, app_model: App, include_secret: bool = False) -> str:
"""
Export app
:param app_model: App instance
@ -168,14 +168,16 @@ class AppDslService:
"app": {
"name": app_model.name,
"mode": app_model.mode,
"icon": '🤖' if app_model.icon_type == 'image' else app_model.icon,
"icon_background": '#FFEAD5' if app_model.icon_type == 'image' else app_model.icon_background,
"description": app_model.description
}
"icon": "🤖" if app_model.icon_type == "image" else app_model.icon,
"icon_background": "#FFEAD5" if app_model.icon_type == "image" else app_model.icon_background,
"description": app_model.description,
},
}
if app_mode in [AppMode.ADVANCED_CHAT, AppMode.WORKFLOW]:
cls._append_workflow_export_data(export_data=export_data, app_model=app_model, include_secret=include_secret)
cls._append_workflow_export_data(
export_data=export_data, app_model=app_model, include_secret=include_secret
)
else:
cls._append_model_config_export_data(export_data, app_model)
@ -188,31 +190,35 @@ class AppDslService:
:param import_data: import data
"""
if not import_data.get('version'):
import_data['version'] = "0.1.0"
if not import_data.get("version"):
import_data["version"] = "0.1.0"
if not import_data.get('kind') or import_data.get('kind') != "app":
import_data['kind'] = "app"
if not import_data.get("kind") or import_data.get("kind") != "app":
import_data["kind"] = "app"
if import_data.get('version') != current_dsl_version:
if import_data.get("version") != current_dsl_version:
# Currently only one DSL version, so no difference checks or compatibility fixes will be performed.
logger.warning(f"DSL version {import_data.get('version')} is not compatible "
f"with current version {current_dsl_version}, related to "
f"Dify version {dsl_to_dify_version_mapping.get(current_dsl_version)}.")
logger.warning(
f"DSL version {import_data.get('version')} is not compatible "
f"with current version {current_dsl_version}, related to "
f"Dify version {dsl_to_dify_version_mapping.get(current_dsl_version)}."
)
return import_data
@classmethod
def _import_and_create_new_workflow_based_app(cls,
tenant_id: str,
app_mode: AppMode,
workflow_data: dict,
account: Account,
name: str,
description: str,
icon_type: str,
icon: str,
icon_background: str) -> App:
def _import_and_create_new_workflow_based_app(
cls,
tenant_id: str,
app_mode: AppMode,
workflow_data: dict,
account: Account,
name: str,
description: str,
icon_type: str,
icon: str,
icon_background: str,
) -> App:
"""
Import app dsl and create new workflow based app
@ -227,8 +233,7 @@ class AppDslService:
:param icon_background: app icon background
"""
if not workflow_data:
raise ValueError("Missing workflow in data argument "
"when app mode is advanced-chat or workflow")
raise ValueError("Missing workflow in data argument " "when app mode is advanced-chat or workflow")
app = cls._create_app(
tenant_id=tenant_id,
@ -238,37 +243,32 @@ class AppDslService:
description=description,
icon_type=icon_type,
icon=icon,
icon_background=icon_background
icon_background=icon_background,
)
# init draft workflow
environment_variables_list = workflow_data.get('environment_variables') or []
environment_variables_list = workflow_data.get("environment_variables") or []
environment_variables = [factory.build_variable_from_mapping(obj) for obj in environment_variables_list]
conversation_variables_list = workflow_data.get('conversation_variables') or []
conversation_variables_list = workflow_data.get("conversation_variables") or []
conversation_variables = [factory.build_variable_from_mapping(obj) for obj in conversation_variables_list]
workflow_service = WorkflowService()
draft_workflow = workflow_service.sync_draft_workflow(
app_model=app,
graph=workflow_data.get('graph', {}),
features=workflow_data.get('../core/app/features', {}),
graph=workflow_data.get("graph", {}),
features=workflow_data.get("../core/app/features", {}),
unique_hash=None,
account=account,
environment_variables=environment_variables,
conversation_variables=conversation_variables,
)
workflow_service.publish_workflow(
app_model=app,
account=account,
draft_workflow=draft_workflow
)
workflow_service.publish_workflow(app_model=app, account=account, draft_workflow=draft_workflow)
return app
@classmethod
def _import_and_overwrite_workflow_based_app(cls,
app_model: App,
workflow_data: dict,
account: Account) -> Workflow:
def _import_and_overwrite_workflow_based_app(
cls, app_model: App, workflow_data: dict, account: Account
) -> Workflow:
"""
Import app dsl and overwrite workflow based app
@ -277,8 +277,7 @@ class AppDslService:
:param account: Account instance
"""
if not workflow_data:
raise ValueError("Missing workflow in data argument "
"when app mode is advanced-chat or workflow")
raise ValueError("Missing workflow in data argument " "when app mode is advanced-chat or workflow")
# fetch draft workflow by app_model
workflow_service = WorkflowService()
@ -289,14 +288,14 @@ class AppDslService:
unique_hash = None
# sync draft workflow
environment_variables_list = workflow_data.get('environment_variables') or []
environment_variables_list = workflow_data.get("environment_variables") or []
environment_variables = [factory.build_variable_from_mapping(obj) for obj in environment_variables_list]
conversation_variables_list = workflow_data.get('conversation_variables') or []
conversation_variables_list = workflow_data.get("conversation_variables") or []
conversation_variables = [factory.build_variable_from_mapping(obj) for obj in conversation_variables_list]
draft_workflow = workflow_service.sync_draft_workflow(
app_model=app_model,
graph=workflow_data.get('graph', {}),
features=workflow_data.get('features', {}),
graph=workflow_data.get("graph", {}),
features=workflow_data.get("features", {}),
unique_hash=unique_hash,
account=account,
environment_variables=environment_variables,
@ -306,16 +305,18 @@ class AppDslService:
return draft_workflow
@classmethod
def _import_and_create_new_model_config_based_app(cls,
tenant_id: str,
app_mode: AppMode,
model_config_data: dict,
account: Account,
name: str,
description: str,
icon_type: str,
icon: str,
icon_background: str) -> App:
def _import_and_create_new_model_config_based_app(
cls,
tenant_id: str,
app_mode: AppMode,
model_config_data: dict,
account: Account,
name: str,
description: str,
icon_type: str,
icon: str,
icon_background: str,
) -> App:
"""
Import app dsl and create new model config based app
@ -329,8 +330,7 @@ class AppDslService:
:param icon_background: app icon background
"""
if not model_config_data:
raise ValueError("Missing model_config in data argument "
"when app mode is chat, agent-chat or completion")
raise ValueError("Missing model_config in data argument " "when app mode is chat, agent-chat or completion")
app = cls._create_app(
tenant_id=tenant_id,
@ -340,7 +340,7 @@ class AppDslService:
description=description,
icon_type=icon_type,
icon=icon,
icon_background=icon_background
icon_background=icon_background,
)
app_model_config = AppModelConfig()
@ -352,23 +352,22 @@ class AppDslService:
app.app_model_config_id = app_model_config.id
app_model_config_was_updated.send(
app,
app_model_config=app_model_config
)
app_model_config_was_updated.send(app, app_model_config=app_model_config)
return app
@classmethod
def _create_app(cls,
tenant_id: str,
app_mode: AppMode,
account: Account,
name: str,
description: str,
icon_type: str,
icon: str,
icon_background: str) -> App:
def _create_app(
cls,
tenant_id: str,
app_mode: AppMode,
account: Account,
name: str,
description: str,
icon_type: str,
icon: str,
icon_background: str,
) -> App:
"""
Create new app
@ -390,7 +389,7 @@ class AppDslService:
icon=icon,
icon_background=icon_background,
enable_site=True,
enable_api=True
enable_api=True,
)
db.session.add(app)
@ -412,7 +411,7 @@ class AppDslService:
if not workflow:
raise ValueError("Missing draft workflow configuration, please check.")
export_data['workflow'] = workflow.to_dict(include_secret=include_secret)
export_data["workflow"] = workflow.to_dict(include_secret=include_secret)
@classmethod
def _append_model_config_export_data(cls, export_data: dict, app_model: App) -> None:
@ -425,4 +424,4 @@ class AppDslService:
if not app_model_config:
raise ValueError("Missing app configuration, please check.")
export_data['model_config'] = app_model_config.to_dict()
export_data["model_config"] = app_model_config.to_dict()

139
api/services/app_generate_service.py
View file

@ -14,14 +14,15 @@ from services.workflow_service import WorkflowService
class AppGenerateService:
@classmethod
def generate(cls, app_model: App,
user: Union[Account, EndUser],
args: Any,
invoke_from: InvokeFrom,
streaming: bool = True,
):
def generate(
cls,
app_model: App,
user: Union[Account, EndUser],
args: Any,
invoke_from: InvokeFrom,
streaming: bool = True,
):
"""
App Content Generate
:param app_model: app model
@ -37,51 +38,54 @@ class AppGenerateService:
try:
request_id = rate_limit.enter(request_id)
if app_model.mode == AppMode.COMPLETION.value:
return rate_limit.generate(CompletionAppGenerator().generate(
app_model=app_model,
user=user,
args=args,
invoke_from=invoke_from,
stream=streaming
), request_id)
return rate_limit.generate(
CompletionAppGenerator().generate(
app_model=app_model, user=user, args=args, invoke_from=invoke_from, stream=streaming
),
request_id,
)
elif app_model.mode == AppMode.AGENT_CHAT.value or app_model.is_agent:
return rate_limit.generate(AgentChatAppGenerator().generate(
app_model=app_model,
user=user,
args=args,
invoke_from=invoke_from,
stream=streaming
), request_id)
return rate_limit.generate(
AgentChatAppGenerator().generate(
app_model=app_model, user=user, args=args, invoke_from=invoke_from, stream=streaming
),
request_id,
)
elif app_model.mode == AppMode.CHAT.value:
return rate_limit.generate(ChatAppGenerator().generate(
app_model=app_model,
user=user,
args=args,
invoke_from=invoke_from,
stream=streaming
), request_id)
return rate_limit.generate(
ChatAppGenerator().generate(
app_model=app_model, user=user, args=args, invoke_from=invoke_from, stream=streaming
),
request_id,
)
elif app_model.mode == AppMode.ADVANCED_CHAT.value:
workflow = cls._get_workflow(app_model, invoke_from)
return rate_limit.generate(AdvancedChatAppGenerator().generate(
app_model=app_model,
workflow=workflow,
user=user,
args=args,
invoke_from=invoke_from,
stream=streaming
), request_id)
return rate_limit.generate(
AdvancedChatAppGenerator().generate(
app_model=app_model,
workflow=workflow,
user=user,
args=args,
invoke_from=invoke_from,
stream=streaming,
),
request_id,
)
elif app_model.mode == AppMode.WORKFLOW.value:
workflow = cls._get_workflow(app_model, invoke_from)
return rate_limit.generate(WorkflowAppGenerator().generate(
app_model=app_model,
workflow=workflow,
user=user,
args=args,
invoke_from=invoke_from,
stream=streaming
), request_id)
return rate_limit.generate(
WorkflowAppGenerator().generate(
app_model=app_model,
workflow=workflow,
user=user,
args=args,
invoke_from=invoke_from,
stream=streaming,
),
request_id,
)
else:
raise ValueError(f'Invalid app mode {app_model.mode}')
raise ValueError(f"Invalid app mode {app_model.mode}")
finally:
if not streaming:
rate_limit.exit(request_id)
@ -94,38 +98,31 @@ class AppGenerateService:
return max_active_requests
@classmethod
def generate_single_iteration(cls, app_model: App,
user: Union[Account, EndUser],
node_id: str,
args: Any,
streaming: bool = True):
def generate_single_iteration(
cls, app_model: App, user: Union[Account, EndUser], node_id: str, args: Any, streaming: bool = True
):
if app_model.mode == AppMode.ADVANCED_CHAT.value:
workflow = cls._get_workflow(app_model, InvokeFrom.DEBUGGER)
return AdvancedChatAppGenerator().single_iteration_generate(
app_model=app_model,
workflow=workflow,
node_id=node_id,
user=user,
args=args,
stream=streaming
app_model=app_model, workflow=workflow, node_id=node_id, user=user, args=args, stream=streaming
)
elif app_model.mode == AppMode.WORKFLOW.value:
workflow = cls._get_workflow(app_model, InvokeFrom.DEBUGGER)
return WorkflowAppGenerator().single_iteration_generate(
app_model=app_model,
workflow=workflow,
node_id=node_id,
user=user,
args=args,
stream=streaming
app_model=app_model, workflow=workflow, node_id=node_id, user=user, args=args, stream=streaming
)
else:
raise ValueError(f'Invalid app mode {app_model.mode}')
raise ValueError(f"Invalid app mode {app_model.mode}")
@classmethod
def generate_more_like_this(cls, app_model: App, user: Union[Account, EndUser],
message_id: str, invoke_from: InvokeFrom, streaming: bool = True) \
-> Union[dict, Generator]:
def generate_more_like_this(
cls,
app_model: App,
user: Union[Account, EndUser],
message_id: str,
invoke_from: InvokeFrom,
streaming: bool = True,
) -> Union[dict, Generator]:
"""
Generate more like this
:param app_model: app model
@ -136,11 +133,7 @@ class AppGenerateService:
:return:
"""
return CompletionAppGenerator().generate_more_like_this(
app_model=app_model,
message_id=message_id,
user=user,
invoke_from=invoke_from,
stream=streaming
app_model=app_model, message_id=message_id, user=user, invoke_from=invoke_from, stream=streaming
)
@classmethod
@ -157,12 +150,12 @@ class AppGenerateService:
workflow = workflow_service.get_draft_workflow(app_model=app_model)
if not workflow:
raise ValueError('Workflow not initialized')
raise ValueError("Workflow not initialized")
else:
# fetch published workflow by app_model
workflow = workflow_service.get_published_workflow(app_model=app_model)
if not workflow:
raise ValueError('Workflow not published')
raise ValueError("Workflow not published")
return workflow

1
api/services/app_model_config_service.py
View file

@ -5,7 +5,6 @@ from models.model import AppMode
class AppModelConfigService:
@classmethod
def validate_configuration(cls, tenant_id: str, config: dict, app_mode: AppMode) -> dict:
if app_mode == AppMode.CHAT:

161
api/services/app_service.py
View file

@ -33,27 +33,22 @@ class AppService:
:param args: request args
:return:
"""
filters = [
App.tenant_id == tenant_id,
App.is_universal == False
]
filters = [App.tenant_id == tenant_id, App.is_universal == False]
if args['mode'] == 'workflow':
if args["mode"] == "workflow":
filters.append(App.mode.in_([AppMode.WORKFLOW.value, AppMode.COMPLETION.value]))
elif args['mode'] == 'chat':
elif args["mode"] == "chat":
filters.append(App.mode.in_([AppMode.CHAT.value, AppMode.ADVANCED_CHAT.value]))
elif args['mode'] == 'agent-chat':
elif args["mode"] == "agent-chat":
filters.append(App.mode == AppMode.AGENT_CHAT.value)
elif args['mode'] == 'channel':
elif args["mode"] == "channel":
filters.append(App.mode == AppMode.CHANNEL.value)
if args.get('name'):
name = args['name'][:30]
filters.append(App.name.ilike(f'%{name}%'))
if args.get('tag_ids'):
target_ids = TagService.get_target_ids_by_tag_ids('app',
tenant_id,
args['tag_ids'])
if args.get("name"):
name = args["name"][:30]
filters.append(App.name.ilike(f"%{name}%"))
if args.get("tag_ids"):
target_ids = TagService.get_target_ids_by_tag_ids("app", tenant_id, args["tag_ids"])
if target_ids:
filters.append(App.id.in_(target_ids))
else:
@ -61,9 +56,9 @@ class AppService:
app_models = db.paginate(
db.select(App).where(*filters).order_by(App.created_at.desc()),
page=args['page'],
per_page=args['limit'],
error_out=False
page=args["page"],
per_page=args["limit"],
error_out=False,
)
return app_models
@ -75,21 +70,20 @@ class AppService:
:param args: request args
:param account: Account instance
"""
app_mode = AppMode.value_of(args['mode'])
app_mode = AppMode.value_of(args["mode"])
app_template = default_app_templates[app_mode]
# get model config
default_model_config = app_template.get('model_config')
default_model_config = app_template.get("model_config")
default_model_config = default_model_config.copy() if default_model_config else None
if default_model_config and 'model' in default_model_config:
if default_model_config and "model" in default_model_config:
# get model provider
model_manager = ModelManager()
# get default model instance
try:
model_instance = model_manager.get_default_model_instance(
tenant_id=account.current_tenant_id,
model_type=ModelType.LLM
tenant_id=account.current_tenant_id, model_type=ModelType.LLM
)
except (ProviderTokenNotInitError, LLMBadRequestError):
model_instance = None
@ -98,39 +92,41 @@ class AppService:
model_instance = None
if model_instance:
if model_instance.model == default_model_config['model']['name'] and model_instance.provider == default_model_config['model']['provider']:
default_model_dict = default_model_config['model']
if (
model_instance.model == default_model_config["model"]["name"]
and model_instance.provider == default_model_config["model"]["provider"]
):
default_model_dict = default_model_config["model"]
else:
llm_model = cast(LargeLanguageModel, model_instance.model_type_instance)
model_schema = llm_model.get_model_schema(model_instance.model, model_instance.credentials)
default_model_dict = {
'provider': model_instance.provider,
'name': model_instance.model,
'mode': model_schema.model_properties.get(ModelPropertyKey.MODE),
'completion_params': {}
"provider": model_instance.provider,
"name": model_instance.model,
"mode": model_schema.model_properties.get(ModelPropertyKey.MODE),
"completion_params": {},
}
else:
provider, model = model_manager.get_default_provider_model_name(
tenant_id=account.current_tenant_id,
model_type=ModelType.LLM
tenant_id=account.current_tenant_id, model_type=ModelType.LLM
)
default_model_config['model']['provider'] = provider
default_model_config['model']['name'] = model
default_model_dict = default_model_config['model']
default_model_config["model"]["provider"] = provider
default_model_config["model"]["name"] = model
default_model_dict = default_model_config["model"]
default_model_config['model'] = json.dumps(default_model_dict)
default_model_config["model"] = json.dumps(default_model_dict)
app = App(**app_template['app'])
app.name = args['name']
app.description = args.get('description', '')
app.mode = args['mode']
app.icon_type = args.get('icon_type', 'emoji')
app.icon = args['icon']
app.icon_background = args['icon_background']
app = App(**app_template["app"])
app.name = args["name"]
app.description = args.get("description", "")
app.mode = args["mode"]
app.icon_type = args.get("icon_type", "emoji")
app.icon = args["icon"]
app.icon_background = args["icon_background"]
app.tenant_id = tenant_id
app.api_rph = args.get('api_rph', 0)
app.api_rpm = args.get('api_rpm', 0)
app.api_rph = args.get("api_rph", 0)
app.api_rpm = args.get("api_rpm", 0)
db.session.add(app)
db.session.flush()
@ -158,7 +154,7 @@ class AppService:
model_config: AppModelConfig = app.app_model_config
agent_mode = model_config.agent_mode_dict
# decrypt agent tool parameters if it's secret-input
for tool in agent_mode.get('tools') or []:
for tool in agent_mode.get("tools") or []:
if not isinstance(tool, dict) or len(tool.keys()) <= 3:
continue
agent_tool_entity = AgentToolEntity(**tool)
@ -174,7 +170,7 @@ class AppService:
tool_runtime=tool_runtime,
provider_name=agent_tool_entity.provider_id,
provider_type=agent_tool_entity.provider_type,
identity_id=f'AGENT.{app.id}'
identity_id=f"AGENT.{app.id}",
)
# get decrypted parameters
@ -185,7 +181,7 @@ class AppService:
masked_parameter = {}
# override tool parameters
tool['tool_parameters'] = masked_parameter
tool["tool_parameters"] = masked_parameter
except Exception as e:
pass
@ -215,12 +211,12 @@ class AppService:
:param args: request args
:return: App instance
"""
app.name = args.get('name')
app.description = args.get('description', '')
app.max_active_requests = args.get('max_active_requests')
app.icon_type = args.get('icon_type', 'emoji')
app.icon = args.get('icon')
app.icon_background = args.get('icon_background')
app.name = args.get("name")
app.description = args.get("description", "")
app.max_active_requests = args.get("max_active_requests")
app.icon_type = args.get("icon_type", "emoji")
app.icon = args.get("icon")
app.icon_background = args.get("icon_background")
app.updated_at = datetime.now(timezone.utc).replace(tzinfo=None)
db.session.commit()
@ -298,10 +294,7 @@ class AppService:
db.session.commit()
# Trigger asynchronous deletion of app and related data
remove_app_and_related_data_task.delay(
tenant_id=app.tenant_id,
app_id=app.id
)
remove_app_and_related_data_task.delay(tenant_id=app.tenant_id, app_id=app.id)
def get_app_meta(self, app_model: App) -> dict:
"""
@ -311,9 +304,7 @@ class AppService:
"""
app_mode = AppMode.value_of(app_model.mode)
meta = {
'tool_icons': {}
}
meta = {"tool_icons": {}}
if app_mode in [AppMode.ADVANCED_CHAT, AppMode.WORKFLOW]:
workflow = app_model.workflow
@ -321,17 +312,19 @@ class AppService:
return meta
graph = workflow.graph_dict
nodes = graph.get('nodes', [])
nodes = graph.get("nodes", [])
tools = []
for node in nodes:
if node.get('data', {}).get('type') == 'tool':
node_data = node.get('data', {})
tools.append({
'provider_type': node_data.get('provider_type'),
'provider_id': node_data.get('provider_id'),
'tool_name': node_data.get('tool_name'),
'tool_parameters': {}
})
if node.get("data", {}).get("type") == "tool":
node_data = node.get("data", {})
tools.append(
{
"provider_type": node_data.get("provider_type"),
"provider_id": node_data.get("provider_id"),
"tool_name": node_data.get("tool_name"),
"tool_parameters": {},
}
)
else:
app_model_config: AppModelConfig = app_model.app_model_config
@ -341,30 +334,26 @@ class AppService:
agent_config = app_model_config.agent_mode_dict or {}
# get all tools
tools = agent_config.get('tools', [])
tools = agent_config.get("tools", [])
url_prefix = (dify_config.CONSOLE_API_URL
+ "/console/api/workspaces/current/tool-provider/builtin/")
url_prefix = dify_config.CONSOLE_API_URL + "/console/api/workspaces/current/tool-provider/builtin/"
for tool in tools:
keys = list(tool.keys())
if len(keys) >= 4:
# current tool standard
provider_type = tool.get('provider_type')
provider_id = tool.get('provider_id')
tool_name = tool.get('tool_name')
if provider_type == 'builtin':
meta['tool_icons'][tool_name] = url_prefix + provider_id + '/icon'
elif provider_type == 'api':
provider_type = tool.get("provider_type")
provider_id = tool.get("provider_id")
tool_name = tool.get("tool_name")
if provider_type == "builtin":
meta["tool_icons"][tool_name] = url_prefix + provider_id + "/icon"
elif provider_type == "api":
try:
provider: ApiToolProvider = db.session.query(ApiToolProvider).filter(
ApiToolProvider.id == provider_id
).first()
meta['tool_icons'][tool_name] = json.loads(provider.icon)
provider: ApiToolProvider = (
db.session.query(ApiToolProvider).filter(ApiToolProvider.id == provider_id).first()
)
meta["tool_icons"][tool_name] = json.loads(provider.icon)
except:
meta['tool_icons'][tool_name] = {
"background": "#252525",
"content": "\ud83d\ude01"
}
meta["tool_icons"][tool_name] = {"background": "#252525", "content": "\ud83d\ude01"}
return meta

56
api/services/audio_service.py
View file

@ -17,7 +17,7 @@ from services.errors.audio import (
FILE_SIZE = 30
FILE_SIZE_LIMIT = FILE_SIZE * 1024 * 1024
ALLOWED_EXTENSIONS = ['mp3', 'mp4', 'mpeg', 'mpga', 'm4a', 'wav', 'webm', 'amr']
ALLOWED_EXTENSIONS = ["mp3", "mp4", "mpeg", "mpga", "m4a", "wav", "webm", "amr"]
logger = logging.getLogger(__name__)
@ -31,19 +31,19 @@ class AudioService:
raise ValueError("Speech to text is not enabled")
features_dict = workflow.features_dict
if 'speech_to_text' not in features_dict or not features_dict['speech_to_text'].get('enabled'):
if "speech_to_text" not in features_dict or not features_dict["speech_to_text"].get("enabled"):
raise ValueError("Speech to text is not enabled")
else:
app_model_config: AppModelConfig = app_model.app_model_config
if not app_model_config.speech_to_text_dict['enabled']:
if not app_model_config.speech_to_text_dict["enabled"]:
raise ValueError("Speech to text is not enabled")
if file is None:
raise NoAudioUploadedServiceError()
extension = file.mimetype
if extension not in [f'audio/{ext}' for ext in ALLOWED_EXTENSIONS]:
if extension not in [f"audio/{ext}" for ext in ALLOWED_EXTENSIONS]:
raise UnsupportedAudioTypeServiceError()
file_content = file.read()
@ -55,20 +55,25 @@ class AudioService:
model_manager = ModelManager()
model_instance = model_manager.get_default_model_instance(
tenant_id=app_model.tenant_id,
model_type=ModelType.SPEECH2TEXT
tenant_id=app_model.tenant_id, model_type=ModelType.SPEECH2TEXT
)
if model_instance is None:
raise ProviderNotSupportSpeechToTextServiceError()
buffer = io.BytesIO(file_content)
buffer.name = 'temp.mp3'
buffer.name = "temp.mp3"
return {"text": model_instance.invoke_speech2text(file=buffer, user=end_user)}
@classmethod
def transcript_tts(cls, app_model: App, text: Optional[str] = None,
voice: Optional[str] = None, end_user: Optional[str] = None, message_id: Optional[str] = None):
def transcript_tts(
cls,
app_model: App,
text: Optional[str] = None,
voice: Optional[str] = None,
end_user: Optional[str] = None,
message_id: Optional[str] = None,
):
from collections.abc import Generator
from flask import Response, stream_with_context
@ -84,65 +89,56 @@ class AudioService:
raise ValueError("TTS is not enabled")
features_dict = workflow.features_dict
if 'text_to_speech' not in features_dict or not features_dict['text_to_speech'].get('enabled'):
if "text_to_speech" not in features_dict or not features_dict["text_to_speech"].get("enabled"):
raise ValueError("TTS is not enabled")
voice = features_dict['text_to_speech'].get('voice') if voice is None else voice
voice = features_dict["text_to_speech"].get("voice") if voice is None else voice
else:
text_to_speech_dict = app_model.app_model_config.text_to_speech_dict
if not text_to_speech_dict.get('enabled'):
if not text_to_speech_dict.get("enabled"):
raise ValueError("TTS is not enabled")
voice = text_to_speech_dict.get('voice') if voice is None else voice
voice = text_to_speech_dict.get("voice") if voice is None else voice
model_manager = ModelManager()
model_instance = model_manager.get_default_model_instance(
tenant_id=app_model.tenant_id,
model_type=ModelType.TTS
tenant_id=app_model.tenant_id, model_type=ModelType.TTS
)
try:
if not voice:
voices = model_instance.get_tts_voices()
if voices:
voice = voices[0].get('value')
voice = voices[0].get("value")
else:
raise ValueError("Sorry, no voice available.")
return model_instance.invoke_tts(
content_text=text_content.strip(),
user=end_user,
tenant_id=app_model.tenant_id,
voice=voice
content_text=text_content.strip(), user=end_user, tenant_id=app_model.tenant_id, voice=voice
)
except Exception as e:
raise e
if message_id:
message = db.session.query(Message).filter(
Message.id == message_id
).first()
if message.answer == '' and message.status == 'normal':
message = db.session.query(Message).filter(Message.id == message_id).first()
if message.answer == "" and message.status == "normal":
return None
else:
response = invoke_tts(message.answer, app_model=app_model, voice=voice)
if isinstance(response, Generator):
return Response(stream_with_context(response), content_type='audio/mpeg')
return Response(stream_with_context(response), content_type="audio/mpeg")
return response
else:
response = invoke_tts(text, app_model, voice)
if isinstance(response, Generator):
return Response(stream_with_context(response), content_type='audio/mpeg')
return Response(stream_with_context(response), content_type="audio/mpeg")
return response
@classmethod
def transcript_tts_voices(cls, tenant_id: str, language: str):
model_manager = ModelManager()
model_instance = model_manager.get_default_model_instance(
tenant_id=tenant_id,
model_type=ModelType.TTS
)
model_instance = model_manager.get_default_model_instance(tenant_id=tenant_id, model_type=ModelType.TTS)
if model_instance is None:
raise ProviderNotSupportTextToSpeechServiceError()

6
api/services/auth/api_key_auth_factory.py
View file

@ -1,14 +1,12 @@
from services.auth.firecrawl import FirecrawlAuth
class ApiKeyAuthFactory:
def __init__(self, provider: str, credentials: dict):
if provider == 'firecrawl':
if provider == "firecrawl":
self.auth = FirecrawlAuth(credentials)
else:
raise ValueError('Invalid provider')
raise ValueError("Invalid provider")
def validate_credentials(self):
return self.auth.validate_credentials()

68
api/services/auth/api_key_auth_service.py
View file

@ -7,39 +7,43 @@ from services.auth.api_key_auth_factory import ApiKeyAuthFactory
class ApiKeyAuthService:
@staticmethod
def get_provider_auth_list(tenant_id: str) -> list:
data_source_api_key_bindings = db.session.query(DataSourceApiKeyAuthBinding).filter(
DataSourceApiKeyAuthBinding.tenant_id == tenant_id,
DataSourceApiKeyAuthBinding.disabled.is_(False)
).all()
data_source_api_key_bindings = (
db.session.query(DataSourceApiKeyAuthBinding)
.filter(DataSourceApiKeyAuthBinding.tenant_id == tenant_id, DataSourceApiKeyAuthBinding.disabled.is_(False))
.all()
)
return data_source_api_key_bindings
@staticmethod
def create_provider_auth(tenant_id: str, args: dict):
auth_result = ApiKeyAuthFactory(args['provider'], args['credentials']).validate_credentials()
auth_result = ApiKeyAuthFactory(args["provider"], args["credentials"]).validate_credentials()
if auth_result:
# Encrypt the api key
api_key = encrypter.encrypt_token(tenant_id, args['credentials']['config']['api_key'])
args['credentials']['config']['api_key'] = api_key
api_key = encrypter.encrypt_token(tenant_id, args["credentials"]["config"]["api_key"])
args["credentials"]["config"]["api_key"] = api_key
data_source_api_key_binding = DataSourceApiKeyAuthBinding()
data_source_api_key_binding.tenant_id = tenant_id
data_source_api_key_binding.category = args['category']
data_source_api_key_binding.provider = args['provider']
data_source_api_key_binding.credentials = json.dumps(args['credentials'], ensure_ascii=False)
data_source_api_key_binding.category = args["category"]
data_source_api_key_binding.provider = args["provider"]
data_source_api_key_binding.credentials = json.dumps(args["credentials"], ensure_ascii=False)
db.session.add(data_source_api_key_binding)
db.session.commit()
@staticmethod
def get_auth_credentials(tenant_id: str, category: str, provider: str):
data_source_api_key_bindings = db.session.query(DataSourceApiKeyAuthBinding).filter(
DataSourceApiKeyAuthBinding.tenant_id == tenant_id,
DataSourceApiKeyAuthBinding.category == category,
DataSourceApiKeyAuthBinding.provider == provider,
DataSourceApiKeyAuthBinding.disabled.is_(False)
).first()
data_source_api_key_bindings = (
db.session.query(DataSourceApiKeyAuthBinding)
.filter(
DataSourceApiKeyAuthBinding.tenant_id == tenant_id,
DataSourceApiKeyAuthBinding.category == category,
DataSourceApiKeyAuthBinding.provider == provider,
DataSourceApiKeyAuthBinding.disabled.is_(False),
)
.first()
)
if not data_source_api_key_bindings:
return None
credentials = json.loads(data_source_api_key_bindings.credentials)
@ -47,24 +51,24 @@ class ApiKeyAuthService:
@staticmethod
def delete_provider_auth(tenant_id: str, binding_id: str):
data_source_api_key_binding = db.session.query(DataSourceApiKeyAuthBinding).filter(
DataSourceApiKeyAuthBinding.tenant_id == tenant_id,
DataSourceApiKeyAuthBinding.id == binding_id
).first()
data_source_api_key_binding = (
db.session.query(DataSourceApiKeyAuthBinding)
.filter(DataSourceApiKeyAuthBinding.tenant_id == tenant_id, DataSourceApiKeyAuthBinding.id == binding_id)
.first()
)
if data_source_api_key_binding:
db.session.delete(data_source_api_key_binding)
db.session.commit()
@classmethod
def validate_api_key_auth_args(cls, args):
if 'category' not in args or not args['category']:
raise ValueError('category is required')
if 'provider' not in args or not args['provider']:
raise ValueError('provider is required')
if 'credentials' not in args or not args['credentials']:
raise ValueError('credentials is required')
if not isinstance(args['credentials'], dict):
raise ValueError('credentials must be a dictionary')
if 'auth_type' not in args['credentials'] or not args['credentials']['auth_type']:
raise ValueError('auth_type is required')
if "category" not in args or not args["category"]:
raise ValueError("category is required")
if "provider" not in args or not args["provider"]:
raise ValueError("provider is required")
if "credentials" not in args or not args["credentials"]:
raise ValueError("credentials is required")
if not isinstance(args["credentials"], dict):
raise ValueError("credentials must be a dictionary")
if "auth_type" not in args["credentials"] or not args["credentials"]["auth_type"]:
raise ValueError("auth_type is required")

41
api/services/auth/firecrawl.py
View file

@ -8,49 +8,40 @@ from services.auth.api_key_auth_base import ApiKeyAuthBase
class FirecrawlAuth(ApiKeyAuthBase):
def __init__(self, credentials: dict):
super().__init__(credentials)
auth_type = credentials.get('auth_type')
if auth_type != 'bearer':
raise ValueError('Invalid auth type, Firecrawl auth type must be Bearer')
self.api_key = credentials.get('config').get('api_key', None)
self.base_url = credentials.get('config').get('base_url', 'https://api.firecrawl.dev')
auth_type = credentials.get("auth_type")
if auth_type != "bearer":
raise ValueError("Invalid auth type, Firecrawl auth type must be Bearer")
self.api_key = credentials.get("config").get("api_key", None)
self.base_url = credentials.get("config").get("base_url", "https://api.firecrawl.dev")
if not self.api_key:
raise ValueError('No API key provided')
raise ValueError("No API key provided")
def validate_credentials(self):
headers = self._prepare_headers()
options = {
'url': 'https://example.com',
'crawlerOptions': {
'excludes': [],
'includes': [],
'limit': 1
},
'pageOptions': {
'onlyMainContent': True
}
"url": "https://example.com",
"crawlerOptions": {"excludes": [], "includes": [], "limit": 1},
"pageOptions": {"onlyMainContent": True},
}
response = self._post_request(f'{self.base_url}/v0/crawl', options, headers)
response = self._post_request(f"{self.base_url}/v0/crawl", options, headers)
if response.status_code == 200:
return True
else:
self._handle_error(response)
def _prepare_headers(self):
return {
'Content-Type': 'application/json',
'Authorization': f'Bearer {self.api_key}'
}
return {"Content-Type": "application/json", "Authorization": f"Bearer {self.api_key}"}
def _post_request(self, url, data, headers):
return requests.post(url, headers=headers, json=data)
def _handle_error(self, response):
if response.status_code in [402, 409, 500]:
error_message = response.json().get('error', 'Unknown error occurred')
raise Exception(f'Failed to authorize. Status code: {response.status_code}. Error: {error_message}')
error_message = response.json().get("error", "Unknown error occurred")
raise Exception(f"Failed to authorize. Status code: {response.status_code}. Error: {error_message}")
else:
if response.text:
error_message = json.loads(response.text).get('error', 'Unknown error occurred')
raise Exception(f'Failed to authorize. Status code: {response.status_code}. Error: {error_message}')
raise Exception(f'Unexpected error occurred while trying to authorize. Status code: {response.status_code}')
error_message = json.loads(response.text).get("error", "Unknown error occurred")
raise Exception(f"Failed to authorize. Status code: {response.status_code}. Error: {error_message}")
raise Exception(f"Unexpected error occurred while trying to authorize. Status code: {response.status_code}")

63
api/services/billing_service.py
View file

@ -7,58 +7,40 @@ from models.account import TenantAccountJoin, TenantAccountRole
class BillingService:
base_url = os.environ.get('BILLING_API_URL', 'BILLING_API_URL')
secret_key = os.environ.get('BILLING_API_SECRET_KEY', 'BILLING_API_SECRET_KEY')
base_url = os.environ.get("BILLING_API_URL", "BILLING_API_URL")
secret_key = os.environ.get("BILLING_API_SECRET_KEY", "BILLING_API_SECRET_KEY")
@classmethod
def get_info(cls, tenant_id: str):
params = {'tenant_id': tenant_id}
params = {"tenant_id": tenant_id}
billing_info = cls._send_request('GET', '/subscription/info', params=params)
billing_info = cls._send_request("GET", "/subscription/info", params=params)
return billing_info
@classmethod
def get_subscription(cls, plan: str,
interval: str,
prefilled_email: str = '',
tenant_id: str = ''):
params = {
'plan': plan,
'interval': interval,
'prefilled_email': prefilled_email,
'tenant_id': tenant_id
}
return cls._send_request('GET', '/subscription/payment-link', params=params)
def get_subscription(cls, plan: str, interval: str, prefilled_email: str = "", tenant_id: str = ""):
params = {"plan": plan, "interval": interval, "prefilled_email": prefilled_email, "tenant_id": tenant_id}
return cls._send_request("GET", "/subscription/payment-link", params=params)
@classmethod
def get_model_provider_payment_link(cls,
provider_name: str,
tenant_id: str,
account_id: str,
prefilled_email: str):
def get_model_provider_payment_link(cls, provider_name: str, tenant_id: str, account_id: str, prefilled_email: str):
params = {
'provider_name': provider_name,
'tenant_id': tenant_id,
'account_id': account_id,
'prefilled_email': prefilled_email
"provider_name": provider_name,
"tenant_id": tenant_id,
"account_id": account_id,
"prefilled_email": prefilled_email,
}
return cls._send_request('GET', '/model-provider/payment-link', params=params)
return cls._send_request("GET", "/model-provider/payment-link", params=params)
@classmethod
def get_invoices(cls, prefilled_email: str = '', tenant_id: str = ''):
params = {
'prefilled_email': prefilled_email,
'tenant_id': tenant_id
}
return cls._send_request('GET', '/invoices', params=params)
def get_invoices(cls, prefilled_email: str = "", tenant_id: str = ""):
params = {"prefilled_email": prefilled_email, "tenant_id": tenant_id}
return cls._send_request("GET", "/invoices", params=params)
@classmethod
def _send_request(cls, method, endpoint, json=None, params=None):
headers = {
"Content-Type": "application/json",
"Billing-Api-Secret-Key": cls.secret_key
}
headers = {"Content-Type": "application/json", "Billing-Api-Secret-Key": cls.secret_key}
url = f"{cls.base_url}{endpoint}"
response = requests.request(method, url, json=json, params=params, headers=headers)
@ -69,10 +51,11 @@ class BillingService:
def is_tenant_owner_or_admin(current_user):
tenant_id = current_user.current_tenant_id
join = db.session.query(TenantAccountJoin).filter(
TenantAccountJoin.tenant_id == tenant_id,
TenantAccountJoin.account_id == current_user.id
).first()
join = (
db.session.query(TenantAccountJoin)
.filter(TenantAccountJoin.tenant_id == tenant_id, TenantAccountJoin.account_id == current_user.id)
.first()
)
if not TenantAccountRole.is_privileged_role(join.role):
raise ValueError('Only team owner or team admin can perform this action')
raise ValueError("Only team owner or team admin can perform this action")

15
api/services/code_based_extension_service.py
View file

@ -2,12 +2,15 @@ from extensions.ext_code_based_extension import code_based_extension
class CodeBasedExtensionService:
@staticmethod
def get_code_based_extension(module: str) -> list[dict]:
module_extensions = code_based_extension.module_extensions(module)
return [{
'name': module_extension.name,
'label': module_extension.label,
'form_schema': module_extension.form_schema
} for module_extension in module_extensions if not module_extension.builtin]
return [
{
"name": module_extension.name,
"label": module_extension.label,
"form_schema": module_extension.form_schema,
}
for module_extension in module_extensions
if not module_extension.builtin
]

79
api/services/conversation_service.py
View file

@ -15,22 +15,27 @@ from services.errors.message import MessageNotExistsError
class ConversationService:
@classmethod
def pagination_by_last_id(cls, app_model: App, user: Optional[Union[Account, EndUser]],
last_id: Optional[str], limit: int,
invoke_from: InvokeFrom,
include_ids: Optional[list] = None,
exclude_ids: Optional[list] = None,
sort_by: str = '-updated_at') -> InfiniteScrollPagination:
def pagination_by_last_id(
cls,
app_model: App,
user: Optional[Union[Account, EndUser]],
last_id: Optional[str],
limit: int,
invoke_from: InvokeFrom,
include_ids: Optional[list] = None,
exclude_ids: Optional[list] = None,
sort_by: str = "-updated_at",
) -> InfiniteScrollPagination:
if not user:
return InfiniteScrollPagination(data=[], limit=limit, has_more=False)
base_query = db.session.query(Conversation).filter(
Conversation.is_deleted == False,
Conversation.app_id == app_model.id,
Conversation.from_source == ('api' if isinstance(user, EndUser) else 'console'),
Conversation.from_source == ("api" if isinstance(user, EndUser) else "console"),
Conversation.from_end_user_id == (user.id if isinstance(user, EndUser) else None),
Conversation.from_account_id == (user.id if isinstance(user, Account) else None),
or_(Conversation.invoke_from.is_(None), Conversation.invoke_from == invoke_from.value)
or_(Conversation.invoke_from.is_(None), Conversation.invoke_from == invoke_from.value),
)
if include_ids is not None:
@ -58,28 +63,26 @@ class ConversationService:
has_more = False
if len(conversations) == limit:
current_page_last_conversation = conversations[-1]
rest_filter_condition = cls._build_filter_condition(sort_field, sort_direction,
current_page_last_conversation, is_next_page=True)
rest_filter_condition = cls._build_filter_condition(
sort_field, sort_direction, current_page_last_conversation, is_next_page=True
)
rest_count = base_query.filter(rest_filter_condition).count()
if rest_count > 0:
has_more = True
return InfiniteScrollPagination(
data=conversations,
limit=limit,
has_more=has_more
)
return InfiniteScrollPagination(data=conversations, limit=limit, has_more=has_more)
@classmethod
def _get_sort_params(cls, sort_by: str) -> tuple[str, callable]:
if sort_by.startswith('-'):
if sort_by.startswith("-"):
return sort_by[1:], desc
return sort_by, asc
@classmethod
def _build_filter_condition(cls, sort_field: str, sort_direction: callable, reference_conversation: Conversation,
is_next_page: bool = False):
def _build_filter_condition(
cls, sort_field: str, sort_direction: callable, reference_conversation: Conversation, is_next_page: bool = False
):
field_value = getattr(reference_conversation, sort_field)
if (sort_direction == desc and not is_next_page) or (sort_direction == asc and is_next_page):
return getattr(Conversation, sort_field) < field_value
@ -87,8 +90,14 @@ class ConversationService:
return getattr(Conversation, sort_field) > field_value
@classmethod
def rename(cls, app_model: App, conversation_id: str,
user: Optional[Union[Account, EndUser]], name: str, auto_generate: bool):
def rename(
cls,
app_model: App,
conversation_id: str,
user: Optional[Union[Account, EndUser]],
name: str,
auto_generate: bool,
):
conversation = cls.get_conversation(app_model, conversation_id, user)
if auto_generate:
@ -103,11 +112,12 @@ class ConversationService:
@classmethod
def auto_generate_name(cls, app_model: App, conversation: Conversation):
# get conversation first message
message = db.session.query(Message) \
.filter(
Message.app_id == app_model.id,
Message.conversation_id == conversation.id
).order_by(Message.created_at.asc()).first()
message = (
db.session.query(Message)
.filter(Message.app_id == app_model.id, Message.conversation_id == conversation.id)
.order_by(Message.created_at.asc())
.first()
)
if not message:
raise MessageNotExistsError()
@ -127,15 +137,18 @@ class ConversationService:
@classmethod
def get_conversation(cls, app_model: App, conversation_id: str, user: Optional[Union[Account, EndUser]]):
conversation = db.session.query(Conversation) \
conversation = (
db.session.query(Conversation)
.filter(
Conversation.id == conversation_id,
Conversation.app_id == app_model.id,
Conversation.from_source == ('api' if isinstance(user, EndUser) else 'console'),
Conversation.from_end_user_id == (user.id if isinstance(user, EndUser) else None),
Conversation.from_account_id == (user.id if isinstance(user, Account) else None),
Conversation.is_deleted == False
).first()
Conversation.id == conversation_id,
Conversation.app_id == app_model.id,
Conversation.from_source == ("api" if isinstance(user, EndUser) else "console"),
Conversation.from_end_user_id == (user.id if isinstance(user, EndUser) else None),
Conversation.from_account_id == (user.id if isinstance(user, Account) else None),
Conversation.is_deleted == False,
)
.first()
)
if not conversation:
raise ConversationNotExistsError()

952
api/services/dataset_service.py
View file

File diff suppressed because it is too large Load diff

9
api/services/enterprise/base.py
View file

@ -4,15 +4,12 @@ import requests
class EnterpriseRequest:
base_url = os.environ.get('ENTERPRISE_API_URL', 'ENTERPRISE_API_URL')
secret_key = os.environ.get('ENTERPRISE_API_SECRET_KEY', 'ENTERPRISE_API_SECRET_KEY')
base_url = os.environ.get("ENTERPRISE_API_URL", "ENTERPRISE_API_URL")
secret_key = os.environ.get("ENTERPRISE_API_SECRET_KEY", "ENTERPRISE_API_SECRET_KEY")
@classmethod
def send_request(cls, method, endpoint, json=None, params=None):
headers = {
"Content-Type": "application/json",
"Enterprise-Api-Secret-Key": cls.secret_key
}
headers = {"Content-Type": "application/json", "Enterprise-Api-Secret-Key": cls.secret_key}
url = f"{cls.base_url}{endpoint}"
response = requests.request(method, url, json=json, params=params, headers=headers)

Some files were not shown because too many files have changed in this diff Show more