mirror of
https://github.com/langgenius/dify.git
synced 2026-08-31 04:53:19 +08:00
949 lines
39 KiB
Python
949 lines
39 KiB
Python
"""SQLite-backed unit tests for :mod:`services.annotation_service`."""
|
|
|
|
# Test functions explicitly request the identity fixture to make the patched
|
|
# request context visible at each service call, even when they do not inspect it.
|
|
# ruff: noqa: ARG002
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from datetime import datetime, timedelta
|
|
from decimal import Decimal
|
|
from io import BytesIO
|
|
from types import SimpleNamespace
|
|
from typing import Any, cast
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pandas as pd
|
|
import pytest
|
|
from sqlalchemy import func, select
|
|
from sqlalchemy.orm import Session, sessionmaker
|
|
from werkzeug.datastructures import FileStorage
|
|
from werkzeug.exceptions import NotFound
|
|
|
|
import services.annotation_service as annotation_service_module
|
|
from models.account import Account
|
|
from models.dataset import DatasetCollectionBinding
|
|
from models.enums import CollectionBindingType
|
|
from models.model import (
|
|
App,
|
|
AppAnnotationHitHistory,
|
|
AppAnnotationSetting,
|
|
AppMode,
|
|
ConversationFromSource,
|
|
IconType,
|
|
Message,
|
|
MessageAnnotation,
|
|
)
|
|
from services.annotation_service import AppAnnotationService
|
|
from services.app_ref_service import AnnotationRef, AppRef
|
|
|
|
TENANT_ID = "tenant-1"
|
|
OTHER_TENANT_ID = "tenant-2"
|
|
|
|
|
|
@pytest.fixture
|
|
def current_user(monkeypatch: pytest.MonkeyPatch) -> Account:
|
|
"""Install a real account model as the request identity."""
|
|
|
|
account = Account(name="Annotation Tester", email="annotation@example.com")
|
|
account.id = "account-1"
|
|
monkeypatch.setattr(
|
|
annotation_service_module,
|
|
"current_account_with_tenant",
|
|
lambda: (account, TENANT_ID),
|
|
)
|
|
return account
|
|
|
|
|
|
def _persist_app(
|
|
session: Session,
|
|
*,
|
|
app_id: str = "app-1",
|
|
tenant_id: str = TENANT_ID,
|
|
status: str = "normal",
|
|
) -> App:
|
|
app = App(
|
|
id=app_id,
|
|
tenant_id=tenant_id,
|
|
name=f"Annotation App {app_id}",
|
|
description="",
|
|
mode=AppMode.CHAT,
|
|
icon_type=IconType.EMOJI,
|
|
icon="chat",
|
|
icon_background="#FFFFFF",
|
|
status=status,
|
|
enable_site=False,
|
|
enable_api=False,
|
|
)
|
|
session.add(app)
|
|
session.commit()
|
|
return app
|
|
|
|
|
|
def _persist_message(
|
|
session: Session,
|
|
app: App,
|
|
*,
|
|
message_id: str = "msg-1",
|
|
query: str = "default-question",
|
|
) -> Message:
|
|
message = Message(
|
|
id=message_id,
|
|
app_id=app.id,
|
|
conversation_id="conv-1",
|
|
_inputs={},
|
|
query=query,
|
|
message={},
|
|
message_unit_price=Decimal(0),
|
|
answer="answer",
|
|
answer_unit_price=Decimal(0),
|
|
currency="USD",
|
|
from_source=ConversationFromSource.API,
|
|
)
|
|
session.add(message)
|
|
session.commit()
|
|
return message
|
|
|
|
|
|
def _persist_annotation(
|
|
session: Session,
|
|
app: App,
|
|
*,
|
|
annotation_id: str = "ann-1",
|
|
question: str = "question",
|
|
content: str = "answer",
|
|
message_id: str | None = None,
|
|
created_at: datetime | None = None,
|
|
) -> MessageAnnotation:
|
|
annotation = MessageAnnotation(
|
|
app_id=app.id,
|
|
question=question,
|
|
content=content,
|
|
account_id="account-1",
|
|
message_id=message_id,
|
|
)
|
|
annotation.id = annotation_id
|
|
if created_at is not None:
|
|
annotation.created_at = created_at
|
|
session.add(annotation)
|
|
session.commit()
|
|
return annotation
|
|
|
|
|
|
def _persist_binding(
|
|
session: Session,
|
|
*,
|
|
binding_id: str = "collection-1",
|
|
provider_name: str = "provider-a",
|
|
model_name: str = "model-a",
|
|
) -> DatasetCollectionBinding:
|
|
binding = DatasetCollectionBinding(
|
|
provider_name=provider_name,
|
|
model_name=model_name,
|
|
type=CollectionBindingType.ANNOTATION,
|
|
collection_name=f"collection-{binding_id}",
|
|
)
|
|
binding.id = binding_id
|
|
session.add(binding)
|
|
session.commit()
|
|
return binding
|
|
|
|
|
|
def _persist_setting(
|
|
session: Session,
|
|
app: App,
|
|
*,
|
|
setting_id: str = "setting-1",
|
|
binding_id: str = "collection-1",
|
|
score_threshold: float = 0.5,
|
|
) -> AppAnnotationSetting:
|
|
setting = AppAnnotationSetting(
|
|
app_id=app.id,
|
|
score_threshold=score_threshold,
|
|
collection_binding_id=binding_id,
|
|
created_user_id="account-1",
|
|
updated_user_id="account-1",
|
|
)
|
|
setting.id = setting_id
|
|
session.add(setting)
|
|
session.commit()
|
|
return setting
|
|
|
|
|
|
def _persist_history(
|
|
session: Session,
|
|
app: App,
|
|
annotation: MessageAnnotation,
|
|
*,
|
|
history_id: str = "history-1",
|
|
created_at: datetime | None = None,
|
|
) -> AppAnnotationHitHistory:
|
|
history = AppAnnotationHitHistory(
|
|
app_id=app.id,
|
|
annotation_id=annotation.id,
|
|
source="hit-testing",
|
|
question="query",
|
|
account_id="account-1",
|
|
score=0.8,
|
|
message_id="message-1",
|
|
annotation_question=annotation.question,
|
|
annotation_content=annotation.content,
|
|
)
|
|
history.id = history_id
|
|
if created_at is not None:
|
|
history.created_at = created_at
|
|
session.add(history)
|
|
session.commit()
|
|
return history
|
|
|
|
|
|
def _app_ref(app: App) -> AppRef:
|
|
return AppRef(tenant_id=app.tenant_id, app_id=app.id)
|
|
|
|
|
|
def _annotation_ref(app: App, annotation_id: str) -> AnnotationRef:
|
|
return AnnotationRef(app=_app_ref(app), annotation_id=annotation_id)
|
|
|
|
|
|
def _file(content: bytes) -> FileStorage:
|
|
return FileStorage(stream=BytesIO(content))
|
|
|
|
|
|
def _observer_get(factory: sessionmaker[Session], model: type[Any], identifier: str) -> Any:
|
|
with factory() as observer:
|
|
return observer.get(model, identifier)
|
|
|
|
|
|
class TestAppAnnotationServiceUpsert:
|
|
def test_rejects_missing_or_cross_tenant_app(self, sqlite_session: Session, current_user: Account) -> None:
|
|
_persist_app(sqlite_session, app_id="other-app", tenant_id=OTHER_TENANT_ID)
|
|
|
|
with pytest.raises(NotFound):
|
|
AppAnnotationService.up_insert_app_annotation_from_message(
|
|
{"answer": "hello", "question": "q"}, "other-app", sqlite_session
|
|
)
|
|
|
|
def test_validates_answer_and_question(self, sqlite_session: Session, current_user: Account) -> None:
|
|
app = _persist_app(sqlite_session)
|
|
|
|
with pytest.raises(ValueError, match="answer.*content"):
|
|
AppAnnotationService.up_insert_app_annotation_from_message({"question": "q"}, app.id, sqlite_session)
|
|
with pytest.raises(ValueError, match="question"):
|
|
AppAnnotationService.up_insert_app_annotation_from_message({"answer": "a"}, app.id, sqlite_session)
|
|
|
|
def test_rejects_message_from_another_app(self, sqlite_session: Session, current_user: Account) -> None:
|
|
app = _persist_app(sqlite_session)
|
|
other_app = _persist_app(sqlite_session, app_id="app-2")
|
|
message = _persist_message(sqlite_session, other_app)
|
|
|
|
with pytest.raises(NotFound, match="Message"):
|
|
AppAnnotationService.up_insert_app_annotation_from_message(
|
|
{"answer": "hello", "message_id": message.id}, app.id, sqlite_session
|
|
)
|
|
|
|
def test_updates_existing_message_annotation_and_enqueues_index(
|
|
self,
|
|
sqlite_session: Session,
|
|
sqlite_session_factory: sessionmaker[Session],
|
|
current_user: Account,
|
|
) -> None:
|
|
app = _persist_app(sqlite_session)
|
|
message = _persist_message(sqlite_session, app)
|
|
annotation = _persist_annotation(sqlite_session, app, message_id=message.id, content="old")
|
|
setting = _persist_setting(sqlite_session, app)
|
|
|
|
with patch.object(annotation_service_module, "add_annotation_to_index_task") as task:
|
|
result = AppAnnotationService.up_insert_app_annotation_from_message(
|
|
{"answer": "updated", "message_id": message.id}, app.id, sqlite_session
|
|
)
|
|
|
|
assert result.id == annotation.id
|
|
stored = _observer_get(sqlite_session_factory, MessageAnnotation, annotation.id)
|
|
assert stored is not None
|
|
assert (stored.question, stored.content) == (message.query, "updated")
|
|
task.delay.assert_called_once_with(
|
|
annotation.id, message.query, TENANT_ID, app.id, setting.collection_binding_id
|
|
)
|
|
|
|
def test_creates_message_annotation_without_setting(
|
|
self,
|
|
sqlite_session: Session,
|
|
sqlite_session_factory: sessionmaker[Session],
|
|
current_user: Account,
|
|
) -> None:
|
|
app = _persist_app(sqlite_session)
|
|
message = _persist_message(sqlite_session, app)
|
|
|
|
with patch.object(annotation_service_module, "add_annotation_to_index_task") as task:
|
|
result = AppAnnotationService.up_insert_app_annotation_from_message(
|
|
{"answer": "hello", "question": "override", "message_id": message.id}, app.id, sqlite_session
|
|
)
|
|
|
|
stored = _observer_get(sqlite_session_factory, MessageAnnotation, result.id)
|
|
assert stored is not None
|
|
assert (stored.app_id, stored.message_id, stored.question, stored.content) == (
|
|
app.id,
|
|
message.id,
|
|
"override",
|
|
"hello",
|
|
)
|
|
assert stored.account_id == current_user.id
|
|
task.delay.assert_not_called()
|
|
|
|
def test_creates_direct_annotation_and_enqueues_index(
|
|
self,
|
|
sqlite_session: Session,
|
|
sqlite_session_factory: sessionmaker[Session],
|
|
current_user: Account,
|
|
) -> None:
|
|
app = _persist_app(sqlite_session)
|
|
setting = _persist_setting(sqlite_session, app)
|
|
|
|
with patch.object(annotation_service_module, "add_annotation_to_index_task") as task:
|
|
result = AppAnnotationService.up_insert_app_annotation_from_message(
|
|
{"answer": "hello", "question": "q1"}, app.id, sqlite_session
|
|
)
|
|
|
|
stored = _observer_get(sqlite_session_factory, MessageAnnotation, result.id)
|
|
assert stored is not None
|
|
assert (stored.conversation_id, stored.message_id, stored.question, stored.content) == (
|
|
None,
|
|
None,
|
|
"q1",
|
|
"hello",
|
|
)
|
|
task.delay.assert_called_once_with(result.id, "q1", TENANT_ID, app.id, setting.collection_binding_id)
|
|
|
|
|
|
class TestAppAnnotationServiceEnableDisable:
|
|
def test_enable_returns_processing_on_cache_hit(self, current_user: Account) -> None:
|
|
args = {"score_threshold": 0.5, "embedding_provider_name": "p", "embedding_model_name": "m"}
|
|
with (
|
|
patch.object(annotation_service_module, "redis_client") as redis,
|
|
patch.object(annotation_service_module, "enable_annotation_reply_task") as task,
|
|
):
|
|
redis.get.return_value = "job-1"
|
|
result = AppAnnotationService.enable_app_annotation(args, "app-1")
|
|
|
|
assert result == {"job_id": "job-1", "job_status": "processing"}
|
|
task.delay.assert_not_called()
|
|
|
|
def test_enable_enqueues_on_cache_miss(self, current_user: Account) -> None:
|
|
args = {"score_threshold": 0.5, "embedding_provider_name": "p", "embedding_model_name": "m"}
|
|
with (
|
|
patch.object(annotation_service_module, "redis_client") as redis,
|
|
patch.object(annotation_service_module.uuid, "uuid4", return_value="uuid-1"),
|
|
patch.object(annotation_service_module, "enable_annotation_reply_task") as task,
|
|
):
|
|
redis.get.return_value = None
|
|
result = AppAnnotationService.enable_app_annotation(args, "app-1")
|
|
|
|
assert result == {"job_id": "uuid-1", "job_status": "waiting"}
|
|
redis.setnx.assert_called_once_with("enable_app_annotation_job_uuid-1", "waiting")
|
|
task.delay.assert_called_once_with("uuid-1", "app-1", current_user.id, TENANT_ID, 0.5, "p", "m")
|
|
|
|
def test_disable_returns_processing_on_cache_hit(self, current_user: Account) -> None:
|
|
with (
|
|
patch.object(annotation_service_module, "redis_client") as redis,
|
|
patch.object(annotation_service_module, "disable_annotation_reply_task") as task,
|
|
):
|
|
redis.get.return_value = "job-2"
|
|
result = AppAnnotationService.disable_app_annotation("app-1")
|
|
|
|
assert result == {"job_id": "job-2", "job_status": "processing"}
|
|
task.delay.assert_not_called()
|
|
|
|
def test_disable_enqueues_on_cache_miss(self, current_user: Account) -> None:
|
|
with (
|
|
patch.object(annotation_service_module, "redis_client") as redis,
|
|
patch.object(annotation_service_module.uuid, "uuid4", return_value="uuid-2"),
|
|
patch.object(annotation_service_module, "disable_annotation_reply_task") as task,
|
|
):
|
|
redis.get.return_value = None
|
|
result = AppAnnotationService.disable_app_annotation("app-1")
|
|
|
|
assert result == {"job_id": "uuid-2", "job_status": "waiting"}
|
|
redis.setnx.assert_called_once_with("disable_app_annotation_job_uuid-2", "waiting")
|
|
task.delay.assert_called_once_with("uuid-2", "app-1", TENANT_ID)
|
|
|
|
|
|
class TestAppAnnotationServiceListAndExport:
|
|
def test_list_rejects_cross_tenant_app(self, sqlite_session: Session, current_user: Account) -> None:
|
|
app = _persist_app(sqlite_session, tenant_id=OTHER_TENANT_ID)
|
|
|
|
with pytest.raises(NotFound):
|
|
AppAnnotationService.get_annotation_list_by_app_id(app.id, 1, 10, "", sqlite_session)
|
|
|
|
def test_list_filters_orders_and_paginates(self, sqlite_session: Session, current_user: Account) -> None:
|
|
app = _persist_app(sqlite_session)
|
|
other_app = _persist_app(sqlite_session, app_id="app-2")
|
|
now = datetime(2026, 1, 1)
|
|
first = _persist_annotation(sqlite_session, app, annotation_id="ann-1", question="needle first", created_at=now)
|
|
second = _persist_annotation(
|
|
sqlite_session,
|
|
app,
|
|
annotation_id="ann-2",
|
|
question="other",
|
|
content="needle second",
|
|
created_at=now + timedelta(seconds=1),
|
|
)
|
|
_persist_annotation(sqlite_session, app, annotation_id="ann-3", question="not matched")
|
|
_persist_annotation(sqlite_session, other_app, annotation_id="decoy", question="needle decoy")
|
|
|
|
items, total = AppAnnotationService.get_annotation_list_by_app_id(app.id, 1, 1, "needle", sqlite_session)
|
|
|
|
assert total == 2
|
|
assert [item.id for item in items] == [second.id]
|
|
items, total = AppAnnotationService.get_annotation_list_by_app_id(app.id, 2, 1, "needle", sqlite_session)
|
|
assert total == 2
|
|
assert [item.id for item in items] == [first.id]
|
|
|
|
def test_list_without_keyword_is_app_scoped(self, sqlite_session: Session, current_user: Account) -> None:
|
|
app = _persist_app(sqlite_session)
|
|
other_app = _persist_app(sqlite_session, app_id="app-2")
|
|
expected = _persist_annotation(sqlite_session, app)
|
|
_persist_annotation(sqlite_session, other_app, annotation_id="decoy")
|
|
|
|
items, total = AppAnnotationService.get_annotation_list_by_app_id(app.id, 1, 10, "", sqlite_session)
|
|
|
|
assert total == 1
|
|
assert [item.id for item in items] == [expected.id]
|
|
|
|
def test_export_sanitizes_and_scopes_rows(self, sqlite_session: Session, current_user: Account) -> None:
|
|
app = _persist_app(sqlite_session)
|
|
other_app = _persist_app(sqlite_session, app_id="app-2")
|
|
first = _persist_annotation(sqlite_session, app, annotation_id="ann-1", question="=cmd", content="+1")
|
|
second = _persist_annotation(sqlite_session, app, annotation_id="ann-2", question="@bad", content="-2")
|
|
_persist_annotation(sqlite_session, other_app, annotation_id="decoy", question="=decoy")
|
|
|
|
result = AppAnnotationService.export_annotation_list_by_app_id(app.id, sqlite_session)
|
|
|
|
assert {annotation.id for annotation in result} == {first.id, second.id}
|
|
assert {(annotation.question, annotation.content) for annotation in result} == {
|
|
("'=cmd", "'+1"),
|
|
("'@bad", "'-2"),
|
|
}
|
|
|
|
|
|
class TestAppAnnotationServiceDirectManipulation:
|
|
def test_insert_rejects_cross_tenant_app_and_missing_question(
|
|
self, sqlite_session: Session, current_user: Account
|
|
) -> None:
|
|
other_app = _persist_app(sqlite_session, tenant_id=OTHER_TENANT_ID)
|
|
with pytest.raises(NotFound):
|
|
AppAnnotationService.insert_app_annotation_directly(
|
|
{"answer": "hello", "question": "q"}, other_app.id, sqlite_session
|
|
)
|
|
|
|
app = _persist_app(sqlite_session, app_id="app-2")
|
|
with pytest.raises(ValueError, match="question"):
|
|
AppAnnotationService.insert_app_annotation_directly({"answer": "hello"}, app.id, sqlite_session)
|
|
|
|
def test_insert_persists_and_enqueues_index(
|
|
self,
|
|
sqlite_session: Session,
|
|
sqlite_session_factory: sessionmaker[Session],
|
|
current_user: Account,
|
|
) -> None:
|
|
app = _persist_app(sqlite_session)
|
|
setting = _persist_setting(sqlite_session, app)
|
|
|
|
with patch.object(annotation_service_module, "add_annotation_to_index_task") as task:
|
|
result = AppAnnotationService.insert_app_annotation_directly(
|
|
{"answer": "hello", "question": "q1"}, app.id, sqlite_session
|
|
)
|
|
|
|
stored = _observer_get(sqlite_session_factory, MessageAnnotation, result.id)
|
|
assert stored is not None
|
|
assert (stored.app_id, stored.question, stored.content, stored.account_id) == (
|
|
app.id,
|
|
"q1",
|
|
"hello",
|
|
current_user.id,
|
|
)
|
|
task.delay.assert_called_once_with(result.id, "q1", TENANT_ID, app.id, setting.collection_binding_id)
|
|
|
|
def test_update_is_app_scoped_and_validates_fields(self, sqlite_session: Session, current_user: Account) -> None:
|
|
app = _persist_app(sqlite_session)
|
|
other_app = _persist_app(sqlite_session, app_id="app-2")
|
|
annotation = _persist_annotation(sqlite_session, other_app)
|
|
|
|
with pytest.raises(NotFound):
|
|
AppAnnotationService.update_app_annotation_directly(
|
|
{"answer": "a", "question": "q"}, _annotation_ref(app, annotation.id), sqlite_session
|
|
)
|
|
|
|
own_annotation = _persist_annotation(sqlite_session, app, annotation_id="own-ann")
|
|
with pytest.raises(ValueError, match="question"):
|
|
AppAnnotationService.update_app_annotation_directly(
|
|
{"answer": "a"}, _annotation_ref(app, own_annotation.id), sqlite_session
|
|
)
|
|
with pytest.raises(ValueError, match="answer"):
|
|
AppAnnotationService.update_app_annotation_directly(
|
|
{"question": "q"}, _annotation_ref(app, own_annotation.id), sqlite_session
|
|
)
|
|
|
|
def test_update_persists_and_enqueues_index(
|
|
self,
|
|
sqlite_session: Session,
|
|
sqlite_session_factory: sessionmaker[Session],
|
|
current_user: Account,
|
|
) -> None:
|
|
app = _persist_app(sqlite_session)
|
|
annotation = _persist_annotation(sqlite_session, app, content="old")
|
|
setting = _persist_setting(sqlite_session, app)
|
|
|
|
with patch.object(annotation_service_module, "update_annotation_to_index_task") as task:
|
|
result = AppAnnotationService.update_app_annotation_directly(
|
|
{"answer": "new", "question": "new q"}, _annotation_ref(app, annotation.id), sqlite_session
|
|
)
|
|
|
|
stored = _observer_get(sqlite_session_factory, MessageAnnotation, annotation.id)
|
|
assert stored is not None
|
|
assert (stored.question, stored.content) == ("new q", "new")
|
|
assert result.id == stored.id
|
|
task.delay.assert_called_once_with(annotation.id, "new q", TENANT_ID, app.id, setting.collection_binding_id)
|
|
|
|
def test_delete_is_app_scoped(self, sqlite_session: Session, current_user: Account) -> None:
|
|
app = _persist_app(sqlite_session)
|
|
other_app = _persist_app(sqlite_session, app_id="app-2")
|
|
annotation = _persist_annotation(sqlite_session, other_app)
|
|
|
|
with pytest.raises(NotFound):
|
|
AppAnnotationService.delete_app_annotation(_annotation_ref(app, annotation.id), sqlite_session)
|
|
|
|
assert sqlite_session.get(MessageAnnotation, annotation.id) is not None
|
|
|
|
def test_delete_removes_annotation_and_histories_and_enqueues_index(
|
|
self,
|
|
sqlite_session: Session,
|
|
sqlite_session_factory: sessionmaker[Session],
|
|
current_user: Account,
|
|
) -> None:
|
|
app = _persist_app(sqlite_session)
|
|
annotation = _persist_annotation(sqlite_session, app)
|
|
histories = [
|
|
_persist_history(sqlite_session, app, annotation, history_id="history-1"),
|
|
_persist_history(sqlite_session, app, annotation, history_id="history-2"),
|
|
]
|
|
setting = _persist_setting(sqlite_session, app)
|
|
|
|
with patch.object(annotation_service_module, "delete_annotation_index_task") as task:
|
|
AppAnnotationService.delete_app_annotation(_annotation_ref(app, annotation.id), sqlite_session)
|
|
|
|
with sqlite_session_factory() as observer:
|
|
assert observer.get(MessageAnnotation, annotation.id) is None
|
|
assert [observer.get(AppAnnotationHitHistory, history.id) for history in histories] == [None, None]
|
|
task.delay.assert_called_once_with(annotation.id, app.id, TENANT_ID, setting.collection_binding_id)
|
|
|
|
def test_batch_delete_returns_zero_without_matching_rows(
|
|
self, sqlite_session: Session, current_user: Account
|
|
) -> None:
|
|
app = _persist_app(sqlite_session)
|
|
other_app = _persist_app(sqlite_session, app_id="app-2")
|
|
_persist_annotation(sqlite_session, other_app, annotation_id="ann-1")
|
|
|
|
result = AppAnnotationService.delete_app_annotations_in_batch(_app_ref(app), ["ann-1"], sqlite_session)
|
|
|
|
assert result == {"deleted_count": 0}
|
|
|
|
def test_batch_delete_scopes_rows_and_histories(
|
|
self,
|
|
sqlite_session: Session,
|
|
sqlite_session_factory: sessionmaker[Session],
|
|
current_user: Account,
|
|
) -> None:
|
|
app = _persist_app(sqlite_session)
|
|
other_app = _persist_app(sqlite_session, app_id="app-2")
|
|
annotation1 = _persist_annotation(sqlite_session, app, annotation_id="ann-1")
|
|
annotation2 = _persist_annotation(sqlite_session, app, annotation_id="ann-2")
|
|
decoy = _persist_annotation(sqlite_session, other_app, annotation_id="ann-3")
|
|
history1 = _persist_history(sqlite_session, app, annotation1, history_id="history-1")
|
|
history2 = _persist_history(sqlite_session, app, annotation2, history_id="history-2")
|
|
decoy_history = _persist_history(sqlite_session, other_app, decoy, history_id="history-3")
|
|
setting = _persist_setting(sqlite_session, app)
|
|
|
|
with patch.object(annotation_service_module, "delete_annotation_index_task") as task:
|
|
result = AppAnnotationService.delete_app_annotations_in_batch(
|
|
_app_ref(app), [annotation1.id, annotation2.id, decoy.id], sqlite_session
|
|
)
|
|
|
|
assert result == {"deleted_count": 2}
|
|
with sqlite_session_factory() as observer:
|
|
assert [observer.get(MessageAnnotation, item.id) for item in (annotation1, annotation2)] == [None, None]
|
|
assert [observer.get(AppAnnotationHitHistory, item.id) for item in (history1, history2)] == [None, None]
|
|
assert observer.get(MessageAnnotation, decoy.id) is not None
|
|
assert observer.get(AppAnnotationHitHistory, decoy_history.id) is not None
|
|
assert task.delay.call_count == 2
|
|
task.delay.assert_any_call(annotation1.id, app.id, TENANT_ID, setting.collection_binding_id)
|
|
task.delay.assert_any_call(annotation2.id, app.id, TENANT_ID, setting.collection_binding_id)
|
|
|
|
|
|
class TestAppAnnotationServiceBatchImport:
|
|
@staticmethod
|
|
def _invoke(
|
|
sqlite_session: Session,
|
|
app: App,
|
|
*,
|
|
dataframe: Any,
|
|
content: bytes = b"question,answer\nq,a\n",
|
|
maximum: int = 5,
|
|
minimum: int = 1,
|
|
features: Any | None = None,
|
|
) -> dict[str, Any]:
|
|
if features is None:
|
|
features = SimpleNamespace(billing=SimpleNamespace(enabled=False), annotation_quota_limit=None)
|
|
with (
|
|
patch.object(annotation_service_module.pd, "read_csv", return_value=dataframe),
|
|
patch.object(annotation_service_module.FeatureService, "get_features", return_value=features),
|
|
patch(
|
|
"configs.dify_config",
|
|
new=SimpleNamespace(
|
|
ANNOTATION_IMPORT_MAX_RECORDS=maximum,
|
|
ANNOTATION_IMPORT_MIN_RECORDS=minimum,
|
|
),
|
|
),
|
|
):
|
|
return AppAnnotationService.batch_import_app_annotations(app.id, _file(content), sqlite_session)
|
|
|
|
def test_rejects_cross_tenant_app(self, sqlite_session: Session, current_user: Account) -> None:
|
|
app = _persist_app(sqlite_session, tenant_id=OTHER_TENANT_ID)
|
|
|
|
with pytest.raises(NotFound):
|
|
AppAnnotationService.batch_import_app_annotations(app.id, _file(b"question,answer\nq,a\n"), sqlite_session)
|
|
|
|
@pytest.mark.parametrize(
|
|
("dataframe", "content", "maximum", "minimum", "expected"),
|
|
[
|
|
(pd.DataFrame({"q": ["only"]}), b"question\nq\n", 5, 1, "Invalid CSV format"),
|
|
(pd.DataFrame({"q": ["q"], "a": ["a"]}), b"", 5, 1, "empty or invalid"),
|
|
(pd.DataFrame({"q": ["q"], "a": ["a"]}), b"question,answer\nq,a\n", 5, 2, "at least"),
|
|
(
|
|
pd.DataFrame({"q": ["q1", "q2"], "a": ["a1", "a2"]}),
|
|
b"question,answer\nq1,a1\nq2,a2\n",
|
|
1,
|
|
1,
|
|
"too many records",
|
|
),
|
|
(pd.DataFrame({"q": ["nan"], "a": ["nan"]}), b"question,answer\nnan,nan\n", 5, 1, "at least"),
|
|
(
|
|
pd.DataFrame({"q": ["q" * 2001], "a": ["a"]}),
|
|
b"question,answer\nq,a\n",
|
|
5,
|
|
1,
|
|
"Question at row",
|
|
),
|
|
(
|
|
pd.DataFrame({"q": ["q"], "a": ["a" * 10001]}),
|
|
b"question,answer\nq,a\n",
|
|
5,
|
|
1,
|
|
"Answer at row",
|
|
),
|
|
],
|
|
)
|
|
def test_validation_errors(
|
|
self,
|
|
sqlite_session: Session,
|
|
current_user: Account,
|
|
dataframe: pd.DataFrame,
|
|
content: bytes,
|
|
maximum: int,
|
|
minimum: int,
|
|
expected: str,
|
|
) -> None:
|
|
app = _persist_app(sqlite_session)
|
|
|
|
result = self._invoke(
|
|
sqlite_session,
|
|
app,
|
|
dataframe=dataframe,
|
|
content=content,
|
|
maximum=maximum,
|
|
minimum=minimum,
|
|
)
|
|
|
|
assert expected in cast(str, result["error_msg"])
|
|
|
|
def test_skips_malformed_rows(self, sqlite_session: Session, current_user: Account) -> None:
|
|
app = _persist_app(sqlite_session)
|
|
malformed_row = MagicMock()
|
|
malformed_row.iloc.__getitem__.side_effect = IndexError()
|
|
dataframe = MagicMock()
|
|
dataframe.columns = ["q", "a"]
|
|
dataframe.iterrows.return_value = [(0, malformed_row)]
|
|
|
|
result = self._invoke(sqlite_session, app, dataframe=dataframe)
|
|
|
|
assert "at least" in cast(str, result["error_msg"])
|
|
|
|
def test_rejects_subscription_quota_overflow(self, sqlite_session: Session, current_user: Account) -> None:
|
|
app = _persist_app(sqlite_session)
|
|
features = SimpleNamespace(
|
|
billing=SimpleNamespace(enabled=True),
|
|
annotation_quota_limit=SimpleNamespace(limit=1, size=1),
|
|
)
|
|
|
|
result = self._invoke(
|
|
sqlite_session,
|
|
app,
|
|
dataframe=pd.DataFrame({"q": ["q1"], "a": ["a1"]}),
|
|
features=features,
|
|
)
|
|
|
|
assert "exceeds the limit" in cast(str, result["error_msg"])
|
|
|
|
def test_valid_import_enqueues_job(self, sqlite_session: Session, current_user: Account) -> None:
|
|
app = _persist_app(sqlite_session)
|
|
dataframe = pd.DataFrame({"q": ["q1"], "a": ["a1"]})
|
|
features = SimpleNamespace(billing=SimpleNamespace(enabled=False), annotation_quota_limit=None)
|
|
with (
|
|
patch.object(annotation_service_module.pd, "read_csv", return_value=dataframe),
|
|
patch.object(annotation_service_module.FeatureService, "get_features", return_value=features),
|
|
patch.object(annotation_service_module, "batch_import_annotations_task") as task,
|
|
patch.object(annotation_service_module, "redis_client") as redis,
|
|
patch.object(annotation_service_module.uuid, "uuid4", return_value="uuid-3"),
|
|
patch.object(annotation_service_module, "naive_utc_now", return_value=datetime.fromtimestamp(1)),
|
|
patch(
|
|
"configs.dify_config",
|
|
new=SimpleNamespace(ANNOTATION_IMPORT_MAX_RECORDS=5, ANNOTATION_IMPORT_MIN_RECORDS=1),
|
|
),
|
|
):
|
|
result = AppAnnotationService.batch_import_app_annotations(
|
|
app.id, _file(b"question,answer\nq,a\n"), sqlite_session
|
|
)
|
|
|
|
assert result == {"job_id": "uuid-3", "job_status": "waiting", "record_count": 1}
|
|
redis.zadd.assert_called_once_with(f"annotation_import_active:{TENANT_ID}", {"uuid-3": 1000})
|
|
redis.expire.assert_called_once_with(f"annotation_import_active:{TENANT_ID}", 7200)
|
|
redis.setnx.assert_called_once_with("app_annotation_batch_import_uuid-3", "waiting")
|
|
task.delay.assert_called_once_with(
|
|
"uuid-3", [{"question": "q1", "answer": "a1"}], app.id, TENANT_ID, current_user.id
|
|
)
|
|
|
|
def test_unexpected_error_cleans_active_job(
|
|
self, sqlite_session: Session, current_user: Account, caplog: pytest.LogCaptureFixture
|
|
) -> None:
|
|
app = _persist_app(sqlite_session)
|
|
dataframe = pd.DataFrame({"q": ["q1"], "a": ["a1"]})
|
|
features = SimpleNamespace(billing=SimpleNamespace(enabled=False), annotation_quota_limit=None)
|
|
with (
|
|
patch.object(annotation_service_module.pd, "read_csv", return_value=dataframe),
|
|
patch.object(annotation_service_module.FeatureService, "get_features", return_value=features),
|
|
patch.object(annotation_service_module, "redis_client") as redis,
|
|
patch.object(annotation_service_module.uuid, "uuid4", return_value="uuid-4"),
|
|
patch.object(annotation_service_module, "naive_utc_now", return_value=datetime.fromtimestamp(1)),
|
|
patch(
|
|
"configs.dify_config",
|
|
new=SimpleNamespace(ANNOTATION_IMPORT_MAX_RECORDS=5, ANNOTATION_IMPORT_MIN_RECORDS=1),
|
|
),
|
|
):
|
|
redis.zadd.side_effect = RuntimeError("boom")
|
|
redis.zrem.side_effect = RuntimeError("cleanup-failed")
|
|
with caplog.at_level(logging.DEBUG):
|
|
result = AppAnnotationService.batch_import_app_annotations(
|
|
app.id, _file(b"question,answer\nq,a\n"), sqlite_session
|
|
)
|
|
|
|
assert result["error_msg"] == "An error occurred while processing the file: boom"
|
|
redis.zrem.assert_called_once_with(f"annotation_import_active:{TENANT_ID}", "uuid-4")
|
|
assert "Failed to clean up active job tracking" in caplog.text
|
|
|
|
|
|
class TestAppAnnotationServiceHitHistoryAndSettings:
|
|
def test_hit_histories_are_annotation_and_app_scoped(self, sqlite_session: Session, current_user: Account) -> None:
|
|
app = _persist_app(sqlite_session)
|
|
other_app = _persist_app(sqlite_session, app_id="app-2")
|
|
other_annotation = _persist_annotation(sqlite_session, other_app)
|
|
|
|
with pytest.raises(NotFound):
|
|
AppAnnotationService.get_annotation_hit_histories(
|
|
_annotation_ref(app, other_annotation.id), 1, 10, sqlite_session
|
|
)
|
|
|
|
annotation = _persist_annotation(sqlite_session, app, annotation_id="own-ann")
|
|
now = datetime(2026, 1, 1)
|
|
old = _persist_history(sqlite_session, app, annotation, history_id="old", created_at=now)
|
|
new = _persist_history(sqlite_session, app, annotation, history_id="new", created_at=now + timedelta(seconds=1))
|
|
_persist_history(sqlite_session, other_app, other_annotation, history_id="decoy")
|
|
|
|
items, total = AppAnnotationService.get_annotation_hit_histories(
|
|
_annotation_ref(app, annotation.id), 1, 1, sqlite_session
|
|
)
|
|
assert total == 2
|
|
assert [item.id for item in items] == [new.id]
|
|
items, total = AppAnnotationService.get_annotation_hit_histories(
|
|
_annotation_ref(app, annotation.id), 2, 1, sqlite_session
|
|
)
|
|
assert total == 2
|
|
assert [item.id for item in items] == [old.id]
|
|
|
|
def test_get_annotation_by_id_uses_real_identity_lookup(
|
|
self, sqlite_session: Session, current_user: Account
|
|
) -> None:
|
|
app = _persist_app(sqlite_session)
|
|
annotation = _persist_annotation(sqlite_session, app)
|
|
|
|
assert AppAnnotationService.get_annotation_by_id("missing", sqlite_session) is None
|
|
assert AppAnnotationService.get_annotation_by_id(annotation.id, sqlite_session) is annotation
|
|
|
|
def test_add_history_increments_count_and_flushes_row(self, sqlite_session: Session, current_user: Account) -> None:
|
|
app = _persist_app(sqlite_session)
|
|
annotation = _persist_annotation(sqlite_session, app)
|
|
|
|
AppAnnotationService.add_annotation_history(
|
|
annotation_id=annotation.id,
|
|
app_id=app.id,
|
|
annotation_question="q",
|
|
annotation_content="a",
|
|
query="user q",
|
|
user_id=current_user.id,
|
|
message_id="msg-1",
|
|
from_source="chat",
|
|
score=0.8,
|
|
session=sqlite_session,
|
|
)
|
|
|
|
sqlite_session.refresh(annotation)
|
|
assert annotation.hit_count == 1
|
|
history = sqlite_session.scalar(
|
|
select(AppAnnotationHitHistory).where(AppAnnotationHitHistory.annotation_id == annotation.id)
|
|
)
|
|
assert history is not None
|
|
assert (history.question, history.annotation_question, history.annotation_content, history.score) == (
|
|
"user q",
|
|
"q",
|
|
"a",
|
|
0.8,
|
|
)
|
|
|
|
def test_get_setting_rejects_cross_tenant_app(self, sqlite_session: Session, current_user: Account) -> None:
|
|
app = _persist_app(sqlite_session, tenant_id=OTHER_TENANT_ID)
|
|
|
|
with pytest.raises(NotFound):
|
|
AppAnnotationService.get_app_annotation_setting_by_app_id(app.id, sqlite_session)
|
|
|
|
def test_get_setting_returns_disabled_without_row(self, sqlite_session: Session, current_user: Account) -> None:
|
|
app = _persist_app(sqlite_session)
|
|
|
|
assert AppAnnotationService.get_app_annotation_setting_by_app_id(app.id, sqlite_session) == {"enabled": False}
|
|
|
|
def test_get_setting_returns_binding_detail(self, sqlite_session: Session, current_user: Account) -> None:
|
|
app = _persist_app(sqlite_session)
|
|
binding = _persist_binding(sqlite_session)
|
|
setting = _persist_setting(sqlite_session, app, binding_id=binding.id)
|
|
|
|
result = AppAnnotationService.get_app_annotation_setting_by_app_id(app.id, sqlite_session)
|
|
|
|
assert result == {
|
|
"id": setting.id,
|
|
"enabled": True,
|
|
"score_threshold": 0.5,
|
|
"embedding_model": {
|
|
"embedding_provider_name": binding.provider_name,
|
|
"embedding_model_name": binding.model_name,
|
|
},
|
|
}
|
|
|
|
def test_get_setting_returns_empty_detail_for_missing_binding(
|
|
self, sqlite_session: Session, current_user: Account
|
|
) -> None:
|
|
app = _persist_app(sqlite_session)
|
|
setting = _persist_setting(sqlite_session, app, binding_id="missing-binding")
|
|
|
|
result = AppAnnotationService.get_app_annotation_setting_by_app_id(app.id, sqlite_session)
|
|
|
|
assert result == {
|
|
"id": setting.id,
|
|
"enabled": True,
|
|
"score_threshold": 0.5,
|
|
"embedding_model": {},
|
|
}
|
|
|
|
def test_update_setting_is_app_scoped(self, sqlite_session: Session, current_user: Account) -> None:
|
|
app = _persist_app(sqlite_session)
|
|
other_app = _persist_app(sqlite_session, app_id="app-2")
|
|
other_setting = _persist_setting(sqlite_session, other_app)
|
|
|
|
with pytest.raises(NotFound):
|
|
AppAnnotationService.update_app_annotation_setting(
|
|
app.id, other_setting.id, {"score_threshold": 0.8}, sqlite_session
|
|
)
|
|
|
|
def test_update_setting_flushes_changes_and_returns_binding(
|
|
self, sqlite_session: Session, current_user: Account
|
|
) -> None:
|
|
app = _persist_app(sqlite_session)
|
|
binding = _persist_binding(sqlite_session)
|
|
setting = _persist_setting(sqlite_session, app, binding_id=binding.id)
|
|
|
|
result = AppAnnotationService.update_app_annotation_setting(
|
|
app.id, setting.id, {"score_threshold": 0.8}, sqlite_session
|
|
)
|
|
|
|
assert result["enabled"] is True
|
|
assert result["score_threshold"] == 0.8
|
|
assert result["embedding_model"] == {
|
|
"embedding_provider_name": binding.provider_name,
|
|
"embedding_model_name": binding.model_name,
|
|
}
|
|
assert setting.score_threshold == 0.8
|
|
assert setting.updated_user_id == current_user.id
|
|
|
|
def test_update_setting_returns_empty_detail_for_missing_binding(
|
|
self, sqlite_session: Session, current_user: Account
|
|
) -> None:
|
|
app = _persist_app(sqlite_session)
|
|
setting = _persist_setting(sqlite_session, app, binding_id="missing-binding")
|
|
|
|
result = AppAnnotationService.update_app_annotation_setting(
|
|
app.id, setting.id, {"score_threshold": 0.7}, sqlite_session
|
|
)
|
|
|
|
assert result["score_threshold"] == 0.7
|
|
assert result["embedding_model"] == {}
|
|
|
|
|
|
class TestAppAnnotationServiceClearAll:
|
|
def test_clear_all_deletes_only_app_rows_and_enqueues_indexes(
|
|
self,
|
|
sqlite_session: Session,
|
|
sqlite_session_factory: sessionmaker[Session],
|
|
current_user: Account,
|
|
) -> None:
|
|
app = _persist_app(sqlite_session)
|
|
other_app = _persist_app(sqlite_session, app_id="app-2")
|
|
annotations = [
|
|
_persist_annotation(sqlite_session, app, annotation_id="ann-1"),
|
|
_persist_annotation(sqlite_session, app, annotation_id="ann-2"),
|
|
]
|
|
histories = [
|
|
_persist_history(sqlite_session, app, annotations[0], history_id="history-1"),
|
|
_persist_history(sqlite_session, app, annotations[1], history_id="history-2"),
|
|
]
|
|
decoy = _persist_annotation(sqlite_session, other_app, annotation_id="decoy")
|
|
decoy_history = _persist_history(sqlite_session, other_app, decoy, history_id="decoy-history")
|
|
setting = _persist_setting(sqlite_session, app)
|
|
|
|
with patch.object(annotation_service_module, "delete_annotation_index_task") as task:
|
|
result = AppAnnotationService.clear_all_annotations(app.id, sqlite_session)
|
|
|
|
assert result == {"result": "success"}
|
|
with sqlite_session_factory() as observer:
|
|
assert [observer.get(MessageAnnotation, item.id) for item in annotations] == [None, None]
|
|
assert [observer.get(AppAnnotationHitHistory, item.id) for item in histories] == [None, None]
|
|
assert observer.get(MessageAnnotation, decoy.id) is not None
|
|
assert observer.get(AppAnnotationHitHistory, decoy_history.id) is not None
|
|
assert task.delay.call_count == 2
|
|
for annotation in annotations:
|
|
task.delay.assert_any_call(annotation.id, app.id, TENANT_ID, setting.collection_binding_id)
|
|
|
|
def test_clear_all_rejects_cross_tenant_app(self, sqlite_session: Session, current_user: Account) -> None:
|
|
app = _persist_app(sqlite_session, tenant_id=OTHER_TENANT_ID)
|
|
|
|
with pytest.raises(NotFound):
|
|
AppAnnotationService.clear_all_annotations(app.id, sqlite_session)
|
|
|
|
assert sqlite_session.scalar(select(func.count()).select_from(MessageAnnotation)) == 0
|