klas/t6_mem0
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Upgrade the chromadb version to 0.4.8 and open its settings configuration. (#517)

Browse Source
This commit is contained in:
wangJm
2023-09-04 14:31:08 +08:00
committed by GitHub
parent 433c4157e0
commit eecdbc5e06
6 changed files with 80 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
tests/embedchain/test_embedchain.py
View File

@@ -3,7 +3,8 @@ import unittest
from unittest.mock import patch
from embedchain import App
from embedchain.config import AppConfig
from embedchain.config import AppConfig, CustomAppConfig
from embedchain.models import EmbeddingFunctions, Providers
class TestChromaDbHostsLoglevel(unittest.TestCase):
@@ -42,7 +43,11 @@ class TestChromaDbHostsLoglevel(unittest.TestCase):
"""
Test if the `App` instance is correctly reconstructed after a reset.
"""
app = App()
app = App(
CustomAppConfig(
provider=Providers.OPENAI, embedding_fn=EmbeddingFunctions.OPENAI, chroma_settings={"allow_reset": True}
)
)
app.reset()
# Make sure the client is still healthy

58
tests/vectordb/test_chroma_db.py
View File

@@ -4,7 +4,8 @@ import unittest
from unittest.mock import patch
from embedchain import App
from embedchain.config import AppConfig
from embedchain.config import AppConfig, CustomAppConfig
from embedchain.models import EmbeddingFunctions, Providers
from embedchain.vectordb.chroma_db import ChromaDB
@@ -21,6 +22,24 @@ class TestChromaDbHosts(unittest.TestCase):
self.assertEqual(settings.chroma_server_host, host)
self.assertEqual(settings.chroma_server_http_port, port)
def test_init_with_basic_auth(self):
host = "test-host"
port = "1234"
chroma_auth_settings = {
"chroma_client_auth_provider": "chromadb.auth.basic.BasicAuthClientProvider",
"chroma_client_auth_credentials": "admin:admin",
}
db = ChromaDB(host=host, port=port, embedding_fn=len, chroma_settings=chroma_auth_settings)
settings = db.client.get_settings()
self.assertEqual(settings.chroma_server_host, host)
self.assertEqual(settings.chroma_server_http_port, port)
self.assertEqual(settings.chroma_client_auth_provider, chroma_auth_settings["chroma_client_auth_provider"])
self.assertEqual(
settings.chroma_client_auth_credentials, chroma_auth_settings["chroma_client_auth_credentials"]
)
# Review this test
class TestChromaDbHostsInit(unittest.TestCase):
@@ -68,12 +87,18 @@ class TestChromaDbHostsLoglevel(unittest.TestCase):
class TestChromaDbDuplicateHandling:
app_with_settings = App(
CustomAppConfig(
provider=Providers.OPENAI, embedding_fn=EmbeddingFunctions.OPENAI, chroma_settings={"allow_reset": True}
)
)
def test_duplicates_throw_warning(self, caplog):
"""
Test that add duplicates throws an error.
"""
# Start with a clean app
App().reset()
self.app_with_settings.reset()
app = App(config=AppConfig(collect_metrics=False))
app.collection.add(embeddings=[[0, 0, 0]], ids=["0"])
@@ -88,7 +113,7 @@ class TestChromaDbDuplicateHandling:
# NOTE: Not part of the TestChromaDbCollection because `unittest.TestCase` doesn't have caplog.
# Start with a clean app
App().reset()
self.app_with_settings.reset()
app = App(config=AppConfig(collect_metrics=False))
app.set_collection("test_collection_1")
@@ -100,6 +125,12 @@ class TestChromaDbDuplicateHandling:
class TestChromaDbCollection(unittest.TestCase):
app_with_settings = App(
CustomAppConfig(
provider=Providers.OPENAI, embedding_fn=EmbeddingFunctions.OPENAI, chroma_settings={"allow_reset": True}
)
)
def test_init_with_default_collection(self):
"""
Test if the `App` instance is initialized with the correct default collection name.
@@ -131,7 +162,7 @@ class TestChromaDbCollection(unittest.TestCase):
Test that changes to one collection do not affect the other collection
"""
# Start with a clean app
App().reset()
self.app_with_settings.reset()
app = App(config=AppConfig(collect_metrics=False))
app.set_collection("test_collection_1")
@@ -157,7 +188,7 @@ class TestChromaDbCollection(unittest.TestCase):
Test that a collection can be picked up later.
"""
# Start with a clean app
App().reset()
self.app_with_settings.reset()
app = App(config=AppConfig(collect_metrics=False))
app.set_collection("test_collection_1")
@@ -175,7 +206,7 @@ class TestChromaDbCollection(unittest.TestCase):
the other app.
"""
# Start clean
App().reset()
self.app_with_settings.reset()
# Create two apps
app1 = App(AppConfig(collection_name="test_collection_1", collect_metrics=False))
@@ -201,7 +232,7 @@ class TestChromaDbCollection(unittest.TestCase):
Different ids should still share collections.
"""
# Start clean
App().reset()
self.app_with_settings.reset()
# Create two apps
app1 = App(AppConfig(collection_name="one_collection", id="new_app_id_1", collect_metrics=False))
@@ -220,11 +251,20 @@ class TestChromaDbCollection(unittest.TestCase):
Resetting should hit all collections and ids.
"""
# Start clean
App().reset()
self.app_with_settings.reset()
# Create four apps.
# app1, which we are about to reset, shares an app with one, and an id with the other, none with the last.
app1 = App(AppConfig(collection_name="one_collection", id="new_app_id_1", collect_metrics=False))
app1 = App(
CustomAppConfig(
collection_name="one_collection",
id="new_app_id_1",
collect_metrics=False,
provider=Providers.OPENAI,
embedding_fn=EmbeddingFunctions.OPENAI,
chroma_settings={"allow_reset": True},
)
)
app2 = App(AppConfig(collection_name="one_collection", id="new_app_id_2", collect_metrics=False))
app3 = App(AppConfig(collection_name="three_collection", id="new_app_id_1", collect_metrics=False))
app4 = App(AppConfig(collection_name="four_collection", id="new_app_id_4", collect_metrics=False))