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

[Feat]: Add support for running chromadb in server mode with embedchain (#220)

Browse Source
This commit is contained in:
Deshraj Yadav
2023-07-11 03:17:40 -07:00
committed by GitHub
parent 9ca836520f
commit 73dd7151cb
4 changed files with 130 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
embedchain/config/InitConfig.py
View File

@@ -9,7 +9,7 @@ class InitConfig(BaseConfig):
Config to initialize an embedchain `App` instance.
"""
def __init__(self, log_level=None, ef=None, db=None):
def __init__(self, log_level=None, ef=None, db=None, host=None, port=None):
"""
:param log_level: Optional. (String) Debug level
['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'].
@@ -33,7 +33,7 @@ class InitConfig(BaseConfig):
if db is None:
from embedchain.vectordb.chroma_db import ChromaDB
self.db = ChromaDB(ef=self.ef)
self.db = ChromaDB(ef=self.ef, host=host, port=port)
else:
self.db = db

24
embedchain/vectordb/chroma_db.py
View File

@@ -9,7 +9,7 @@ from embedchain.vectordb.base_vector_db import BaseVectorDB
class ChromaDB(BaseVectorDB):
"""Vector database using ChromaDB."""
def __init__(self, db_dir=None, ef=None):
def __init__(self, db_dir=None, ef=None, host=None, port=None):
if ef:
self.ef = ef
else:
@@ -18,13 +18,21 @@ class ChromaDB(BaseVectorDB):
organization_id=os.getenv("OPENAI_ORGANIZATION"),
model_name="text-embedding-ada-002",
)
if db_dir is None:
db_dir = "db"
self.client_settings = chromadb.config.Settings(
chroma_db_impl="duckdb+parquet",
persist_directory=db_dir,
anonymized_telemetry=False,
)
if host and port:
self.client_settings = chromadb.config.Settings(
chroma_api_impl="rest",
chroma_server_host=host,
chroma_server_http_port=port,
)
else:
if db_dir is None:
db_dir = "db"
self.client_settings = chromadb.config.Settings(
chroma_db_impl="duckdb+parquet",
persist_directory=db_dir,
anonymized_telemetry=False,
)
super().__init__()
def _get_or_create_db(self):