scripts/validate_system.py

#!/usr/bin/env python3
"""
Comprehensive system validation script
Tests all components: parser, vector DB, hybrid search, evidence grounding
"""

import asyncio
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent))

from services.neo4j_service import Neo4jService
from services.mongo_service import MongoService
from services.lancedb_service import LanceDBService
from services.vector_store_service import VectorStoreService
from services.llm_service import LLMService
from memory.graph_memory import GraphMemory
from utils.logger import get_logger

logger = get_logger("validation")


class SystemValidator:
    """Validates all system components."""

    def __init__(self):
        self.results = {
            "passed": [],
            "failed": [],
            "warnings": [],
        }

    def log_pass(self, test_name: str, message: str = ""):
        """Log passing test."""
        self.results["passed"].append(f"{test_name}: {message}".strip())
        logger.info(f"✓ {test_name}")

    def log_fail(self, test_name: str, error: str):
        """Log failing test."""
        self.results["failed"].append(f"{test_name}: {error}")
        logger.error(f"✗ {test_name}: {error}")

    def log_warning(self, test_name: str, warning: str):
        """Log warning."""
        self.results["warnings"].append(f"{test_name}: {warning}")
        logger.warning(f"⚠ {test_name}: {warning}")

    async def validate_all(self):
        """Run all validation tests."""
        logger.info("=" * 70)
        logger.info("CITATIONEDGE SYSTEM VALIDATION")
        logger.info("=" * 70)

        # Test individual components
        await self.validate_neo4j()
        await self.validate_mongodb()
        await self.validate_lancedb()
        await self.validate_embeddings()
        await self.validate_llm()
        await self.validate_memory()

        # Print results
        self.print_results()

    async def validate_neo4j(self):
        """Validate Neo4j connection and schema."""
        logger.info("\n📊 Testing Neo4j...")

        try:
            neo4j = Neo4jService()
            neo4j.connect()
            self.log_pass("Neo4j Connection")

            # Test basic query
            result = neo4j.run("MATCH (n) RETURN count(n) AS count")
            if result:
                count = result[0].get("count", 0)
                self.log_pass("Neo4j Query", f"{count} nodes in database")
            else:
                self.log_fail("Neo4j Query", "No result returned")

            # Check indexes
            try:
                neo4j.create_indexes()
                self.log_pass("Neo4j Indexes", "Created/verified")
            except Exception as e:
                self.log_warning("Neo4j Indexes", f"Creation had issue: {e}")

            # Check node types
            node_types = neo4j.run(
                "MATCH (n) RETURN DISTINCT labels(n)[0] AS type, count(*) AS count"
            )
            if node_types:
                types_str = ", ".join(
                    [
                        f"{row.get('type', '?')}({row.get('count', 0)})"
                        for row in node_types[:5]
                    ]
                )
                self.log_pass("Neo4j Node Types", types_str)

        except Exception as e:
            self.log_fail("Neo4j", str(e))

    async def validate_mongodb(self):
        """Validate MongoDB connection."""
        logger.info("\n💾 Testing MongoDB...")

        try:
            mongo = MongoService()
            await mongo.ping()
            self.log_pass("MongoDB Connection")

            # Test collection access
            collections = await mongo.db.list_collection_names()
            self.log_pass("MongoDB Collections", f"{len(collections)} collections")

        except Exception as e:
            self.log_warning("MongoDB", f"Connection failed (optional): {e}")

    async def validate_lancedb(self):
        """Validate LanceDB vector database."""
        logger.info("\n🔍 Testing LanceDB...")

        try:
            vector_db = LanceDBService()

            # Check directory exists
            db_path = Path(vector_db.db_path)
            if db_path.exists():
                self.log_pass("LanceDB Directory", str(db_path))
            else:
                self.log_warning("LanceDB Directory", f"Creating {db_path}")
                db_path.mkdir(parents=True, exist_ok=True)

            # List tables
            tables = vector_db.list_tables()
            self.log_pass("LanceDB Tables", f"{len(tables)} tables")

            # Test create/search
            test_data = [
                {"id": "test_1", "text": "test text", "vector": [0.1] * 768},
                {"id": "test_2", "text": "another test", "vector": [0.2] * 768},
            ]

            success = vector_db.create_table("test_validation", test_data)
            if success:
                self.log_pass("LanceDB Create Table", "test_validation created")

                # Test search
                results = vector_db.search("test_validation", [0.15] * 768, top_k=1)
                if results:
                    self.log_pass("LanceDB Search", f"Found {len(results)} results")
                else:
                    self.log_warning("LanceDB Search", "No results returned")
            else:
                self.log_fail("LanceDB Create Table", "Failed to create test table")

        except Exception as e:
            self.log_fail("LanceDB", str(e))

    async def validate_embeddings(self):
        """Validate embedding model."""
        logger.info("\n🧬 Testing Embeddings...")

        try:
            embedder = VectorStoreService()

            # Test single embedding
            text = "test sentence for embedding"
            embedding = embedder.embed(text)

            if embedding and len(embedding) == 768:
                self.log_pass("Embedding Model", f"768-dim vector generated")
            else:
                self.log_fail(
                    "Embedding Model",
                    f"Unexpected dimension: {len(embedding) if embedding else 0}",
                )

            # Test batch embedding
            texts = ["text1", "text2", "text3"]
            embeddings = embedder.embed_batch(texts)

            if embeddings and len(embeddings) == 3:
                self.log_pass("Batch Embedding", f"Generated 3 embeddings")
            else:
                self.log_fail(
                    "Batch Embedding",
                    f"Expected 3, got {len(embeddings) if embeddings else 0}",
                )

            # Test similarity
            if embedding:
                sim = embedder.cosine_similarity(embedding, embedding)
                if abs(sim - 1.0) < 0.01:
                    self.log_pass("Cosine Similarity", "Self-similarity ≈ 1.0")
                else:
                    self.log_fail("Cosine Similarity", f"Expected 1.0, got {sim}")

        except Exception as e:
            self.log_fail("Embeddings", str(e))

    async def validate_llm(self):
        """Validate LLM service."""
        logger.info("\n🤖 Testing LLM...")

        try:
            llm = LLMService()

            # Test basic completion
            response = await llm.complete(
                system="You are a helpful assistant.",
                user="What is 2+2?",
            )

            if response and "4" in response:
                self.log_pass("LLM Completion", "Got valid response")
            else:
                self.log_warning(
                    "LLM Completion", f"Unexpected response: {response[:50]}"
                )

        except Exception as e:
            self.log_fail("LLM", str(e))

    async def validate_memory(self):
        """Validate memory/retrieval layer."""
        logger.info("\n🧠 Testing Memory Layer...")

        try:
            neo4j = Neo4jService()
            embedder = VectorStoreService()
            vector_db = LanceDBService()

            graph_memory = GraphMemory(neo4j, embedder, vector_db)
            self.log_pass("GraphMemory Initialization")

            # Test basic retrieval (if data exists)
            try:
                docs = neo4j.run("MATCH (d:Document) RETURN d.doc_id LIMIT 1")
                if docs:
                    doc_id = docs[0].get("doc_id")
                    context = graph_memory.get_sections_context(doc_id)
                    if context:
                        token_count = len(context.split())
                        self.log_pass(
                            "Sections Context", f"Retrieved {token_count} tokens"
                        )
                    else:
                        self.log_warning(
                            "Sections Context", "Empty context (expected if no data)"
                        )
                else:
                    self.log_warning(
                        "Sections Context", "No documents in database (expected)"
                    )
            except Exception as e:
                self.log_warning("Sections Context", f"Retrieval test failed: {e}")

        except Exception as e:
            self.log_fail("Memory Layer", str(e))

    def print_results(self):
        """Print validation results."""
        logger.info("\n" + "=" * 70)
        logger.info("VALIDATION RESULTS")
        logger.info("=" * 70)

        # Passed tests
        if self.results["passed"]:
            logger.info(f"\n✓ PASSED ({len(self.results['passed'])})")
            for test in self.results["passed"]:
                logger.info(f"  • {test}")

        # Failed tests
        if self.results["failed"]:
            logger.info(f"\n✗ FAILED ({len(self.results['failed'])})")
            for test in self.results["failed"]:
                logger.error(f"  • {test}")

        # Warnings
        if self.results["warnings"]:
            logger.info(f"\n⚠ WARNINGS ({len(self.results['warnings'])})")
            for test in self.results["warnings"]:
                logger.warning(f"  • {test}")

        # Summary
        total = len(self.results["passed"]) + len(self.results["failed"])
        passed = len(self.results["passed"])

        logger.info("\n" + "=" * 70)
        logger.info(f"SUMMARY: {passed}/{total} tests passed")

        if self.results["failed"]:
            logger.info("Status: ✗ VALIDATION FAILED")
            return False
        else:
            logger.info("Status: ✓ VALIDATION PASSED")
            return True


async def main():
    """Run validation."""
    validator = SystemValidator()
    success = await validator.validate_all()

    sys.exit(0 if success else 1)


if __name__ == "__main__":
    asyncio.run(main())