Spaces:
Runtime error
Runtime error
File size: 3,393 Bytes
7ab7df1 ed005f8 | 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 | """Pytest configuration and shared fixtures."""
import os
from unittest.mock import MagicMock, Mock
import pytest
from langchain_core.documents import Document
from langchain_core.messages import AIMessage, HumanMessage
@pytest.fixture
def mock_embeddings():
"""Create a mock embeddings model."""
mock = MagicMock()
mock.embed_query.return_value = [0.1] * 384
return mock
@pytest.fixture
def mock_vectorstore():
"""Create a mock vectorstore."""
mock = MagicMock()
mock.get.return_value = {
"documents": ["Document 1", "Document 2"],
"metadatas": [
{"source": "pdf/test1.pdf", "page": 1},
{"source": "pdf/test2.pdf", "page": 2},
],
}
return mock
@pytest.fixture
def sample_documents():
"""Create sample document objects."""
return [
Document(
page_content="This is a test document about machine learning.",
metadata={"source": "pdf/test1.pdf", "page": 1},
),
Document(
page_content="Another document about neural networks.",
metadata={"source": "pdf/test2.pdf", "page": 2},
),
Document(
page_content="Document in references section.",
metadata={"source": "pdf/test3.pdf", "page": 3, "section": "references"},
),
]
@pytest.fixture
def mock_llm():
"""Create a mock LLM."""
mock = MagicMock()
mock.invoke.return_value = MagicMock(content="Mocked response")
mock.stream.return_value = [
MagicMock(content="Chunk "),
MagicMock(content="1 "),
MagicMock(content="2"),
]
return mock
@pytest.fixture
def mock_reranker():
"""Create a mock reranker."""
mock = MagicMock()
mock.predict.return_value = [0.9, 0.8, 0.7]
return mock
@pytest.fixture
def sample_chat_history():
"""Create sample chat history."""
return [
{"role": "user", "content": "What is RAG?"},
{"role": "assistant", "content": "RAG is Retrieval-Augmented Generation."},
{"role": "user", "content": "How does it work?"},
]
@pytest.fixture
def sample_chat_history_tuples():
"""Create sample chat history as tuples."""
return [
("What is RAG?", "RAG is Retrieval-Augmented Generation."),
("How does it work?", "It retrieves documents and generates answers."),
]
@pytest.fixture
def temp_pdf_dir(tmp_path):
"""Create a temporary directory for PDF files."""
pdf_dir = tmp_path / "pdf"
pdf_dir.mkdir()
return str(pdf_dir)
@pytest.fixture
def temp_vectorstore_dir(tmp_path):
"""Create a temporary directory for vectorstore."""
vs_dir = tmp_path / "vectorstore"
vs_dir.mkdir()
return str(vs_dir)
@pytest.fixture
def mock_hybrid_results():
"""Create mock hybrid search results."""
doc1 = Document(
page_content="Test content 1",
metadata={"source": "pdf/test1.pdf", "page": 1},
)
doc2 = Document(
page_content="Test content 2",
metadata={"source": "pdf/test2.pdf", "page": 2},
)
return [
{
"doc": doc1,
"fused_score": 0.9,
"semantic_score": 0.85,
"keyword_score": 0.95,
},
{
"doc": doc2,
"fused_score": 0.8,
"semantic_score": 0.75,
"keyword_score": 0.85,
},
]
|