Spaces:
Runtime error
Runtime error
File size: 6,580 Bytes
7ab7df1 1521cd2 7ab7df1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 | """Tests for empty vectorstore and initial state scenarios."""
from unittest.mock import MagicMock, Mock, patch
import pytest
from qa_chain import QAChainWrapper, create_qa_chain
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.documents import Document
from ui.handlers import create_stream_chat_response
@pytest.fixture
def empty_vectorstore():
"""Create an empty vectorstore mock."""
mock_vs = MagicMock()
mock_vs.get.return_value = {
"documents": [],
"metadatas": [],
"ids": [],
}
mock_retriever = MagicMock()
mock_retriever.invoke.return_value = [] # No documents
mock_vs.as_retriever.return_value = mock_retriever
return mock_vs
@patch("qa_chain.create_llm")
@patch("qa_chain.format_chat_history")
def test_rag_query_with_empty_vectorstore(mock_format_history, mock_create_llm, empty_vectorstore):
"""Test RAG query when vectorstore is empty."""
prompt = ChatPromptTemplate.from_template("Test: {question}")
qa_chain = QAChainWrapper(empty_vectorstore, prompt)
mock_format_history.return_value = ""
mock_llm = MagicMock()
mock_chunk = MagicMock()
mock_chunk.content = "I don't have any documents to search."
mock_llm.stream.return_value = [mock_chunk]
mock_create_llm.return_value = mock_llm
mock_chain = MagicMock()
mock_chain.stream.return_value = [mock_chunk]
qa_chain._prompt.__or__ = MagicMock(return_value=mock_chain)
inputs = {
"question": "What is RAG?",
"chat_history": [],
}
results = list(qa_chain.stream(inputs))
# Should handle empty vectorstore gracefully
assert len(results) > 0
# Should have empty source documents
assert len(results[0]["source_documents"]) == 0
# Context should be empty
context = "\n\n".join(doc.page_content for doc in results[0]["source_documents"])
assert context == ""
@patch("qa_chain.create_llm")
@patch("qa_chain.format_chat_history")
def test_rag_query_with_no_retrieved_documents(mock_format_history, mock_create_llm, mock_vectorstore):
"""Test RAG query when retrieval returns no documents."""
prompt = ChatPromptTemplate.from_template("Test: {question}")
qa_chain = QAChainWrapper(mock_vectorstore, prompt)
mock_format_history.return_value = ""
mock_llm = MagicMock()
mock_chunk = MagicMock()
mock_chunk.content = "No relevant documents found."
mock_llm.stream.return_value = [mock_chunk]
mock_create_llm.return_value = mock_llm
# Retriever returns empty list
mock_retriever = MagicMock()
mock_retriever.invoke.return_value = [] # No documents retrieved
qa_chain._retriever = mock_retriever
mock_chain = MagicMock()
mock_chain.stream.return_value = [mock_chunk]
qa_chain._prompt.__or__ = MagicMock(return_value=mock_chain)
inputs = {
"question": "What is RAG?",
"chat_history": [],
}
results = list(qa_chain.stream(inputs))
# Should handle no documents gracefully
assert len(results) > 0
assert len(results[0]["source_documents"]) == 0
def test_initialize_chain_with_empty_vectorstore():
"""Test chain initialization with empty vectorstore."""
from ui.app import initialize_chain
with patch("ui.app.create_embeddings") as mock_embeddings, \
patch("ui.app.load_or_create_vectorstore") as mock_load_vs, \
patch("ui.app.create_qa_chain") as mock_create_chain:
mock_embeddings.return_value = MagicMock()
mock_vectorstore = MagicMock()
mock_vectorstore.get.return_value = {
"metadatas": [], # Empty metadatas
}
mock_load_vs.return_value = mock_vectorstore
mock_chain = MagicMock()
mock_create_chain.return_value = mock_chain
chain, sources = initialize_chain()
assert chain == mock_chain
assert sources == [] # Should return empty list
@patch("ui.handlers.create_llm")
def test_handlers_with_empty_vectorstore(mock_create_llm, empty_vectorstore):
"""Test handlers with empty vectorstore."""
mock_qa_chain = MagicMock()
mock_qa_chain.stream.return_value = [
{
"chunk": "No documents available.",
"source_documents": [],
"docs_with_scores": None,
"rewritten_query": None,
"hybrid_scores": None,
}
]
stream_fn = create_stream_chat_response(mock_qa_chain)
results = list(
stream_fn(
"What is RAG?",
[],
"RAG",
doc_filter=None,
search_type="mmr",
)
)
assert len(results) > 0
# Should handle empty results gracefully
assert len(results[0][1]) == 0 # source_documents is empty
@patch("retrievers.BM25Okapi")
def test_hybrid_search_with_empty_documents(mock_bm25, empty_vectorstore):
"""Test hybrid search with empty document collection."""
from retrievers import HybridRetriever
empty_vectorstore.get.return_value = {
"documents": [],
"metadatas": [],
}
empty_vectorstore.similarity_search_with_score.return_value = []
retriever = HybridRetriever(empty_vectorstore)
results = retriever.hybrid_search("test query", k=5)
# Should return empty list
assert results == []
@patch("vectorstore.Chroma")
@patch("vectorstore.get_pdf_files")
def test_create_new_vectorstore_no_pdfs(mock_get_pdfs, mock_chroma):
"""Test creating new vectorstore when no PDFs exist."""
from vectorstore import create_new_vectorstore
mock_get_pdfs.return_value = [] # No PDF files
with pytest.raises(FileNotFoundError):
create_new_vectorstore(MagicMock())
def test_rerank_with_empty_documents(mock_vectorstore):
"""Test re-ranking with empty document list."""
from qa_chain import QAChainWrapper
from langchain_core.prompts import ChatPromptTemplate
prompt = ChatPromptTemplate.from_template("Test: {question}")
qa_chain_wrapper = QAChainWrapper(mock_vectorstore, prompt)
result = qa_chain_wrapper.rerank_documents("query", [], top_k=5)
# Should return empty list
assert result == []
def test_hybrid_retriever_with_empty_collection(empty_vectorstore):
"""Test hybrid retriever initialization with empty collection."""
from retrievers import HybridRetriever
retriever = HybridRetriever(empty_vectorstore)
retriever._build_bm25_index()
# BM25 should be None when no documents
assert retriever._bm25 is None
assert retriever._documents == []
|