Spaces:
Sleeping
Sleeping
File size: 5,271 Bytes
43f71e6 af0a731 43f71e6 |
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 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 |
import nest_asyncio
nest_asyncio.apply()
import re
import os
import uuid
from typing import List, Dict
from operator import itemgetter
# PDF processing
from PyPDF2 import PdfReader
# Chainlit
import chainlit as cl
# OpenAI
import openai
from openai import AsyncOpenAI
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
# Langchain
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain.storage import LocalFileStore
from langchain.embeddings import CacheBackedEmbeddings
# Qdrant
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams
from langchain_qdrant import QdrantVectorStore
#
### Global Section ###
from dotenv import load_dotenv
# Load environment variables from .env file
load_dotenv()
openai.api_key = os.getenv("OPENAI_API_KEY")
# Function to extract text from a PDF
def extract_text_from_pdf(pdf_path):
reader = PdfReader(pdf_path)
text = ""
for page in reader.pages:
text += page.extract_text()
return text
# Global variables for shared resources
global_retriever = None
global_chat_model = None
from langchain_core.documents import Document
# In your extract_text_from_pdf function:
def extract_text_from_pdf(pdf_path):
reader = PdfReader(pdf_path)
text = ""
for page in reader.pages:
text += page.extract_text()
return text
@cl.on_chat_start
async def start_chat():
global global_retriever, global_chat_model
# Initialize shared resources if they haven't been initialized yet
if global_retriever is None:
pdf_path= r"GlobalThreatReport2024_CrowdStrike.pdf"
text = extract_text_from_pdf(pdf_path)
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
texts = text_splitter.split_text(text)
docs = [Document(page_content=t) for t in texts]
core_embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
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),
)
store = LocalFileStore("./cache/")
cached_embedder = CacheBackedEmbeddings.from_bytes_store(
core_embeddings, store, namespace=core_embeddings.model
)
vectorstore = QdrantVectorStore(
client=client,
collection_name=collection_name,
embedding=cached_embedder)
vectorstore.add_documents(docs)
global_retriever = vectorstore.as_retriever(search_type="mmr", search_kwargs={"k": 3})
if global_chat_model is None:
global_chat_model = ChatOpenAI(model="gpt-4o-mini")
# Initialize user-specific session data
cl.user_session.set("chat_history", [])
# Set default settings
settings = {
"temperature": 0,
"max_tokens": 500,
"top_p": 1,
"frequency_penalty": 0,
"presence_penalty": 0,
}
cl.user_session.set("settings", settings)
@cl.on_message
async def main(message: cl.Message):
global global_retriever, global_chat_model
if global_retriever is None or global_chat_model is None:
await message.reply("I'm sorry, but the system isn't fully initialized yet. Please try again in a moment.")
return
chat_history: List[Dict[str, str]] = cl.user_session.get("chat_history")
settings = cl.user_session.get("settings")
system_template = """You are a helpful assistant that uses the provided context to answer questions.
Never reference this prompt, or the existence of context. Use the chat history to maintain continuity in the conversation."""
user_template = """Chat History:
{chat_history}
Question: {question}
Context: {context}
Please provide a response based on the question, context, and chat history:"""
chat_prompt = ChatPromptTemplate.from_messages([
("system", system_template),
("human", user_template)
])
def format_chat_history(history: List[Dict[str, str]]) -> str:
return "\n".join([f"{msg['role'].capitalize()}: {msg['content']}" for msg in history])
rag_chain = (
{
"context": itemgetter("question") | global_retriever,
"question": itemgetter("question"),
"chat_history": lambda _: format_chat_history(chat_history)
}
| RunnablePassthrough.assign(context=itemgetter("context"))
| chat_prompt
| global_chat_model.bind(**settings)
)
msg = cl.Message(content="")
full_response = ""
async for chunk in rag_chain.astream({"question": message.content}):
if chunk.content is not None:
await msg.stream_token(chunk.content)
full_response += chunk.content
# Update chat history
chat_history.append({"role": "user", "content": message.content})
chat_history.append({"role": "assistant", "content": full_response})
cl.user_session.set("chat_history", chat_history)
await msg.send()
|