library-desk/tests/test_hybrid_rag.py

"""
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.clients.content_extractor import ContentExtractor
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, neo4j_test_uri) -> AsyncGenerator[Neo4jClient, None]:
    """Get connected Neo4j client."""
    client = Neo4jClient(
        uri=neo4j_test_uri,
        user=settings.neo4j_user,
        password=settings.neo4j_password
    )
    await client.connect()
    yield client
    await client.close()


@pytest.fixture
def qdrant_client(qdrant_test_url) -> QdrantClientWrapper:
    """Get Qdrant client."""
    return QdrantClientWrapper(url=qdrant_test_url)


@pytest_asyncio.fixture
async def wiki_client(wikijs_test_config) -> AsyncGenerator[WikiJSClient, None]:
    """Get Wiki.js client."""
    client = WikiJSClient(
        base_url=wikijs_test_config["base_url"],
        api_token=wikijs_test_config["api_token"]
    )
    yield client


@pytest.fixture
def searxng_client(searxng_test_url) -> SearXNGClient:
    """Get SearXNG client."""
    return SearXNGClient(base_url=searxng_test_url)


@pytest.fixture
def ollama_client(ollama_test_config) -> OllamaClient:
    """Get Ollama client."""
    return OllamaClient(
        base_url=ollama_test_config["base_url"],
        model=ollama_test_config["model"]
    )


@pytest.fixture
def content_extractor(settings) -> ContentExtractor:
    """Get ContentExtractor client."""
    return ContentExtractor(
        timeout=settings.content_extraction_timeout,
        max_length=settings.content_max_length
    )


@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,
    content_extractor,
    settings
):
    """Get HybridRAGService instance."""
    return HybridRAGService(
        vector_service=vector_service,
        graph_service=graph_service,
        searxng_client=searxng_client,
        ollama_client=ollama_client,
        content_extractor=content_extractor,
        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 two-stage Reciprocal Rank Fusion algorithm."""

    def test_wiki_merge_single_source(self, hybrid_rag_service):
        """Test wiki merge with single source (vector only)."""
        vector_results = [
            {"page_id": 1, "title": "Doc 1", "content": "test"},
            {"page_id": 2, "title": "Doc 2", "content": "test"}
        ]

        merged = hybrid_rag_service._merge_wiki_sources(vector_results, [], k=60)

        assert len(merged) == 2
        assert merged[0]["wiki_rrf_score"] > merged[1]["wiki_rrf_score"]  # Rank 1 > Rank 2
        assert merged[0]["found_by"] == ["vector"]

    def test_wiki_merge_multiple_sources_same_doc(self, hybrid_rag_service):
        """Test wiki merge with same document from vector and graph."""
        vector_results = [{"page_id": 1, "title": "Doc 1", "content": "test"}]
        graph_results = [{"page_id": 1, "title": "Doc 1", "content": ""}]

        merged = hybrid_rag_service._merge_wiki_sources(vector_results, graph_results, k=60)

        assert len(merged) == 1  # Deduplicated
        assert len(merged[0]["found_by"]) == 2  # Both sources
        assert "vector" in merged[0]["found_by"]
        assert "graph" in merged[0]["found_by"]
        # Wiki RRF score should be sum: 1/(60+1) + 1/(60+1)
        expected_score = 1/61 + 1/61
        assert abs(merged[0]["wiki_rrf_score"] - expected_score) < 0.001

    def test_final_rrf_wiki_and_web(self, hybrid_rag_service):
        """Test final RRF between wiki and web results."""
        # Pre-merged wiki results
        wiki_results = [
            {"page_id": 1, "title": "Wiki 1", "content": "test", "found_by": ["vector"]}
        ]
        web_results = [
            {"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(wiki_results, web_results, k=60)

        assert len(fused) == 3
        # Wiki rank 1 and web rank 1 should have same RRF score
        wiki_score = next(r["rrf_score"] for r in fused if r["source_type"] == "wiki")
        web_score = next(r["rrf_score"] for r in fused if r["source_type"] == "web")
        assert abs(wiki_score - web_score) < 0.001  # Equal footing


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}

        import uuid as _uuid
        search_id = await hybrid_rag_service._persist_search_for_librarian(
            search_id=str(_uuid.uuid4()),
            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}")