AWS-Guard-Bot / vectorstore.py
SSK-14's picture
Upload 17 files
1049895 verified
import os
import PyPDF2
from langchain_text_splitters import RecursiveCharacterTextSplitter
from qdrant_client import QdrantClient
from dotenv import load_dotenv
load_dotenv()
PATH_TO_KNOWLEDGE_BASE = "knowledge_base" # Path where the PDFs are stored
COLLECTION_NAME = "aws_faq" # Name of the collection
VECTOR_DB_PATH = "./qdrant" # Change this to your own path
# qdrant_client = QdrantClient(path=VECTOR_DB_PATH)
# If using qdrant cloud, use the following code
qdrant_client = QdrantClient(
os.getenv("QDRANT_URL"),
api_key=os.getenv("QDRANT_API_KEY"),
)
def ingest_embeddings():
metadatas = []
text = []
for file in os.listdir(PATH_TO_KNOWLEDGE_BASE):
if file.endswith('.pdf'):
pdf_path = os.path.join(PATH_TO_KNOWLEDGE_BASE, file)
pdf_reader = PyPDF2.PdfReader(pdf_path)
page_number = 1
for page in pdf_reader.pages:
text.append(page.extract_text())
metadatas.append({"page": page_number, "file": file})
page_number += 1
text_splitter = RecursiveCharacterTextSplitter(separators=["\n\n"], chunk_size=400, chunk_overlap=50)
chunked_documents = text_splitter.create_documents(text, metadatas=metadatas)
chunks, metadata, ids = zip(*[(chunk.page_content, chunk.metadata, i+1) for i, chunk in enumerate(chunked_documents)])
try:
qdrant_client.add(
collection_name=COLLECTION_NAME,
documents=chunks,
metadata=metadata,
ids=ids
)
print("Collection created and persisted")
except Exception as error:
print(f"Error: {error}")
if __name__ == "__main__":
ingest_embeddings()