replypier/backend/tests/core/test_knowledge_base.py

"""
Copyright 2024-2026 ChatterMate

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
from unittest.mock import patch, MagicMock, AsyncMock
from app.knowledge.knowledge_base import KnowledgeManager
from app.models.knowledge import Knowledge, SourceType
from app.models.knowledge_to_agent import KnowledgeToAgent
from uuid import UUID, uuid4
import os
from agno.vectordb.pgvector import PgVector, SearchType

@pytest.fixture
def mock_db():
    """Create a mock database session"""
    db = MagicMock()
    return db

@pytest.fixture
def mock_ai_config():
    """Create a mock AI config"""
    config = MagicMock()
    config.encrypted_api_key = "encrypted_key"
    return config

@pytest.fixture
def mock_knowledge():
    """Create a mock knowledge entry"""
    knowledge_id = uuid4()
    org_id = uuid4()
    return Knowledge(
        id=knowledge_id,
        organization_id=org_id,
        source="test_source",
        source_type=SourceType.FILE,
        table_name="test_table",
        schema="test_schema"
    )

@pytest.fixture
def mock_vector_db():
    """Create a mock vector database"""
    vector_db = MagicMock(spec=PgVector)
    vector_db.table_name = "test_table"
    vector_db.schema = "test_schema"
    return vector_db

@pytest.fixture
def knowledge_manager(mock_db, mock_ai_config, mock_vector_db):
    """Create a KnowledgeManager instance with mocked dependencies"""
    org_id = uuid4()
    agent_id = uuid4()
    
    with patch('app.knowledge.knowledge_base.SessionLocal') as mock_session_local, \
         patch('app.knowledge.knowledge_base.AIConfigRepository') as mock_ai_config_repo, \
         patch('app.knowledge.knowledge_base.OptimizedPgVector') as mock_pg_vector, \
         patch('app.knowledge.knowledge_base.KnowledgeRepository') as mock_knowledge_repo, \
         patch('app.knowledge.knowledge_base.KnowledgeToAgentRepository') as mock_link_repo:
        
        # Configure basic mocks
        mock_session_local.return_value.__enter__.return_value = mock_db
        mock_ai_config_repo.return_value.get_active_config.return_value = mock_ai_config
        mock_pg_vector.return_value = mock_vector_db
        
        # Configure repository mocks
        mock_knowledge_repo_instance = MagicMock()
        mock_link_repo_instance = MagicMock()
        mock_knowledge_repo.return_value = mock_knowledge_repo_instance
        mock_link_repo.return_value = mock_link_repo_instance
        
        manager = KnowledgeManager(org_id=org_id, agent_id=str(agent_id))
        
        # Replace repository instances with mocks
        manager.knowledge_repo = mock_knowledge_repo_instance
        manager.link_repo = mock_link_repo_instance
        
        return manager

@pytest.mark.asyncio
async def test_add_pdf_urls_success(knowledge_manager):
    """A PDF URL is downloaded through the guarded fetcher, then indexed."""
    urls = ["https://example.com/docs/test.pdf"]

    with patch('app.knowledge.knowledge_base.download_pdf') as mock_download, \
         patch.object(knowledge_manager, 'add_pdf_files', AsyncMock(return_value=True)) as mock_add_files, \
         patch('app.knowledge.knowledge_base.delete_file_from_s3', AsyncMock()) as mock_delete:
        result = await knowledge_manager.add_pdf_urls(urls)

    assert result is True
    url, temp_path = mock_download.call_args.args
    assert url == urls[0]
    assert os.path.basename(temp_path) == "test.pdf"
    mock_add_files.assert_awaited_once_with([temp_path], filename="test")
    mock_delete.assert_not_called()  # not one of our S3 uploads
    assert not os.path.exists(os.path.dirname(temp_path))  # temp dir cleaned up

@pytest.mark.asyncio
async def test_add_pdf_urls_failure(knowledge_manager):
    """A refused or failed download is reported, and nothing is indexed."""
    from app.knowledge.url_safety import BlockedHostError

    with patch('app.knowledge.knowledge_base.download_pdf',
               side_effect=BlockedHostError("internal")), \
         patch.object(knowledge_manager, 'add_pdf_files', AsyncMock()) as mock_add_files:
        result = await knowledge_manager.add_pdf_urls(["http://169.254.169.254/latest"])

    assert result is False
    mock_add_files.assert_not_called()

@pytest.mark.parametrize("url, name", [
    ("https://example.com/a/guide.pdf", "guide"),
    ("https://example.com/a/My%20Guide.pdf?x=1", "My Guide"),
    # An encoded traversal decodes to a separator; only the last part survives.
    ("https://b.s3.amazonaws.com/knowledge/..%2F..%2Fapp%2Fmain.py", "main"),
    ("https://example.com/", "example.com"),
])
def test_pdf_name_from_url_never_carries_a_path(url, name):
    from app.knowledge.knowledge_base import _pdf_name_from_url
    assert _pdf_name_from_url(url) == name

@pytest.mark.asyncio
async def test_only_this_orgs_uploaded_pdf_is_deleted_from_s3(knowledge_manager):
    from app.knowledge.knowledge_base import settings
    own = f"https://{settings.S3_BUCKET}.s3.us-east-1.amazonaws.com/knowledge/{knowledge_manager.org_id}/f.pdf?X-Amz-Signature=x"
    other_org = f"https://{settings.S3_BUCKET}.s3.us-east-1.amazonaws.com/knowledge/{uuid4()}/f.pdf"
    # Another bucket, laid out so a path-style read lands on our org's key.
    foreign = f"https://evil.s3.amazonaws.com/x/knowledge/{knowledge_manager.org_id}/f.pdf"

    with patch('app.knowledge.knowledge_base.delete_file_from_s3', AsyncMock(return_value=True)) as mock_delete:
        for url in (own, other_org, foreign):
            await knowledge_manager._delete_uploaded_pdf(url)

    mock_delete.assert_awaited_once_with(own)

@pytest.mark.asyncio
async def test_add_websites_success(knowledge_manager, mock_knowledge):
    """Test successful addition of websites"""
    # Setup
    urls = ["http://example.com"]
    max_links = 5
    
    with patch('app.knowledge.knowledge_base.EnhancedWebsiteKnowledgeBase') as mock_web_kb, \
         patch.object(knowledge_manager, '_add_knowledge_source', return_value=mock_knowledge) as mock_add_source:
        # Configure mock
        mock_web_kb_instance = MagicMock()
        mock_web_kb.return_value = mock_web_kb_instance
        
        # Execute
        result = await knowledge_manager.add_websites(urls, max_links)
        
        # Assert
        assert result is True
        mock_web_kb.assert_called_once()
        mock_web_kb_instance.load.assert_called_once()
        mock_add_source.assert_called_once()

@pytest.mark.asyncio
async def test_add_pdf_files_success(knowledge_manager, mock_knowledge):
    """Test successful addition of PDF files"""
    # Setup
    files = ["/path/to/test.pdf"]

    with patch('app.knowledge.knowledge_base.EnhancedPDFKnowledgeBase') as mock_pdf_kb, \
         patch('app.knowledge.knowledge_base.PDFImageReader') as mock_pdf_reader, \
         patch('os.path.exists') as mock_exists, \
         patch('os.path.getsize') as mock_getsize, \
         patch('os.path.basename') as mock_basename, \
         patch('os.path.splitext') as mock_splitext, \
         patch.object(knowledge_manager, '_add_knowledge_source', return_value=mock_knowledge) as mock_add_source:
        # Configure mocks
        mock_pdf_kb_instance = MagicMock()
        mock_pdf_kb.return_value = mock_pdf_kb_instance
        mock_exists.return_value = True
        mock_getsize.return_value = 1024
        mock_basename.return_value = "test.pdf"
        mock_splitext.return_value = ("test", ".pdf")

        # Execute
        result = await knowledge_manager.add_pdf_files(files)

        # Assert
        assert result is True
        mock_pdf_kb.assert_called_once()
        mock_pdf_kb_instance.load.assert_called_once()
        mock_add_source.assert_called_once()

def test_get_knowledge_base_for_agent(knowledge_manager, mock_knowledge):
    """Test getting knowledge base for a specific agent"""
    # Setup
    mock_knowledge.agent_links = [
        MagicMock(agent_id=knowledge_manager.agent_id)
    ]
    
    with patch('app.knowledge.knowledge_base.KnowledgeRepository') as mock_knowledge_repo_class:
        mock_knowledge_repo_instance = MagicMock()
        mock_knowledge_repo_class.return_value = mock_knowledge_repo_instance
        mock_knowledge_repo_instance.get_by_agent.return_value = [mock_knowledge]
        
        # Execute
        result = knowledge_manager.get_knowledge_base()
        
        # Assert
        assert len(result) == 1
        assert result[0]['id'] == mock_knowledge.id
        assert result[0]['source'] == mock_knowledge.source
        assert result[0]['source_type'] == mock_knowledge.source_type
        mock_knowledge_repo_instance.get_by_agent.assert_called_once_with(knowledge_manager.agent_id)

def test_get_knowledge_base_for_org(knowledge_manager, mock_knowledge):
    """Test getting knowledge base for an organization"""
    # Setup
    knowledge_manager.agent_id = None
    mock_knowledge.agent_links = []
    
    with patch('app.knowledge.knowledge_base.KnowledgeRepository') as mock_knowledge_repo_class:
        mock_knowledge_repo_instance = MagicMock()
        mock_knowledge_repo_class.return_value = mock_knowledge_repo_instance
        mock_knowledge_repo_instance.get_by_org.return_value = [mock_knowledge]
        
        # Execute
        result = knowledge_manager.get_knowledge_base()
        
        # Assert
        assert len(result) == 1
        assert result[0]['id'] == mock_knowledge.id
        assert result[0]['source'] == mock_knowledge.source
        assert result[0]['source_type'] == mock_knowledge.source_type
        mock_knowledge_repo_instance.get_by_org.assert_called_once_with(knowledge_manager.org_id)

@pytest.mark.asyncio
async def test_process_knowledge_pdf_file(knowledge_manager):
    """Test processing PDF file knowledge"""
    # Setup
    queue_item = MagicMock()
    queue_item.id = 123  # ensure a real integer ID to avoid DB adapter issues
    queue_item.source_type = "pdf_file"
    queue_item.source = "/path/to/test.pdf"
    
    with patch.object(knowledge_manager, 'add_pdf_files') as mock_add_pdf_files, \
         patch('os.path.exists') as mock_exists, \
         patch('os.remove') as mock_remove, \
         patch('app.knowledge.knowledge_base.KnowledgeQueueRepository') as mock_queue_repo_cls:
        # Configure mocks
        mock_add_pdf_files.return_value = True
        mock_exists.return_value = True
        mock_queue_repo = MagicMock()
        mock_queue_repo.update_progress.return_value = True
        mock_queue_repo.update_status.return_value = True
        mock_queue_repo_cls.return_value = mock_queue_repo
        
        # Execute
        result = await knowledge_manager.process_knowledge(queue_item)
        
        # Assert
        assert result is True
        mock_add_pdf_files.assert_called_once_with([queue_item.source])
        mock_exists.assert_called_once_with(queue_item.source)
        mock_remove.assert_called_once_with(queue_item.source)

@pytest.mark.asyncio
async def test_process_knowledge_unsupported_type(knowledge_manager):
    """Test processing unsupported knowledge type"""
    # Setup
    queue_item = MagicMock()
    queue_item.source_type = "unsupported_type"
    
    # Execute and Assert
    with pytest.raises(ValueError, match="Unsupported source type: unsupported_type"):
        await knowledge_manager.process_knowledge(queue_item) 

# ---- PDFs over a limit: refused with the reason, not "failed to process" ----

from app.knowledge.pdf_download import PdfTooLargeError  # noqa: E402
from app.knowledge.pdf_limits import PdfTooManyPages  # noqa: E402
from tests.knowledge.test_pdf_limits import pdf_with_pages  # noqa: E402


@pytest.mark.asyncio
async def test_a_pdf_over_the_page_limit_is_refused_before_it_is_read(knowledge_manager, tmp_path):
    path = tmp_path / "catalogue.pdf"
    path.write_bytes(pdf_with_pages(4))

    with patch('app.core.config.settings.KB_PDF_MAX_PAGES', 3), \
         patch('app.knowledge.knowledge_base.EnhancedPDFKnowledgeBase') as reader, \
         pytest.raises(PdfTooManyPages, match="4 pages; the limit is 3"):
        await knowledge_manager.add_pdf_files([str(path)])

    reader.assert_not_called()


@pytest.mark.asyncio
async def test_a_pdf_link_over_the_size_limit_says_so_without_the_link(knowledge_manager):
    signed = "https://bucket.s3.amazonaws.com/knowledge/x/big.pdf?X-Amz-Signature=secret"
    with patch('app.knowledge.knowledge_base.download_pdf', side_effect=PdfTooLargeError(25 * 1024 * 1024)), \
         pytest.raises(PdfTooLargeError) as refused:
        await knowledge_manager.add_pdf_urls([signed])

    assert "25 MB" in str(refused.value)
    assert "amazonaws" not in str(refused.value)


@pytest.mark.asyncio
async def test_a_refused_upload_is_reported_and_its_copy_removed(knowledge_manager, tmp_path):
    path = tmp_path / "0e9f_catalogue.pdf"
    path.write_bytes(pdf_with_pages(4))
    queue_item = MagicMock(id=123, source_type="pdf_file", source=str(path))

    with patch('app.core.config.settings.KB_PDF_MAX_PAGES', 3), \
         patch('app.knowledge.knowledge_base.KnowledgeQueueRepository') as queue_repo, \
         pytest.raises(PdfTooManyPages):
        await knowledge_manager.process_knowledge(queue_item)

    assert not path.exists()
    queue_repo.return_value.update_status.assert_not_called()  # never reported as completed


@pytest.mark.asyncio
@pytest.mark.parametrize("source_type,method", [("pdf_file", "add_pdf_files"), ("pdf_url", "add_pdf_urls")])
async def test_a_pdf_that_fails_is_never_marked_completed_first(knowledge_manager, source_type, method):
    queue_item = MagicMock(id=123, source_type=source_type, source="https://example.com/x.pdf")

    with patch.object(knowledge_manager, method, AsyncMock(return_value=False)), \
         patch('app.knowledge.knowledge_base.KnowledgeQueueRepository') as queue_repo, \
         pytest.raises(Exception, match="Failed to process knowledge source"):
        await knowledge_manager.process_knowledge(queue_item)

    queue_repo.return_value.update_status.assert_not_called()