File size: 1,033 Bytes
dbb2933
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
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})