rag / app.py
dfasd
Update app.py
620b6be verified
from dotenv import load_dotenv
import os
import gradio as gr
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import CharacterTextSplitter
from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_core.runnables import RunnablePassthrough
from langchain_openai import ChatOpenAI
from langchain import hub
from langchain_core.output_parsers import StrOutputParser
from langchain.chains import create_history_aware_retriever
from langchain.prompts import PromptTemplate
from langchain.chains.question_answering import load_qa_chain
import pydantic
# Load environment variables
load_dotenv()
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
# Initialize components
text_splitter = CharacterTextSplitter(separator="\n", chunk_size=1000, chunk_overlap=200, length_function=len)
embeddings = OpenAIEmbeddings(openai_api_key=OPENAI_API_KEY)
llm = ChatOpenAI(model="gpt-4-1106-preview", api_key=OPENAI_API_KEY)
vectordb_path = './vector_db'
dbname = 'vector_db'
uploaded_files = ['airbus.pdf', 'annualreport2223.pdf']
vectorstore = None
def create_vectordb():
for file in uploaded_files:
loader = PyPDFLoader(file)
data = loader.load()
texts = text_splitter.split_documents(data)
if vectorstore is None:
vectorstore = Chroma.from_documents(documents=texts, embedding=embeddings, persist_directory=os.path.join(vectordb_path, dbname))
else:
vectorstore.add_documents(texts)
def rag_bot(query, chat_history):
print(f"Received query: {query}")
template = """Please answer to human's input based on context. If the input is not mentioned in context, output something like 'I don't know'.
Context: {context}
Human: {human_input}
Your Response as Chatbot:"""
prompt_s = PromptTemplate(
input_variables=["human_input", "context"],
template=template
)
# Initialize vector store
vectorstore = Chroma(persist_directory=os.path.join(vectordb_path), embedding_function=embeddings)
# prompt = hub.pull("langchain-ai/chat-langchain-rephrase")
docs = vectorstore.similarity_search(query)
try:
stuff_chain = load_qa_chain(llm, chain_type="stuff", prompt=prompt_s)
except pydantic.ValidationError as e:
print(f"Validation error: {e}")
output = stuff_chain({"input_documents": docs, "human_input": query}, return_only_outputs=False)
final_answer = output["output_text"]
print(f"Final Answer ---> {final_answer}")
return final_answer
def chat(query, chat_history):
response = rag_bot(query, chat_history)
# chat_history.append((query, response))
return response
chatbot = gr.Chatbot(avatar_images=["user.jpg", "bot.png"], height=600)
clear_but = gr.Button(value="Clear Chat")
demo = gr.ChatInterface(fn=chat, title="RAG Chatbot Prototype", multimodal=False, retry_btn=None, undo_btn=None, clear_btn=clear_but, chatbot=chatbot)
if __name__ == '__main__':
demo.launch(debug=True, share=True)