Add default db loader, fix import

This commits builds on DumoeDss's PR. It

- adds a default db directory name.
- adds a default db instance (Chroma).

Both points offers flexibility for users who
want to use default and users who want to
customize

Lastly, it fixes an import
This commit is contained in:
Taranjeet Singh
2023-06-23 11:51:31 +05:30
parent 85a6a0c161
commit cff244b894
2 changed files with 11 additions and 4 deletions

View File

@@ -1,15 +1,19 @@
import os
import chromadb
from base_vector_db import BaseVectorDB
import os
from chromadb.utils import embedding_functions
from embedchain.vectordb.base_vector_db import BaseVectorDB
openai_ef = embedding_functions.OpenAIEmbeddingFunction(
api_key=os.getenv("OPENAI_API_KEY"),
model_name="text-embedding-ada-002"
)
class ChromaDB(BaseVectorDB):
def __init__(self, db_dir):
def __init__(self, db_dir=None):
if db_dir is None:
db_dir = "db"
self.client_settings = chromadb.config.Settings(
chroma_db_impl="duckdb+parquet",
persist_directory=db_dir,