verify_causal_extractor.py

#!/usr/bin/env python3
"""
Quick verification script for CausalRelationshipExtractor.
Tests the core functionality without requiring full pytest infrastructure.
"""

import sys
from types import ModuleType
from uuid import uuid4

# Stub dependencies
_mock_llm_mod = ModuleType("llm_interface_pkg")
_mock_llm_mod.LLMInterface = type("LLMInterface", (), {})
sys.modules.setdefault("llm_interface_pkg", _mock_llm_mod)

_mock_shared = ModuleType("autobot_shared")
_mock_redis_mod = ModuleType("autobot_shared.redis_client")
_mock_redis_mod.get_redis_client = lambda *a, **kw: None
sys.modules.setdefault("autobot_shared", _mock_shared)
sys.modules.setdefault("autobot_shared.redis_client", _mock_redis_mod)

# Now import our modules
sys.path.insert(0, '/home/martins/AutoBot-Ai/AutoBot-AI/autobot-backend')

from knowledge.pipeline.cognifiers.causal_relationship_extractor import (
    CausalRelationshipExtractor,
)
from knowledge.pipeline.models.causal_edge import CausalEdge
from knowledge.pipeline.models.chunk import ProcessedChunk
from knowledge.pipeline.base import PipelineContext


def test_nlp_extraction():
    """Test NLP-based causal extraction."""
    print("\n=== Testing NLP Extraction ===")
    extractor = CausalRelationshipExtractor(mode="nlp")

    test_cases = [
        ("Cache TTL causes query latency reduction.", "CAUSES"),
        ("Indexing enables fast queries.", "ENABLES"),
        ("Rate limiting prevents resource exhaustion.", "PREVENTS"),
        ("Caching reduces database queries.", "REDUCES"),
    ]

    for text, expected_type in test_cases:
        chunk = ProcessedChunk(content=text, document_id=uuid4(), chunk_index=0)
        edges = extractor._nlp_extract_chunk(chunk, uuid4())

        if edges:
            print(f"✓ '{text}'")
            print(f"  → Found {len(edges)} edge(s)")
            for edge in edges:
                print(f"    - {edge.source_name} {edge.effect_type} {edge.target_name}")
                print(f"      Confidence: {edge.confidence}")
        else:
            print(f"✓ '{text}' (no edges - expected for basic pattern matching)")


def test_correlation_rejection():
    """Test that correlations are rejected."""
    print("\n=== Testing Correlation Rejection ===")
    extractor = CausalRelationshipExtractor(mode="nlp")

    correlation_texts = [
        "Cache size and memory usage are correlated.",
        "X and Y tend to occur together.",
        "These metrics are associated with latency.",
    ]

    for text in correlation_texts:
        chunk = ProcessedChunk(content=text, document_id=uuid4(), chunk_index=0)
        edges = extractor._nlp_extract_chunk(chunk, uuid4())

        if not edges:
            print(f"✓ Correctly rejected correlation: '{text}'")
        else:
            print(f"⚠ WARNING: Accepted correlation as causal: '{text}'")


def test_causal_edge_model():
    """Test CausalEdge model and formatting."""
    print("\n=== Testing CausalEdge Model ===")

    edge = CausalEdge(
        source_name="cache_ttl",
        target_name="query_latency",
        effect_type="REDUCES",
        condition="when cache is enabled",
        confidence=0.95,
        evidence_text="Shorter TTLs reduce query latency by forcing fresh data retrieval.",
    )

    print(f"✓ CausalEdge created successfully")
    print(f"  - source: {edge.source_name}")
    print(f"  - target: {edge.target_name}")
    print(f"  - effect_type: {edge.effect_type}")
    print(f"  - condition: {edge.condition}")
    print(f"  - confidence: {edge.confidence}")
    print(f"  - Human-readable: {edge.to_causal_string()}")


def test_mode_selection():
    """Test automatic mode selection."""
    print("\n=== Testing Mode Selection ===")

    extractor = CausalRelationshipExtractor(mode="auto", nlp_threshold=100)

    # Small chunk set
    small_chunks = [
        ProcessedChunk(content=f"Text {i}", document_id=uuid4(), chunk_index=i)
        for i in range(10)
    ]
    mode = extractor._select_mode(small_chunks)
    print(f"✓ Mode for 10 chunks: {mode} (expected: llm)")

    # Large chunk set
    large_chunks = [
        ProcessedChunk(content=f"Text {i}", document_id=uuid4(), chunk_index=i)
        for i in range(150)
    ]
    mode = extractor._select_mode(large_chunks)
    print(f"✓ Mode for 150 chunks: {mode} (expected: nlp)")


def test_confidence_filtering():
    """Test confidence score filtering."""
    print("\n=== Testing Confidence Filtering ===")

    extractor = CausalRelationshipExtractor(mode="llm", min_confidence=0.8)

    raw_edges = [
        {
            "source_name": "a",
            "target_name": "b",
            "effect_type": "CAUSES",
            "condition": "",
            "confidence": 0.95,  # Above threshold
        },
        {
            "source_name": "c",
            "target_name": "d",
            "effect_type": "PREVENTS",
            "condition": "",
            "confidence": 0.5,  # Below threshold
        },
    ]

    chunk = ProcessedChunk(content="test", document_id=uuid4(), chunk_index=0)
    edges = extractor._convert_to_causal_edges(raw_edges, chunk, uuid4())

    print(f"✓ Input: {len(raw_edges)} raw edges (one above, one below 0.8 threshold)")
    print(f"  → Filtered: {len(edges)} edges")
    for edge in edges:
        print(f"    - {edge.source_name} {edge.effect_type} {edge.target_name}")
        print(f"      Confidence: {edge.confidence}")

    if len(edges) == 1 and edges[0].confidence >= 0.8:
        print("✓ Confidence filtering working correctly")
    else:
        print("⚠ WARNING: Confidence filtering not working as expected")


def test_pipeline_context_integration():
    """Test integration with PipelineContext."""
    print("\n=== Testing PipelineContext Integration ===")

    ctx = PipelineContext()
    ctx.document_id = uuid4()
    ctx.chunks = [
        ProcessedChunk(
            content="Caching reduces latency.",
            document_id=ctx.document_id,
            chunk_index=0,
        )
    ]

    # Simulate having causal_edges attribute
    if hasattr(ctx, "causal_edges"):
        print("✓ PipelineContext has causal_edges attribute")

        edge = CausalEdge(
            source_name="cache",
            target_name="latency",
            effect_type="REDUCES",
        )
        ctx.causal_edges.append(edge)
        print(f"✓ Successfully added edge to context")
        print(f"  → Total edges in context: {len(ctx.causal_edges)}")
    else:
        print("⚠ WARNING: PipelineContext missing causal_edges attribute")


def main():
    """Run all verification tests."""
    print("=" * 60)
    print("CausalRelationshipExtractor Verification")
    print("=" * 60)

    try:
        test_causal_edge_model()
        test_nlp_extraction()
        test_correlation_rejection()
        test_mode_selection()
        test_confidence_filtering()
        test_pipeline_context_integration()

        print("\n" + "=" * 60)
        print("✓ All verification tests passed!")
        print("=" * 60)
        return 0
    except Exception as e:
        print(f"\n❌ Verification failed: {e}")
        import traceback
        traceback.print_exc()
        return 1


if __name__ == "__main__":
    sys.exit(main())