Richard Hein
8 months ago
parent
de8edca6d4
commit
d30f5f8259
@ -1,276 +0,0 @@
|
||||
""" Vector storage with RAG (Retrieval Augmented Generation) support for Markdown."""
|
||||
|
||||
import asyncio
|
||||
import glob
|
||||
import json
|
||||
import os
|
||||
from datetime import datetime
|
||||
from typing import Dict
|
||||
|
||||
from chromadb.config import Settings
|
||||
from langchain_community.document_loaders import UnstructuredMarkdownLoader
|
||||
from langchain_community.embeddings import HuggingFaceBgeEmbeddings
|
||||
from langchain.schema import BaseRetriever
|
||||
from langchain.storage import LocalFileStore
|
||||
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
||||
from langchain_chroma import Chroma
|
||||
|
||||
from playground.demos.chatbot.server.async_parent_document_retriever import \
|
||||
AsyncParentDocumentRetriever
|
||||
|
||||
STORE_TYPE = "local" # "redis" or "local"
|
||||
|
||||
|
||||
class VectorStorage:
|
||||
"""Vector storage class handles loading documents from a given directory."""
|
||||
|
||||
def __init__(self, directory, use_gpu=False):
|
||||
self.embeddings = HuggingFaceBgeEmbeddings(
|
||||
cache_folder="./.embeddings",
|
||||
model_name="BAAI/bge-large-en",
|
||||
model_kwargs={"device": "cuda" if use_gpu else "cpu"},
|
||||
encode_kwargs={"normalize_embeddings": True},
|
||||
query_instruction="Represent this sentence for searching relevant passages: ",
|
||||
)
|
||||
self.directory = directory
|
||||
self.child_splitter = RecursiveCharacterTextSplitter(
|
||||
chunk_size=200, chunk_overlap=20
|
||||
)
|
||||
self.parent_splitter = RecursiveCharacterTextSplitter(
|
||||
chunk_size=2000, chunk_overlap=200
|
||||
)
|
||||
if STORE_TYPE == "redis":
|
||||
from langchain_community.storage import RedisStore
|
||||
from langchain_community.storage.redis import get_client
|
||||
|
||||
username = r"username"
|
||||
password = r"password"
|
||||
client = get_client(
|
||||
redis_url=f"redis://{username}:{password}@localhost:6239"
|
||||
)
|
||||
self.store = RedisStore(client=client)
|
||||
else:
|
||||
self.store = LocalFileStore(root_path="./.parent_documents")
|
||||
self.settings = Settings(
|
||||
persist_directory="./.chroma_db",
|
||||
is_persistent=True,
|
||||
anonymized_telemetry=False,
|
||||
)
|
||||
# create a new vectorstore or get an existing one, with default collection
|
||||
self.vectorstore = self.get_vector_store()
|
||||
self.client = self.vectorstore._client
|
||||
self.retrievers: Dict[str, BaseRetriever] = {}
|
||||
# default retriever for when no collection title is specified
|
||||
self.retrievers["swarms"] = self.vectorstore.as_retriever()
|
||||
|
||||
async def init_retrievers(self, directories: list[str] | None = None):
|
||||
"""Initializes the vector storage retrievers."""
|
||||
start_time = datetime.now()
|
||||
print(f"Start vectorstore initialization time: {start_time}")
|
||||
|
||||
# for each subdirectory in the directory, create a new collection if it doesn't exist
|
||||
dirs = directories or os.listdir(self.directory)
|
||||
# make sure the subdir is not a file on MacOS (which has a hidden .DS_Store file)
|
||||
dirs = [
|
||||
subdir
|
||||
for subdir in dirs
|
||||
if not os.path.isfile(f"{self.directory}/{subdir}")
|
||||
]
|
||||
print(f"{len(dirs)} subdirectories to load: {dirs}")
|
||||
|
||||
self.retrievers[self.directory] = await self.init_retriever(
|
||||
self.directory
|
||||
)
|
||||
|
||||
end_time = datetime.now()
|
||||
print("Vectorstore initialization complete.")
|
||||
print(f"Vectorstore initialization end time: {end_time}")
|
||||
print(f"Total time taken: {end_time - start_time}")
|
||||
|
||||
return self.retrievers
|
||||
|
||||
async def init_retriever(self, subdir: str) -> BaseRetriever:
|
||||
""" Initialize each retriever. """
|
||||
# Ensure only one process/thread is executing this method at a time
|
||||
lock = asyncio.Lock()
|
||||
async with lock:
|
||||
subdir_start_time = datetime.now()
|
||||
print(f"Start {subdir} processing time: {subdir_start_time}")
|
||||
|
||||
# get all existing collections
|
||||
collections = self.client.list_collections()
|
||||
print(f"Existing collections: {collections}")
|
||||
|
||||
# Initialize an empty list to hold the documents
|
||||
documents = []
|
||||
# Define the maximum number of files to load at a time
|
||||
max_files = 1000
|
||||
|
||||
# Load existing metadata
|
||||
metadata_file = f"{self.directory}/metadata.json"
|
||||
metadata = {
|
||||
"processDate": str(datetime.now()),
|
||||
"processed_files": [],
|
||||
}
|
||||
processed_files = set() # Track processed files
|
||||
if os.path.isfile(metadata_file):
|
||||
with open(
|
||||
metadata_file, "r",
|
||||
) as metadata_file_handle:
|
||||
metadata = dict[str, str](json.load(metadata_file_handle))
|
||||
processed_files = {
|
||||
entry["file"]
|
||||
for entry in metadata.get("processed_files", [])
|
||||
}
|
||||
|
||||
# Get a list of all files in the directory and exclude processed files
|
||||
all_files = [
|
||||
file
|
||||
for file in glob.glob(
|
||||
f"{self.directory}/**/*.md", recursive=True
|
||||
)
|
||||
if file not in processed_files
|
||||
]
|
||||
|
||||
print(
|
||||
f"Loading {len(all_files)} documents for title version {subdir}."
|
||||
)
|
||||
# Load files in chunks of max_files
|
||||
for i in range(0, len(all_files), max_files):
|
||||
chunks_start_time = datetime.now()
|
||||
chunk_files = all_files[i : i + max_files]
|
||||
for file in chunk_files:
|
||||
loader = UnstructuredMarkdownLoader(
|
||||
file, mode="single", strategy="fast"
|
||||
)
|
||||
print(f"Loaded {file} in {subdir} ...")
|
||||
documents.extend(loader.load())
|
||||
|
||||
# Record the file as processed in metadata
|
||||
metadata["processed_files"].append(
|
||||
{"file": file, "processed_at": str(datetime.now())}
|
||||
)
|
||||
|
||||
print(
|
||||
f"Creating new collection for {self.directory}..."
|
||||
)
|
||||
# Create or get the collection
|
||||
collection = self.client.create_collection(
|
||||
name=self.directory,
|
||||
get_or_create=True,
|
||||
metadata={"processDate": metadata["processDate"]},
|
||||
)
|
||||
|
||||
# Reload vectorstore based on collection
|
||||
vectorstore = self.get_vector_store(
|
||||
collection_name=collection.name
|
||||
)
|
||||
|
||||
# Create a new parent document retriever
|
||||
retriever = AsyncParentDocumentRetriever(
|
||||
docstore=self.store,
|
||||
vectorstore=vectorstore,
|
||||
child_splitter=self.child_splitter,
|
||||
parent_splitter=self.parent_splitter,
|
||||
)
|
||||
|
||||
# force reload of collection to make sure we don't have
|
||||
# the default langchain collection
|
||||
collection = self.client.get_collection(
|
||||
name=self.directory
|
||||
)
|
||||
vectorstore = self.get_vector_store(
|
||||
collection_name=self.directory
|
||||
)
|
||||
|
||||
# Add documents to the collection and docstore
|
||||
print(
|
||||
f"Adding {len(documents)} documents to collection..."
|
||||
)
|
||||
add_docs_start_time = datetime.now()
|
||||
await retriever.aadd_documents(
|
||||
documents=documents, add_to_docstore=True
|
||||
)
|
||||
add_docs_end_time = datetime.now()
|
||||
total_time = add_docs_end_time - add_docs_start_time
|
||||
print(
|
||||
f"Adding {len(documents)} documents to collection took: {total_time}"
|
||||
)
|
||||
|
||||
documents = [] # clear documents list for next chunk
|
||||
|
||||
# Save metadata to the metadata.json file
|
||||
with open(
|
||||
metadata_file, "w"
|
||||
) as metadata_file_handle: # Changed variable name here
|
||||
json.dump(metadata, metadata_file_handle, indent=4)
|
||||
|
||||
print(
|
||||
f"Loaded {len(documents)} documents for directory '{subdir}'."
|
||||
)
|
||||
chunks_end_time = datetime.now()
|
||||
chunk_time = chunks_end_time - chunks_start_time
|
||||
print(
|
||||
f"{max_files} markdown file chunks processing time: {chunk_time}"
|
||||
)
|
||||
|
||||
subdir_end_time = datetime.now()
|
||||
print(
|
||||
f"Subdir {subdir} processing end time: {subdir_end_time}"
|
||||
)
|
||||
print(f"Time taken: {subdir_end_time - subdir_start_time}")
|
||||
|
||||
# Reload vectorstore based on collection to pass to parent doc retriever
|
||||
# collection = self.client.get_collection(name=self.directory)
|
||||
vectorstore = self.get_vector_store()
|
||||
retriever = AsyncParentDocumentRetriever(
|
||||
docstore=self.store,
|
||||
vectorstore=vectorstore,
|
||||
child_splitter=self.child_splitter,
|
||||
parent_splitter=self.parent_splitter,
|
||||
)
|
||||
return retriever
|
||||
|
||||
def get_vector_store(self, collection_name: str | None = None) -> Chroma:
|
||||
""" get a specific vector store for a collection """
|
||||
if collection_name is None or "" or "None":
|
||||
collection_name = "swarms"
|
||||
print("collection_name: " + collection_name)
|
||||
vectorstore = Chroma(
|
||||
client_settings=self.settings,
|
||||
embedding_function=self.embeddings,
|
||||
collection_name=collection_name,
|
||||
)
|
||||
return vectorstore
|
||||
|
||||
def list_collections(self):
|
||||
""" Get a list of all collections in the vectorstore """
|
||||
vectorstore = Chroma(
|
||||
client_settings=self.settings,
|
||||
embedding_function=self.embeddings,
|
||||
)
|
||||
return vectorstore._client.list_collections()
|
||||
|
||||
async def get_retriever(self, collection_name: str | None = None):
|
||||
""" get a specific retriever for a collection in the vectorstore """
|
||||
if self.retrievers is None:
|
||||
self.retrievers = await self.init_retrievers()
|
||||
|
||||
if (
|
||||
collection_name is None
|
||||
or collection_name == ""
|
||||
or collection_name == "None"
|
||||
):
|
||||
name = "swarms"
|
||||
else:
|
||||
name = collection_name
|
||||
|
||||
try:
|
||||
retriever = self.retrievers[name]
|
||||
except KeyError:
|
||||
print(f"Retriever for {name} not found, using default...")
|
||||
retriever = self.retrievers[
|
||||
"swarms"
|
||||
]
|
||||
|
||||
return retriever
|
||||
Loading…
Reference in new issue