replypier/backend/tests/knowledge/test_enhanced_website_kb.py

"""
Copyright 2024-2026 ReplyPier

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import pytest
import os
from unittest.mock import MagicMock, patch, call, ANY
from agno.document.base import Document
from app.core.config import settings
from app.knowledge.crawl_scope import DEFAULT_CRAWL_SCOPE
from app.knowledge.deadline import Deadline, SourceTimedOut
from app.knowledge.enhanced_website_kb import EnhancedWebsiteKnowledgeBase
from app.knowledge.failure_advice import FAILURE_ADVICE
from app.knowledge.enhanced_website_reader import (
    BotProtectionError,
    EmptyCrawlError,
    EnhancedWebsiteReader,
    FailureCause,
)

# Test Data
TEST_URLS = [
    "https://example.com",
    "https://test.com"
]

def _embedded(content, url):
    """A document as it looks once a real run has embedded it.

    Documents without an embedding are skipped by the vector store, so a fixture
    without one describes a source that stores nothing — which the guard now
    (correctly) fails.
    """
    document = Document(content=content, meta_data={"url": url})
    document.embedding = [0.1, 0.2, 0.3]
    return document


TEST_DOCUMENTS = [
    _embedded("Test content 1", "https://example.com/page1"),
    _embedded("Test content 2", "https://example.com/page2"),
]

@pytest.fixture
def mock_vector_db():
    """Create a mock vector database"""
    mock_db = MagicMock()
    mock_db.upsert_available.return_value = True
    mock_db.name_exists.return_value = False
    mock_db.doc_exists.return_value = False
    return mock_db

@pytest.fixture
def mock_reader():
    """Create a mock website reader"""
    mock = MagicMock(spec=EnhancedWebsiteReader)
    mock.read.return_value = TEST_DOCUMENTS
    # The real reader declares this as an int defaulting to 0; left as a Mock it
    # is truthy, which would read as "bot protection blocked every page".
    mock._challenge_blocked = 0
    return mock

class TestEnhancedWebsiteKnowledgeBase:
    def test_initialization_with_defaults(self):
        """Test initialization with default values"""
        kb = EnhancedWebsiteKnowledgeBase(urls=TEST_URLS)

        # Against settings, not literals: every one of these is configurable per
        # install (KB_MAX_DEPTH and friends), so hardcoding the defaults here
        # fails on any machine whose .env tunes a crawl.
        assert kb.urls == TEST_URLS
        assert isinstance(kb.reader, EnhancedWebsiteReader)
        assert kb.max_depth == settings.KB_MAX_DEPTH
        assert kb.max_links == settings.KB_MAX_LINKS
        assert kb.min_content_length == settings.KB_MIN_CONTENT_LENGTH
        assert kb.timeout == settings.KB_TIMEOUT
        assert kb.max_retries == settings.KB_MAX_RETRIES
        assert kb.max_workers == settings.KB_MAX_WORKERS
        assert kb.crawl_scope == DEFAULT_CRAWL_SCOPE
    
    def test_initialization_with_custom_params(self):
        """Test initialization with custom parameters"""
        kb = EnhancedWebsiteKnowledgeBase(
            urls=TEST_URLS,
            max_depth=5,
            max_links=20,
            min_content_length=200,
            timeout=60,
            max_retries=5,
            max_workers=8
        )
        
        assert kb.urls == TEST_URLS
        assert isinstance(kb.reader, EnhancedWebsiteReader)
        assert kb.max_depth == 5
        assert kb.max_links == 20
        assert kb.min_content_length == 200
        assert kb.timeout == 60
        assert kb.max_retries == 5
        assert kb.max_workers == 8
        
        # Verify reader was initialized with custom params
        assert kb.reader.max_depth == 5
        assert kb.reader.max_links == 20
        assert kb.reader.min_content_length == 200
        assert kb.reader.timeout == 60
        assert kb.reader.max_retries == 5
        assert kb.reader.max_workers == 8
    
    def test_initialization_with_custom_reader(self):
        """Test initialization with a custom reader"""
        custom_reader = EnhancedWebsiteReader(max_depth=4, max_links=15)
        kb = EnhancedWebsiteKnowledgeBase(urls=TEST_URLS, reader=custom_reader)

        assert kb.reader is custom_reader
        assert kb.reader.max_depth == 4
        assert kb.reader.max_links == 15

    def test_raise_if_nothing_stored(self):
        """A run that stored no pages must raise so the queue item fails.

        Previously only the bot-blocked case raised; an empty crawl for any other
        reason was marked COMPLETED, making a broken source look identical to a
        healthy one with no sub-pages.
        """
        reader = EnhancedWebsiteReader(max_links=5)
        kb = EnhancedWebsiteKnowledgeBase(urls=TEST_URLS, reader=reader)

        # Anything stored → never raises, whatever happened along the way.
        reader._challenge_blocked = 0
        kb._raise_if_nothing_stored(5)
        reader._challenge_blocked = 3
        kb._raise_if_nothing_stored(2)

        # Zero stored because bot protection blocked every page → the specific,
        # actionable error, which takes precedence.
        reader._challenge_blocked = 3
        with pytest.raises(BotProtectionError):
            kb._raise_if_nothing_stored(0)

        # Zero stored for any other reason → still a failure, not a success.
        reader._challenge_blocked = 0
        with pytest.raises(EmptyCrawlError) as excinfo:
            kb._raise_if_nothing_stored(0)
        # The message names the source so the dashboard error is actionable.
        assert TEST_URLS[0] in str(excinfo.value)

    def test_a_crawled_page_that_never_reached_the_vector_store_is_a_failure(self):
        """Crawled and stored are different numbers.

        chatring.ai crawled one document, failed to embed it, stored nothing,
        and was still marked COMPLETED — the org's vector table was empty while
        the dashboard showed a healthy source. The guard counts what survived
        the upsert, so a batch that stored nothing still fails the run.
        """
        reader = EnhancedWebsiteReader(max_links=5)
        reader._challenge_blocked = 0
        kb = EnhancedWebsiteKnowledgeBase(urls=TEST_URLS, reader=reader)

        with pytest.raises(EmptyCrawlError):
            kb._raise_if_nothing_stored(0)

    def test_failed_batch_reports_zero_stored(self, mock_vector_db):
        """An upsert that raises must not be counted as progress."""
        kb = EnhancedWebsiteKnowledgeBase(urls=TEST_URLS)
        kb.vector_db = mock_vector_db

        assert kb._process_document_batch(TEST_DOCUMENTS) == len(TEST_DOCUMENTS)

        mock_vector_db.upsert.side_effect = RuntimeError("connection reset")
        assert kb._process_document_batch(TEST_DOCUMENTS) == 0

    def test_a_run_whose_documents_never_embedded_fails(self, mock_vector_db, mock_reader):
        """Crawling is not storing.

        chatring.ai crawled one document, failed to embed it, stored nothing, and
        was still marked COMPLETED — the org's vector table was empty while the
        dashboard showed a healthy source. A document with no embedding is
        skipped by the vector store, so the run has produced nothing retrievable
        and must fail.
        """
        unembedded = Document(content="crawled but never embedded", meta_data={})
        unembedded.embedding = None
        mock_reader.read.return_value = [unembedded]

        kb = EnhancedWebsiteKnowledgeBase(urls=TEST_URLS[:1], reader=mock_reader)
        kb.vector_db = mock_vector_db
        mock_vector_db.embedder = None

        with pytest.raises(EmptyCrawlError):
            kb.load()

    @pytest.mark.parametrize("cause,expected", [
        (FailureCause.DNS, "does not resolve"),
        (FailureCause.TIMEOUT, "did not respond in time"),
        (FailureCause.BOT_BLOCKED, "refused our crawler"),
        (FailureCause.EMPTY_RESPONSE, "returned an empty page"),
        (FailureCause.THIN_CONTENT, "JavaScript"),
        (FailureCause.HTTP_ERROR, "error response"),
        (FailureCause.BLOCKED_HOST, "not allowed to fetch"),
    ])
    def test_the_message_names_the_cause_the_crawl_actually_hit(self, cause, expected):
        """Every distinct cause gets its own advice.

        One string covering "empty, require JavaScript, or be unreachable" was
        undiagnosable for the user and for us; the reader already knew which of
        the three it hit.
        """
        reader = EnhancedWebsiteReader(max_links=5)
        reader._challenge_blocked = 0
        reader._failure_causes = {cause: 2}
        kb = EnhancedWebsiteKnowledgeBase(urls=TEST_URLS, reader=reader)

        with pytest.raises(EmptyCrawlError) as excinfo:
            kb._raise_if_nothing_stored(0)

        message = str(excinfo.value)
        assert expected in message
        assert TEST_URLS[0] in message

    def test_message_falls_back_when_no_cause_was_recorded(self):
        """An unexplained empty crawl still says something sensible."""
        reader = EnhancedWebsiteReader(max_links=5)
        reader._challenge_blocked = 0
        reader._failure_causes = {}
        kb = EnhancedWebsiteKnowledgeBase(urls=TEST_URLS, reader=reader)

        with pytest.raises(EmptyCrawlError) as excinfo:
            kb._raise_if_nothing_stored(0)
        assert "No content could be read" in str(excinfo.value)

    def test_nothing_stored_because_everything_was_already_indexed(self):
        """Skipping URLs we already hold is a no-op, not an unreadable source.

        load() drops URLs already in the vector store, so a re-add reaches this
        guard having stored nothing. Raising there failed the whole source and
        blamed the site — "may require JavaScript" — for content already indexed.
        """
        reader = EnhancedWebsiteReader(max_links=5)
        reader._challenge_blocked = 0
        kb = EnhancedWebsiteKnowledgeBase(urls=TEST_URLS, reader=reader)

        # Every URL skipped as already present → nothing to do, no error.
        kb._raise_if_nothing_stored(0, skipped_existing=len(TEST_URLS))

        # Some URLs were still meant to be read and yielded nothing → real failure.
        with pytest.raises(EmptyCrawlError):
            kb._raise_if_nothing_stored(0, skipped_existing=len(TEST_URLS) - 1)

        # Bot protection keeps precedence over the skip shortcut: its message is
        # the actionable one, and a blocked sitemap must still fail loudly.
        reader._challenge_blocked = 2
        with pytest.raises(BotProtectionError):
            kb._raise_if_nothing_stored(0, skipped_existing=len(TEST_URLS))
    
    def test_document_lists_property(self, mock_reader):
        """Test the document_lists property"""
        kb = EnhancedWebsiteKnowledgeBase(urls=TEST_URLS, reader=mock_reader)
        
        # Convert iterator to list to verify contents
        doc_lists = list(kb.document_lists)
        
        assert len(doc_lists) == len(TEST_URLS)
        assert all(docs == TEST_DOCUMENTS for docs in doc_lists)
        assert mock_reader.read.call_count == len(TEST_URLS)
        
        # Verify read was called with correct URLs and new callback args
        for url in TEST_URLS:
            mock_reader.read.assert_any_call(
                url=url,
                vector_db_callback=ANY,
                url_crawled_callback=ANY
            )
    
    def test_load_without_vector_db(self, mock_reader):
        """Test load method when no vector db is provided"""
        kb = EnhancedWebsiteKnowledgeBase(urls=TEST_URLS, reader=mock_reader)
        
        with patch('app.knowledge.enhanced_website_kb.logger') as mock_logger:
            kb.load()
            mock_logger.warning.assert_called_with("No vector db provided")
    
    def test_load_with_recreate(self, mock_vector_db, mock_reader):
        """Test load method with recreate=True"""
        kb = EnhancedWebsiteKnowledgeBase(urls=TEST_URLS, reader=mock_reader)
        kb.vector_db = mock_vector_db
        
        # Set up mock reader to return different documents for each URL
        doc1 = _embedded("Test content 1", "https://example.com/page1")
        doc2 = _embedded("Test content 2", "https://test.com/page1")
        mock_reader.read.side_effect = [[doc1], [doc2]]
        
        kb.load(recreate=True)
        
        # Verify vector db operations
        mock_vector_db.drop.assert_called_once()
        mock_vector_db.create.assert_called_once()
        
        # Verify reader operations
        assert mock_reader.read.call_count == len(TEST_URLS)
        mock_reader.read.assert_has_calls([
            call(url=TEST_URLS[0], vector_db_callback=ANY, url_crawled_callback=ANY),
            call(url=TEST_URLS[1], vector_db_callback=ANY, url_crawled_callback=ANY)
        ], any_order=True)
        
        # Verify batch upsert was called with combined documents
        mock_vector_db.upsert.assert_called_once()
        # Get the actual documents passed to upsert
        actual_docs = mock_vector_db.upsert.call_args[1]['documents']
        assert len(actual_docs) == 2
        assert doc1 in actual_docs
        assert doc2 in actual_docs

    def test_load_with_existing_urls(self, mock_vector_db, mock_reader):
        """Test load method with existing URLs in vector db"""
        kb = EnhancedWebsiteKnowledgeBase(urls=TEST_URLS, reader=mock_reader)
        kb.vector_db = mock_vector_db
        
        # Simulate first URL already exists
        mock_vector_db.name_exists.side_effect = lambda name: name == TEST_URLS[0]
        
        kb.load(recreate=False)
        
        # Should only process second URL
        assert mock_reader.read.call_count == 1
        mock_reader.read.assert_called_once_with(
            url=TEST_URLS[1],
            vector_db_callback=ANY,
            url_crawled_callback=ANY
        )
        
        # Verify upsert was called once with the documents
        mock_vector_db.upsert.assert_called_once()

    def test_load_when_every_url_is_already_indexed(self, mock_vector_db, mock_reader):
        """Re-adding a source that is fully indexed must not fail it.

        Nothing is crawled and nothing is stored, which used to surface on the
        queue item as "No content could be read ... may require JavaScript" even
        though the content was sitting in the vector store the whole time.
        """
        kb = EnhancedWebsiteKnowledgeBase(urls=TEST_URLS, reader=mock_reader)
        kb.vector_db = mock_vector_db
        mock_vector_db.name_exists.return_value = True

        kb.load(recreate=False)

        mock_reader.read.assert_not_called()

    def test_load_with_existing_documents(self, mock_vector_db, mock_reader):
        """Test load method with existing documents"""
        kb = EnhancedWebsiteKnowledgeBase(urls=TEST_URLS, reader=mock_reader)
        kb.vector_db = mock_vector_db
        
        # Simulate first document already exists
        mock_vector_db.doc_exists.side_effect = lambda doc: doc == TEST_DOCUMENTS[0]
        
        kb.load(recreate=False)
        
        # Verify that batch upsert is used with all documents
        mock_vector_db.upsert.assert_called_once()
        # Extract the documents passed to upsert
        actual_docs = mock_vector_db.upsert.call_args[1]['documents']
        assert len(actual_docs) >= 1  # We should have at least one document
    
    def test_load_without_upsert(self, mock_vector_db, mock_reader):
        """Test load method when upsert is not available"""
        kb = EnhancedWebsiteKnowledgeBase(urls=TEST_URLS, reader=mock_reader)
        kb.vector_db = mock_vector_db
        
        # Simulate upsert not available
        mock_vector_db.upsert_available.return_value = False
        
        kb.load(upsert=False)
        
        # Verify reader was called
        assert mock_reader.read.call_count > 0
        
        # When upsert is False, no upsert should be called
        mock_vector_db.upsert.assert_not_called()
    
    def test_load_never_builds_an_approximate_index(self, mock_vector_db, mock_reader):
        """Agent searches filter by agent_id; an HNSW/IVFFlat index would filter
        after choosing candidates and could return fewer results. Exact search only."""
        kb = EnhancedWebsiteKnowledgeBase(urls=TEST_URLS, reader=mock_reader)
        kb.vector_db = mock_vector_db

        kb.load()

        mock_vector_db.optimize.assert_not_called()
    
    def test_load_with_error_handling(self, mock_vector_db, mock_reader):
        """Test load method error handling"""
        kb = EnhancedWebsiteKnowledgeBase(urls=TEST_URLS, reader=mock_reader)
        kb.vector_db = mock_vector_db
        
        # Simulate error for first URL
        mock_reader.read.side_effect = [Exception("Test error"), TEST_DOCUMENTS]
        
        # Disable embedding for this test to avoid extra error logging
        mock_vector_db.embedder = None
        
        with patch('app.knowledge.enhanced_website_kb.logger') as mock_logger:
            kb.load()
            
            # Should log error and continue with second URL
            # Only check for the specific error we're expecting
            mock_logger.error.assert_any_call("Error reading URL https://example.com: Test error")
            
            # Should have still processed the second URL
            assert mock_reader.read.call_count == 2
            
            # Should still have called upsert with the successful documents
            mock_vector_db.upsert.assert_called_once() 

class ExpiringDeadline:
    expired = False


class CountingEmbedder:
    def __init__(self):
        self.calls = 0

    def get_embedding_and_usage(self, text):
        self.calls += 1
        return [0.1, 0.2, 0.3], None


class TestASourceThatRunsOutOfTime:
    """Seen in production: one source ran for three hours until the worker hit
    its memory limit. Past its time limit a source fails, and stores nothing:
    a partial copy would make adding it again skip the crawl."""

    def test_it_fails_with_advice_and_stores_nothing(self, mock_vector_db, mock_reader):
        deadline = ExpiringDeadline()

        def slow_read(**kwargs):
            deadline.expired = True
            return TEST_DOCUMENTS
        mock_reader.read.side_effect = slow_read
        kb = EnhancedWebsiteKnowledgeBase(urls=TEST_URLS[:1], reader=mock_reader)
        kb.vector_db = mock_vector_db

        with patch('app.knowledge.enhanced_website_kb.Deadline', return_value=deadline), \
             pytest.raises(SourceTimedOut) as raised:
            kb.load()

        assert str(raised.value) == FAILURE_ADVICE[FailureCause.TOO_LONG].format(url=TEST_URLS[0])
        mock_vector_db.upsert.assert_not_called()
        assert mock_reader.deadline is deadline  # the crawl stops at the same time

    def test_its_queued_chunks_are_not_embedded(self, mock_vector_db, mock_reader):
        deadline = ExpiringDeadline()
        embedder = CountingEmbedder()
        mock_vector_db.embedder = embedder
        chunks = [Document(content=f"chunk {i}", meta_data={"url": "https://example.com"}) for i in range(200)]

        def slow_read(url, vector_db_callback=None, url_crawled_callback=None):
            deadline.expired = True
            for chunk in chunks:
                vector_db_callback(chunk)
            return chunks
        mock_reader.read.side_effect = slow_read
        kb = EnhancedWebsiteKnowledgeBase(urls=TEST_URLS[:1], reader=mock_reader)
        kb.vector_db = mock_vector_db

        with patch('app.knowledge.enhanced_website_kb.Deadline', return_value=deadline), \
             pytest.raises(SourceTimedOut):
            kb.load()

        assert embedder.calls == 0

    def test_a_source_read_in_time_is_stored(self, mock_vector_db, mock_reader):
        mock_vector_db.embedder = CountingEmbedder()
        kb = EnhancedWebsiteKnowledgeBase(urls=TEST_URLS[:1], reader=mock_reader)
        kb.vector_db = mock_vector_db

        kb.load()

        mock_vector_db.upsert.assert_called_once()
        assert isinstance(mock_reader.deadline, Deadline)

    def test_the_limit_comes_from_settings(self, mock_vector_db, mock_reader):
        kb = EnhancedWebsiteKnowledgeBase(urls=TEST_URLS[:1], reader=mock_reader)
        kb.vector_db = mock_vector_db
        with patch('app.knowledge.enhanced_website_kb.Deadline', wraps=Deadline) as made:
            kb.load()
        made.assert_called_once_with(settings.KB_SOURCE_TIME_LIMIT_SECONDS)