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:
Asuka Minato
2025-11-24 11:50:20 +09:00
committed by GitHub
parent 2c9e435558
commit 6241b87f90
8 changed files with 179 additions and 161 deletions

View File

@@ -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,

View File

@@ -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()

View File

@@ -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)

View File

@@ -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"}