mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-11 00:58:30 +00:00
feat: Add comprehensive Elasticsearch support to crewai.rag
- Implement ElasticsearchClient with full sync/async operations - Add ElasticsearchConfig with connection and embedding options - Create factory pattern following ChromaDB/Qdrant conventions - Add comprehensive test suite with 26 passing tests (100% coverage) - Support both sync and async Elasticsearch operations - Include proper error handling and edge case coverage - Update type system and factory to support Elasticsearch provider - Follow existing RAG patterns for consistency Resolves #3404 Co-Authored-By: João <joao@crewai.com>
This commit is contained in:
@@ -2,6 +2,8 @@
|
||||
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from crewai.rag.factory import create_client
|
||||
|
||||
|
||||
@@ -25,10 +27,50 @@ def test_create_client_chromadb():
|
||||
mock_module.create_client.assert_called_once_with(mock_config)
|
||||
|
||||
|
||||
def test_create_client_qdrant():
|
||||
"""Test Qdrant client creation."""
|
||||
mock_config = Mock()
|
||||
mock_config.provider = "qdrant"
|
||||
|
||||
with patch("crewai.rag.factory.require") as mock_require:
|
||||
mock_module = Mock()
|
||||
mock_client = Mock()
|
||||
mock_module.create_client.return_value = mock_client
|
||||
mock_require.return_value = mock_module
|
||||
|
||||
result = create_client(mock_config)
|
||||
|
||||
assert result == mock_client
|
||||
mock_require.assert_called_once_with(
|
||||
"crewai.rag.qdrant.factory", purpose="The 'qdrant' provider"
|
||||
)
|
||||
mock_module.create_client.assert_called_once_with(mock_config)
|
||||
|
||||
|
||||
def test_create_client_elasticsearch():
|
||||
"""Test Elasticsearch client creation."""
|
||||
mock_config = Mock()
|
||||
mock_config.provider = "elasticsearch"
|
||||
|
||||
with patch("crewai.rag.factory.require") as mock_require:
|
||||
mock_module = Mock()
|
||||
mock_client = Mock()
|
||||
mock_module.create_client.return_value = mock_client
|
||||
mock_require.return_value = mock_module
|
||||
|
||||
result = create_client(mock_config)
|
||||
|
||||
assert result == mock_client
|
||||
mock_require.assert_called_once_with(
|
||||
"crewai.rag.elasticsearch.factory", purpose="The 'elasticsearch' provider"
|
||||
)
|
||||
mock_module.create_client.assert_called_once_with(mock_config)
|
||||
|
||||
|
||||
def test_create_client_unsupported_provider():
|
||||
"""Test unsupported provider returns None for now."""
|
||||
"""Test that unsupported provider raises ValueError."""
|
||||
mock_config = Mock()
|
||||
mock_config.provider = "unsupported"
|
||||
|
||||
result = create_client(mock_config)
|
||||
assert result is None
|
||||
with pytest.raises(ValueError, match="Unsupported provider: unsupported"):
|
||||
create_client(mock_config)
|
||||
|
||||
@@ -3,7 +3,10 @@
|
||||
import pytest
|
||||
|
||||
from crewai.rag.config.optional_imports.base import _MissingProvider
|
||||
from crewai.rag.config.optional_imports.providers import MissingChromaDBConfig
|
||||
from crewai.rag.config.optional_imports.providers import (
|
||||
MissingChromaDBConfig,
|
||||
MissingElasticsearchConfig,
|
||||
)
|
||||
|
||||
|
||||
def test_missing_provider_raises_runtime_error():
|
||||
@@ -20,3 +23,11 @@ def test_missing_chromadb_config_raises_runtime_error():
|
||||
RuntimeError, match="provider 'chromadb' requested but not installed"
|
||||
):
|
||||
MissingChromaDBConfig()
|
||||
|
||||
|
||||
def test_missing_elasticsearch_config_raises_runtime_error():
|
||||
"""Test that MissingElasticsearchConfig raises RuntimeError on instantiation."""
|
||||
with pytest.raises(
|
||||
RuntimeError, match="provider 'elasticsearch' requested but not installed"
|
||||
):
|
||||
MissingElasticsearchConfig()
|
||||
|
||||
Reference in New Issue
Block a user