|
|
|
|
|
|
|
import os |
|
import pinecone |
|
from langchain.vectorstores import Chroma, Pinecone |
|
|
|
PERSIST_DIRECTORY = "./chroma_db" |
|
|
|
try: |
|
from dotenv import load_dotenv |
|
|
|
load_dotenv() |
|
except: |
|
pass |
|
|
|
|
|
def get_vectorstore(embeddings_function): |
|
if has_pinecone_config(): |
|
return get_pinecone_vectorstore(embeddings_function) |
|
return get_chroma_vectore_store(embeddings_function) |
|
|
|
|
|
def get_chroma_vectore_store(embedding_function): |
|
return Chroma( |
|
persist_directory=PERSIST_DIRECTORY, embedding_function=embedding_function |
|
) |
|
|
|
|
|
def has_pinecone_config(): |
|
return all( |
|
key in os.environ |
|
for key in [ |
|
"PINECONE_API_KEY", |
|
"PINECONE_API_ENVIRONMENT", |
|
"PINECONE_API_INDEX", |
|
] |
|
) |
|
|
|
|
|
def get_pinecone_vectorstore(embeddings_function, text_key="content"): |
|
|
|
pinecone.init( |
|
api_key=os.getenv("PINECONE_API_KEY"), |
|
environment=os.getenv("PINECONE_API_ENVIRONMENT"), |
|
) |
|
|
|
index_name = os.getenv("PINECONE_API_INDEX") |
|
vectorstore = Pinecone.from_existing_index( |
|
index_name, embeddings_function, text_key=text_key |
|
) |
|
return vectorstore |
|
|