more typed orm (#28519)
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
@@ -852,6 +852,7 @@ class TestAgentService:
|
||||
# Add files to message
|
||||
from models.model import MessageFile
|
||||
|
||||
assert message.from_account_id is not None
|
||||
message_file1 = MessageFile(
|
||||
message_id=message.id,
|
||||
type=FileType.IMAGE,
|
||||
|
||||
@@ -860,22 +860,24 @@ class TestAnnotationService:
|
||||
from models.model import AppAnnotationSetting
|
||||
|
||||
# Create a collection binding first
|
||||
collection_binding = DatasetCollectionBinding()
|
||||
collection_binding.id = fake.uuid4()
|
||||
collection_binding.provider_name = "openai"
|
||||
collection_binding.model_name = "text-embedding-ada-002"
|
||||
collection_binding.type = "annotation"
|
||||
collection_binding.collection_name = f"annotation_collection_{fake.uuid4()}"
|
||||
collection_binding = DatasetCollectionBinding(
|
||||
provider_name="openai",
|
||||
model_name="text-embedding-ada-002",
|
||||
type="annotation",
|
||||
collection_name=f"annotation_collection_{fake.uuid4()}",
|
||||
)
|
||||
collection_binding.id = str(fake.uuid4())
|
||||
db.session.add(collection_binding)
|
||||
db.session.flush()
|
||||
|
||||
# Create annotation setting
|
||||
annotation_setting = AppAnnotationSetting()
|
||||
annotation_setting.app_id = app.id
|
||||
annotation_setting.score_threshold = 0.8
|
||||
annotation_setting.collection_binding_id = collection_binding.id
|
||||
annotation_setting.created_user_id = account.id
|
||||
annotation_setting.updated_user_id = account.id
|
||||
annotation_setting = AppAnnotationSetting(
|
||||
app_id=app.id,
|
||||
score_threshold=0.8,
|
||||
collection_binding_id=collection_binding.id,
|
||||
created_user_id=account.id,
|
||||
updated_user_id=account.id,
|
||||
)
|
||||
db.session.add(annotation_setting)
|
||||
db.session.commit()
|
||||
|
||||
@@ -919,22 +921,24 @@ class TestAnnotationService:
|
||||
from models.model import AppAnnotationSetting
|
||||
|
||||
# Create a collection binding first
|
||||
collection_binding = DatasetCollectionBinding()
|
||||
collection_binding.id = fake.uuid4()
|
||||
collection_binding.provider_name = "openai"
|
||||
collection_binding.model_name = "text-embedding-ada-002"
|
||||
collection_binding.type = "annotation"
|
||||
collection_binding.collection_name = f"annotation_collection_{fake.uuid4()}"
|
||||
collection_binding = DatasetCollectionBinding(
|
||||
provider_name="openai",
|
||||
model_name="text-embedding-ada-002",
|
||||
type="annotation",
|
||||
collection_name=f"annotation_collection_{fake.uuid4()}",
|
||||
)
|
||||
collection_binding.id = str(fake.uuid4())
|
||||
db.session.add(collection_binding)
|
||||
db.session.flush()
|
||||
|
||||
# Create annotation setting
|
||||
annotation_setting = AppAnnotationSetting()
|
||||
annotation_setting.app_id = app.id
|
||||
annotation_setting.score_threshold = 0.8
|
||||
annotation_setting.collection_binding_id = collection_binding.id
|
||||
annotation_setting.created_user_id = account.id
|
||||
annotation_setting.updated_user_id = account.id
|
||||
annotation_setting = AppAnnotationSetting(
|
||||
app_id=app.id,
|
||||
score_threshold=0.8,
|
||||
collection_binding_id=collection_binding.id,
|
||||
created_user_id=account.id,
|
||||
updated_user_id=account.id,
|
||||
)
|
||||
db.session.add(annotation_setting)
|
||||
db.session.commit()
|
||||
|
||||
@@ -1020,22 +1024,24 @@ class TestAnnotationService:
|
||||
from models.model import AppAnnotationSetting
|
||||
|
||||
# Create a collection binding first
|
||||
collection_binding = DatasetCollectionBinding()
|
||||
collection_binding.id = fake.uuid4()
|
||||
collection_binding.provider_name = "openai"
|
||||
collection_binding.model_name = "text-embedding-ada-002"
|
||||
collection_binding.type = "annotation"
|
||||
collection_binding.collection_name = f"annotation_collection_{fake.uuid4()}"
|
||||
collection_binding = DatasetCollectionBinding(
|
||||
provider_name="openai",
|
||||
model_name="text-embedding-ada-002",
|
||||
type="annotation",
|
||||
collection_name=f"annotation_collection_{fake.uuid4()}",
|
||||
)
|
||||
collection_binding.id = str(fake.uuid4())
|
||||
db.session.add(collection_binding)
|
||||
db.session.flush()
|
||||
|
||||
# Create annotation setting
|
||||
annotation_setting = AppAnnotationSetting()
|
||||
annotation_setting.app_id = app.id
|
||||
annotation_setting.score_threshold = 0.8
|
||||
annotation_setting.collection_binding_id = collection_binding.id
|
||||
annotation_setting.created_user_id = account.id
|
||||
annotation_setting.updated_user_id = account.id
|
||||
annotation_setting = AppAnnotationSetting(
|
||||
app_id=app.id,
|
||||
score_threshold=0.8,
|
||||
collection_binding_id=collection_binding.id,
|
||||
created_user_id=account.id,
|
||||
updated_user_id=account.id,
|
||||
)
|
||||
db.session.add(annotation_setting)
|
||||
db.session.commit()
|
||||
|
||||
@@ -1080,22 +1086,24 @@ class TestAnnotationService:
|
||||
from models.model import AppAnnotationSetting
|
||||
|
||||
# Create a collection binding first
|
||||
collection_binding = DatasetCollectionBinding()
|
||||
collection_binding.id = fake.uuid4()
|
||||
collection_binding.provider_name = "openai"
|
||||
collection_binding.model_name = "text-embedding-ada-002"
|
||||
collection_binding.type = "annotation"
|
||||
collection_binding.collection_name = f"annotation_collection_{fake.uuid4()}"
|
||||
collection_binding = DatasetCollectionBinding(
|
||||
provider_name="openai",
|
||||
model_name="text-embedding-ada-002",
|
||||
type="annotation",
|
||||
collection_name=f"annotation_collection_{fake.uuid4()}",
|
||||
)
|
||||
collection_binding.id = str(fake.uuid4())
|
||||
db.session.add(collection_binding)
|
||||
db.session.flush()
|
||||
|
||||
# Create annotation setting
|
||||
annotation_setting = AppAnnotationSetting()
|
||||
annotation_setting.app_id = app.id
|
||||
annotation_setting.score_threshold = 0.8
|
||||
annotation_setting.collection_binding_id = collection_binding.id
|
||||
annotation_setting.created_user_id = account.id
|
||||
annotation_setting.updated_user_id = account.id
|
||||
annotation_setting = AppAnnotationSetting(
|
||||
app_id=app.id,
|
||||
score_threshold=0.8,
|
||||
collection_binding_id=collection_binding.id,
|
||||
created_user_id=account.id,
|
||||
updated_user_id=account.id,
|
||||
)
|
||||
db.session.add(annotation_setting)
|
||||
db.session.commit()
|
||||
|
||||
@@ -1151,22 +1159,25 @@ class TestAnnotationService:
|
||||
from models.model import AppAnnotationSetting
|
||||
|
||||
# Create a collection binding first
|
||||
collection_binding = DatasetCollectionBinding()
|
||||
collection_binding.id = fake.uuid4()
|
||||
collection_binding.provider_name = "openai"
|
||||
collection_binding.model_name = "text-embedding-ada-002"
|
||||
collection_binding.type = "annotation"
|
||||
collection_binding.collection_name = f"annotation_collection_{fake.uuid4()}"
|
||||
collection_binding = DatasetCollectionBinding(
|
||||
provider_name="openai",
|
||||
model_name="text-embedding-ada-002",
|
||||
type="annotation",
|
||||
collection_name=f"annotation_collection_{fake.uuid4()}",
|
||||
)
|
||||
collection_binding.id = str(fake.uuid4())
|
||||
db.session.add(collection_binding)
|
||||
db.session.flush()
|
||||
|
||||
# Create annotation setting
|
||||
annotation_setting = AppAnnotationSetting()
|
||||
annotation_setting.app_id = app.id
|
||||
annotation_setting.score_threshold = 0.8
|
||||
annotation_setting.collection_binding_id = collection_binding.id
|
||||
annotation_setting.created_user_id = account.id
|
||||
annotation_setting.updated_user_id = account.id
|
||||
annotation_setting = AppAnnotationSetting(
|
||||
app_id=app.id,
|
||||
score_threshold=0.8,
|
||||
collection_binding_id=collection_binding.id,
|
||||
created_user_id=account.id,
|
||||
updated_user_id=account.id,
|
||||
)
|
||||
|
||||
db.session.add(annotation_setting)
|
||||
db.session.commit()
|
||||
|
||||
@@ -1211,22 +1222,24 @@ class TestAnnotationService:
|
||||
from models.model import AppAnnotationSetting
|
||||
|
||||
# Create a collection binding first
|
||||
collection_binding = DatasetCollectionBinding()
|
||||
collection_binding.id = fake.uuid4()
|
||||
collection_binding.provider_name = "openai"
|
||||
collection_binding.model_name = "text-embedding-ada-002"
|
||||
collection_binding.type = "annotation"
|
||||
collection_binding.collection_name = f"annotation_collection_{fake.uuid4()}"
|
||||
collection_binding = DatasetCollectionBinding(
|
||||
provider_name="openai",
|
||||
model_name="text-embedding-ada-002",
|
||||
type="annotation",
|
||||
collection_name=f"annotation_collection_{fake.uuid4()}",
|
||||
)
|
||||
collection_binding.id = str(fake.uuid4())
|
||||
db.session.add(collection_binding)
|
||||
db.session.flush()
|
||||
|
||||
# Create annotation setting
|
||||
annotation_setting = AppAnnotationSetting()
|
||||
annotation_setting.app_id = app.id
|
||||
annotation_setting.score_threshold = 0.8
|
||||
annotation_setting.collection_binding_id = collection_binding.id
|
||||
annotation_setting.created_user_id = account.id
|
||||
annotation_setting.updated_user_id = account.id
|
||||
annotation_setting = AppAnnotationSetting(
|
||||
app_id=app.id,
|
||||
score_threshold=0.8,
|
||||
collection_binding_id=collection_binding.id,
|
||||
created_user_id=account.id,
|
||||
updated_user_id=account.id,
|
||||
)
|
||||
db.session.add(annotation_setting)
|
||||
db.session.commit()
|
||||
|
||||
|
||||
@@ -502,11 +502,11 @@ class TestAddDocumentToIndexTask:
|
||||
auto_disable_logs = []
|
||||
for _ in range(2):
|
||||
log_entry = DatasetAutoDisableLog(
|
||||
id=fake.uuid4(),
|
||||
tenant_id=document.tenant_id,
|
||||
dataset_id=dataset.id,
|
||||
document_id=document.id,
|
||||
)
|
||||
log_entry.id = str(fake.uuid4())
|
||||
db.session.add(log_entry)
|
||||
auto_disable_logs.append(log_entry)
|
||||
|
||||
|
||||
@@ -39,9 +39,9 @@ def test__to_model_settings(mocker: MockerFixture, mock_provider_entity):
|
||||
ps.id = "id"
|
||||
|
||||
provider_model_settings = [ps]
|
||||
|
||||
load_balancing_model_configs = [
|
||||
LoadBalancingModelConfig(
|
||||
id="id1",
|
||||
tenant_id="tenant_id",
|
||||
provider_name="openai",
|
||||
model_name="gpt-4",
|
||||
@@ -51,7 +51,6 @@ def test__to_model_settings(mocker: MockerFixture, mock_provider_entity):
|
||||
enabled=True,
|
||||
),
|
||||
LoadBalancingModelConfig(
|
||||
id="id2",
|
||||
tenant_id="tenant_id",
|
||||
provider_name="openai",
|
||||
model_name="gpt-4",
|
||||
@@ -61,6 +60,8 @@ def test__to_model_settings(mocker: MockerFixture, mock_provider_entity):
|
||||
enabled=True,
|
||||
),
|
||||
]
|
||||
load_balancing_model_configs[0].id = "id1"
|
||||
load_balancing_model_configs[1].id = "id2"
|
||||
|
||||
mocker.patch(
|
||||
"core.helper.model_provider_cache.ProviderCredentialsCache.get", return_value={"openai_api_key": "fake_key"}
|
||||
@@ -101,7 +102,6 @@ def test__to_model_settings_only_one_lb(mocker: MockerFixture, mock_provider_ent
|
||||
provider_model_settings = [ps]
|
||||
load_balancing_model_configs = [
|
||||
LoadBalancingModelConfig(
|
||||
id="id1",
|
||||
tenant_id="tenant_id",
|
||||
provider_name="openai",
|
||||
model_name="gpt-4",
|
||||
@@ -111,6 +111,7 @@ def test__to_model_settings_only_one_lb(mocker: MockerFixture, mock_provider_ent
|
||||
enabled=True,
|
||||
)
|
||||
]
|
||||
load_balancing_model_configs[0].id = "id1"
|
||||
|
||||
mocker.patch(
|
||||
"core.helper.model_provider_cache.ProviderCredentialsCache.get", return_value={"openai_api_key": "fake_key"}
|
||||
@@ -148,7 +149,6 @@ def test__to_model_settings_lb_disabled(mocker: MockerFixture, mock_provider_ent
|
||||
provider_model_settings = [ps]
|
||||
load_balancing_model_configs = [
|
||||
LoadBalancingModelConfig(
|
||||
id="id1",
|
||||
tenant_id="tenant_id",
|
||||
provider_name="openai",
|
||||
model_name="gpt-4",
|
||||
@@ -158,7 +158,6 @@ def test__to_model_settings_lb_disabled(mocker: MockerFixture, mock_provider_ent
|
||||
enabled=True,
|
||||
),
|
||||
LoadBalancingModelConfig(
|
||||
id="id2",
|
||||
tenant_id="tenant_id",
|
||||
provider_name="openai",
|
||||
model_name="gpt-4",
|
||||
@@ -168,6 +167,8 @@ def test__to_model_settings_lb_disabled(mocker: MockerFixture, mock_provider_ent
|
||||
enabled=True,
|
||||
),
|
||||
]
|
||||
load_balancing_model_configs[0].id = "id1"
|
||||
load_balancing_model_configs[1].id = "id2"
|
||||
|
||||
mocker.patch(
|
||||
"core.helper.model_provider_cache.ProviderCredentialsCache.get", return_value={"openai_api_key": "fake_key"}
|
||||
|
||||
Reference in New Issue
Block a user