feat(library-desk): implement HybridRAG query system
HybridRAG Service: - Combine vector (Qdrant), graph (Neo4j), and web (SearXNG) search - Reciprocal Rank Fusion (RRF) for result merging - LLM re-ranking with mistral-nemo - Graph enrichment with related dossiers - Query enhancement with keyword/synonym extraction - Search result persistence for offline processing Router: - POST /query/hybrid endpoint - Configurable search limits per source - Enable/disable individual sources - Timing breakdown for performance monitoring Models: - HybridRAGRequest, HybridRAGResponse - HybridRAGResult with source tracking - KeywordExtraction for query analysis - TimingBreakdown for performance metrics Tests: - End-to-end HybridRAG query tests - RRF fusion algorithm validation - Multi-source result merging
This commit is contained in:
@@ -0,0 +1,711 @@
|
||||
"""
|
||||
Comprehensive tests for HybridRAG system.
|
||||
|
||||
Tests cover all 6 phases:
|
||||
- Phase 0: Query Enhancement (keyword/synonym extraction)
|
||||
- Phase 1: Parallel Retrieval (vector + graph + web)
|
||||
- Phase 2: RRF Fusion
|
||||
- Phase 3: Enrichment (related dossiers)
|
||||
- Phase 4: LLM Re-ranking
|
||||
- Phase 5: Context Formatting
|
||||
- Phase 6: Persistence (search storage)
|
||||
|
||||
Uses 'llm-tester' user to avoid contaminating production data.
|
||||
|
||||
Run with: pytest tests/test_hybrid_rag.py -v -s
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from typing import AsyncGenerator
|
||||
import json
|
||||
|
||||
from src.clients.neo4j_client import Neo4jClient
|
||||
from src.clients.qdrant_client import QdrantClientWrapper
|
||||
from src.clients.wikijs_client import WikiJSClient
|
||||
from src.clients.searxng_client import SearXNGClient
|
||||
from src.clients.ollama_client import OllamaClient
|
||||
from src.services.hybrid_rag_service import HybridRAGService
|
||||
from src.services.vector_service import VectorService
|
||||
from src.services.graph_service import GraphService
|
||||
from src.models.hybrid_rag import HybridRAGConfig, HybridRAGRequest
|
||||
from src.config import get_settings
|
||||
|
||||
# Test user to isolate test data
|
||||
TEST_USER = "llm-tester"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def settings():
|
||||
"""Get application settings."""
|
||||
return get_settings()
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def neo4j_client(settings) -> AsyncGenerator[Neo4jClient, None]:
|
||||
"""Get connected Neo4j client."""
|
||||
client = Neo4jClient(
|
||||
uri=settings.neo4j_uri,
|
||||
user=settings.neo4j_user,
|
||||
password=settings.neo4j_password
|
||||
)
|
||||
await client.connect()
|
||||
yield client
|
||||
await client.close()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def qdrant_client(settings) -> QdrantClientWrapper:
|
||||
"""Get Qdrant client."""
|
||||
return QdrantClientWrapper(url=settings.qdrant_url)
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def wiki_client(settings) -> AsyncGenerator[WikiJSClient, None]:
|
||||
"""Get Wiki.js client."""
|
||||
client = WikiJSClient(
|
||||
base_url=settings.wikijs_url,
|
||||
username=settings.wikijs_username,
|
||||
password=settings.wikijs_password
|
||||
)
|
||||
yield client
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def searxng_client(settings) -> SearXNGClient:
|
||||
"""Get SearXNG client."""
|
||||
return SearXNGClient(base_url=settings.searxng_url)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def ollama_client(settings) -> OllamaClient:
|
||||
"""Get Ollama client."""
|
||||
return OllamaClient(base_url=settings.ollama_url)
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def vector_service(qdrant_client, wiki_client, ollama_client):
|
||||
"""Get VectorService instance."""
|
||||
return VectorService(qdrant_client, wiki_client, ollama_client)
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def graph_service(neo4j_client, wiki_client):
|
||||
"""Get GraphService instance."""
|
||||
return GraphService(neo4j_client, wiki_client)
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def hybrid_rag_service(
|
||||
vector_service,
|
||||
graph_service,
|
||||
searxng_client,
|
||||
ollama_client,
|
||||
settings
|
||||
):
|
||||
"""Get HybridRAGService instance."""
|
||||
return HybridRAGService(
|
||||
vector_service=vector_service,
|
||||
graph_service=graph_service,
|
||||
searxng_client=searxng_client,
|
||||
ollama_client=ollama_client,
|
||||
settings=settings
|
||||
)
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def test_wiki_page(wiki_client):
|
||||
"""
|
||||
Create test wiki page for llm-tester user.
|
||||
|
||||
Creates a page about Docker and Kubernetes for testing.
|
||||
"""
|
||||
from src.core.multi_tenancy import get_wikijs_namespace
|
||||
|
||||
namespace = get_wikijs_namespace(TEST_USER)
|
||||
path = f"{namespace}/testing/docker-kubernetes"
|
||||
|
||||
# Create test page
|
||||
page_data = {
|
||||
"title": "Docker and Kubernetes Testing",
|
||||
"path": path,
|
||||
"content": """# Docker and Kubernetes
|
||||
|
||||
Docker is a containerization platform that packages applications into containers.
|
||||
Kubernetes (k8s) is an orchestration platform for managing Docker containers at scale.
|
||||
|
||||
## Key Technologies
|
||||
- Docker: Container runtime
|
||||
- Kubernetes: Orchestration platform
|
||||
- Helm: Package manager for Kubernetes
|
||||
- kubectl: Command-line tool for k8s
|
||||
|
||||
## Use Cases
|
||||
Our infrastructure uses Docker containers orchestrated by Kubernetes clusters.
|
||||
We deploy microservices using Helm charts and manage them with kubectl.
|
||||
""",
|
||||
"description": "Test page for HybridRAG testing",
|
||||
"tags": ["testing", "infrastructure", "docker"]
|
||||
}
|
||||
|
||||
try:
|
||||
# Delete if exists
|
||||
existing = await wiki_client.search_pages(query="Docker and Kubernetes Testing")
|
||||
for page in existing:
|
||||
if page.get("path") == path:
|
||||
await wiki_client.delete_page(page["id"])
|
||||
|
||||
# Create new
|
||||
page = await wiki_client.create_page(**page_data)
|
||||
yield page
|
||||
|
||||
# Cleanup
|
||||
try:
|
||||
await wiki_client.delete_page(page["id"])
|
||||
except:
|
||||
pass
|
||||
except Exception as e:
|
||||
pytest.skip(f"Could not create test page: {e}")
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def test_graph_data(graph_service, test_wiki_page):
|
||||
"""
|
||||
Populate graph with test data for llm-tester.
|
||||
|
||||
Extracts entities from test page.
|
||||
"""
|
||||
try:
|
||||
summary = await graph_service.update_from_page(
|
||||
page_id=test_wiki_page["id"],
|
||||
user=TEST_USER
|
||||
)
|
||||
yield summary
|
||||
except Exception as e:
|
||||
pytest.skip(f"Could not populate graph: {e}")
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def test_vector_data(vector_service, test_wiki_page):
|
||||
"""
|
||||
Populate vector DB with test data for llm-tester.
|
||||
|
||||
Creates embeddings from test page.
|
||||
"""
|
||||
try:
|
||||
summary = await vector_service.update_from_page(
|
||||
page_id=test_wiki_page["id"],
|
||||
user=TEST_USER
|
||||
)
|
||||
yield summary
|
||||
except Exception as e:
|
||||
pytest.skip(f"Could not populate vectors: {e}")
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Unit Tests - Individual Components
|
||||
# ============================================================================
|
||||
|
||||
class TestRRFFusion:
|
||||
"""Test Reciprocal Rank Fusion algorithm."""
|
||||
|
||||
def test_rrf_single_source(self, hybrid_rag_service):
|
||||
"""Test RRF with single source."""
|
||||
results_by_source = {
|
||||
"vector": [
|
||||
{"page_id": 1, "title": "Doc 1", "content": "test"},
|
||||
{"page_id": 2, "title": "Doc 2", "content": "test"}
|
||||
]
|
||||
}
|
||||
|
||||
fused = hybrid_rag_service._reciprocal_rank_fusion(results_by_source, k=60)
|
||||
|
||||
assert len(fused) == 2
|
||||
assert fused[0]["rrf_score"] > fused[1]["rrf_score"] # Rank 1 > Rank 2
|
||||
assert fused[0]["sources"] == ["vector"]
|
||||
|
||||
def test_rrf_multiple_sources_same_doc(self, hybrid_rag_service):
|
||||
"""Test RRF with same document from multiple sources."""
|
||||
results_by_source = {
|
||||
"vector": [{"page_id": 1, "title": "Doc 1", "content": "test"}],
|
||||
"graph": [{"page_id": 1, "title": "Doc 1", "content": ""}],
|
||||
}
|
||||
|
||||
fused = hybrid_rag_service._reciprocal_rank_fusion(results_by_source, k=60)
|
||||
|
||||
assert len(fused) == 1 # Deduplicated
|
||||
assert len(fused[0]["sources"]) == 2 # Both sources
|
||||
assert "vector" in fused[0]["sources"]
|
||||
assert "graph" in fused[0]["sources"]
|
||||
# RRF score should be sum: 1/(60+1) + 1/(60+1)
|
||||
expected_score = 1/61 + 1/61
|
||||
assert abs(fused[0]["rrf_score"] - expected_score) < 0.001
|
||||
|
||||
def test_rrf_web_results(self, hybrid_rag_service):
|
||||
"""Test RRF with web results (URL-based)."""
|
||||
results_by_source = {
|
||||
"web": [
|
||||
{"url": "https://example.com/1", "title": "Web 1", "content": "test"},
|
||||
{"url": "https://example.com/2", "title": "Web 2", "content": "test"}
|
||||
]
|
||||
}
|
||||
|
||||
fused = hybrid_rag_service._reciprocal_rank_fusion(results_by_source, k=60)
|
||||
|
||||
assert len(fused) == 2
|
||||
assert fused[0]["result"]["url"] == "https://example.com/1"
|
||||
|
||||
|
||||
class TestContextFormatting:
|
||||
"""Test context formatting for LLM."""
|
||||
|
||||
def test_format_basic(self, hybrid_rag_service):
|
||||
"""Test basic context formatting."""
|
||||
from src.models.hybrid_rag import HybridRAGResult
|
||||
|
||||
results = [
|
||||
HybridRAGResult(
|
||||
source_type="vector",
|
||||
title="Test Document",
|
||||
content="This is test content for formatting",
|
||||
page_id=1,
|
||||
rrf_score=0.5,
|
||||
final_rank=1,
|
||||
sources=["vector"]
|
||||
)
|
||||
]
|
||||
|
||||
context = hybrid_rag_service._format_context_for_llm(results)
|
||||
|
||||
assert "Test Document" in context
|
||||
assert "[VECTOR]" in context
|
||||
assert "test content" in context
|
||||
|
||||
def test_format_with_related_dossiers(self, hybrid_rag_service):
|
||||
"""Test context formatting with related dossiers."""
|
||||
from src.models.hybrid_rag import HybridRAGResult, RelatedDossier
|
||||
|
||||
results = [
|
||||
HybridRAGResult(
|
||||
source_type="vector+graph",
|
||||
title="Test Document",
|
||||
content="Content",
|
||||
page_id=1,
|
||||
rrf_score=0.5,
|
||||
final_rank=1,
|
||||
sources=["vector", "graph"],
|
||||
related_dossiers=[
|
||||
RelatedDossier(
|
||||
page_id=2,
|
||||
title="Related Doc",
|
||||
path="/test/related",
|
||||
tag="infrastructure",
|
||||
shared_entities=5
|
||||
)
|
||||
]
|
||||
)
|
||||
]
|
||||
|
||||
context = hybrid_rag_service._format_context_for_llm(results)
|
||||
|
||||
assert "Related research: infrastructure" in context
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Integration Tests - Phase Testing
|
||||
# ============================================================================
|
||||
|
||||
class TestPhase0_QueryEnhancement:
|
||||
"""Test Phase 0: Query Enhancement (keyword/synonym extraction)."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_extract_keywords_basic(self, hybrid_rag_service):
|
||||
"""Test basic keyword extraction."""
|
||||
query = "Docker container orchestration with Kubernetes"
|
||||
|
||||
keywords_data = await hybrid_rag_service._extract_keywords_and_synonyms(query)
|
||||
|
||||
assert "core_keywords" in keywords_data
|
||||
assert "entities" in keywords_data
|
||||
assert "synonyms" in keywords_data
|
||||
assert "expansions" in keywords_data
|
||||
|
||||
# Should extract Docker and Kubernetes
|
||||
all_terms = (
|
||||
keywords_data["core_keywords"] +
|
||||
keywords_data["entities"]
|
||||
)
|
||||
assert any("docker" in term.lower() for term in all_terms)
|
||||
assert any("kubernetes" in term.lower() or "k8s" in term.lower() for term in all_terms)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_extract_keywords_with_abbreviations(self, hybrid_rag_service):
|
||||
"""Test keyword extraction handles abbreviations."""
|
||||
query = "k8s cluster management"
|
||||
|
||||
keywords_data = await hybrid_rag_service._extract_keywords_and_synonyms(query)
|
||||
|
||||
# Should expand k8s to kubernetes
|
||||
all_data = json.dumps(keywords_data).lower()
|
||||
assert "k8s" in all_data or "kubernetes" in all_data
|
||||
|
||||
|
||||
class TestPhase1_ParallelRetrieval:
|
||||
"""Test Phase 1: Parallel Retrieval."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parallel_retrieval_all_sources(
|
||||
self,
|
||||
hybrid_rag_service,
|
||||
test_wiki_page,
|
||||
test_graph_data,
|
||||
test_vector_data
|
||||
):
|
||||
"""Test parallel retrieval from all sources."""
|
||||
config = HybridRAGConfig(
|
||||
enable_vector=True,
|
||||
enable_graph=True,
|
||||
enable_web=True,
|
||||
vector_limit=5,
|
||||
graph_limit=5,
|
||||
web_limit=3
|
||||
)
|
||||
|
||||
keywords_data = {
|
||||
"core_keywords": ["docker", "kubernetes"],
|
||||
"entities": ["Docker", "Kubernetes"],
|
||||
"synonyms": {"docker": ["container"], "kubernetes": ["k8s"]},
|
||||
"expansions": {"k8s": ["kubernetes"]}
|
||||
}
|
||||
|
||||
results = await hybrid_rag_service._retrieve_parallel(
|
||||
query="docker kubernetes",
|
||||
user=TEST_USER,
|
||||
config=config,
|
||||
keywords_data=keywords_data
|
||||
)
|
||||
|
||||
assert "vector" in results
|
||||
assert "graph" in results
|
||||
assert "web" in results
|
||||
assert "timing" in results
|
||||
|
||||
# Should have timing for each source
|
||||
assert results["timing"]["vector_ms"] >= 0
|
||||
assert results["timing"]["graph_ms"] >= 0
|
||||
assert results["timing"]["web_ms"] >= 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parallel_retrieval_graceful_degradation(self, hybrid_rag_service):
|
||||
"""Test graceful degradation when sources fail."""
|
||||
config = HybridRAGConfig(
|
||||
enable_vector=True,
|
||||
enable_graph=True,
|
||||
enable_web=True
|
||||
)
|
||||
|
||||
keywords_data = {"core_keywords": ["test"], "entities": [], "synonyms": {}, "expansions": {}}
|
||||
|
||||
# Even if some sources fail, should return results from working sources
|
||||
results = await hybrid_rag_service._retrieve_parallel(
|
||||
query="test query",
|
||||
user=TEST_USER,
|
||||
config=config,
|
||||
keywords_data=keywords_data
|
||||
)
|
||||
|
||||
# Should have all keys even if empty
|
||||
assert "vector" in results
|
||||
assert "graph" in results
|
||||
assert "web" in results
|
||||
|
||||
|
||||
class TestPhase3_Enrichment:
|
||||
"""Test Phase 3: Graph Enrichment."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_enrich_with_related_dossiers(
|
||||
self,
|
||||
hybrid_rag_service,
|
||||
graph_service,
|
||||
test_wiki_page,
|
||||
test_graph_data
|
||||
):
|
||||
"""Test enriching results with related dossiers."""
|
||||
# Create mock fused results
|
||||
fused_results = [
|
||||
{
|
||||
"result": {
|
||||
"page_id": test_wiki_page["id"],
|
||||
"title": test_wiki_page["title"],
|
||||
"content": "test"
|
||||
},
|
||||
"rrf_score": 0.5,
|
||||
"sources": ["vector"]
|
||||
}
|
||||
]
|
||||
|
||||
enriched = await hybrid_rag_service._enrich_with_related_dossiers(
|
||||
fused_results,
|
||||
user=TEST_USER
|
||||
)
|
||||
|
||||
assert len(enriched) == 1
|
||||
assert "related_dossiers" in enriched[0]
|
||||
# May or may not have related docs depending on graph state
|
||||
assert isinstance(enriched[0]["related_dossiers"], list)
|
||||
|
||||
|
||||
class TestPhase6_Persistence:
|
||||
"""Test Phase 6: Search Persistence."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_persist_search_creates_node(
|
||||
self,
|
||||
hybrid_rag_service,
|
||||
neo4j_client,
|
||||
test_wiki_page
|
||||
):
|
||||
"""Test that search persistence creates SearchQuery node."""
|
||||
keywords_data = {
|
||||
"core_keywords": ["docker", "kubernetes"],
|
||||
"entities": [],
|
||||
"synonyms": {},
|
||||
"expansions": {}
|
||||
}
|
||||
|
||||
raw_results = {
|
||||
"vector": [{"page_id": test_wiki_page["id"], "title": "Test", "content": "test"}],
|
||||
"graph": [],
|
||||
"web": []
|
||||
}
|
||||
|
||||
final_results = [
|
||||
{
|
||||
"result": {"page_id": test_wiki_page["id"], "title": "Test"},
|
||||
"rrf_score": 0.5,
|
||||
"final_rank": 1,
|
||||
"sources": ["vector"]
|
||||
}
|
||||
]
|
||||
|
||||
timing = {"total_ms": 1000}
|
||||
|
||||
search_id = await hybrid_rag_service._persist_search_for_librarian(
|
||||
query="test query",
|
||||
user=TEST_USER,
|
||||
keywords_data=keywords_data,
|
||||
raw_results=raw_results,
|
||||
final_results=final_results,
|
||||
timing=timing
|
||||
)
|
||||
|
||||
assert search_id is not None
|
||||
|
||||
# Verify SearchQuery node was created
|
||||
from src.core.multi_tenancy import get_neo4j_user_base_label
|
||||
user_label = get_neo4j_user_base_label(TEST_USER)
|
||||
|
||||
query = f"""
|
||||
MATCH (sq:{user_label}_SearchQuery:SearchQuery {{id: $search_id}})
|
||||
RETURN sq.query as query, sq.processed as processed
|
||||
"""
|
||||
|
||||
result = await neo4j_client.execute_query(query, {"search_id": search_id})
|
||||
assert len(result) == 1
|
||||
assert result[0]["query"] == "test query"
|
||||
assert result[0]["processed"] == False
|
||||
|
||||
# Cleanup
|
||||
cleanup_query = f"""
|
||||
MATCH (sq:{user_label}_SearchQuery:SearchQuery {{id: $search_id}})
|
||||
DETACH DELETE sq
|
||||
"""
|
||||
await neo4j_client.execute_query(cleanup_query, {"search_id": search_id})
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# End-to-End Tests
|
||||
# ============================================================================
|
||||
|
||||
class TestHybridRAG_EndToEnd:
|
||||
"""End-to-end tests for complete HybridRAG flow."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_full_search_pipeline(
|
||||
self,
|
||||
hybrid_rag_service,
|
||||
test_wiki_page,
|
||||
test_graph_data,
|
||||
test_vector_data
|
||||
):
|
||||
"""
|
||||
Test complete HybridRAG search pipeline with all 6 phases.
|
||||
|
||||
This is the main end-to-end test that validates:
|
||||
- Phase 0: Query enhancement
|
||||
- Phase 1: Parallel retrieval
|
||||
- Phase 2: RRF fusion
|
||||
- Phase 3: Enrichment
|
||||
- Phase 4: Re-ranking
|
||||
- Phase 5: Context formatting
|
||||
- Phase 6: Persistence
|
||||
"""
|
||||
query = "How does Docker work with Kubernetes?"
|
||||
config = HybridRAGConfig(
|
||||
vector_limit=5,
|
||||
graph_limit=5,
|
||||
web_limit=3,
|
||||
enable_reranking=True,
|
||||
enable_enrichment=True,
|
||||
final_result_count=10
|
||||
)
|
||||
|
||||
# Execute full search
|
||||
response = await hybrid_rag_service.search(
|
||||
query=query,
|
||||
user=TEST_USER,
|
||||
config=config
|
||||
)
|
||||
|
||||
# Validate response structure
|
||||
assert response.query == query
|
||||
assert response.keywords is not None
|
||||
assert response.results is not None
|
||||
assert response.context is not None
|
||||
assert response.source_counts is not None
|
||||
assert response.total_results >= 0
|
||||
assert response.timing is not None
|
||||
assert response.config_used == config
|
||||
assert response.search_id is not None
|
||||
|
||||
# Validate timing breakdown
|
||||
assert response.timing.query_enhancement_ms >= 0
|
||||
assert response.timing.vector_ms >= 0
|
||||
assert response.timing.graph_ms >= 0
|
||||
assert response.timing.web_ms >= 0
|
||||
assert response.timing.fusion_ms >= 0
|
||||
assert response.timing.enrichment_ms >= 0
|
||||
assert response.timing.reranking_ms >= 0
|
||||
assert response.timing.persistence_ms >= 0
|
||||
assert response.timing.total_ms >= 0
|
||||
|
||||
# Validate keywords extraction
|
||||
assert len(response.keywords.core_keywords) > 0
|
||||
|
||||
# Validate context is formatted
|
||||
assert len(response.context) > 0
|
||||
|
||||
# Log results for inspection
|
||||
print(f"\n=== HybridRAG E2E Test Results ===")
|
||||
print(f"Query: {response.query}")
|
||||
print(f"Total Results: {response.total_results}")
|
||||
print(f"Source Counts: {response.source_counts}")
|
||||
print(f"Keywords: {response.keywords.core_keywords}")
|
||||
print(f"Total Time: {response.timing.total_ms:.0f}ms")
|
||||
print(f"Search ID: {response.search_id}")
|
||||
|
||||
if response.results:
|
||||
print(f"\nTop Result:")
|
||||
top = response.results[0]
|
||||
print(f" Title: {top.title}")
|
||||
print(f" Source: {top.source_type}")
|
||||
print(f" RRF Score: {top.rrf_score:.4f}")
|
||||
print(f" Rank: {top.final_rank}")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_search_with_disabled_sources(
|
||||
self,
|
||||
hybrid_rag_service,
|
||||
test_wiki_page,
|
||||
test_vector_data
|
||||
):
|
||||
"""Test HybridRAG with some sources disabled."""
|
||||
config = HybridRAGConfig(
|
||||
enable_vector=True,
|
||||
enable_graph=False, # Disabled
|
||||
enable_web=False, # Disabled
|
||||
enable_reranking=False,
|
||||
final_result_count=5
|
||||
)
|
||||
|
||||
response = await hybrid_rag_service.search(
|
||||
query="docker containers",
|
||||
user=TEST_USER,
|
||||
config=config
|
||||
)
|
||||
|
||||
# Should only have vector results
|
||||
assert response.total_results >= 0
|
||||
if response.total_results > 0:
|
||||
assert all(
|
||||
"vector" in result.sources
|
||||
for result in response.results
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_search_performance_target(
|
||||
self,
|
||||
hybrid_rag_service,
|
||||
test_wiki_page,
|
||||
test_graph_data,
|
||||
test_vector_data
|
||||
):
|
||||
"""Test that search completes within performance target (<3.5s)."""
|
||||
import time
|
||||
|
||||
config = HybridRAGConfig()
|
||||
|
||||
start = time.time()
|
||||
response = await hybrid_rag_service.search(
|
||||
query="kubernetes orchestration",
|
||||
user=TEST_USER,
|
||||
config=config
|
||||
)
|
||||
duration_ms = (time.time() - start) * 1000
|
||||
|
||||
print(f"\nPerformance: {duration_ms:.0f}ms (target: <3500ms)")
|
||||
|
||||
# Soft assertion - warn if exceeds target
|
||||
if duration_ms > 3500:
|
||||
print(f"WARNING: Search exceeded 3.5s target ({duration_ms:.0f}ms)")
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Cleanup Tests
|
||||
# ============================================================================
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cleanup_test_data(neo4j_client, qdrant_client):
|
||||
"""
|
||||
Cleanup test data for llm-tester user.
|
||||
|
||||
Run this to clean up test data:
|
||||
pytest tests/test_hybrid_rag.py::test_cleanup_test_data -v -s
|
||||
"""
|
||||
from src.core.multi_tenancy import (
|
||||
get_neo4j_user_base_label,
|
||||
get_neo4j_user_label,
|
||||
get_qdrant_collection_name
|
||||
)
|
||||
|
||||
# Clean Neo4j
|
||||
user_base_label = get_neo4j_user_base_label(TEST_USER)
|
||||
user_doc_label = get_neo4j_user_label(TEST_USER)
|
||||
|
||||
# Delete all test user nodes
|
||||
delete_query = f"""
|
||||
MATCH (n)
|
||||
WHERE n:{user_base_label} OR n:{user_doc_label}
|
||||
DETACH DELETE n
|
||||
"""
|
||||
await neo4j_client.execute_query(delete_query, {})
|
||||
|
||||
# Clean Qdrant
|
||||
collection_name = get_qdrant_collection_name(TEST_USER)
|
||||
try:
|
||||
await qdrant_client.delete_collection(collection_name)
|
||||
except:
|
||||
pass
|
||||
|
||||
print(f"\n✓ Cleaned up test data for user: {TEST_USER}")
|
||||
Reference in New Issue
Block a user