mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-10 16:48:30 +00:00
e0de166592
* INTPYTHON-580 Design and Implement MongoDBVectorSearchTool * add implementation * wip * wip * finish tests * add todo * refactor to wrap langchain-mongodb * cleanup * address review * Fix usage of EnvVar class * inline code * lint * lint * fix usage of SearchIndexModel * Refactor: Update EnvVar import path and remove unused tests.utils module - Changed import of EnvVar from tests.utils to crewai.tools in multiple files. - Updated README.md for MongoDB vector search tool with additional context. - Modified subprocess command in vector_search.py for package installation. - Cleaned up test_generate_tool_specs.py to improve mock patching syntax. - Deleted unused tests/utils.py file. * update the crewai dep and the lockfile * chore: update package versions and dependencies in uv.lock - Removed `auth0-python` package. - Updated `crewai` version to 0.140.0 and adjusted its dependencies. - Changed `json-repair` version to 0.25.2. - Updated `litellm` version to 1.72.6. - Modified dependency markers for several packages to improve compatibility with Python versions. * refactor: improve MongoDB vector search tool with enhanced error handling and new dimensions field - Added logging for error handling in the _run method and during client cleanup. - Introduced a new 'dimensions' field in the MongoDBVectorSearchConfig for embedding vector size. - Refactored the _run method to return JSON formatted results and handle exceptions gracefully. - Cleaned up import statements and improved code readability. * address review * update tests * debug * fix test * fix test * fix test * support azure openai --------- Co-authored-by: lorenzejay <lorenzejaytech@gmail.com>
76 lines
2.5 KiB
Python
76 lines
2.5 KiB
Python
import json
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
from crewai_tools import MongoDBVectorSearchConfig, MongoDBVectorSearchTool
|
|
|
|
|
|
# Unit Test Fixtures
|
|
@pytest.fixture
|
|
def mongodb_vector_search_tool():
|
|
tool = MongoDBVectorSearchTool(
|
|
connection_string="foo", database_name="bar", collection_name="test"
|
|
)
|
|
tool._embed_texts = lambda x: [[0.1]]
|
|
yield tool
|
|
|
|
|
|
# Unit Tests
|
|
def test_successful_query_execution(mongodb_vector_search_tool):
|
|
# Enable embedding
|
|
with patch.object(mongodb_vector_search_tool._coll, "aggregate") as mock_aggregate:
|
|
mock_aggregate.return_value = [dict(text="foo", score=0.1, _id=1)]
|
|
|
|
results = json.loads(mongodb_vector_search_tool._run(query="sandwiches"))
|
|
|
|
assert len(results) == 1
|
|
assert results[0]["text"] == "foo"
|
|
assert results[0]["_id"] == 1
|
|
|
|
|
|
def test_provide_config():
|
|
query_config = MongoDBVectorSearchConfig(limit=10)
|
|
tool = MongoDBVectorSearchTool(
|
|
connection_string="foo",
|
|
database_name="bar",
|
|
collection_name="test",
|
|
query_config=query_config,
|
|
vector_index_name="foo",
|
|
embedding_model="bar",
|
|
)
|
|
tool._embed_texts = lambda x: [[0.1]]
|
|
with patch.object(tool._coll, "aggregate") as mock_aggregate:
|
|
mock_aggregate.return_value = [dict(text="foo", score=0.1, _id=1)]
|
|
|
|
tool._run(query="sandwiches")
|
|
assert mock_aggregate.mock_calls[-1].args[0][0]["$vectorSearch"]["limit"] == 10
|
|
|
|
mock_aggregate.return_value = [dict(text="foo", score=0.1, _id=1)]
|
|
|
|
|
|
def test_cleanup_on_deletion(mongodb_vector_search_tool):
|
|
with patch.object(mongodb_vector_search_tool, "_client") as mock_client:
|
|
# Trigger cleanup
|
|
mongodb_vector_search_tool.__del__()
|
|
|
|
mock_client.close.assert_called_once()
|
|
|
|
|
|
def test_create_search_index(mongodb_vector_search_tool):
|
|
with patch(
|
|
"crewai_tools.tools.mongodb_vector_search_tool.vector_search.create_vector_search_index"
|
|
) as mock_create_search_index:
|
|
mongodb_vector_search_tool.create_vector_search_index(dimensions=10)
|
|
kwargs = mock_create_search_index.mock_calls[0].kwargs
|
|
assert kwargs["dimensions"] == 10
|
|
assert kwargs["similarity"] == "cosine"
|
|
|
|
|
|
def test_add_texts(mongodb_vector_search_tool):
|
|
with patch.object(mongodb_vector_search_tool._coll, "bulk_write") as bulk_write:
|
|
mongodb_vector_search_tool.add_texts(["foo"])
|
|
args = bulk_write.mock_calls[0].args
|
|
assert "ReplaceOne" in str(args[0][0])
|
|
assert "foo" in str(args[0][0])
|