ultimate-rag / backend /src /directory_reader.py
Ashad001's picture
pdf search
cd7bbf0
Raw
History Blame Contribute Delete
2.91 kB
import os
import time
from watchdog.observers import Observer
from watchdog.events import FileSystemEventHandler
from dotenv import load_dotenv
from langchain_groq import ChatGroq
from langchain import hub
from langchain_chroma import Chroma
from langchain_community.document_loaders import DirectoryLoader
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain_huggingface.embeddings import HuggingFaceEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
class Watcher:
def __init__(self, directory_to_watch, callback):
self.DIRECTORY_TO_WATCH = directory_to_watch
self.observer = Observer()
self.callback = callback
def run(self):
event_handler = Handler(self.callback)
self.observer.schedule(event_handler, self.DIRECTORY_TO_WATCH, recursive=True)
self.observer.start()
try:
while True:
time.sleep(5)
except KeyboardInterrupt:
self.observer.stop()
self.observer.join()
class Handler(FileSystemEventHandler):
def __init__(self, callback):
self.callback = callback
def on_any_event(self, event):
if event.is_directory:
return None
elif event.event_type == 'created' or event.event_type == 'modified' or event.event_type == 'deleted':
self.callback()
def update_database():
docs = DirectoryLoader(
path="./data/files",
silent_errors=True,
show_progress=True,
use_multithreading=True,
).load_documents()
if os.path.exists("./chroma_db"):
vectorstore = Chroma.from_directory("./chroma_db")
else:
text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
splits = text_splitter.split_documents(docs)
vectorstore = Chroma.from_documents(
documents=splits,
embedding=embeddings,
persist_directory="./chroma_db",
)
retriever = vectorstore.as_retriever()
prompt = hub.pull("rlm/rag-prompt")
def format_docs(docs):
return "\n\n".join(doc.page_content for doc in docs)
rag_chain = (
{"context": retriever | format_docs, "question": RunnablePassthrough()}
| prompt
| llm
| StrOutputParser()
)
response = rag_chain.invoke("Summarize the abstract of Importance of AI in evaluatiing climate change and food safety risk paper")
print(response)
model_name = "BAAI/bge-small-en"
model_kwargs = {"device": "cpu"}
encode_kwargs = {"normalize_embeddings": True}
embeddings = HuggingFaceEmbeddings()
load_dotenv()
os.environ["LANGCHAIN_TRACING_V2"] = "true"
llm = ChatGroq(model="llama3-8b-8192")
if __name__ == "__main__":
watcher = Watcher(directory_to_watch="./data/files", callback=update_database)
watcher.run()