rag_datategy / app.py
Warblade's picture
minor change
175870c
raw
history blame
No virus
2.91 kB
import gradio as gr
from langchain.vectorstores import Qdrant
import qdrant_client
from langchain.schema.runnable import RunnablePassthrough
from langchain.schema.output_parser import StrOutputParser
from langchain.prompts import ChatPromptTemplate
from langchain.chat_models import ChatOpenAI
from langchain.embeddings import OpenAIEmbeddings
import dotenv
import os
from utils import template
import time
# Load environment variables and validate
dotenv.load_dotenv()
QDRANT_URL = os.getenv("QDRANT_URL")
QDRANT_API_KEY = os.getenv("QDRANT_API_KEY")
if not QDRANT_URL or not QDRANT_API_KEY:
raise ValueError("QDRANT_URL and QDRANT_API_KEY must be set in the environment")
# Initialize the vector store
def initiliaze_vector_store():
"""
Initialize and return the vector store.
Only run this on launch.
"""
embeddings = OpenAIEmbeddings()
client = qdrant_client.QdrantClient(url=QDRANT_URL, api_key=QDRANT_API_KEY)
vectorstore = Qdrant(client=client, collection_name="doc_datategy", embeddings=embeddings)
return vectorstore
# Initialize the retriever
def initiliaze_retriever(vectorstore):
"""
Initialize and return the retriever using the given vectorstore.
"""
return vectorstore.as_retriever()
# Initialize the chatbot
def initiliaze_chatbot(template, model_name="gpt-3.5-turbo-1106", temperature=0):
"""
Initialize and return the chatbot components: prompt and language model.
"""
prompt = ChatPromptTemplate.from_template(template)
llm = ChatOpenAI(model_name=model_name, temperature=temperature)
return prompt, llm
# Initialize the RAG chain
def initiliaze_RAG(retriever, prompt, llm):
"""
Initialize and return the RAG chain.
"""
context_function = {"context": retriever, "question": RunnablePassthrough()}
rag_chain = context_function | prompt | llm | StrOutputParser()
return rag_chain
# Launch Gradio app
vectorstore = initiliaze_vector_store()
retriever = initiliaze_retriever(vectorstore)
with gr.Blocks() as demo:
chatbot = gr.Chatbot(label="PapAI custom chatbot")
msg = gr.Textbox(label="Prompt", value='What is the minimum configuration for running papAI ?', interactive=True)
clear = gr.Button("Clear")
template_user = gr.Textbox(label="Template", value=template, interactive=True)
def change_template(template_user_str):
prompt, llm = initiliaze_chatbot(template_user_str)
return initiliaze_RAG(retriever, prompt, llm)
def RAG_answer(query, chat_history, template_user_str):
rag_chain = change_template(template_user_str)
answer = rag_chain.invoke(query)
chat_history.append((query, answer))
time.sleep(1.3) # Consider optimizing or dynamic handling
return "", chat_history
msg.submit(RAG_answer, [msg, chatbot, template_user], [msg, chatbot])
demo.queue()
demo.launch(share=False, debug=True)