170 lines
6.7 KiB
Python
170 lines
6.7 KiB
Python
"""
|
|
Tests for RAG Engine.
|
|
"""
|
|
|
|
import pytest
|
|
from unittest.mock import Mock, patch, MagicMock
|
|
|
|
|
|
@pytest.mark.integration
|
|
@pytest.mark.rag
|
|
class TestRAGEngine:
|
|
"""Test RAG engine functionality."""
|
|
|
|
@patch('app.core.rag_engine.Settings')
|
|
@patch('app.core.rag_engine.Ollama')
|
|
@patch('app.core.rag_engine.OllamaEmbedding')
|
|
@patch('app.core.rag_engine.get_sqlite_manager')
|
|
def test_rag_engine_initialization(self, mock_get_db, mock_embed, mock_llm, mock_settings, sample_kb, clean_db):
|
|
"""Test that RAG engine initializes correctly."""
|
|
# Configure mocks
|
|
mock_db = MagicMock()
|
|
mock_db.collection_exists.return_value = True
|
|
mock_db.get_collection_metadata.return_value = {"kb_type": "generic", "id": sample_kb["id"]}
|
|
mock_get_db.return_value = mock_db
|
|
|
|
mock_llm_instance = MagicMock()
|
|
mock_llm.return_value = mock_llm_instance
|
|
|
|
mock_embed_instance = MagicMock()
|
|
mock_embed.return_value = mock_embed_instance
|
|
|
|
from app.core.rag_engine import RAGEngine
|
|
|
|
engine = RAGEngine(sample_kb["name"])
|
|
|
|
assert engine.collection_name == sample_kb["name"]
|
|
assert engine.llm == mock_llm_instance
|
|
assert engine.embed_model == mock_embed_instance
|
|
mock_db.collection_exists.assert_called_with(sample_kb["name"])
|
|
|
|
@patch('app.core.rag_engine.Settings')
|
|
@patch('app.core.rag_engine.Ollama')
|
|
@patch('app.core.rag_engine.OllamaEmbedding')
|
|
@patch('app.core.rag_engine.get_sqlite_manager')
|
|
def test_rag_engine_collection_not_found(self, mock_get_db, mock_embed, mock_llm, mock_settings, sample_kb):
|
|
"""Test that RAG engine raises error for non-existent collection."""
|
|
# Configure mocks
|
|
mock_db = MagicMock()
|
|
mock_db.collection_exists.return_value = False
|
|
mock_get_db.return_value = mock_db
|
|
|
|
from app.core.rag_engine import RAGEngine
|
|
|
|
with pytest.raises(ValueError, match="Collection .* not found"):
|
|
RAGEngine("nonexistent_collection")
|
|
|
|
@patch('app.core.rag_engine.Settings')
|
|
@patch('app.core.rag_engine.Ollama')
|
|
@patch('app.core.rag_engine.OllamaEmbedding')
|
|
@patch('app.core.rag_engine.get_sqlite_manager')
|
|
def test_query_with_no_documents(self, mock_get_db, mock_embed, mock_llm, mock_settings, sample_kb):
|
|
"""Test querying when no documents exist."""
|
|
# Configure mocks
|
|
mock_db = MagicMock()
|
|
mock_db.collection_exists.return_value = True
|
|
mock_db.get_collection_metadata.return_value = {"kb_type": "generic", "id": sample_kb["id"]}
|
|
mock_db.query_documents.return_value = {
|
|
"ids": [[]],
|
|
"documents": [[]],
|
|
"metadatas": [[]],
|
|
"distances": [[]]
|
|
}
|
|
mock_get_db.return_value = mock_db
|
|
|
|
mock_llm_instance = MagicMock()
|
|
mock_llm.return_value = mock_llm_instance
|
|
|
|
mock_embed_instance = MagicMock()
|
|
mock_embed_instance.get_text_embedding.return_value = [0.1] * 768
|
|
mock_embed.return_value = mock_embed_instance
|
|
|
|
from app.core.rag_engine import RAGEngine
|
|
|
|
engine = RAGEngine(sample_kb["name"])
|
|
result = engine.query("test question")
|
|
|
|
# When no documents found, returns "No relevant information found"
|
|
assert "No relevant information found" in result["answer"]
|
|
assert len(result["sources"]) == 0
|
|
|
|
@patch('app.core.rag_engine.Settings')
|
|
@patch('app.core.rag_engine.Ollama')
|
|
@patch('app.core.rag_engine.OllamaEmbedding')
|
|
@patch('app.core.rag_engine.get_sqlite_manager')
|
|
def test_query_with_documents(self, mock_get_db, mock_embed, mock_llm, mock_settings, sample_kb, sample_embedding):
|
|
"""Test querying with documents returns sources."""
|
|
# Configure mocks
|
|
mock_db = MagicMock()
|
|
mock_db.collection_exists.return_value = True
|
|
mock_db.get_collection_metadata.return_value = {"kb_type": "generic", "id": sample_kb["id"]}
|
|
mock_db.query_documents.return_value = {
|
|
"ids": [["doc1"]],
|
|
"documents": [["This is test content about the topic."]],
|
|
"metadatas": [[{"source": "test.md"}]],
|
|
"distances": [[0.2]]
|
|
}
|
|
mock_get_db.return_value = mock_db
|
|
|
|
mock_llm_instance = MagicMock()
|
|
mock_response = MagicMock()
|
|
mock_response.message.content = "This is a test answer."
|
|
mock_llm_instance.chat.return_value = mock_response
|
|
mock_llm.return_value = mock_llm_instance
|
|
|
|
mock_embed_instance = MagicMock()
|
|
mock_embed_instance.get_text_embedding.return_value = sample_embedding
|
|
mock_embed.return_value = mock_embed_instance
|
|
|
|
from app.core.rag_engine import RAGEngine
|
|
|
|
engine = RAGEngine(sample_kb["name"])
|
|
result = engine.query("test question")
|
|
|
|
# Result should contain answer and sources
|
|
assert result["answer"] is not None
|
|
assert len(result["sources"]) == 1
|
|
assert result["sources"][0]["text"] == "This is test content about the topic."
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestSimpleVectorRetriever:
|
|
"""Test SimpleVectorRetriever functionality."""
|
|
|
|
def test_retriever_initialization(self, sample_kb, clean_db):
|
|
"""Test that retriever initializes correctly."""
|
|
from app.core.rag_engine import SimpleVectorRetriever
|
|
|
|
retriever = SimpleVectorRetriever(sample_kb["name"], clean_db, similarity_top_k=5)
|
|
|
|
assert retriever.kb_name == sample_kb["name"]
|
|
assert retriever.db_manager == clean_db
|
|
assert retriever.similarity_top_k == 5
|
|
|
|
def test_retriever_with_results(self, sample_kb, clean_db, sample_embedding):
|
|
"""Test retriever returns results correctly."""
|
|
from app.core.rag_engine import SimpleVectorRetriever
|
|
|
|
# add_documents signature: (kb_name, documents, embeddings, metadatas, ids)
|
|
clean_db.add_documents(
|
|
sample_kb["name"],
|
|
["Test content 1", "Test content 2"], # documents
|
|
[sample_embedding, sample_embedding], # embeddings
|
|
[{"source": "file1.md"}, {"source": "file2.md"}], # metadatas
|
|
["doc1", "doc2"] # ids
|
|
)
|
|
|
|
retriever = SimpleVectorRetriever(sample_kb["name"], clean_db, similarity_top_k=5)
|
|
nodes = retriever.retrieve("test query", sample_embedding)
|
|
|
|
assert len(nodes) >= 1
|
|
assert nodes[0].text in ["Test content 1", "Test content 2"]
|
|
|
|
def test_retriever_with_no_results(self, sample_kb, clean_db, sample_embedding):
|
|
"""Test retriever with empty KB."""
|
|
from app.core.rag_engine import SimpleVectorRetriever
|
|
|
|
retriever = SimpleVectorRetriever(sample_kb["name"], clean_db, similarity_top_k=5)
|
|
nodes = retriever.retrieve("test query", sample_embedding)
|
|
|
|
assert len(nodes) == 0
|