From 2ffa9e12b951361d0ef7833493e564397c15a4e6 Mon Sep 17 00:00:00 2001 From: betterwang Date: Thu, 7 Mar 2024 13:54:09 +0800 Subject: [PATCH] fix chromadb ut --- metagpt/document_store/chromadb_store.py | 4 ++-- tests/metagpt/document_store/test_faiss_store.py | 6 +++--- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/metagpt/document_store/chromadb_store.py b/metagpt/document_store/chromadb_store.py index d7344d41b..8c2114f66 100644 --- a/metagpt/document_store/chromadb_store.py +++ b/metagpt/document_store/chromadb_store.py @@ -11,9 +11,9 @@ import chromadb class ChromaStore: """If inherited from BaseStore, or importing other modules from metagpt, a Python exception occurs, which is strange.""" - def __init__(self, name): + def __init__(self, name, get_or_create: bool = True): client = chromadb.Client() - collection = client.create_collection(name) + collection = client.create_collection(name, get_or_create=get_or_create) self.client = client self.collection = collection diff --git a/tests/metagpt/document_store/test_faiss_store.py b/tests/metagpt/document_store/test_faiss_store.py index 0c5a55e0f..7c712294e 100644 --- a/tests/metagpt/document_store/test_faiss_store.py +++ b/tests/metagpt/document_store/test_faiss_store.py @@ -28,7 +28,7 @@ def mock_openai_embed_documents(self, texts: list[str], chunk_size: Optional[int async def test_search_json(mocker): mocker.patch("langchain_community.embeddings.openai.OpenAIEmbeddings.embed_documents", mock_openai_embed_documents) - store = FaissStore(EXAMPLE_PATH / "data/example.json") + store = FaissStore(EXAMPLE_PATH / "data/search_kb/example.json") role = Sales(profile="Sales", store=store) query = "Which facial cleanser is good for oily skin?" result = await role.run(query) @@ -39,7 +39,7 @@ async def test_search_json(mocker): async def test_search_xlsx(mocker): mocker.patch("langchain_community.embeddings.openai.OpenAIEmbeddings.embed_documents", mock_openai_embed_documents) - store = FaissStore(EXAMPLE_PATH / "data/example.xlsx") + store = FaissStore(EXAMPLE_PATH / "data/search_kb/example.xlsx") role = Sales(profile="Sales", store=store) query = "Which facial cleanser is good for oily skin?" result = await role.run(query) @@ -50,7 +50,7 @@ async def test_search_xlsx(mocker): async def test_write(mocker): mocker.patch("langchain_community.embeddings.openai.OpenAIEmbeddings.embed_documents", mock_openai_embed_documents) - store = FaissStore(EXAMPLE_PATH / "data/example.xlsx", meta_col="Answer", content_col="Question") + store = FaissStore(EXAMPLE_PATH / "data/search_kb/example.xlsx", meta_col="Answer", content_col="Question") _faiss_store = store.write() assert _faiss_store.storage_context.docstore assert _faiss_store.storage_context.vector_store.client