Spaces:
Paused
Paused
File size: 2,709 Bytes
1e2602b f119af7 31f3de1 8669df3 31f3de1 2d7499f 31f3de1 2d7499f 8669df3 31f3de1 1e2602b 2d7499f 31f3de1 1e2602b 31f3de1 1e2602b 31f3de1 8669df3 31f3de1 8669df3 31f3de1 8669df3 2d7499f 31f3de1 8669df3 31f3de1 b29e219 31f3de1 b29e219 1e2602b 31f3de1 |
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 |
"""
IMPORTS HERE
"""
import chainlit as cl
from qdrant_client import QdrantClient
from qdrant_client.http.models import Distance, VectorParams
from langchain_qdrant import QdrantVectorStore
from operator import itemgetter
from langchain_core.runnables.passthrough import RunnablePassthrough
from langchain_core.runnables.config import RunnableConfig
import uuid
from prompts import chat_prompt
from handle_files import split_file
from models import chat_model, cached_embedder
"""
GLOBAL CODE HERE
"""
# Typical QDrant Client Set-up
collection_name = f"pdf_to_parse_{uuid.uuid4()}"
client = QdrantClient(":memory:")
client.create_collection(
collection_name=collection_name,
vectors_config=VectorParams(size=1536, distance=Distance.COSINE),
)
# Typical QDrant Vector Store Set-up
vectorstore = QdrantVectorStore(
client=client,
collection_name=collection_name,
embedding=cached_embedder)
### On Chat Start (Session Start) Section ###
@cl.on_chat_start
async def on_chat_start():
""" SESSION SPECIFIC CODE HERE """
files = None
# Wait for the user to upload a file
while files == None:
files = await cl.AskFileMessage(
content="Please upload a PDF File file to begin!",
accept=["application/pdf"],
max_size_mb=20,
timeout=180,
).send()
file = files[0]
msg = cl.Message(
content=f"Processing `{file.name}`..."
)
await msg.send()
docs = split_file(file)
vectorstore.add_documents(docs)
retriever = vectorstore.as_retriever(search_type="mmr", search_kwargs={"k": 15})
retrieval_augmented_qa_chain = (
{"context": itemgetter("question") | retriever, "question": itemgetter("question")}
| RunnablePassthrough.assign(context=itemgetter("context"))
| chat_prompt | chat_model
)
msg.content = f"Processing `{file.name}` done. You can now ask questions!"
await msg.send()
cl.user_session.set("chain", retrieval_augmented_qa_chain)
# ### Rename Chains ###
@cl.author_rename
def rename(orig_author: str):
""" RENAME CODE HERE """
rename_dict = {"ChatOpenAI": "the Generator ...", "VectorStoreRetriever" : "the Retriever"}
return rename_dict.get(orig_author, orig_author)
### On Message Section ###
@cl.on_message
async def main(message: cl.Message):
"""
MESSAGE CODE HERE
"""
chain = cl.user_session.get("chain")
msg = cl.Message(content="")
async for stream_response in chain.astream(
{"question":message.content},
config=RunnableConfig(callbacks=[cl.LangchainCallbackHandler()])
):
await msg.stream_token(stream_response.content)
await msg.send() |