opensource
/
dify
mirror of https://github.com/langgenius/dify.git
2
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

more typed orm (#28494) 

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
This commit is contained in:
Asuka Minato 2025-11-21 16:42:27 +09:00 committed by GitHub
parent 4486b54680
commit 237bb4595b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
8 changed files with 177 additions and 141 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
api/core/app/apps/workflow/generate_task_pipeline.py
View File

@ -644,14 +644,15 @@ class WorkflowAppGenerateTaskPipeline(GraphRuntimeStateSupport):
if not workflow_run_id:
return
workflow_app_log = WorkflowAppLog()
workflow_app_log.tenant_id = self._application_generate_entity.app_config.tenant_id
workflow_app_log.app_id = self._application_generate_entity.app_config.app_id
workflow_app_log.workflow_id = self._workflow.id
workflow_app_log.workflow_run_id = workflow_run_id
workflow_app_log.created_from = created_from.value
workflow_app_log.created_by_role = self._created_by_role
workflow_app_log.created_by = self._user_id
workflow_app_log = WorkflowAppLog(
tenant_id=self._application_generate_entity.app_config.tenant_id,
app_id=self._application_generate_entity.app_config.app_id,
workflow_id=self._workflow.id,
workflow_run_id=workflow_run_id,
created_from=created_from.value,
created_by_role=self._created_by_role,
created_by=self._user_id,
)
session.add(workflow_app_log)
session.commit()

66
api/models/model.py
View File

@ -24,7 +24,7 @@ from libs.helper import generate_string # type: ignore[import-not-found]
from libs.uuid_utils import uuidv7
from .account import Account, Tenant
from .base import Base
from .base import Base, TypeBase
from .engine import db
from .enums import CreatorUserRole
from .provider_ids import GenericProviderID
@ -34,12 +34,14 @@ if TYPE_CHECKING:
from models.workflow import Workflow
class DifySetup(Base):
class DifySetup(TypeBase):
__tablename__ = "dify_setups"
__table_args__ = (sa.PrimaryKeyConstraint("version", name="dify_setup_pkey"),)
version: Mapped[str] = mapped_column(String(255), nullable=False)
setup_at = mapped_column(sa.DateTime, nullable=False, server_default=func.current_timestamp())
setup_at: Mapped[datetime] = mapped_column(
sa.DateTime, nullable=False, server_default=func.current_timestamp(), init=False
)
class AppMode(StrEnum):
@ -561,7 +563,7 @@ class RecommendedApp(Base):
return app
class InstalledApp(Base):
class InstalledApp(TypeBase):
__tablename__ = "installed_apps"
__table_args__ = (
sa.PrimaryKeyConstraint("id", name="installed_app_pkey"),
@ -570,14 +572,16 @@ class InstalledApp(Base):
sa.UniqueConstraint("tenant_id", "app_id", name="unique_tenant_app"),
)
id = mapped_column(StringUUID, default=lambda: str(uuid4()))
tenant_id = mapped_column(StringUUID, nullable=False)
app_id = mapped_column(StringUUID, nullable=False)
app_owner_tenant_id = mapped_column(StringUUID, nullable=False)
id: Mapped[str] = mapped_column(StringUUID, default=lambda: str(uuid4()), init=False)
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
app_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
app_owner_tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
position: Mapped[int] = mapped_column(sa.Integer, nullable=False, default=0)
is_pinned: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, server_default=sa.text("false"))
last_used_at = mapped_column(sa.DateTime, nullable=True)
created_at = mapped_column(sa.DateTime, nullable=False, server_default=func.current_timestamp())
is_pinned: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, server_default=sa.text("false"), default=False)
last_used_at: Mapped[datetime | None] = mapped_column(sa.DateTime, nullable=True, default=None)
created_at: Mapped[datetime] = mapped_column(
sa.DateTime, nullable=False, server_default=func.current_timestamp(), init=False
)
@property
def app(self) -> App | None:
@ -1533,25 +1537,31 @@ class EndUser(Base, UserMixin):
)
class AppMCPServer(Base):
class AppMCPServer(TypeBase):
__tablename__ = "app_mcp_servers"
__table_args__ = (
sa.PrimaryKeyConstraint("id", name="app_mcp_server_pkey"),
sa.UniqueConstraint("tenant_id", "app_id", name="unique_app_mcp_server_tenant_app_id"),
sa.UniqueConstraint("server_code", name="unique_app_mcp_server_server_code"),
)
id = mapped_column(StringUUID, default=lambda: str(uuid4()))
tenant_id = mapped_column(StringUUID, nullable=False)
app_id = mapped_column(StringUUID, nullable=False)
id: Mapped[str] = mapped_column(StringUUID, default=lambda: str(uuid4()), init=False)
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
app_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
name: Mapped[str] = mapped_column(String(255), nullable=False)
description: Mapped[str] = mapped_column(String(255), nullable=False)
server_code: Mapped[str] = mapped_column(String(255), nullable=False)
status = mapped_column(String(255), nullable=False, server_default=sa.text("'normal'"))
parameters = mapped_column(LongText, nullable=False)
status: Mapped[str] = mapped_column(String(255), nullable=False, server_default=sa.text("'normal'"))
parameters: Mapped[str] = mapped_column(LongText, nullable=False)
created_at = mapped_column(sa.DateTime, nullable=False, server_default=func.current_timestamp())
updated_at = mapped_column(
sa.DateTime, nullable=False, server_default=func.current_timestamp(), onupdate=func.current_timestamp()
created_at: Mapped[datetime] = mapped_column(
sa.DateTime, nullable=False, server_default=func.current_timestamp(), init=False
)
updated_at: Mapped[datetime] = mapped_column(
sa.DateTime,
nullable=False,
server_default=func.current_timestamp(),
onupdate=func.current_timestamp(),
init=False,
)
@staticmethod
@ -1928,7 +1938,7 @@ class Tag(Base):
created_at = mapped_column(sa.DateTime, nullable=False, server_default=func.current_timestamp())
class TagBinding(Base):
class TagBinding(TypeBase):
__tablename__ = "tag_bindings"
__table_args__ = (
sa.PrimaryKeyConstraint("id", name="tag_binding_pkey"),
@ -1936,12 +1946,14 @@ class TagBinding(Base):
sa.Index("tag_bind_tag_id_idx", "tag_id"),
)
id = mapped_column(StringUUID, default=lambda: str(uuid4()))
tenant_id = mapped_column(StringUUID, nullable=True)
tag_id = mapped_column(StringUUID, nullable=True)
target_id = mapped_column(StringUUID, nullable=True)
created_by = mapped_column(StringUUID, nullable=False)
created_at = mapped_column(sa.DateTime, nullable=False, server_default=func.current_timestamp())
id: Mapped[str] = mapped_column(StringUUID, default=lambda: str(uuid4()), init=False)
tenant_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True)
tag_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True)
target_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True)
created_by: Mapped[str] = mapped_column(StringUUID, nullable=False)
created_at: Mapped[datetime] = mapped_column(
sa.DateTime, nullable=False, server_default=func.current_timestamp(), init=False
)
class TraceAppConfig(Base):

50
api/models/provider.py
View File

@ -118,7 +118,7 @@ class Provider(TypeBase):
return self.is_valid and self.token_is_set
class ProviderModel(Base):
class ProviderModel(TypeBase):
"""
Provider model representing the API provider_models and their configurations.
"""
@ -132,16 +132,18 @@ class ProviderModel(Base):
),
)
id: Mapped[str] = mapped_column(StringUUID, default=lambda: str(uuid4()))
id: Mapped[str] = mapped_column(StringUUID, default=lambda: str(uuid4()), init=False)
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
provider_name: Mapped[str] = mapped_column(String(255), nullable=False)
model_name: Mapped[str] = mapped_column(String(255), nullable=False)
model_type: Mapped[str] = mapped_column(String(40), nullable=False)
credential_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True)
is_valid: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, server_default=text("false"))
created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp())
credential_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True, default=None)
is_valid: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, server_default=text("false"), default=False)
created_at: Mapped[datetime] = mapped_column(
DateTime, nullable=False, server_default=func.current_timestamp(), init=False
)
updated_at: Mapped[datetime] = mapped_column(
DateTime, nullable=False, server_default=func.current_timestamp(), onupdate=func.current_timestamp()
DateTime, nullable=False, server_default=func.current_timestamp(), onupdate=func.current_timestamp(), init=False
)
@cached_property
@ -182,31 +184,33 @@ class TenantDefaultModel(Base):
)
class TenantPreferredModelProvider(Base):
class TenantPreferredModelProvider(TypeBase):
__tablename__ = "tenant_preferred_model_providers"
__table_args__ = (
sa.PrimaryKeyConstraint("id", name="tenant_preferred_model_provider_pkey"),
sa.Index("tenant_preferred_model_provider_tenant_provider_idx", "tenant_id", "provider_name"),
)
id: Mapped[str] = mapped_column(StringUUID, default=lambda: str(uuid4()))
id: Mapped[str] = mapped_column(StringUUID, default=lambda: str(uuid4()), init=False)
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
provider_name: Mapped[str] = mapped_column(String(255), nullable=False)
preferred_provider_type: Mapped[str] = mapped_column(String(40), nullable=False)
created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp())
created_at: Mapped[datetime] = mapped_column(
DateTime, nullable=False, server_default=func.current_timestamp(), init=False
)
updated_at: Mapped[datetime] = mapped_column(
DateTime, nullable=False, server_default=func.current_timestamp(), onupdate=func.current_timestamp()
DateTime, nullable=False, server_default=func.current_timestamp(), onupdate=func.current_timestamp(), init=False
)
class ProviderOrder(Base):
class ProviderOrder(TypeBase):
__tablename__ = "provider_orders"
__table_args__ = (
sa.PrimaryKeyConstraint("id", name="provider_order_pkey"),
sa.Index("provider_order_tenant_provider_idx", "tenant_id", "provider_name"),
)
id: Mapped[str] = mapped_column(StringUUID, default=lambda: str(uuid4()))
id: Mapped[str] = mapped_column(StringUUID, default=lambda: str(uuid4()), init=False)
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
provider_name: Mapped[str] = mapped_column(String(255), nullable=False)
account_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
@ -220,13 +224,15 @@ class ProviderOrder(Base):
paid_at: Mapped[datetime | None] = mapped_column(DateTime)
pay_failed_at: Mapped[datetime | None] = mapped_column(DateTime)
refunded_at: Mapped[datetime | None] = mapped_column(DateTime)
created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp())
created_at: Mapped[datetime] = mapped_column(
DateTime, nullable=False, server_default=func.current_timestamp(), init=False
)
updated_at: Mapped[datetime] = mapped_column(
DateTime, nullable=False, server_default=func.current_timestamp(), onupdate=func.current_timestamp()
DateTime, nullable=False, server_default=func.current_timestamp(), onupdate=func.current_timestamp(), init=False
)
class ProviderModelSetting(Base):
class ProviderModelSetting(TypeBase):
"""
Provider model settings for record the model enabled status and load balancing status.
"""
@ -237,16 +243,20 @@ class ProviderModelSetting(Base):
sa.Index("provider_model_setting_tenant_provider_model_idx", "tenant_id", "provider_name", "model_type"),
)
id: Mapped[str] = mapped_column(StringUUID, default=lambda: str(uuid4()))
id: Mapped[str] = mapped_column(StringUUID, default=lambda: str(uuid4()), init=False)
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
provider_name: Mapped[str] = mapped_column(String(255), nullable=False)
model_name: Mapped[str] = mapped_column(String(255), nullable=False)
model_type: Mapped[str] = mapped_column(String(40), nullable=False)
enabled: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, server_default=text("true"))
load_balancing_enabled: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, server_default=text("false"))
created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp())
enabled: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, server_default=text("true"), default=True)
load_balancing_enabled: Mapped[bool] = mapped_column(
sa.Boolean, nullable=False, server_default=text("false"), default=False
)
created_at: Mapped[datetime] = mapped_column(
DateTime, nullable=False, server_default=func.current_timestamp(), init=False
)
updated_at: Mapped[datetime] = mapped_column(
DateTime, nullable=False, server_default=func.current_timestamp(), onupdate=func.current_timestamp()
DateTime, nullable=False, server_default=func.current_timestamp(), onupdate=func.current_timestamp(), init=False
)

29
api/models/trigger.py
View File

@ -271,7 +271,7 @@ class WorkflowTriggerLog(Base):
}
class WorkflowWebhookTrigger(Base):
class WorkflowWebhookTrigger(TypeBase):
"""
Workflow Webhook Trigger
@ -294,18 +294,21 @@ class WorkflowWebhookTrigger(Base):
sa.UniqueConstraint("webhook_id", name="uniq_webhook_id"),
)
id: Mapped[str] = mapped_column(StringUUID, default=lambda: str(uuidv7()))
id: Mapped[str] = mapped_column(StringUUID, default=lambda: str(uuidv7()), init=False)
app_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
node_id: Mapped[str] = mapped_column(String(64), nullable=False)
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
webhook_id: Mapped[str] = mapped_column(String(24), nullable=False)
created_by: Mapped[str] = mapped_column(StringUUID, nullable=False)
created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp())
created_at: Mapped[datetime] = mapped_column(
DateTime, nullable=False, server_default=func.current_timestamp(), init=False
)
updated_at: Mapped[datetime] = mapped_column(
DateTime,
nullable=False,
server_default=func.current_timestamp(),
server_onupdate=func.current_timestamp(),
init=False,
)
@cached_property
@ -323,7 +326,7 @@ class WorkflowWebhookTrigger(Base):
return generate_webhook_trigger_endpoint(self.webhook_id, True)
class WorkflowPluginTrigger(Base):
class WorkflowPluginTrigger(TypeBase):
"""
Workflow Plugin Trigger
@ -348,23 +351,26 @@ class WorkflowPluginTrigger(Base):
sa.UniqueConstraint("app_id", "node_id", name="uniq_app_node_subscription"),
)
id: Mapped[str] = mapped_column(StringUUID, default=lambda: str(uuid4()))
id: Mapped[str] = mapped_column(StringUUID, default=lambda: str(uuid4()), init=False)
app_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
node_id: Mapped[str] = mapped_column(String(64), nullable=False)
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
provider_id: Mapped[str] = mapped_column(String(512), nullable=False)
event_name: Mapped[str] = mapped_column(String(255), nullable=False)
subscription_id: Mapped[str] = mapped_column(String(255), nullable=False)
created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp())
created_at: Mapped[datetime] = mapped_column(
DateTime, nullable=False, server_default=func.current_timestamp(), init=False
)
updated_at: Mapped[datetime] = mapped_column(
DateTime,
nullable=False,
server_default=func.current_timestamp(),
server_onupdate=func.current_timestamp(),
init=False,
)
class AppTrigger(Base):
class AppTrigger(TypeBase):
"""
App Trigger
@ -389,22 +395,25 @@ class AppTrigger(Base):
sa.Index("app_trigger_tenant_app_idx", "tenant_id", "app_id"),
)
id: Mapped[str] = mapped_column(StringUUID, default=lambda: str(uuidv7()))
id: Mapped[str] = mapped_column(StringUUID, default=lambda: str(uuidv7()), init=False)
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
app_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
node_id: Mapped[str | None] = mapped_column(String(64), nullable=False)
trigger_type: Mapped[str] = mapped_column(EnumText(AppTriggerType, length=50), nullable=False)
title: Mapped[str] = mapped_column(String(255), nullable=False)
provider_name: Mapped[str] = mapped_column(String(255), server_default="", nullable=True)
provider_name: Mapped[str] = mapped_column(String(255), server_default="", default="") # why it is nullable?
status: Mapped[str] = mapped_column(
EnumText(AppTriggerStatus, length=50), nullable=False, default=AppTriggerStatus.ENABLED
)
created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp())
created_at: Mapped[datetime] = mapped_column(
DateTime, nullable=False, server_default=func.current_timestamp(), init=False
)
updated_at: Mapped[datetime] = mapped_column(
DateTime,
nullable=False,
default=naive_utc_now(),
server_onupdate=func.current_timestamp(),
init=False,
)

43
api/models/workflow.py
View File

@ -7,7 +7,19 @@ from typing import TYPE_CHECKING, Any, Optional, Union, cast
from uuid import uuid4
import sqlalchemy as sa
from sqlalchemy import DateTime, Select, exists, orm, select
from sqlalchemy import (
DateTime,
Index,
PrimaryKeyConstraint,
Select,
String,
UniqueConstraint,
exists,
func,
orm,
select,
)
from sqlalchemy.orm import Mapped, declared_attr, mapped_column
from core.file.constants import maybe_file_object
from core.file.models import File
@ -26,10 +38,8 @@ from libs.uuid_utils import uuidv7
from ._workflow_exc import NodeNotFoundError, WorkflowDataError
if TYPE_CHECKING:
from models.model import AppMode, UploadFile
from .model import AppMode, UploadFile
from sqlalchemy import Index, PrimaryKeyConstraint, String, UniqueConstraint, func
from sqlalchemy.orm import Mapped, declared_attr, mapped_column
from constants import DEFAULT_FILE_NUMBER_LIMITS, HIDDEN_VALUE
from core.helper import encrypter
@ -38,7 +48,7 @@ from factories import variable_factory
from libs import helper
from .account import Account
from .base import Base, DefaultFieldsMixin
from .base import Base, DefaultFieldsMixin, TypeBase
from .engine import db
from .enums import CreatorUserRole, DraftVariableType, ExecutionOffLoadType
from .types import EnumText, LongText, StringUUID
@ -1052,7 +1062,7 @@ class WorkflowAppLogCreatedFrom(StrEnum):
raise ValueError(f"invalid workflow app log created from value {value}")
class WorkflowAppLog(Base):
class WorkflowAppLog(TypeBase):
"""
Workflow App execution log, excluding workflow debugging records.
@ -1088,7 +1098,7 @@ class WorkflowAppLog(Base):
sa.Index("workflow_app_log_workflow_run_id_idx", "workflow_run_id"),
)
id: Mapped[str] = mapped_column(StringUUID, default=lambda: str(uuid4()))
id: Mapped[str] = mapped_column(StringUUID, default=lambda: str(uuid4()), init=False)
tenant_id: Mapped[str] = mapped_column(StringUUID)
app_id: Mapped[str] = mapped_column(StringUUID)
workflow_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
@ -1096,7 +1106,9 @@ class WorkflowAppLog(Base):
created_from: Mapped[str] = mapped_column(String(255), nullable=False)
created_by_role: Mapped[str] = mapped_column(String(255), nullable=False)
created_by: Mapped[str] = mapped_column(StringUUID, nullable=False)
created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp())
created_at: Mapped[datetime] = mapped_column(
DateTime, nullable=False, server_default=func.current_timestamp(), init=False
)
@property
def workflow_run(self):
@ -1137,7 +1149,7 @@ class WorkflowAppLog(Base):
}
class ConversationVariable(Base):
class ConversationVariable(TypeBase):
__tablename__ = "workflow_conversation_variables"
id: Mapped[str] = mapped_column(StringUUID, primary_key=True)
@ -1145,21 +1157,12 @@ class ConversationVariable(Base):
app_id: Mapped[str] = mapped_column(StringUUID, nullable=False, index=True)
data: Mapped[str] = mapped_column(LongText, nullable=False)
created_at: Mapped[datetime] = mapped_column(
DateTime, nullable=False, server_default=func.current_timestamp(), index=True
DateTime, nullable=False, server_default=func.current_timestamp(), index=True, init=False
)
updated_at: Mapped[datetime] = mapped_column(
DateTime,
nullable=False,
server_default=func.current_timestamp(),
onupdate=func.current_timestamp(),
DateTime, nullable=False, server_default=func.current_timestamp(), onupdate=func.current_timestamp(), init=False
)
def __init__(self, *, id: str, app_id: str, conversation_id: str, data: str):
self.id = id
self.app_id = app_id
self.conversation_id = conversation_id
self.data = data
@classmethod
def from_variable(cls, *, app_id: str, conversation_id: str, variable: Variable) -> "ConversationVariable":
obj = cls(

4
api/tests/test_containers_integration_tests/services/test_webhook_service.py
View File

@ -67,6 +67,7 @@ class TestWebhookService:
)
TenantService.create_owner_tenant_if_not_exist(account, name=fake.company())
tenant = account.current_tenant
assert tenant is not None
# Create app
app = App(
@ -131,7 +132,7 @@ class TestWebhookService:
app_id=app.id,
node_id="webhook_node",
tenant_id=tenant.id,
webhook_id=webhook_id,
webhook_id=str(webhook_id),
created_by=account.id,
)
db_session_with_containers.add(webhook_trigger)
@ -143,6 +144,7 @@ class TestWebhookService:
app_id=app.id,
node_id="webhook_node",
trigger_type=AppTriggerType.TRIGGER_WEBHOOK,
provider_name="webhook",
title="Test Webhook",
status=AppTriggerStatus.ENABLED,
)

44
api/tests/test_containers_integration_tests/services/test_workflow_app_service.py
View File

@ -209,7 +209,6 @@ class TestWorkflowAppService:
# Create workflow app log
workflow_app_log = WorkflowAppLog(
id=str(uuid.uuid4()),
tenant_id=app.tenant_id,
app_id=app.id,
workflow_id=workflow.id,
@ -217,8 +216,9 @@ class TestWorkflowAppService:
created_from="service-api",
created_by_role=CreatorUserRole.ACCOUNT,
created_by=account.id,
created_at=datetime.now(UTC),
)
workflow_app_log.id = str(uuid.uuid4())
workflow_app_log.created_at = datetime.now(UTC)
db.session.add(workflow_app_log)
db.session.commit()
@ -365,7 +365,6 @@ class TestWorkflowAppService:
db.session.commit()
workflow_app_log = WorkflowAppLog(
id=str(uuid.uuid4()),
tenant_id=app.tenant_id,
app_id=app.id,
workflow_id=workflow.id,
@ -373,8 +372,9 @@ class TestWorkflowAppService:
created_from="service-api",
created_by_role=CreatorUserRole.ACCOUNT,
created_by=account.id,
created_at=datetime.now(UTC) + timedelta(minutes=i),
)
workflow_app_log.id = str(uuid.uuid4())
workflow_app_log.created_at = datetime.now(UTC) + timedelta(minutes=i)
db.session.add(workflow_app_log)
db.session.commit()
@ -473,7 +473,6 @@ class TestWorkflowAppService:
db.session.commit()
workflow_app_log = WorkflowAppLog(
id=str(uuid.uuid4()),
tenant_id=app.tenant_id,
app_id=app.id,
workflow_id=workflow.id,
@ -481,8 +480,9 @@ class TestWorkflowAppService:
created_from="service-api",
created_by_role=CreatorUserRole.ACCOUNT,
created_by=account.id,
created_at=timestamp,
)
workflow_app_log.id = str(uuid.uuid4())
workflow_app_log.created_at = timestamp
db.session.add(workflow_app_log)
db.session.commit()
@ -580,7 +580,6 @@ class TestWorkflowAppService:
db.session.commit()
workflow_app_log = WorkflowAppLog(
id=str(uuid.uuid4()),
tenant_id=app.tenant_id,
app_id=app.id,
workflow_id=workflow.id,
@ -588,8 +587,9 @@ class TestWorkflowAppService:
created_from="service-api",
created_by_role=CreatorUserRole.ACCOUNT,
created_by=account.id,
created_at=datetime.now(UTC) + timedelta(minutes=i),
)
workflow_app_log.id = str(uuid.uuid4())
workflow_app_log.created_at = datetime.now(UTC) + timedelta(minutes=i)
db.session.add(workflow_app_log)
db.session.commit()
@ -710,7 +710,6 @@ class TestWorkflowAppService:
db.session.commit()
workflow_app_log = WorkflowAppLog(
id=str(uuid.uuid4()),
tenant_id=app.tenant_id,
app_id=app.id,
workflow_id=workflow.id,
@ -718,8 +717,9 @@ class TestWorkflowAppService:
created_from="service-api",
created_by_role=CreatorUserRole.ACCOUNT,
created_by=account.id,
created_at=datetime.now(UTC) + timedelta(minutes=i),
)
workflow_app_log.id = str(uuid.uuid4())
workflow_app_log.created_at = datetime.now(UTC) + timedelta(minutes=i)
db.session.add(workflow_app_log)
db.session.commit()
@ -752,7 +752,6 @@ class TestWorkflowAppService:
db.session.commit()
workflow_app_log = WorkflowAppLog(
id=str(uuid.uuid4()),
tenant_id=app.tenant_id,
app_id=app.id,
workflow_id=workflow.id,
@ -760,8 +759,9 @@ class TestWorkflowAppService:
created_from="web-app",
created_by_role=CreatorUserRole.END_USER,
created_by=end_user.id,
created_at=datetime.now(UTC) + timedelta(minutes=i + 10),
)
workflow_app_log.id = str(uuid.uuid4())
workflow_app_log.created_at = datetime.now(UTC) + timedelta(minutes=i + 10)
db.session.add(workflow_app_log)
db.session.commit()
@ -889,7 +889,6 @@ class TestWorkflowAppService:
# Create workflow app log
workflow_app_log = WorkflowAppLog(
id=str(uuid.uuid4()),
tenant_id=app.tenant_id,
app_id=app.id,
workflow_id=workflow.id,
@ -897,8 +896,9 @@ class TestWorkflowAppService:
created_from="service-api",
created_by_role=CreatorUserRole.ACCOUNT,
created_by=account.id,
created_at=datetime.now(UTC),
)
workflow_app_log.id = str(uuid.uuid4())
workflow_app_log.created_at = datetime.now(UTC)
db.session.add(workflow_app_log)
db.session.commit()
@ -979,7 +979,6 @@ class TestWorkflowAppService:
# Create workflow app log
workflow_app_log = WorkflowAppLog(
id=str(uuid.uuid4()),
tenant_id=app.tenant_id,
app_id=app.id,
workflow_id=workflow.id,
@ -987,8 +986,9 @@ class TestWorkflowAppService:
created_from="service-api",
created_by_role=CreatorUserRole.ACCOUNT,
created_by=account.id,
created_at=datetime.now(UTC),
)
workflow_app_log.id = str(uuid.uuid4())
workflow_app_log.created_at = datetime.now(UTC)
db.session.add(workflow_app_log)
db.session.commit()
@ -1133,7 +1133,6 @@ class TestWorkflowAppService:
db_session_with_containers.flush()
log = WorkflowAppLog(
id=str(uuid.uuid4()),
tenant_id=app.tenant_id,
app_id=app.id,
workflow_id=workflow.id,
@ -1141,8 +1140,9 @@ class TestWorkflowAppService:
created_from="service-api",
created_by_role=CreatorUserRole.ACCOUNT,
created_by=account.id,
created_at=datetime.now(UTC) + timedelta(minutes=i),
)
log.id = str(uuid.uuid4())
log.created_at = datetime.now(UTC) + timedelta(minutes=i)
db_session_with_containers.add(log)
logs_data.append((log, workflow_run))
@ -1233,7 +1233,6 @@ class TestWorkflowAppService:
db_session_with_containers.flush()
log = WorkflowAppLog(
id=str(uuid.uuid4()),
tenant_id=app.tenant_id,
app_id=app.id,
workflow_id=workflow.id,
@ -1241,8 +1240,9 @@ class TestWorkflowAppService:
created_from="service-api",
created_by_role=CreatorUserRole.ACCOUNT,
created_by=account.id,
created_at=datetime.now(UTC) + timedelta(minutes=i),
)
log.id = str(uuid.uuid4())
log.created_at = datetime.now(UTC) + timedelta(minutes=i)
db_session_with_containers.add(log)
logs_data.append((log, workflow_run))
@ -1335,7 +1335,6 @@ class TestWorkflowAppService:
db_session_with_containers.flush()
log = WorkflowAppLog(
id=str(uuid.uuid4()),
tenant_id=app.tenant_id,
app_id=app.id,
workflow_id=workflow.id,
@ -1343,8 +1342,9 @@ class TestWorkflowAppService:
created_from="service-api",
created_by_role=CreatorUserRole.ACCOUNT,
created_by=account.id,
created_at=datetime.now(UTC) + timedelta(minutes=i * 10 + j),
)
log.id = str(uuid.uuid4())
log.created_at = datetime.now(UTC) + timedelta(minutes=i * 10 + j)
db_session_with_containers.add(log)
db_session_with_containers.commit()

65
api/tests/unit_tests/core/test_provider_manager.py
View File

@ -28,17 +28,17 @@ def mock_provider_entity(mocker: MockerFixture):
def test__to_model_settings(mocker: MockerFixture, mock_provider_entity):
# Mocking the inputs
provider_model_settings = [
ProviderModelSetting(
id="id",
tenant_id="tenant_id",
provider_name="openai",
model_name="gpt-4",
model_type="text-generation",
enabled=True,
load_balancing_enabled=True,
)
]
ps = ProviderModelSetting(
tenant_id="tenant_id",
provider_name="openai",
model_name="gpt-4",
model_type="text-generation",
enabled=True,
load_balancing_enabled=True,
)
ps.id = "id"
provider_model_settings = [ps]
load_balancing_model_configs = [
LoadBalancingModelConfig(
id="id1",
@ -88,17 +88,17 @@ def test__to_model_settings(mocker: MockerFixture, mock_provider_entity):
def test__to_model_settings_only_one_lb(mocker: MockerFixture, mock_provider_entity):
# Mocking the inputs
provider_model_settings = [
ProviderModelSetting(
id="id",
tenant_id="tenant_id",
provider_name="openai",
model_name="gpt-4",
model_type="text-generation",
enabled=True,
load_balancing_enabled=True,
)
]
ps = ProviderModelSetting(
tenant_id="tenant_id",
provider_name="openai",
model_name="gpt-4",
model_type="text-generation",
enabled=True,
load_balancing_enabled=True,
)
ps.id = "id"
provider_model_settings = [ps]
load_balancing_model_configs = [
LoadBalancingModelConfig(
id="id1",
@ -136,17 +136,16 @@ def test__to_model_settings_only_one_lb(mocker: MockerFixture, mock_provider_ent
def test__to_model_settings_lb_disabled(mocker: MockerFixture, mock_provider_entity):
# Mocking the inputs
provider_model_settings = [
ProviderModelSetting(
id="id",
tenant_id="tenant_id",
provider_name="openai",
model_name="gpt-4",
model_type="text-generation",
enabled=True,
load_balancing_enabled=False,
)
]
ps = ProviderModelSetting(
tenant_id="tenant_id",
provider_name="openai",
model_name="gpt-4",
model_type="text-generation",
enabled=True,
load_balancing_enabled=False,
)
ps.id = "id"
provider_model_settings = [ps]
load_balancing_model_configs = [
LoadBalancingModelConfig(
id="id1",