tests/integration/episodes/test_lancedb_integration.py

"""
Integration tests for LanceDB vector search and embedding generation

Tests cover:
1. LanceDB client connection and table creation
2. Vector embedding generation for episode content
3. Semantic search with cosine similarity
4. Hybrid search (temporal + semantic)
5. Batch embedding operations
6. Query performance and result ranking

These tests use actual LanceDB in-memory database for realistic testing.
External embedding APIs are mocked for reliability and speed.
"""

import pytest
import tempfile
import shutil
from datetime import datetime, timedelta
from pathlib import Path
from unittest.mock import Mock, patch, AsyncMock
from sqlalchemy.orm import Session
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

from core.episode_retrieval_service import EpisodeRetrievalService
from core.models import (
    AgentRegistry,
    AgentStatus,
    Episode,
    EpisodeSegment,
    CanvasAudit,
    AgentFeedback,
    ChatSession,
    ChatMessage,
    User,
)


# ============================================================================
# Test Configuration
# ============================================================================

@pytest.fixture(scope="module")
def temp_lancedb_dir():
    """Create temporary directory for LanceDB data."""
    temp_dir = tempfile.mkdtemp()
    yield temp_dir
    # Cleanup
    try:
        shutil.rmtree(temp_dir)
    except Exception:
        pass


@pytest.fixture(scope="function")
def test_db():
    """Create in-memory SQLite database for testing."""
    engine = create_engine("sqlite:///:memory:")
    from core.models import Base
    Base.metadata.create_all(engine)
    SessionLocal = sessionmaker(bind=engine)
    session = SessionLocal()
    yield session
    session.close()


@pytest.fixture
def mock_governance():
    """Mock governance service to always allow access."""
    gov = Mock()
    gov.can_perform_action = Mock(return_value={
        "allowed": True,
        "agent_maturity": "SUPERVISED"
    })
    return gov


@pytest.fixture
def sample_user(test_db):
    """Create sample user."""
    user = User(
        id="user-123",
        email="test@example.com",
        name="Test User"
    )
    test_db.add(user)
    test_db.commit()
    return user


@pytest.fixture
def sample_agent(test_db):
    """Create sample agent."""
    agent = AgentRegistry(
        id="agent-456",
        name="TestAgent",
        status=AgentStatus.SUPERVISED,
        description="Test agent for LanceDB integration"
    )
    test_db.add(agent)
    test_db.commit()
    return agent


@pytest.fixture
def sample_episodes(test_db, sample_agent, sample_user):
    """Create sample episodes with varied content."""
    now = datetime.now()
    episodes = []

    episode_data = [
        {
            "title": "Data Analysis Dashboard",
            "description": "Created analytics dashboard with sales metrics",
            "summary": "Built comprehensive dashboard showing KPIs",
            "topics": ["analytics", "dashboard", "sales"],
            "importance": 0.8
        },
        {
            "title": "Customer Support Automation",
            "description": "Automated customer support ticket routing",
            "summary": "Implemented AI-based ticket classification",
            "topics": ["automation", "support", "tickets"],
            "importance": 0.9
        },
        {
            "title": "Report Generation",
            "description": "Generated monthly sales reports",
            "summary": "Automated PDF report generation",
            "topics": ["reports", "sales", "automation"],
            "importance": 0.7
        },
        {
            "title": "Database Migration",
            "description": "Migrated customer data to new schema",
            "summary": "Successfully migrated 100k records",
            "topics": ["database", "migration", "data"],
            "importance": 0.85
        },
        {
            "title": "API Integration",
            "description": "Integrated third-party payment API",
            "summary": "Connected Stripe payment processing",
            "topics": ["api", "integration", "payments"],
            "importance": 0.75
        }
    ]

    for i, data in enumerate(episode_data):
        episode = Episode(
            id=f"episode-{i}",
            title=data["title"],
            description=data["description"],
            summary=data["summary"],
            agent_id=sample_agent.id,
            user_id=sample_user.id,
            workspace_id="default",
            topics=data["topics"],
            entities=[],
            importance_score=data["importance"],
            status="completed",
            started_at=now - timedelta(days=i),
            ended_at=now - timedelta(days=i) + timedelta(hours=1),
            duration_seconds=3600,
            maturity_at_time="SUPERVISED",
            human_intervention_count=i % 3,
            constitutional_score=0.8 + (i * 0.02),
            decay_score=1.0 - (i * 0.1),
            access_count=10 - i,
            canvas_action_count=i,
            feedback_ids=[]
        )
        test_db.add(episode)

        # Add segments for each episode
        for j in range(3):
            segment = EpisodeSegment(
                id=f"segment-{i}-{j}",
                episode_id=episode.id,
                segment_type="conversation",
                sequence_order=j,
                content=f"Content segment {j} for {data['title']}",
                content_summary=f"Summary {j}",
                source_type="chat_message",
                source_id=f"msg-{j}"
            )
            test_db.add(segment)

        episodes.append(episode)

    test_db.commit()
    return episodes


# ============================================================================
# LanceDB Connection and Table Creation Tests
# ============================================================================

class TestLanceDBConnection:
    """Test LanceDB client initialization and table operations."""

    @pytest.mark.asyncio
    async def test_lancedb_handler_initialization(self, temp_lancedb_dir):
        """Test LanceDB handler can be initialized with temporary directory."""
        from core.lancedb_handler import LanceDBHandler, LANCEDB_AVAILABLE

        if not LANCEDB_AVAILABLE:
            pytest.skip("LanceDB not available")

        handler = LanceDBHandler(db_path=temp_lancedb_dir)

        # Verify handler initialized
        assert handler is not None
        assert handler.db_path == temp_lancedb_dir

    @pytest.mark.asyncio
    async def test_table_creation(self, test_db, temp_lancedb_dir, sample_agent):
        """Test creating episodes table in LanceDB."""
        from core.lancedb_handler import LanceDBHandler, LANCEDB_AVAILABLE

        if not LANCEDB_AVAILABLE:
            pytest.skip("LanceDB not available")

        handler = LanceDBHandler(db_path=temp_lancedb_dir)

        # Create table
        table_name = "test_episodes"
        handler.create_table(table_name)

        # Verify table exists
        assert table_name in handler.db.table_names()

    @pytest.mark.asyncio
    async def test_add_documents_to_lancedb(self, test_db, temp_lancedb_dir, sample_episodes):
        """Test adding episode documents to LanceDB."""
        from core.lancedb_handler import LanceDBHandler, LANCEDB_AVAILABLE

        if not LANCEDB_AVAILABLE:
            pytest.skip("LanceDB not available")

        handler = LanceDBHandler(db_path=temp_lancedb_dir)
        table_name = "episodes"
        handler.create_table(table_name)

        # Add episodes to LanceDB
        for episode in sample_episodes:
            content = f"""
Title: {episode.title}
Description: {episode.description}
Summary: {episode.summary}
Topics: {', '.join(episode.topics)}
            """.strip()

            metadata = {
                "episode_id": episode.id,
                "agent_id": episode.agent_id,
                "user_id": episode.user_id,
                "status": episode.status,
                "topics": episode.topics,
                "type": "episode"
            }

            handler.add_document(
                table_name=table_name,
                text=content,
                source=f"episode:{episode.id}",
                metadata=metadata,
                user_id=episode.user_id,
                extract_knowledge=False
            )

        # Verify documents added
        table = handler.db.open_table(table_name)
        count = len(table.to_arrow())
        assert count == len(sample_episodes)


# ============================================================================
# Vector Embedding Generation Tests
# ============================================================================

class TestEmbeddingGeneration:
    """Test vector embedding generation for episode content."""

    @pytest.mark.asyncio
    async def test_embed_text_generates_vector(self, temp_lancedb_dir):
        """Test text embedding generates fixed-length vector."""
        from core.lancedb_handler import LanceDBHandler, LANCEDB_AVAILABLE

        if not LANCEDB_AVAILABLE:
            pytest.skip("LanceDB not available")

        handler = LanceDBHandler(db_path=temp_lancedb_dir)

        text = "This is a test episode about data analytics"
        embedding = handler.embed_text(text)

        assert embedding is not None
        # Check it's a list or array
        assert isinstance(embedding, (list, tuple))
        # Check vector dimension (should be 384 for MiniLM-L6-v2)
        assert len(embedding) > 0
        # Check values are floats
        assert all(isinstance(x, (float, int)) for x in embedding)

    @pytest.mark.asyncio
    async def test_embedding_deterministic_for_same_text(self, temp_lancedb_dir):
        """Test embeddings are consistent for identical text."""
        from core.lancedb_handler import LanceDBHandler, LANCEDB_AVAILABLE

        if not LANCEDB_AVAILABLE:
            pytest.skip("LanceDB not available")

        handler = LanceDBHandler(db_path=temp_lancedb_dir)

        text = "Consistent test text"
        embedding1 = handler.embed_text(text)
        embedding2 = handler.embed_text(text)

        # Embeddings should be identical or very similar
        assert len(embedding1) == len(embedding2)

        # Calculate cosine similarity
        dot_product = sum(a * b for a, b in zip(embedding1, embedding2))
        magnitude1 = sum(a ** 2 for a in embedding1) ** 0.5
        magnitude2 = sum(b ** 2 for b in embedding2) ** 0.5
        similarity = dot_product / (magnitude1 * magnitude2) if magnitude1 * magnitude2 > 0 else 0

        # Should be nearly identical (similarity > 0.99)
        assert similarity > 0.99

    @pytest.mark.asyncio
    async def test_embedding_different_for_different_text(self, temp_lancedb_dir):
        """Test embeddings differ for semantically different text."""
        from core.lancedb_handler import LanceDBHandler, LANCEDB_AVAILABLE

        if not LANCEDB_AVAILABLE:
            pytest.skip("LanceDB not available")

        handler = LanceDBHandler(db_path=temp_lancedb_dir)

        text1 = "Data analytics and business intelligence"
        text2 = "Customer support and ticket management"

        embedding1 = handler.embed_text(text1)
        embedding2 = handler.embed_text(text2)

        # Calculate cosine similarity
        dot_product = sum(a * b for a, b in zip(embedding1, embedding2))
        magnitude1 = sum(a ** 2 for a in embedding1) ** 0.5
        magnitude2 = sum(b ** 2 for b in embedding2) ** 0.5
        similarity = dot_product / (magnitude1 * magnitude2) if magnitude1 * magnitude2 > 0 else 0

        # Should be different (similarity < 0.9 for unrelated topics)
        assert similarity < 0.9

    @pytest.mark.asyncio
    async def test_batch_embedding_performance(self, temp_lancedb_dir):
        """Test batch embedding multiple texts."""
        from core.lancedb_handler import LanceDBHandler, LANCEDB_AVAILABLE

        if not LANCEDB_AVAILABLE:
            pytest.skip("LanceDB not available")

        handler = LanceDBHandler(db_path=temp_lancedb_dir)

        texts = [
            "First episode about analytics",
            "Second episode about automation",
            "Third episode about reporting",
            "Fourth episode about database",
            "Fifth episode about API integration"
        ]

        import time
        start = time.time()
        embeddings = [handler.embed_text(text) for text in texts]
        duration = time.time() - start

        # All embeddings generated
        assert len(embeddings) == len(texts)
        assert all(e is not None for e in embeddings)

        # Performance check: should complete in reasonable time
        # (This is a soft check - just ensure it's not excessively slow)
        assert duration < 30  # 5 embeddings should take less than 30 seconds


# ============================================================================
# Semantic Search Tests
# ============================================================================

class TestSemanticSearch:
    """Test semantic similarity search via LanceDB."""

    @pytest.mark.asyncio
    async def test_semantic_search_returns_results(self, test_db, temp_lancedb_dir,
                                                     sample_episodes, sample_agent):
        """Test semantic search returns relevant episodes."""
        from core.lancedb_handler import LanceDBHandler, LANCEDB_AVAILABLE

        if not LANCEDB_AVAILABLE:
            pytest.skip("LanceDB not available")

        # Setup: Add episodes to LanceDB
        handler = LanceDBHandler(db_path=temp_lancedb_dir)
        table_name = "episodes"
        handler.create_table(table_name)

        for episode in sample_episodes:
            content = f"{episode.title} {episode.description} {episode.summary}"
            metadata = {
                "episode_id": episode.id,
                "agent_id": episode.agent_id,
                "type": "episode"
            }
            handler.add_document(
                table_name=table_name,
                text=content,
                source=f"episode:{episode.id}",
                metadata=metadata,
                user_id=episode.user_id,
                extract_knowledge=False
            )

        # Test search
        results = handler.search(
            table_name=table_name,
            query="data analytics dashboard",
            filter_str=f"agent_id == '{sample_agent.id}'",
            limit=3
        )

        assert len(results) > 0
        assert len(results) <= 3

    @pytest.mark.asyncio
    async def test_semantic_search_relevance_ranking(self, test_db, temp_lancedb_dir,
                                                       sample_episodes, sample_agent):
        """Test semantic search ranks results by relevance."""
        from core.lancedb_handler import LanceDBHandler, LANCEDB_AVAILABLE

        if not LANCEDB_AVAILABLE:
            pytest.skip("LanceDB not available")

        # Setup
        handler = LanceDBHandler(db_path=temp_lancedb_dir)
        table_name = "episodes"
        handler.create_table(table_name)

        for episode in sample_episodes:
            content = f"{episode.title} {episode.description} {episode.summary}"
            metadata = {
                "episode_id": episode.id,
                "agent_id": episode.agent_id,
                "type": "episode"
            }
            handler.add_document(
                table_name=table_name,
                text=content,
                source=f"episode:{episode.id}",
                metadata=metadata,
                user_id=episode.user_id,
                extract_knowledge=False
            )

        # Search for analytics-related content
        results = handler.search(
            table_name=table_name,
            query="analytics dashboard metrics",
            filter_str=f"agent_id == '{sample_agent.id}'",
            limit=5
        )

        # Results should be ranked by distance (lower = more similar)
        if len(results) > 1:
            distances = [r.get("_distance", 1.0) for r in results]
            # Check distances are sorted (non-decreasing)
            assert distances == sorted(distances)

    @pytest.mark.asyncio
    async def test_semantic_search_with_filter(self, test_db, temp_lancedb_dir,
                                                 sample_episodes, sample_agent):
        """Test semantic search with agent filtering."""
        from core.lancedb_handler import LanceDBHandler, LANCEDB_AVAILABLE

        if not LANCEDB_AVAILABLE:
            pytest.skip("LanceDB not available")

        # Setup
        handler = LanceDBHandler(db_path=temp_lancedb_dir)
        table_name = "episodes"
        handler.create_table(table_name)

        for episode in sample_episodes:
            content = f"{episode.title} {episode.description}"
            metadata = {
                "episode_id": episode.id,
                "agent_id": episode.agent_id,
                "type": "episode"
            }
            handler.add_document(
                table_name=table_name,
                text=content,
                source=f"episode:{episode.id}",
                metadata=metadata,
                user_id=episode.user_id,
                extract_knowledge=False
            )

        # Search with agent filter
        results = handler.search(
            table_name=table_name,
            query="automation",
            filter_str=f"agent_id == '{sample_agent.id}'",
            limit=10
        )

        # All results should be from the specified agent
        for result in results:
            metadata = result.get("metadata", {})
            if isinstance(metadata, str):
                import json
                metadata = json.loads(metadata)
            assert metadata.get("agent_id") == sample_agent.id


# ============================================================================
# Hybrid Search Tests
# ============================================================================

class TestHybridSearch:
    """Test hybrid search combining temporal and semantic relevance."""

    @pytest.mark.asyncio
    async def test_hybrid_search_combines_signals(self, test_db, temp_lancedb_dir,
                                                    sample_episodes, sample_agent, mock_governance):
        """Test hybrid search combines temporal and semantic scores."""
        from core.lancedb_handler import LanceDBHandler, LANCEDB_AVAILABLE

        if not LANCEDB_AVAILABLE:
            pytest.skip("LanceDB not available")

        # Setup LanceDB
        handler = LanceDBHandler(db_path=temp_lancedb_dir)
        table_name = "episodes"
        handler.create_table(table_name)

        for episode in sample_episodes:
            content = f"{episode.title} {episode.description}"
            metadata = {
                "episode_id": episode.id,
                "agent_id": episode.agent_id,
                "type": "episode"
            }
            handler.add_document(
                table_name=table_name,
                text=content,
                source=f"episode:{episode.id}",
                metadata=metadata,
                user_id=episode.user_id,
                extract_knowledge=False
            )

        # Create retrieval service with mocked LanceDB
        with patch('core.episode_retrieval_service.get_lancedb_handler', return_value=handler):
            service = EpisodeRetrievalService(test_db)
            service.governance = mock_governance

            # Test contextual retrieval (hybrid)
            result = await service.retrieve_contextual(
                agent_id=sample_agent.id,
                current_task="Create analytics dashboard",
                limit=5
            )

            assert "episodes" in result
            assert "count" in result
            assert isinstance(result["episodes"], list)
            assert result["count"] == len(result["episodes"])

    @pytest.mark.asyncio
    async def test_hybrid_search_weights(self, test_db, temp_lancedb_dir,
                                           sample_episodes, sample_agent, mock_governance):
        """Test hybrid search applies correct weights to temporal and semantic."""
        from core.lancedb_handler import LanceDBHandler, LANCEDB_AVAILABLE

        if not LANCEDB_AVAILABLE:
            pytest.skip("LanceDB not available")

        # Setup
        handler = LanceDBHandler(db_path=temp_lancedb_dir)
        table_name = "episodes"
        handler.create_table(table_name)

        for episode in sample_episodes:
            content = f"{episode.title} {episode.description}"
            metadata = {
                "episode_id": episode.id,
                "agent_id": episode.agent_id,
                "type": "episode"
            }
            handler.add_document(
                table_name=table_name,
                text=content,
                source=f"episode:{episode.id}",
                metadata=metadata,
                user_id=episode.user_id,
                extract_knowledge=False
            )

        with patch('core.episode_retrieval_service.get_lancedb_handler', return_value=handler):
            service = EpisodeRetrievalService(test_db)
            service.governance = mock_governance

            result = await service.retrieve_contextual(
                agent_id=sample_agent.id,
                current_task="automation and reporting",
                limit=3
            )

            # Results should have relevance scores
            assert "episodes" in result
            for episode in result["episodes"]:
                # Check for relevance_score in results
                assert "relevance_score" in episode or "id" in episode


# ============================================================================
# Performance Tests
# ============================================================================

class TestQueryPerformance:
    """Test LanceDB query performance and optimization."""

    @pytest.mark.asyncio
    async def test_search_performance_small_dataset(self, test_db, temp_lancedb_dir,
                                                      sample_episodes):
        """Test search performance with small dataset."""
        from core.lancedb_handler import LanceDBHandler, LANCEDB_AVAILABLE

        if not LANCEDB_AVAILABLE:
            pytest.skip("LanceDB not available")

        handler = LanceDBHandler(db_path=temp_lancedb_dir)
        table_name = "episodes"
        handler.create_table(table_name)

        # Add episodes
        for episode in sample_episodes:
            content = f"{episode.title} {episode.description}"
            metadata = {
                "episode_id": episode.id,
                "agent_id": episode.agent_id,
                "type": "episode"
            }
            handler.add_document(
                table_name=table_name,
                text=content,
                source=f"episode:{episode.id}",
                metadata=metadata,
                user_id=episode.user_id,
                extract_knowledge=False
            )

        # Measure search performance
        import time
        start = time.time()
        results = handler.search(
            table_name=table_name,
            query="test query",
            limit=5
        )
        duration = time.time() - start

        # Search should be fast (< 5 seconds for small dataset)
        assert duration < 5.0
        assert isinstance(results, list)

    @pytest.mark.asyncio
    async def test_retrieval_service_integration(self, test_db, temp_lancedb_dir,
                                                   sample_episodes, sample_agent, mock_governance):
        """Test EpisodeRetrievalService with real LanceDB operations."""
        from core.lancedb_handler import LanceDBHandler, LANCEDB_AVAILABLE

        if not LANCEDB_AVAILABLE:
            pytest.skip("LanceDB not available")

        # Setup
        handler = LanceDBHandler(db_path=temp_lancedb_dir)
        table_name = "episodes"
        handler.create_table(table_name)

        for episode in sample_episodes:
            content = f"{episode.title} {episode.description}"
            metadata = {
                "episode_id": episode.id,
                "agent_id": episode.agent_id,
                "type": "episode"
            }
            handler.add_document(
                table_name=table_name,
                text=content,
                source=f"episode:{episode.id}",
                metadata=metadata,
                user_id=episode.user_id,
                extract_knowledge=False
            )

        # Test service integration
        with patch('core.episode_retrieval_service.get_lancedb_handler', return_value=handler):
            service = EpisodeRetrievalService(test_db)
            service.governance = mock_governance

            # Test semantic retrieval
            result = await service.retrieve_semantic(
                agent_id=sample_agent.id,
                query="analytics and reporting",
                limit=3
            )

            assert "episodes" in result
            assert "count" in result
            assert isinstance(result["episodes"], list)
            assert result["count"] >= 0


# ============================================================================
# Edge Cases and Error Handling
# ============================================================================

class TestLanceDBEdgeCases:
    """Test edge cases and error handling in LanceDB operations."""

    @pytest.mark.asyncio
    async def test_empty_query_handling(self, temp_lancedb_dir):
        """Test handling of empty or invalid queries."""
        from core.lancedb_handler import LanceDBHandler, LANCEDB_AVAILABLE

        if not LANCEDB_AVAILABLE:
            pytest.skip("LanceDB not available")

        handler = LanceDBHandler(db_path=temp_lancedb_dir)
        table_name = "test_table"
        handler.create_table(table_name)

        # Empty query should return empty results or handle gracefully
        results = handler.search(
            table_name=table_name,
            query="",
            limit=5
        )

        assert isinstance(results, list)

    @pytest.mark.asyncio
    async def test_nonexistent_table_search(self, temp_lancedb_dir):
        """Test searching in non-existent table."""
        from core.lancedb_handler import LanceDBHandler, LANCEDB_AVAILABLE

        if not LANCEDB_AVAILABLE:
            pytest.skip("LanceDB not available")

        handler = LanceDBHandler(db_path=temp_lancedb_dir)

        # Should handle error gracefully
        try:
            results = handler.search(
                table_name="nonexistent_table",
                query="test",
                limit=5
            )
            # Either returns empty list or raises exception
            assert isinstance(results, list) or True
        except Exception:
            # Exception is acceptable
            pass

    @pytest.mark.asyncio
    async def test_special_characters_in_query(self, test_db, temp_lancedb_dir,
                                                 sample_episodes):
        """Test handling special characters in search queries."""
        from core.lancedb_handler import LanceDBHandler, LANCEDB_AVAILABLE

        if not LANCEDB_AVAILABLE:
            pytest.skip("LanceDB not available")

        handler = LanceDBHandler(db_path=temp_lancedb_dir)
        table_name = "episodes"
        handler.create_table(table_name)

        for episode in sample_episodes[:1]:  # Just add one
            content = f"{episode.title} {episode.description}"
            metadata = {
                "episode_id": episode.id,
                "agent_id": episode.agent_id,
                "type": "episode"
            }
            handler.add_document(
                table_name=table_name,
                text=content,
                source=f"episode:{episode.id}",
                metadata=metadata,
                user_id=episode.user_id,
                extract_knowledge=False
            )

        # Query with special characters
        special_queries = [
            "data & analytics",
            "reporting (monthly)",
            "API + integration",
            "automation/support",
        ]

        for query in special_queries:
            try:
                results = handler.search(
                    table_name=table_name,
                    query=query,
                    limit=5
                )
                assert isinstance(results, list)
            except Exception:
                # Some special chars might cause issues - that's acceptable
                pass