AgenticRagNCERT / utils /vector_store.py
Ashvanth.S
Add initial files
dbb2933
raw
history blame
1.03 kB
import os
from dotenv import load_dotenv
from langchain_chroma import Chroma
load_dotenv()
persist_directory = os.getenv("VECTOR_STORE")
def create_vector_store(documents, unique_ids, embeddings):
"""
Creates a new vector store with the given documents, unique IDs, and embeddings.
"""
vector_store = Chroma(
collection_name="NCERT-Chapters",
embedding_function=embeddings,
persist_directory=persist_directory
)
vector_store.add_documents(documents=documents, ids=unique_ids)
vector_store.persist()
return vector_store
def load_vector_store(embeddings):
"""
Loads an existing vector store using the embeddings provided.
"""
return Chroma(
collection_name="NCERT-Chapters",
persist_directory=persist_directory,
embedding_function=embeddings
)
def get_retriever(vector_store, k=5):
"""
Returns a retriever object to search through the vector store.
"""
return vector_store.as_retriever(search_kwargs={"k": k})