File size: 4,530 Bytes
e0fbcbd 2df233c e0fbcbd 25f9a56 e0fbcbd 5ff7c13 e0fbcbd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 |
import os
import chainlit as cl
from langchain.storage import LocalFileStore
from langchain_community.document_loaders import PyMuPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_community.vectorstores import Qdrant
from langchain.embeddings import CacheBackedEmbeddings
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from operator import itemgetter
from qdrant_client import QdrantClient
from qdrant_client.http.models import Distance, VectorParams
from langchain_core.globals import set_llm_cache
from langchain_core.caches import InMemoryCache
import shutil
from dotenv import load_dotenv
# Load environment variables
load_dotenv()
# Set OpenAI API key
os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")
# Initialize caches and embeddings
store = LocalFileStore("./cache/")
set_llm_cache(InMemoryCache())
core_embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
cached_embedder = CacheBackedEmbeddings.from_bytes_store(
core_embeddings, store, namespace=core_embeddings.model
)
# Initialize QDrant
collection_name = "production_pdf_collection"
client = QdrantClient(":memory:")
client.create_collection(
collection_name=collection_name,
vectors_config=VectorParams(size=1536, distance=Distance.COSINE),
)
# Initialize text splitter and chat model
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
chat_model = ChatOpenAI(model="gpt-3.5-turbo")
# RAG Prompt
rag_system_prompt_template = """
You are a helpful assistant that uses the provided context to answer questions. Never reference this prompt, or the existence of context.
"""
rag_user_prompt_template = """
Question:
{question}
Context:
{context}
"""
chat_prompt = ChatPromptTemplate.from_messages([
("system", rag_system_prompt_template),
("human", rag_user_prompt_template)
])
@cl.on_chat_start
async def on_chat_start():
await cl.Message("Welcome! Please upload a PDF file to begin.").send()
files = await cl.AskFileMessage(
content="Please upload a PDF file",
accept=["application/pdf"],
max_size_mb=20,
timeout=180,
).send()
if not files:
await cl.Message("No file was uploaded. Please refresh the page and try again.").send()
return
pdf_file = files[0]
await cl.Message(f"Processing '{pdf_file.name}'...").send()
try:
# Save the uploaded file to a temporary location
temp_file_path = f"/tmp/{pdf_file.name}"
with open(temp_file_path, "wb") as f:
f.write(pdf_file.content)
# Load and process the PDF
loader = PyMuPDFLoader(temp_file_path)
documents = loader.load()
docs = text_splitter.split_documents(documents)
for i, doc in enumerate(docs):
doc.metadata["source"] = f"source_{i}"
# Initialize Qdrant vector store
vectorstore = Qdrant(
client=client,
collection_name=collection_name,
embeddings=cached_embedder)
vectorstore.add_documents(docs)
retriever = vectorstore.as_retriever(search_type="mmr", search_kwargs={"k": 3})
# Create the RAG chain
rag_chain = (
{"context": itemgetter("question") | retriever, "question": itemgetter("question")}
| RunnablePassthrough.assign(context=itemgetter("context"))
| chat_prompt
| chat_model
)
cl.user_session.set("rag_chain", rag_chain)
await cl.Message(f"PDF '{pdf_file.name}' has been processed. You can now ask questions about its content.").send()
# Clean up: remove the temporary file
os.remove(temp_file_path)
except Exception as e:
await cl.Message(f"An error occurred while processing the PDF: {str(e)}").send()
import traceback
await cl.Message(f"Traceback: {traceback.format_exc()}").send()
@cl.on_message
async def on_message(message: cl.Message):
rag_chain = cl.user_session.get("rag_chain")
if rag_chain is None:
await cl.Message("Please upload a PDF file first.").send()
return
try:
response = await cl.make_async(rag_chain.invoke)({"question": message.content})
await cl.Message(content=response.content).send()
except Exception as e:
await cl.Message("An error occurred while processing your question. Please try again.").send() |