mirror of
https://github.com/langgenius/dify.git
synced 2026-07-14 17:07:03 +08:00
test: use sqlite3 session in test_annotation_reply (#38714)
This commit is contained in:
Asuka Minato
GitHub
@ -1,12 +1,13 @@
|
||||
import logging
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from core.app.entities.app_invoke_entities import InvokeFrom
|
||||
from core.rag.datasource.vdb.vector_factory import Vector
|
||||
from core.rag.index_processor.constant.index_type import IndexTechniqueType
|
||||
from extensions.ext_database import db
|
||||
from models.dataset import Dataset
|
||||
from models.dataset import Dataset, DatasetCollectionBinding
|
||||
from models.enums import CollectionBindingType, ConversationFromSource
|
||||
from models.model import App, AppAnnotationSetting, Message, MessageAnnotation
|
||||
from services.annotation_service import AppAnnotationService
|
||||
@ -17,24 +18,33 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
class AnnotationReplyFeature:
|
||||
def query(
|
||||
self, app_record: App, message: Message, query: str, user_id: str, invoke_from: InvokeFrom
|
||||
self,
|
||||
app_record: App,
|
||||
message: Message,
|
||||
query: str,
|
||||
user_id: str,
|
||||
invoke_from: InvokeFrom,
|
||||
*,
|
||||
session: Session | None = None,
|
||||
) -> MessageAnnotation | None:
|
||||
"""Return the closest annotation reply and record a hit in ``session``.
|
||||
|
||||
The caller may provide its transaction so the setting lookup, annotation
|
||||
lookup, and hit-history write share one session. Runtime callers that do
|
||||
not provide one continue to use Flask-SQLAlchemy's scoped session.
|
||||
Vector-search failures are logged and return ``None``; transaction
|
||||
cleanup remains the caller's responsibility.
|
||||
"""
|
||||
Query app annotations to reply
|
||||
:param app_record: app record
|
||||
:param message: message
|
||||
:param query: query
|
||||
:param user_id: user id
|
||||
:param invoke_from: invoke from
|
||||
:return:
|
||||
"""
|
||||
if session is None:
|
||||
session = db.session()
|
||||
|
||||
stmt = select(AppAnnotationSetting).where(AppAnnotationSetting.app_id == app_record.id)
|
||||
annotation_setting = db.session.scalar(stmt)
|
||||
annotation_setting = session.scalar(stmt)
|
||||
|
||||
if not annotation_setting:
|
||||
return None
|
||||
|
||||
collection_binding_detail = annotation_setting.collection_binding_detail
|
||||
collection_binding_detail = session.get(DatasetCollectionBinding, annotation_setting.collection_binding_id)
|
||||
|
||||
if not collection_binding_detail:
|
||||
return None
|
||||
@ -45,7 +55,7 @@ class AnnotationReplyFeature:
|
||||
embedding_model_name = collection_binding_detail.model_name
|
||||
|
||||
dataset_collection_binding = DatasetCollectionBindingService.get_dataset_collection_binding(
|
||||
embedding_provider_name, embedding_model_name, db.session(), CollectionBindingType.ANNOTATION
|
||||
embedding_provider_name, embedding_model_name, session, CollectionBindingType.ANNOTATION
|
||||
)
|
||||
|
||||
dataset = Dataset(
|
||||
@ -66,7 +76,7 @@ class AnnotationReplyFeature:
|
||||
if documents and documents[0].metadata:
|
||||
annotation_id = documents[0].metadata["annotation_id"]
|
||||
score = documents[0].metadata["score"]
|
||||
annotation = AppAnnotationService.get_annotation_by_id(annotation_id, session=db.session())
|
||||
annotation = AppAnnotationService.get_annotation_by_id(annotation_id, session=session)
|
||||
if annotation:
|
||||
if invoke_from in {InvokeFrom.SERVICE_API, InvokeFrom.WEB_APP}:
|
||||
from_source = ConversationFromSource.API
|
||||
@ -84,7 +94,7 @@ class AnnotationReplyFeature:
|
||||
message.id,
|
||||
from_source,
|
||||
score,
|
||||
session=db.session(),
|
||||
session=session,
|
||||
)
|
||||
|
||||
return annotation
|
||||
|
||||
@ -3,163 +3,171 @@ from types import SimpleNamespace
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from core.app.entities.app_invoke_entities import InvokeFrom
|
||||
from core.app.features.annotation_reply.annotation_reply import AnnotationReplyFeature
|
||||
from models.dataset import DatasetCollectionBinding
|
||||
from models.enums import CollectionBindingType, ConversationFromSource
|
||||
from models.model import AppAnnotationHitHistory, AppAnnotationSetting, MessageAnnotation
|
||||
|
||||
TABLES = (AppAnnotationSetting, DatasetCollectionBinding, MessageAnnotation, AppAnnotationHitHistory)
|
||||
|
||||
|
||||
def _persist_binding(session: Session) -> DatasetCollectionBinding:
|
||||
binding = DatasetCollectionBinding(
|
||||
provider_name="prov",
|
||||
model_name="model",
|
||||
type=CollectionBindingType.ANNOTATION,
|
||||
collection_name="annotation-collection",
|
||||
)
|
||||
session.add(binding)
|
||||
session.flush()
|
||||
return binding
|
||||
|
||||
|
||||
def _persist_setting(
|
||||
session: Session,
|
||||
*,
|
||||
app_id: str = "app-1",
|
||||
collection_binding_id: str,
|
||||
score_threshold: float = 0.5,
|
||||
) -> AppAnnotationSetting:
|
||||
setting = AppAnnotationSetting(
|
||||
app_id=app_id,
|
||||
score_threshold=score_threshold,
|
||||
collection_binding_id=collection_binding_id,
|
||||
created_user_id="user-1",
|
||||
updated_user_id="user-1",
|
||||
)
|
||||
session.add(setting)
|
||||
session.flush()
|
||||
return setting
|
||||
|
||||
|
||||
def _persist_annotation(session: Session) -> MessageAnnotation:
|
||||
annotation = MessageAnnotation(
|
||||
app_id="app-1",
|
||||
question="question",
|
||||
content="content",
|
||||
account_id="acct-1",
|
||||
)
|
||||
session.add(annotation)
|
||||
session.flush()
|
||||
return annotation
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True)
|
||||
class TestAnnotationReplyFeature:
|
||||
def test_query_returns_none_when_setting_missing(self):
|
||||
feature = AnnotationReplyFeature()
|
||||
def test_query_returns_none_when_setting_missing(self, sqlite_session: Session):
|
||||
binding = _persist_binding(sqlite_session)
|
||||
_persist_setting(sqlite_session, app_id="other-app", collection_binding_id=binding.id)
|
||||
|
||||
with patch("core.app.features.annotation_reply.annotation_reply.db") as mock_db:
|
||||
mock_db.session.scalar.return_value = None
|
||||
|
||||
result = feature.query(
|
||||
app_record=SimpleNamespace(id="app-1", tenant_id="tenant-1"),
|
||||
message=SimpleNamespace(id="msg-1"),
|
||||
query="hi",
|
||||
user_id="user-1",
|
||||
invoke_from=InvokeFrom.SERVICE_API,
|
||||
)
|
||||
result = AnnotationReplyFeature().query(
|
||||
app_record=SimpleNamespace(id="app-1", tenant_id="tenant-1"),
|
||||
message=SimpleNamespace(id="msg-1"),
|
||||
query="hi",
|
||||
user_id="user-1",
|
||||
invoke_from=InvokeFrom.SERVICE_API,
|
||||
session=sqlite_session,
|
||||
)
|
||||
|
||||
assert result is None
|
||||
|
||||
def test_query_returns_none_when_binding_missing(self):
|
||||
feature = AnnotationReplyFeature()
|
||||
annotation_setting = SimpleNamespace(collection_binding_detail=None)
|
||||
def test_query_returns_none_when_binding_missing(self, sqlite_session: Session):
|
||||
_persist_setting(sqlite_session, collection_binding_id="missing-binding")
|
||||
|
||||
with patch("core.app.features.annotation_reply.annotation_reply.db") as mock_db:
|
||||
mock_db.session.scalar.return_value = annotation_setting
|
||||
|
||||
result = feature.query(
|
||||
app_record=SimpleNamespace(id="app-1", tenant_id="tenant-1"),
|
||||
message=SimpleNamespace(id="msg-1"),
|
||||
query="hi",
|
||||
user_id="user-1",
|
||||
invoke_from=InvokeFrom.SERVICE_API,
|
||||
)
|
||||
result = AnnotationReplyFeature().query(
|
||||
app_record=SimpleNamespace(id="app-1", tenant_id="tenant-1"),
|
||||
message=SimpleNamespace(id="msg-1"),
|
||||
query="hi",
|
||||
user_id="user-1",
|
||||
invoke_from=InvokeFrom.SERVICE_API,
|
||||
session=sqlite_session,
|
||||
)
|
||||
|
||||
assert result is None
|
||||
|
||||
def test_query_returns_annotation_and_records_history_for_api(self):
|
||||
feature = AnnotationReplyFeature()
|
||||
annotation_setting = SimpleNamespace(
|
||||
score_threshold=None,
|
||||
collection_binding_detail=SimpleNamespace(provider_name="prov", model_name="model"),
|
||||
)
|
||||
dataset_binding = SimpleNamespace(id="binding-1")
|
||||
annotation = SimpleNamespace(
|
||||
id="ann-1",
|
||||
question_text="question",
|
||||
content="content",
|
||||
account_id="acct-1",
|
||||
account=SimpleNamespace(name="Alice"),
|
||||
)
|
||||
document = SimpleNamespace(metadata={"annotation_id": "ann-1", "score": 0.8})
|
||||
def test_query_returns_annotation_and_persists_history_for_api(self, sqlite_session: Session):
|
||||
binding = _persist_binding(sqlite_session)
|
||||
_persist_setting(sqlite_session, collection_binding_id=binding.id, score_threshold=0)
|
||||
annotation = _persist_annotation(sqlite_session)
|
||||
document = SimpleNamespace(metadata={"annotation_id": annotation.id, "score": 0.8})
|
||||
vector_instance = Mock()
|
||||
vector_instance.search_by_vector.return_value = [document]
|
||||
|
||||
with (
|
||||
patch("core.app.features.annotation_reply.annotation_reply.db") as mock_db,
|
||||
patch(
|
||||
"core.app.features.annotation_reply.annotation_reply.DatasetCollectionBindingService"
|
||||
) as mock_binding_service,
|
||||
patch("core.app.features.annotation_reply.annotation_reply.Vector") as mock_vector,
|
||||
patch(
|
||||
"core.app.features.annotation_reply.annotation_reply.AppAnnotationService"
|
||||
) as mock_annotation_service,
|
||||
):
|
||||
mock_db.session.scalar.return_value = annotation_setting
|
||||
mock_binding_service.get_dataset_collection_binding.return_value = dataset_binding
|
||||
mock_vector.return_value = vector_instance
|
||||
mock_annotation_service.get_annotation_by_id.return_value = annotation
|
||||
|
||||
result = feature.query(
|
||||
with patch("core.app.features.annotation_reply.annotation_reply.Vector", return_value=vector_instance):
|
||||
result = AnnotationReplyFeature().query(
|
||||
app_record=SimpleNamespace(id="app-1", tenant_id="tenant-1"),
|
||||
message=SimpleNamespace(id="msg-1"),
|
||||
query="hi",
|
||||
user_id="user-1",
|
||||
invoke_from=InvokeFrom.SERVICE_API,
|
||||
session=sqlite_session,
|
||||
)
|
||||
|
||||
assert result == annotation
|
||||
mock_annotation_service.add_annotation_history.assert_called_once()
|
||||
_, _, _, _, _, _, _, from_source, score = mock_annotation_service.add_annotation_history.call_args[0]
|
||||
assert from_source == "api"
|
||||
assert score == 0.8
|
||||
assert result is annotation
|
||||
vector_instance.search_by_vector.assert_called_once_with(
|
||||
query="hi", top_k=1, score_threshold=1, filter={"group_id": ["app-1"]}
|
||||
)
|
||||
sqlite_session.refresh(annotation)
|
||||
assert annotation.hit_count == 1
|
||||
history = sqlite_session.scalar(select(AppAnnotationHitHistory))
|
||||
assert history is not None
|
||||
assert history.annotation_id == annotation.id
|
||||
assert history.app_id == "app-1"
|
||||
assert history.message_id == "msg-1"
|
||||
assert history.account_id == "user-1"
|
||||
assert history.source == ConversationFromSource.API
|
||||
assert history.score == 0.8
|
||||
|
||||
def test_query_returns_annotation_and_records_history_for_console(self):
|
||||
feature = AnnotationReplyFeature()
|
||||
annotation_setting = SimpleNamespace(
|
||||
score_threshold=0.5,
|
||||
collection_binding_detail=SimpleNamespace(provider_name="prov", model_name="model"),
|
||||
)
|
||||
dataset_binding = SimpleNamespace(id="binding-1")
|
||||
annotation = SimpleNamespace(
|
||||
id="ann-1",
|
||||
question_text="question",
|
||||
content="content",
|
||||
account_id="acct-1",
|
||||
account=None,
|
||||
)
|
||||
document = SimpleNamespace(metadata={"annotation_id": "ann-1", "score": 0.6})
|
||||
def test_query_returns_annotation_and_persists_history_for_console(self, sqlite_session: Session):
|
||||
binding = _persist_binding(sqlite_session)
|
||||
_persist_setting(sqlite_session, collection_binding_id=binding.id)
|
||||
annotation = _persist_annotation(sqlite_session)
|
||||
document = SimpleNamespace(metadata={"annotation_id": annotation.id, "score": 0.6})
|
||||
vector_instance = Mock()
|
||||
vector_instance.search_by_vector.return_value = [document]
|
||||
|
||||
with (
|
||||
patch("core.app.features.annotation_reply.annotation_reply.db") as mock_db,
|
||||
patch(
|
||||
"core.app.features.annotation_reply.annotation_reply.DatasetCollectionBindingService"
|
||||
) as mock_binding_service,
|
||||
patch("core.app.features.annotation_reply.annotation_reply.Vector") as mock_vector,
|
||||
patch(
|
||||
"core.app.features.annotation_reply.annotation_reply.AppAnnotationService"
|
||||
) as mock_annotation_service,
|
||||
):
|
||||
mock_db.session.scalar.return_value = annotation_setting
|
||||
mock_binding_service.get_dataset_collection_binding.return_value = dataset_binding
|
||||
mock_vector.return_value = vector_instance
|
||||
mock_annotation_service.get_annotation_by_id.return_value = annotation
|
||||
|
||||
result = feature.query(
|
||||
with patch("core.app.features.annotation_reply.annotation_reply.Vector", return_value=vector_instance):
|
||||
result = AnnotationReplyFeature().query(
|
||||
app_record=SimpleNamespace(id="app-1", tenant_id="tenant-1"),
|
||||
message=SimpleNamespace(id="msg-1"),
|
||||
query="hi",
|
||||
user_id="user-1",
|
||||
invoke_from=InvokeFrom.EXPLORE,
|
||||
session=sqlite_session,
|
||||
)
|
||||
|
||||
assert result == annotation
|
||||
_, _, _, _, _, _, _, from_source, _ = mock_annotation_service.add_annotation_history.call_args[0]
|
||||
assert from_source == "console"
|
||||
assert result is annotation
|
||||
history = sqlite_session.scalar(select(AppAnnotationHitHistory))
|
||||
assert history is not None
|
||||
assert history.source == ConversationFromSource.CONSOLE
|
||||
|
||||
def test_query_logs_and_returns_none_on_exception(self, caplog: pytest.LogCaptureFixture):
|
||||
feature = AnnotationReplyFeature()
|
||||
annotation_setting = SimpleNamespace(
|
||||
score_threshold=None,
|
||||
collection_binding_detail=SimpleNamespace(provider_name="prov", model_name="model"),
|
||||
)
|
||||
def test_query_logs_and_returns_none_on_exception(self, sqlite_session: Session, caplog: pytest.LogCaptureFixture):
|
||||
binding = _persist_binding(sqlite_session)
|
||||
_persist_setting(sqlite_session, collection_binding_id=binding.id)
|
||||
vector_instance = Mock()
|
||||
vector_instance.search_by_vector.side_effect = RuntimeError("boom")
|
||||
|
||||
with (
|
||||
patch("core.app.features.annotation_reply.annotation_reply.db") as mock_db,
|
||||
patch(
|
||||
"core.app.features.annotation_reply.annotation_reply.DatasetCollectionBindingService"
|
||||
) as mock_binding_service,
|
||||
patch("core.app.features.annotation_reply.annotation_reply.Vector") as mock_vector,
|
||||
"core.app.features.annotation_reply.annotation_reply.Vector",
|
||||
return_value=vector_instance,
|
||||
),
|
||||
caplog.at_level(logging.WARNING),
|
||||
):
|
||||
mock_db.session.scalar.return_value = annotation_setting
|
||||
mock_binding_service.get_dataset_collection_binding.return_value = SimpleNamespace(id="binding-1")
|
||||
mock_vector.return_value.search_by_vector.side_effect = RuntimeError("boom")
|
||||
|
||||
with caplog.at_level(logging.WARNING):
|
||||
result = feature.query(
|
||||
app_record=SimpleNamespace(id="app-1", tenant_id="tenant-1"),
|
||||
message=SimpleNamespace(id="msg-1"),
|
||||
query="hi",
|
||||
user_id="user-1",
|
||||
invoke_from=InvokeFrom.SERVICE_API,
|
||||
)
|
||||
result = AnnotationReplyFeature().query(
|
||||
app_record=SimpleNamespace(id="app-1", tenant_id="tenant-1"),
|
||||
message=SimpleNamespace(id="msg-1"),
|
||||
query="hi",
|
||||
user_id="user-1",
|
||||
invoke_from=InvokeFrom.SERVICE_API,
|
||||
session=sqlite_session,
|
||||
)
|
||||
|
||||
assert result is None
|
||||
assert "Query annotation failed" in caplog.text
|
||||
assert sqlite_session.scalar(select(AppAnnotationHitHistory)) is None
|
||||
assert sqlite_session.is_active
|
||||
|
||||
Reference in New Issue
Block a user