update exp_pool manager

This commit is contained in:
seehi 2024-08-19 14:07:13 +08:00
parent 243c7a65d6
commit 665ca6ff97
9 changed files with 31 additions and 22 deletions

View file

@ -3,7 +3,7 @@ import pytest
from metagpt.config2 import Config
from metagpt.configs.exp_pool_config import (
ExperiencePoolConfig,
ExperiencePoolStorageType,
ExperiencePoolRetrievalType,
)
from metagpt.configs.llm_config import LLMConfig
from metagpt.exp_pool.manager import Experience, ExperienceManager
@ -16,7 +16,7 @@ class TestExperienceManager:
return Config(
llm=LLMConfig(),
exp_pool=ExperiencePoolConfig(
enable_write=True, enable_read=True, enabled=True, storage_type=ExperiencePoolStorageType.BM25
enable_write=True, enable_read=True, enabled=True, retrieval_type=ExperiencePoolRetrievalType.BM25
),
)
@ -96,7 +96,7 @@ class TestExperienceManager:
assert exp_manager.get_exps_count() == 10
def test_resolve_storage_bm25(self, mocker, mock_config):
mock_config.exp_pool.storage_type = ExperiencePoolStorageType.BM25
mock_config.exp_pool.retrieval_type = ExperiencePoolRetrievalType.BM25
mocker.patch.object(ExperienceManager, "_create_bm25_storage", return_value=mocker.MagicMock())
manager = ExperienceManager(config=mock_config)
storage = manager._resolve_storage()
@ -104,7 +104,7 @@ class TestExperienceManager:
assert storage is not None
def test_resolve_storage_chroma(self, mocker, mock_config):
mock_config.exp_pool.storage_type = ExperiencePoolStorageType.CHROMA
mock_config.exp_pool.retrieval_type = ExperiencePoolRetrievalType.CHROMA
mocker.patch.object(ExperienceManager, "_create_chroma_storage", return_value=mocker.MagicMock())
manager = ExperienceManager(config=mock_config)
storage = manager._resolve_storage()

View file

@ -75,7 +75,7 @@ class TestSimpleEngine:
)
# Assert
mock_simple_directory_reader.assert_called_once_with(input_dir=input_dir, input_files=input_files)
mock_simple_directory_reader.assert_called_once_with(input_dir=input_dir, input_files=input_files, fs=None)
mock_get_retriever.assert_called_once()
mock_get_rankers.assert_called_once()
mock_get_response_synthesizer.assert_called_once_with(llm=llm)

View file

@ -1,5 +1,6 @@
import pytest
from metagpt.config2 import Config
from metagpt.configs.embedding_config import EmbeddingType
from metagpt.configs.llm_config import LLMType
from metagpt.rag.factories.embedding import RAGEmbeddingFactory
@ -12,7 +13,10 @@ class TestRAGEmbeddingFactory:
@pytest.fixture
def mock_config(self, mocker):
return mocker.patch("metagpt.rag.factories.embedding.config")
config = Config.default().model_copy(deep=True)
default = mocker.patch("metagpt.config2.Config.default")
default.return_value = config
return config
@staticmethod
def mock_openai_embedding(mocker):