apunkt/MetaGPT
1
0
Fork 0
mirror of https://github.com/FoundationAgents/MetaGPT.git synced 2026-06-08 15:05:17 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix chromadb ut

Browse source
This commit is contained in:
betterwang 2024-03-07 13:54:09 +08:00
parent f149007752
commit ac14814d74
2 changed files with 5 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

4
metagpt/document_store/chromadb_store.py
View file

@ -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

6
tests/metagpt/document_store/test_faiss_store.py
View file

@ -16,7 +16,7 @@ from metagpt.roles import Sales
@pytest.mark.asyncio
async def test_search_json():
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)
@ -25,7 +25,7 @@ async def test_search_json():
@pytest.mark.asyncio
async def test_search_xlsx():
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")
role = Sales(profile="Sales", store=store)
query = "Which facial cleanser is good for oily skin?"
result = await role.run(query)
@ -34,7 +34,7 @@ async def test_search_xlsx():
@pytest.mark.asyncio
async def test_write():
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