mirror of
https://github.com/crewAIInc/crewAI.git
synced 2025-12-16 12:28:30 +00:00
82 lines
2.7 KiB
Python
82 lines
2.7 KiB
Python
"""Tests for RAGStorage custom path functionality."""
|
|
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from crewai.memory.storage.rag_storage import RAGStorage
|
|
|
|
|
|
@patch("crewai.memory.storage.rag_storage.create_client")
|
|
@patch("crewai.memory.storage.rag_storage.build_embedder")
|
|
def test_rag_storage_custom_path(
|
|
mock_build_embedder: MagicMock,
|
|
mock_create_client: MagicMock,
|
|
) -> None:
|
|
"""Test RAGStorage uses custom path when provided."""
|
|
mock_build_embedder.return_value = MagicMock(return_value=[[0.1, 0.2, 0.3]])
|
|
mock_create_client.return_value = MagicMock()
|
|
|
|
custom_path = "/custom/memory/path"
|
|
embedder_config = {"provider": "openai", "config": {"model": "text-embedding-3-small"}}
|
|
|
|
RAGStorage(
|
|
type="short_term",
|
|
crew=None,
|
|
path=custom_path,
|
|
embedder_config=embedder_config,
|
|
)
|
|
|
|
mock_create_client.assert_called_once()
|
|
config_arg = mock_create_client.call_args[0][0]
|
|
assert config_arg.settings.persist_directory == custom_path
|
|
|
|
|
|
@patch("crewai.memory.storage.rag_storage.create_client")
|
|
@patch("crewai.memory.storage.rag_storage.build_embedder")
|
|
def test_rag_storage_default_path_when_none(
|
|
mock_build_embedder: MagicMock,
|
|
mock_create_client: MagicMock,
|
|
) -> None:
|
|
"""Test RAGStorage uses default path when no custom path is provided."""
|
|
mock_build_embedder.return_value = MagicMock(return_value=[[0.1, 0.2, 0.3]])
|
|
mock_create_client.return_value = MagicMock()
|
|
|
|
embedder_config = {"provider": "openai", "config": {"model": "text-embedding-3-small"}}
|
|
|
|
storage = RAGStorage(
|
|
type="short_term",
|
|
crew=None,
|
|
path=None,
|
|
embedder_config=embedder_config,
|
|
)
|
|
|
|
mock_create_client.assert_called_once()
|
|
assert storage.path is None
|
|
|
|
|
|
@patch("crewai.memory.storage.rag_storage.create_client")
|
|
@patch("crewai.memory.storage.rag_storage.build_embedder")
|
|
def test_rag_storage_custom_path_with_batch_size(
|
|
mock_build_embedder: MagicMock,
|
|
mock_create_client: MagicMock,
|
|
) -> None:
|
|
"""Test RAGStorage uses custom path with batch_size in config."""
|
|
mock_build_embedder.return_value = MagicMock(return_value=[[0.1, 0.2, 0.3]])
|
|
mock_create_client.return_value = MagicMock()
|
|
|
|
custom_path = "/custom/batch/path"
|
|
embedder_config = {
|
|
"provider": "openai",
|
|
"config": {"model": "text-embedding-3-small", "batch_size": 100},
|
|
}
|
|
|
|
RAGStorage(
|
|
type="long_term",
|
|
crew=None,
|
|
path=custom_path,
|
|
embedder_config=embedder_config,
|
|
)
|
|
|
|
mock_create_client.assert_called_once()
|
|
config_arg = mock_create_client.call_args[0][0]
|
|
assert config_arg.settings.persist_directory == custom_path
|
|
assert config_arg.batch_size == 100 |