Spaces:
Running on Zero
Running on Zero
File size: 12,175 Bytes
0828c2c | 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 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 | """Integration tests for RAG pipeline with main document."""
from pathlib import Path
from unittest.mock import MagicMock, patch
import pytest
from src.config_loader import get_config
from src.main_document_loader import reload_main_document_loader
from src.rag_pipeline import RAGPipeline
@pytest.fixture(autouse=True)
def ensure_fail_silently():
"""Ensure fail_silently is always set in config for all tests in this module."""
import copy
config = get_config()
# Reload singleton at the start to ensure clean state
reload_main_document_loader()
# Save original config with deep copy to avoid reference issues
original_main_doc_config = copy.deepcopy(config._config.get("main_document", {}))
# Ensure fail_silently is set
if "main_document" not in config._config:
config._config["main_document"] = {}
if "fail_silently" not in config._config["main_document"]:
config._config["main_document"]["fail_silently"] = True
yield
# Restore original config after test
if original_main_doc_config:
config._config["main_document"] = original_main_doc_config
elif "main_document" in config._config:
# If there was no original config, remove what we added
del config._config["main_document"]
# Also reload the singleton to ensure it picks up the restored config
reload_main_document_loader()
class TestRAGWithMainDocument:
"""Test RAG pipeline with main document integration."""
def test_rag_pipeline_initialization(self):
"""Test that RAG pipeline initializes with main document loader."""
pipeline = RAGPipeline()
# Should have main_doc_loader attribute
assert hasattr(pipeline, "main_doc_loader")
assert hasattr(pipeline, "main_doc_content")
def test_main_doc_info_disabled(self):
"""Test main document info when feature is disabled."""
# Temporarily disable main document
config = get_config()
# Save the entire main_document config section
original_main_doc_config = config._config.get("main_document", {}).copy()
try:
# Mock disabled state but preserve other config values
config._config["main_document"] = original_main_doc_config.copy()
config._config["main_document"]["enabled"] = False
pipeline = RAGPipeline()
info = pipeline.get_main_document_info()
assert "enabled" in info
assert "loaded" in info
assert info["enabled"] is False
assert info["loaded"] is False
finally:
# Restore original config completely
if original_main_doc_config:
config._config["main_document"] = original_main_doc_config
def test_main_doc_info_enabled(self):
"""Test main document info when feature is enabled and loaded."""
config = get_config()
# Only run if main document is actually configured and exists
main_doc_path = config.get("main_document.path", "")
if not main_doc_path or not Path(main_doc_path).exists():
pytest.skip("Main document not configured or doesn't exist")
pipeline = RAGPipeline()
# If main document is loaded
if pipeline.main_doc_content:
info = pipeline.get_main_document_info()
assert info["enabled"] is True
assert info["loaded"] is True
assert "tokens" in info
assert "path" in info
assert "size_bytes" in info
assert info["tokens"] > 0
assert info["size_bytes"] > 0
def test_token_budget_calculation(self):
"""Test token budget calculation."""
pipeline = RAGPipeline()
budget = pipeline._calculate_context_budget()
# Check all required keys are present
assert "model_context_window" in budget
assert "main_doc_tokens" in budget
assert "max_output_tokens" in budget
assert "buffer_tokens" in budget
assert "estimated_chat_history_tokens" in budget
assert "available_for_retrieval" in budget
assert "total_input_budget" in budget
# Check reasonable values
assert budget["model_context_window"] > 0
assert budget["main_doc_tokens"] >= 0
assert budget["max_output_tokens"] > 0
assert budget["buffer_tokens"] >= 0
assert budget["estimated_chat_history_tokens"] >= 0
assert budget["available_for_retrieval"] >= 0
# Check math (must account for estimated chat history tokens)
expected_available = (
budget["model_context_window"]
- budget["main_doc_tokens"]
- budget["max_output_tokens"]
- budget["buffer_tokens"]
- budget["estimated_chat_history_tokens"]
)
assert budget["available_for_retrieval"] == max(0, expected_available)
def test_format_main_doc_empty(self):
"""Test _format_main_doc returns empty string when no content."""
pipeline = RAGPipeline()
pipeline.main_doc_content = ""
formatted = pipeline._format_main_doc()
assert formatted == ""
def test_format_main_doc_with_content(self):
"""Test _format_main_doc formats content correctly."""
pipeline = RAGPipeline()
pipeline.main_doc_content = "Test profile content"
formatted = pipeline._format_main_doc()
assert "ESSENTIAL PROFILE INFORMATION" in formatted
assert "Test profile content" in formatted
assert "ADDITIONAL CONTEXT FROM DOCUMENTS" in formatted
def test_reload_main_document(self, tmp_path):
"""Test runtime reloading of main document."""
# Create a test main document
test_doc = tmp_path / "test_profile.md"
test_doc.write_text("# Test Profile\n\nOriginal content")
# Temporarily override config
config = get_config()
# Save the entire main_document config section
original_main_doc_config = config._config.get("main_document", {}).copy()
try:
config._config["main_document"] = {
"enabled": True,
"path": str(test_doc),
"max_tokens": 10000,
"summarize_if_exceeds": True,
"cache_enabled": True,
"cache_check_interval": 60,
"fail_silently": True, # Add missing key
"fallback_to_vectordb_only": True,
}
# Reload the singleton to pick up new config
reload_main_document_loader()
pipeline = RAGPipeline()
# Initial content
assert "Original content" in pipeline.main_doc_content
# Modify file
test_doc.write_text("# Test Profile\n\nModified content")
# Reload
success = pipeline.reload_main_document()
assert success is True
assert "Modified content" in pipeline.main_doc_content
finally:
# Restore original config completely
if original_main_doc_config:
config._config["main_document"] = original_main_doc_config
def test_reload_main_document_error_handling(self):
"""Test reload_main_document handles errors gracefully."""
pipeline = RAGPipeline()
# Set invalid path
pipeline.main_doc_loader.path = Path("/nonexistent/path/file.md")
# Reload should succeed (with empty content) when fail_silently is True
success = pipeline.reload_main_document()
assert success is True
# But content should be empty since file doesn't exist
assert pipeline.main_doc_content == ""
def test_query_with_main_doc(self):
"""Test that query works with main document enabled."""
config = get_config()
# Only run if main document is configured
main_doc_enabled = config.get("main_document.enabled", False)
if not main_doc_enabled:
pytest.skip("Main document not enabled in config")
pipeline = RAGPipeline()
# Simple test query
response = pipeline.query("What is the name?")
assert "result" in response
assert isinstance(response["result"], str)
assert len(response["result"]) > 0
def test_format_docs_empty_list(self):
"""Test _format_docs handles empty list."""
pipeline = RAGPipeline()
formatted = pipeline._format_docs([])
assert formatted == ""
def test_context_budget_warning_large_main_doc(self, caplog):
"""Test that warning is logged when main doc uses >50% of context."""
pipeline = RAGPipeline()
# Simulate large main document
original_content = pipeline.main_doc_content
pipeline.main_doc_content = "Test " * 10000 # Very large content
try:
import logging
with caplog.at_level(logging.WARNING):
budget = pipeline._calculate_context_budget()
# Check if warning was logged (only if main doc is actually large)
if budget["main_doc_tokens"] > budget["model_context_window"] * 0.5:
assert any("uses" in record.message for record in caplog.records)
finally:
# Restore original content
pipeline.main_doc_content = original_content
def test_model_context_window_detection(self):
"""Test that different models get correct context windows."""
config = get_config()
original_model = config.get("llm.model", "")
test_cases = {
"llama3.2:3b": 8192,
"llama3.1:8b": 128000,
"phi3:mini": 4096,
"gemma2:2b": 8192,
"unknown:model": 8192, # Default
}
# Mock LLM handler to avoid actual LLM initialization
with patch("src.rag_pipeline.get_llm_handler") as mock_llm_handler:
mock_handler = MagicMock()
mock_handler.get_llm.return_value = MagicMock()
mock_llm_handler.return_value = mock_handler
for model_name, expected_window in test_cases.items():
config._config["llm"]["model"] = model_name
pipeline = RAGPipeline()
budget = pipeline._calculate_context_budget()
assert budget["model_context_window"] == expected_window
# Restore original model
config._config["llm"]["model"] = original_model
class TestMainDocumentContextInsertion:
"""Test that main document content is properly inserted into prompts."""
def test_prompt_template_structure(self):
"""Test that prompt template includes main document placeholder."""
# Mock LLM handler to avoid actual LLM initialization
with patch("src.rag_pipeline.get_llm_handler") as mock_llm_handler:
mock_handler = MagicMock()
mock_handler.get_llm.return_value = MagicMock()
mock_llm_handler.return_value = mock_handler
pipeline = RAGPipeline()
template_str = pipeline.prompt_template.messages[0].prompt.template
# Should contain main_document_section placeholder
assert "{main_document_section}" in template_str
assert "{context}" in template_str
assert "{question}" in template_str
def test_main_doc_position_before_context(self):
"""Test that main document appears BEFORE vectordb context in template."""
# Mock LLM handler to avoid actual LLM initialization
with patch("src.rag_pipeline.get_llm_handler") as mock_llm_handler:
mock_handler = MagicMock()
mock_handler.get_llm.return_value = MagicMock()
mock_llm_handler.return_value = mock_handler
pipeline = RAGPipeline()
template_str = pipeline.prompt_template.messages[0].prompt.template
main_doc_pos = template_str.find("{main_document_section}")
context_pos = template_str.find("{context}")
# Main document should come before context
assert main_doc_pos < context_pos
|