diff --git a/api/core/app/features/annotation_reply/annotation_reply.py b/api/core/app/features/annotation_reply/annotation_reply.py index 9eff9747764..3af2211188a 100644 --- a/api/core/app/features/annotation_reply/annotation_reply.py +++ b/api/core/app/features/annotation_reply/annotation_reply.py @@ -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 diff --git a/api/tests/unit_tests/core/app/features/test_annotation_reply.py b/api/tests/unit_tests/core/app/features/test_annotation_reply.py index ef4ee0dca6b..2c9204e64fa 100644 --- a/api/tests/unit_tests/core/app/features/test_annotation_reply.py +++ b/api/tests/unit_tests/core/app/features/test_annotation_reply.py @@ -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