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

update: chroma v0.4.0 (#330)

Browse Source
This commit is contained in:
Candido Sales Gomes
2023-07-20 17:17:53 -04:00
committed by GitHub
parent cdbf75e0ec
commit d590e4423b
3 changed files with 10 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
embedchain/vectordb/chroma_db.py
View File

@@ -1,6 +1,7 @@
import logging import logging
import chromadb import chromadb
from chromadb.config import Settings
from embedchain.vectordb.base_vector_db import BaseVectorDB from embedchain.vectordb.base_vector_db import BaseVectorDB
@@ -16,24 +17,16 @@ class ChromaDB(BaseVectorDB):
if host and port: if host and port:
logging.info(f"Connecting to ChromaDB server: {host}:{port}") logging.info(f"Connecting to ChromaDB server: {host}:{port}")
self.client_settings = chromadb.config.Settings( self.settings = Settings(chroma_server_host=host, chroma_server_http_port=port)
chroma_api_impl="rest",
chroma_server_host=host,
chroma_server_http_port=port,
)
else: else:
if db_dir is None: if db_dir is None:
db_dir = "db" db_dir = "db"
self.client_settings = chromadb.config.Settings( self.settings = Settings(persist_directory=db_dir, anonymized_telemetry=False, allow_reset=True)
chroma_db_impl="duckdb+parquet",
persist_directory=db_dir,
anonymized_telemetry=False,
)
super().__init__() super().__init__()
def _get_or_create_db(self): def _get_or_create_db(self):
"""Get or create the database.""" """Get or create the database."""
return chromadb.Client(self.client_settings) return chromadb.Client(self.settings)
def _get_or_create_collection(self): def _get_or_create_collection(self):
"""Get or create the collection.""" """Get or create the collection."""

4
pyproject.toml
View File

@@ -82,10 +82,10 @@ color = true
[tool.poetry.dependencies] [tool.poetry.dependencies]
python = ">=3.9,<3.9.7 || >3.9.7,<4.0" python = ">=3.9,<3.9.7 || >3.9.7,<4.0"
python-dotenv = "^1.0.0" python-dotenv = "^1.0.0"
langchain = "^0.0.205" langchain = "^0.0.237"
requests = "^2.31.0" requests = "^2.31.0"
openai = "^0.27.5" openai = "^0.27.5"
chromadb ="^0.3.26" chromadb ="^0.4.2"
youtube-transcript-api = "^0.6.1" youtube-transcript-api = "^0.6.1"
beautifulsoup4 = "^4.12.2" beautifulsoup4 = "^4.12.2"
pypdf = "^3.11.0" pypdf = "^3.11.0"

8
tests/vectordb/test_chroma_db.py
View File

@@ -6,6 +6,7 @@ from unittest.mock import patch
from embedchain import App from embedchain import App
from embedchain.config import AppConfig from embedchain.config import AppConfig
from embedchain.vectordb.chroma_db import ChromaDB, chromadb from embedchain.vectordb.chroma_db import ChromaDB, chromadb
from chromadb.config import Settings
class TestChromaDbHosts(unittest.TestCase): class TestChromaDbHosts(unittest.TestCase):
@@ -19,10 +20,9 @@ class TestChromaDbHosts(unittest.TestCase):
with patch.object(chromadb, "Client") as mock_client: with patch.object(chromadb, "Client") as mock_client:
_db = ChromaDB(host=host, port=port, embedding_fn=len) _db = ChromaDB(host=host, port=port, embedding_fn=len)
expected_settings = chromadb.config.Settings( expected_settings = Settings(
chroma_api_impl="rest", chroma_server_host="test-host",
chroma_server_host=host, chroma_server_http_port="1234",
chroma_server_http_port=port,
) )
mock_client.assert_called_once_with(expected_settings) mock_client.assert_called_once_with(expected_settings)