mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-07-02 13:48:09 +00:00
8919026326
Some checks failed
Add opt-in extension seams so an application can route memory, knowledge, RAG, and flow persistence through a custom backend without subclassing or threading an explicit instance through every construction site -- mirroring the existing crewai_core.lock_store.set_lock_backend seam. - memory: crewai.memory.storage.factory.set_memory_storage_factory - knowledge: crewai.knowledge.storage.factory.set_knowledge_storage_factory - rag: crewai.rag.factory.register_rag_client_factory (provider registry) - flow: crewai.flow.persistence.factory.set_flow_persistence_factory Each construction site consults the registered factory and falls back to the built-in default when none is set; an explicit instance always wins. Widen Knowledge.storage and the knowledge source base classes to BaseKnowledgeStorage (consistent with BaseAgent.knowledge_storage) so any base-interface backend plugs in. Runtime-free tests cover each seam.
67 lines
2.0 KiB
Python
67 lines
2.0 KiB
Python
"""Tests for the RAG client factory registry seam.
|
|
|
|
We verify our own logic: a registered factory is used for its provider,
|
|
factories override the built-in providers, unregister removes them, and an
|
|
unknown provider still raises.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
import crewai.rag.factory as factory
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def reset_registry():
|
|
"""Reset the registry around each test without clobbering preexisting state."""
|
|
original = dict(factory._factories)
|
|
factory._factories.clear()
|
|
yield
|
|
factory._factories.clear()
|
|
factory._factories.update(original)
|
|
|
|
|
|
def test_registered_factory_is_used_for_its_provider():
|
|
sentinel = object()
|
|
factory.register_rag_client_factory("custom", lambda config: sentinel)
|
|
|
|
assert factory.create_client(SimpleNamespace(provider="custom")) is sentinel
|
|
|
|
|
|
def test_factory_receives_the_config():
|
|
seen: list[object] = []
|
|
config = SimpleNamespace(provider="custom")
|
|
factory.register_rag_client_factory("custom", lambda cfg: seen.append(cfg) or object())
|
|
|
|
factory.create_client(config)
|
|
|
|
assert seen == [config]
|
|
|
|
|
|
def test_factory_overrides_builtin_provider():
|
|
sentinel = object()
|
|
factory.register_rag_client_factory("chromadb", lambda config: sentinel)
|
|
|
|
# Resolves via the registry without importing the built-in chromadb factory.
|
|
assert factory.create_client(SimpleNamespace(provider="chromadb")) is sentinel
|
|
|
|
|
|
def test_unregister_removes_factory():
|
|
factory.register_rag_client_factory("custom", lambda config: object())
|
|
factory.unregister_rag_client_factory("custom")
|
|
|
|
with pytest.raises(ValueError, match="Unsupported provider: custom"):
|
|
factory.create_client(SimpleNamespace(provider="custom"))
|
|
|
|
|
|
def test_unregister_unknown_provider_is_noop():
|
|
factory.unregister_rag_client_factory("never-registered")
|
|
|
|
|
|
def test_unknown_provider_still_raises():
|
|
with pytest.raises(ValueError, match="Unsupported provider: nope"):
|
|
factory.create_client(SimpleNamespace(provider="nope"))
|