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

@@ -11,6 +11,7 @@ from embedchain.loaders.web_page import WebPageLoader
from embedchain.chunkers.youtube_video import YoutubeVideoChunker
from embedchain.chunkers.pdf_file import PdfFileChunker
from embedchain.chunkers.web_page import WebPageChunker
from embedchain.vectordb.chroma_db import ChromaDB
load_dotenv()
@@ -21,13 +22,15 @@ DB_DIR = os.path.join(ABS_PATH, "db")
class EmbedChain:
def __init__(self, db):
def __init__(self, db=None):
"""
Initializes the EmbedChain instance, sets up a vector DB client and
creates a collection.
:param db: The instance of the VectorDB subclass.
"""
if db is None:
db = ChromaDB()
self.db_client = db.client
self.collection = db.collection
self.user_asks = []

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,