238 lines
8.4 KiB
Python
238 lines
8.4 KiB
Python
|
"""
|
||
|
|
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 MagicMock, patch, PropertyMock
|
||
|
|
from agno.document import Document
|
||
|
|
|
||
|
|
from app.knowledge.enhanced_pdf_kb import EnhancedPDFKnowledgeBase
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.fixture
|
||
|
|
def mock_vector_db():
|
||
|
|
"""Create a mock vector database"""
|
||
|
|
db = MagicMock()
|
||
|
|
db.exists.return_value = True
|
||
|
|
db.embedder = MagicMock()
|
||
|
|
return db
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.fixture
|
||
|
|
def mock_reader():
|
||
|
|
"""Create a mock PDF reader"""
|
||
|
|
reader = MagicMock()
|
||
|
|
return reader
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.fixture
|
||
|
|
def mock_document():
|
||
|
|
"""Create a mock document"""
|
||
|
|
doc = MagicMock(spec=Document)
|
||
|
|
doc.id = "test_doc_1"
|
||
|
|
doc.embedding = None
|
||
|
|
return doc
|
||
|
|
|
||
|
|
|
||
|
|
class TestEnhancedPDFKnowledgeBase:
|
||
|
|
"""Tests for EnhancedPDFKnowledgeBase class"""
|
||
|
|
|
||
|
|
def test_load_no_vector_db(self):
|
||
|
|
"""Test load when no vector db is provided"""
|
||
|
|
kb = EnhancedPDFKnowledgeBase(path="/test/path")
|
||
|
|
kb.vector_db = None
|
||
|
|
|
||
|
|
# Should return early without error
|
||
|
|
kb.load()
|
||
|
|
|
||
|
|
def test_load_no_reader(self, mock_vector_db):
|
||
|
|
"""Test load when no reader is provided"""
|
||
|
|
kb = EnhancedPDFKnowledgeBase(path="/test/path")
|
||
|
|
kb.vector_db = mock_vector_db
|
||
|
|
kb.reader = None
|
||
|
|
|
||
|
|
# Should return early without error
|
||
|
|
kb.load()
|
||
|
|
|
||
|
|
def test_load_recreate_db(self, mock_vector_db, mock_reader):
|
||
|
|
"""Test load with recreate=True"""
|
||
|
|
with patch.object(EnhancedPDFKnowledgeBase, 'document_lists', new_callable=PropertyMock) as mock_doc_lists:
|
||
|
|
mock_doc_lists.return_value = []
|
||
|
|
|
||
|
|
kb = EnhancedPDFKnowledgeBase(path="/test/path")
|
||
|
|
kb.vector_db = mock_vector_db
|
||
|
|
kb.reader = mock_reader
|
||
|
|
|
||
|
|
kb.load(recreate=True)
|
||
|
|
|
||
|
|
mock_vector_db.drop.assert_called_once()
|
||
|
|
mock_vector_db.create.assert_called_once()
|
||
|
|
|
||
|
|
def test_load_create_if_not_exists(self, mock_vector_db, mock_reader):
|
||
|
|
"""Test load creates db if it doesn't exist"""
|
||
|
|
mock_vector_db.exists.return_value = False
|
||
|
|
|
||
|
|
with patch.object(EnhancedPDFKnowledgeBase, 'document_lists', new_callable=PropertyMock) as mock_doc_lists:
|
||
|
|
mock_doc_lists.return_value = []
|
||
|
|
|
||
|
|
kb = EnhancedPDFKnowledgeBase(path="/test/path")
|
||
|
|
kb.vector_db = mock_vector_db
|
||
|
|
kb.reader = mock_reader
|
||
|
|
|
||
|
|
kb.load(recreate=False)
|
||
|
|
|
||
|
|
mock_vector_db.create.assert_called_once()
|
||
|
|
|
||
|
|
def test_load_with_documents(self, mock_vector_db, mock_reader, mock_document):
|
||
|
|
"""Test load with documents"""
|
||
|
|
with patch.object(EnhancedPDFKnowledgeBase, 'document_lists', new_callable=PropertyMock) as mock_doc_lists:
|
||
|
|
mock_doc_lists.return_value = [[mock_document]]
|
||
|
|
|
||
|
|
kb = EnhancedPDFKnowledgeBase(path="/test/path")
|
||
|
|
kb.vector_db = mock_vector_db
|
||
|
|
kb.reader = mock_reader
|
||
|
|
|
||
|
|
kb.load(upsert=True)
|
||
|
|
|
||
|
|
# Document should be embedded
|
||
|
|
mock_document.embed.assert_called_once_with(embedder=mock_vector_db.embedder)
|
||
|
|
# Documents should be upserted
|
||
|
|
mock_vector_db.upsert.assert_called_once()
|
||
|
|
|
||
|
|
def test_load_with_filters(self, mock_vector_db, mock_reader, mock_document):
|
||
|
|
"""Test load with metadata filters"""
|
||
|
|
with patch.object(EnhancedPDFKnowledgeBase, 'document_lists', new_callable=PropertyMock) as mock_doc_lists:
|
||
|
|
mock_doc_lists.return_value = [[mock_document]]
|
||
|
|
|
||
|
|
kb = EnhancedPDFKnowledgeBase(path="/test/path")
|
||
|
|
kb.vector_db = mock_vector_db
|
||
|
|
kb.reader = mock_reader
|
||
|
|
|
||
|
|
filters = {"organization_id": "test_org", "agent_id": "test_agent"}
|
||
|
|
kb.load(upsert=True, filters=filters)
|
||
|
|
|
||
|
|
mock_vector_db.upsert.assert_called_once_with(
|
||
|
|
documents=[mock_document],
|
||
|
|
filters=filters
|
||
|
|
)
|
||
|
|
|
||
|
|
def test_load_document_already_embedded(self, mock_vector_db, mock_reader):
|
||
|
|
"""Test load when document is already embedded"""
|
||
|
|
doc = MagicMock(spec=Document)
|
||
|
|
doc.id = "test_doc_1"
|
||
|
|
doc.embedding = [0.1, 0.2, 0.3] # Already has embedding
|
||
|
|
|
||
|
|
with patch.object(EnhancedPDFKnowledgeBase, 'document_lists', new_callable=PropertyMock) as mock_doc_lists:
|
||
|
|
mock_doc_lists.return_value = [[doc]]
|
||
|
|
|
||
|
|
kb = EnhancedPDFKnowledgeBase(path="/test/path")
|
||
|
|
kb.vector_db = mock_vector_db
|
||
|
|
kb.reader = mock_reader
|
||
|
|
|
||
|
|
kb.load(upsert=True)
|
||
|
|
|
||
|
|
# embed should not be called since doc already has embedding
|
||
|
|
doc.embed.assert_not_called()
|
||
|
|
|
||
|
|
def test_load_embedding_error(self, mock_vector_db, mock_reader, mock_document):
|
||
|
|
"""Test load when embedding fails"""
|
||
|
|
mock_document.embed.side_effect = Exception("Embedding error")
|
||
|
|
|
||
|
|
with patch.object(EnhancedPDFKnowledgeBase, 'document_lists', new_callable=PropertyMock) as mock_doc_lists:
|
||
|
|
mock_doc_lists.return_value = [[mock_document]]
|
||
|
|
|
||
|
|
kb = EnhancedPDFKnowledgeBase(path="/test/path")
|
||
|
|
kb.vector_db = mock_vector_db
|
||
|
|
kb.reader = mock_reader
|
||
|
|
|
||
|
|
# Should not raise, but log error
|
||
|
|
kb.load(upsert=True)
|
||
|
|
|
||
|
|
def test_load_upsert_error(self, mock_vector_db, mock_reader, mock_document):
|
||
|
|
"""Test load when upsert fails"""
|
||
|
|
mock_vector_db.upsert.side_effect = Exception("Upsert error")
|
||
|
|
|
||
|
|
with patch.object(EnhancedPDFKnowledgeBase, 'document_lists', new_callable=PropertyMock) as mock_doc_lists:
|
||
|
|
mock_doc_lists.return_value = [[mock_document]]
|
||
|
|
|
||
|
|
kb = EnhancedPDFKnowledgeBase(path="/test/path")
|
||
|
|
kb.vector_db = mock_vector_db
|
||
|
|
kb.reader = mock_reader
|
||
|
|
|
||
|
|
# Should not raise, but log error
|
||
|
|
kb.load(upsert=True)
|
||
|
|
|
||
|
|
def test_load_skip_upsert(self, mock_vector_db, mock_reader, mock_document):
|
||
|
|
"""Test load with upsert=False"""
|
||
|
|
with patch.object(EnhancedPDFKnowledgeBase, 'document_lists', new_callable=PropertyMock) as mock_doc_lists:
|
||
|
|
mock_doc_lists.return_value = [[mock_document]]
|
||
|
|
|
||
|
|
kb = EnhancedPDFKnowledgeBase(path="/test/path")
|
||
|
|
kb.vector_db = mock_vector_db
|
||
|
|
kb.reader = mock_reader
|
||
|
|
|
||
|
|
kb.load(upsert=False)
|
||
|
|
|
||
|
|
# upsert should not be called
|
||
|
|
mock_vector_db.upsert.assert_not_called()
|
||
|
|
|
||
|
|
def test_load_multiple_document_lists(self, mock_vector_db, mock_reader):
|
||
|
|
"""Test load with multiple document lists"""
|
||
|
|
doc1 = MagicMock(spec=Document)
|
||
|
|
doc1.id = "doc1"
|
||
|
|
doc1.embedding = None
|
||
|
|
|
||
|
|
doc2 = MagicMock(spec=Document)
|
||
|
|
doc2.id = "doc2"
|
||
|
|
doc2.embedding = None
|
||
|
|
|
||
|
|
with patch.object(EnhancedPDFKnowledgeBase, 'document_lists', new_callable=PropertyMock) as mock_doc_lists:
|
||
|
|
mock_doc_lists.return_value = [[doc1], [doc2]]
|
||
|
|
|
||
|
|
kb = EnhancedPDFKnowledgeBase(path="/test/path")
|
||
|
|
kb.vector_db = mock_vector_db
|
||
|
|
kb.reader = mock_reader
|
||
|
|
|
||
|
|
kb.load(upsert=True)
|
||
|
|
|
||
|
|
# Both documents should be embedded
|
||
|
|
doc1.embed.assert_called_once()
|
||
|
|
doc2.embed.assert_called_once()
|
||
|
|
|
||
|
|
# upsert should be called twice (once per document list)
|
||
|
|
assert mock_vector_db.upsert.call_count == 2
|
||
|
|
|
||
|
|
def test_load_no_embedder(self, mock_reader):
|
||
|
|
"""Test load when vector_db has no embedder"""
|
||
|
|
mock_db = MagicMock()
|
||
|
|
mock_db.exists.return_value = True
|
||
|
|
# No embedder attribute
|
||
|
|
del mock_db.embedder
|
||
|
|
|
||
|
|
doc = MagicMock(spec=Document)
|
||
|
|
doc.id = "doc1"
|
||
|
|
doc.embedding = None
|
||
|
|
|
||
|
|
with patch.object(EnhancedPDFKnowledgeBase, 'document_lists', new_callable=PropertyMock) as mock_doc_lists:
|
||
|
|
mock_doc_lists.return_value = [[doc]]
|
||
|
|
|
||
|
|
kb = EnhancedPDFKnowledgeBase(path="/test/path")
|
||
|
|
kb.vector_db = mock_db
|
||
|
|
kb.reader = mock_reader
|
||
|
|
|
||
|
|
kb.load(upsert=True)
|
||
|
|
|
||
|
|
# embed should not be called since no embedder
|
||
|
|
doc.embed.assert_not_called()
|