opensource/dify
2
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-07-30 16:59:35 +08:00
Code Issues Packages Projects Releases Wiki Activity

test: use SQLite sessions in core rag (#39092)

Browse Source 
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
Asuka Minato 2026-07-30 12:04:18 +09:00 committed by GitHub
parent 6804dd474d
commit 3453fdce44
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 124 additions and 72 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

196
api/tests/unit_tests/core/rag/embedding/test_cached_embedding.py
View File

@ -8,19 +8,55 @@ This test file covers the methods not fully tested in test_embedding_service.py:
import base64
import logging
from dataclasses import dataclass
from decimal import Decimal
from unittest.mock import Mock, patch
import numpy as np
import pytest
from sqlalchemy.exc import IntegrityError
from sqlalchemy import event, func, select
from sqlalchemy.orm import Session
from core.rag.embedding import cached_embedding as cached_embedding_module
from core.rag.embedding.cached_embedding import CacheEmbedding
from graphon.model_runtime.entities.model_entities import ModelPropertyKey
from graphon.model_runtime.entities.text_embedding_entities import EmbeddingResult, EmbeddingUsage
from models.dataset import Embedding
@dataclass(frozen=True)
class _DatabaseBinding:
session: Session
@pytest.fixture
def embedding_session(sqlite_session: Session, monkeypatch: pytest.MonkeyPatch) -> Session:
"""Bind CacheEmbedding to the shared SQLite session."""
monkeypatch.setattr(cached_embedding_module, "db", _DatabaseBinding(session=sqlite_session))
return sqlite_session
def _persist_embedding(
session: Session,
*,
cache_key: str,
vector: list[float],
model_name: str = "vision-embedding-model",
provider_name: str = "openai",
) -> Embedding:
embedding = Embedding(
model_name=model_name,
hash=cache_key,
provider_name=provider_name,
embedding=b"placeholder",
)
embedding.set_embedding(vector)
session.add(embedding)
session.commit()
return embedding
class TestCacheEmbeddingMultimodalDocuments:
"""Test suite for CacheEmbedding.embed_multimodal_documents method."""
@ -65,27 +101,29 @@ class TestCacheEmbeddingMultimodalDocuments:
)
def test_embed_single_multimodal_document_cache_miss(
self, mock_model_instance, sample_multimodal_result: EmbeddingResult
self,
mock_model_instance,
sample_multimodal_result: EmbeddingResult,
embedding_session: Session,
):
"""Test embedding a single multimodal document when cache is empty."""
cache_embedding = CacheEmbedding(mock_model_instance)
documents = [{"file_id": "file123", "content": "test content"}]
with patch("core.rag.embedding.cached_embedding.db.session") as mock_session:
mock_session.scalar.return_value = None
mock_model_instance.invoke_multimodal_embedding.return_value = sample_multimodal_result
mock_model_instance.invoke_multimodal_embedding.return_value = sample_multimodal_result
result = cache_embedding.embed_multimodal_documents(documents)
result = cache_embedding.embed_multimodal_documents(documents)
assert len(result) == 1
assert isinstance(result[0], list)
assert len(result[0]) == 1536
assert len(result) == 1
assert isinstance(result[0], list)
assert len(result[0]) == 1536
mock_model_instance.invoke_multimodal_embedding.assert_called_once()
mock_session.add.assert_called_once()
mock_session.commit.assert_called_once()
mock_model_instance.invoke_multimodal_embedding.assert_called_once()
persisted = embedding_session.scalar(select(Embedding).where(Embedding.hash == "file123"))
assert persisted is not None
assert persisted.get_embedding() == result[0]
def test_embed_multiple_multimodal_documents_cache_miss(self, mock_model_instance):
def test_embed_multiple_multimodal_documents_cache_miss(self, mock_model_instance, embedding_session: Session):
"""Test embedding multiple multimodal documents when cache is empty."""
cache_embedding = CacheEmbedding(mock_model_instance)
documents = [
@ -116,16 +154,15 @@ class TestCacheEmbeddingMultimodalDocuments:
usage=usage,
)
with patch("core.rag.embedding.cached_embedding.db.session") as mock_session:
mock_session.scalar.return_value = None
mock_model_instance.invoke_multimodal_embedding.return_value = embedding_result
mock_model_instance.invoke_multimodal_embedding.return_value = embedding_result
result = cache_embedding.embed_multimodal_documents(documents)
result = cache_embedding.embed_multimodal_documents(documents)
assert len(result) == 3
assert all(len(emb) == 1536 for emb in result)
assert len(result) == 3
assert all(len(emb) == 1536 for emb in result)
assert embedding_session.scalar(select(func.count()).select_from(Embedding)) == 3
def test_embed_multimodal_documents_cache_hit(self, mock_model_instance):
def test_embed_multimodal_documents_cache_hit(self, mock_model_instance, embedding_session: Session):
"""Test embedding multimodal documents when embeddings are cached."""
cache_embedding = CacheEmbedding(mock_model_instance)
documents = [{"file_id": "file123"}]
@ -133,19 +170,15 @@ class TestCacheEmbeddingMultimodalDocuments:
cached_vector = np.random.randn(1536)
normalized_cached = (cached_vector / np.linalg.norm(cached_vector)).tolist()
mock_cached_embedding = Mock(spec=Embedding)
mock_cached_embedding.get_embedding.return_value = normalized_cached
_persist_embedding(embedding_session, cache_key="file123", vector=normalized_cached)
with patch("core.rag.embedding.cached_embedding.db.session") as mock_session:
mock_session.scalar.return_value = mock_cached_embedding
result = cache_embedding.embed_multimodal_documents(documents)
result = cache_embedding.embed_multimodal_documents(documents)
assert len(result) == 1
assert result[0] == normalized_cached
mock_model_instance.invoke_multimodal_embedding.assert_not_called()
assert len(result) == 1
assert result[0] == normalized_cached
mock_model_instance.invoke_multimodal_embedding.assert_not_called()
def test_embed_multimodal_documents_partial_cache_hit(self, mock_model_instance):
def test_embed_multimodal_documents_partial_cache_hit(self, mock_model_instance, embedding_session: Session):
"""Test embedding multimodal documents with mixed cache hits and misses."""
cache_embedding = CacheEmbedding(mock_model_instance)
documents = [
@ -157,8 +190,7 @@ class TestCacheEmbeddingMultimodalDocuments:
cached_vector = np.random.randn(1536)
normalized_cached = (cached_vector / np.linalg.norm(cached_vector)).tolist()
mock_cached_embedding = Mock(spec=Embedding)
mock_cached_embedding.get_embedding.return_value = normalized_cached
_persist_embedding(embedding_session, cache_key="cached_file", vector=normalized_cached)
new_embeddings = []
for _ in range(2):
@ -182,16 +214,20 @@ class TestCacheEmbeddingMultimodalDocuments:
usage=usage,
)
with patch("core.rag.embedding.cached_embedding.db.session") as mock_session:
mock_session.scalar.side_effect = [mock_cached_embedding, None, None]
mock_model_instance.invoke_multimodal_embedding.return_value = embedding_result
mock_model_instance.invoke_multimodal_embedding.return_value = embedding_result
result = cache_embedding.embed_multimodal_documents(documents)
result = cache_embedding.embed_multimodal_documents(documents)
assert len(result) == 3
assert result[0] == normalized_cached
assert len(result) == 3
assert result[0] == normalized_cached
assert embedding_session.scalar(select(func.count()).select_from(Embedding)) == 3
def test_embed_multimodal_documents_nan_handling(self, mock_model_instance, caplog: pytest.LogCaptureFixture):
def test_embed_multimodal_documents_nan_handling(
self,
mock_model_instance,
embedding_session: Session,
caplog: pytest.LogCaptureFixture,
):
"""Test handling of NaN values in multimodal embeddings."""
cache_embedding = CacheEmbedding(mock_model_instance)
documents = [{"file_id": "valid"}, {"file_id": "nan"}]
@ -215,20 +251,19 @@ class TestCacheEmbeddingMultimodalDocuments:
usage=usage,
)
with patch("core.rag.embedding.cached_embedding.db.session") as mock_session:
mock_session.scalar.return_value = None
mock_model_instance.invoke_multimodal_embedding.return_value = embedding_result
mock_model_instance.invoke_multimodal_embedding.return_value = embedding_result
with caplog.at_level(logging.WARNING, logger="core.rag.embedding.cached_embedding"):
result = cache_embedding.embed_multimodal_documents(documents)
with caplog.at_level(logging.WARNING, logger="core.rag.embedding.cached_embedding"):
result = cache_embedding.embed_multimodal_documents(documents)
assert len(result) == 2
assert result[0] is not None
assert result[1] is None
assert len(result) == 2
assert result[0] is not None
assert result[1] is None
assert embedding_session.scalar(select(func.count()).select_from(Embedding)) == 1
assert any(record.levelno == logging.WARNING for record in caplog.records)
assert any(record.levelno == logging.WARNING for record in caplog.records)
def test_embed_multimodal_documents_large_batch(self, mock_model_instance):
def test_embed_multimodal_documents_large_batch(self, mock_model_instance, embedding_session: Session):
"""Test embedding large batch of multimodal documents respecting MAX_CHUNKS."""
cache_embedding = CacheEmbedding(mock_model_instance)
documents = [{"file_id": f"file{i}"} for i in range(25)]
@ -256,49 +291,66 @@ class TestCacheEmbeddingMultimodalDocuments:
usage=usage,
)
with patch("core.rag.embedding.cached_embedding.db.session") as mock_session:
mock_session.scalar.return_value = None
batch_results = [create_batch_result(10), create_batch_result(10), create_batch_result(5)]
mock_model_instance.invoke_multimodal_embedding.side_effect = batch_results
batch_results = [create_batch_result(10), create_batch_result(10), create_batch_result(5)]
mock_model_instance.invoke_multimodal_embedding.side_effect = batch_results
result = cache_embedding.embed_multimodal_documents(documents)
result = cache_embedding.embed_multimodal_documents(documents)
assert len(result) == 25
assert mock_model_instance.invoke_multimodal_embedding.call_count == 3
assert embedding_session.scalar(select(func.count()).select_from(Embedding)) == 25
assert len(result) == 25
assert mock_model_instance.invoke_multimodal_embedding.call_count == 3
def test_embed_multimodal_documents_api_error(self, mock_model_instance):
def test_embed_multimodal_documents_api_error(self, mock_model_instance, embedding_session: Session):
"""Test handling of API errors during multimodal embedding."""
cache_embedding = CacheEmbedding(mock_model_instance)
documents = [{"file_id": "file123"}]
with patch("core.rag.embedding.cached_embedding.db.session") as mock_session:
mock_session.scalar.return_value = None
mock_model_instance.invoke_multimodal_embedding.side_effect = Exception("API Error")
mock_model_instance.invoke_multimodal_embedding.side_effect = Exception("API Error")
with pytest.raises(Exception) as exc_info:
cache_embedding.embed_multimodal_documents(documents)
with pytest.raises(Exception, match="API Error"):
cache_embedding.embed_multimodal_documents(documents)
assert "API Error" in str(exc_info.value)
mock_session.rollback.assert_called()
assert not embedding_session.in_transaction()
assert embedding_session.scalar(select(func.count()).select_from(Embedding)) == 0
def test_embed_multimodal_documents_integrity_error_during_transform(
self, mock_model_instance, sample_multimodal_result
self,
mock_model_instance,
sample_multimodal_result,
embedding_session: Session,
):
"""Test handling of IntegrityError during embedding transformation."""
cache_embedding = CacheEmbedding(mock_model_instance)
documents = [{"file_id": "file123"}]
with patch("core.rag.embedding.cached_embedding.db.session") as mock_session:
mock_session.scalar.return_value = None
mock_model_instance.invoke_multimodal_embedding.return_value = sample_multimodal_result
mock_model_instance.invoke_multimodal_embedding.return_value = sample_multimodal_result
mock_session.commit.side_effect = IntegrityError("Duplicate key", None, None)
injected = False
def add_competing_row(session: Session, _flush_context: object, _instances: object) -> None:
nonlocal injected
if injected:
return
pending = next(item for item in session.new if isinstance(item, Embedding))
session.add(
Embedding(
model_name=pending.model_name,
hash=pending.hash,
provider_name=pending.provider_name,
embedding=pending.embedding,
)
)
injected = True
event.listen(embedding_session, "before_flush", add_competing_row)
try:
result = cache_embedding.embed_multimodal_documents(documents)
finally:
event.remove(embedding_session, "before_flush", add_competing_row)
assert len(result) == 1
mock_session.rollback.assert_called()
assert len(result) == 1
assert not embedding_session.in_transaction()
assert embedding_session.scalar(select(func.count()).select_from(Embedding)) == 0
class TestCacheEmbeddingMultimodalQuery: