File size: 6,673 Bytes
a2ac738 |
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 |
import os
import dotenv
from time import time
import streamlit as st
from langchain_community.document_loaders.text import TextLoader
from langchain_community.document_loaders import (
WebBaseLoader,
PyPDFLoader,
Docx2txtLoader,
)
from langchain_community.vectorstores import Chroma
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain.chains import create_history_aware_retriever, create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain
dotenv.load_dotenv()
os.environ["USER_AGENT"] = "myagent"
DB_DOCS_LIMIT = 10
# Stream non-RAG LLM response
def stream_llm_response(llm_stream, messages):
response_message = ""
for chunk in llm_stream.stream(messages):
response_message += chunk.content
yield chunk
st.session_state.messages.append({"role": "assistant", "content": response_message})
# --- Document Loading and Indexing ---
def load_doc_to_db():
if "rag_docs" in st.session_state and st.session_state.rag_docs:
docs = []
for doc_file in st.session_state.rag_docs:
if doc_file.name not in st.session_state.rag_sources:
if len(st.session_state.rag_sources) < DB_DOCS_LIMIT:
os.makedirs("source_files", exist_ok=True)
file_path = f"./source_files/{doc_file.name}"
with open(file_path, "wb") as file:
file.write(doc_file.read())
try:
if doc_file.type == "application/pdf":
loader = PyPDFLoader(file_path)
elif doc_file.name.endswith(".docx"):
loader = Docx2txtLoader(file_path)
elif doc_file.type in ["text/plain", "text/markdown"]:
loader = TextLoader(file_path)
else:
st.warning(f"Unsupported document type: {doc_file.type}")
continue
docs.extend(loader.load())
st.session_state.rag_sources.append(doc_file.name)
except Exception as e:
st.toast(f"Error loading document {doc_file.name}: {e}", icon="⚠️")
finally:
os.remove(file_path)
else:
st.error(f"Max documents reached ({DB_DOCS_LIMIT}).")
if docs:
_split_and_load_docs(docs)
st.toast(f"Documents loaded successfully.", icon="✅")
def load_url_to_db():
if "rag_url" in st.session_state and st.session_state.rag_url:
url = st.session_state.rag_url
docs = []
if url not in st.session_state.rag_sources:
if len(st.session_state.rag_sources) < DB_DOCS_LIMIT:
try:
loader = WebBaseLoader(url)
docs.extend(loader.load())
st.session_state.rag_sources.append(url)
except Exception as e:
st.error(f"Error loading from URL {url}: {e}")
if docs:
_split_and_load_docs(docs)
st.toast(f"Loaded content from URL: {url}", icon="✅")
else:
st.error(f"Max documents reached ({DB_DOCS_LIMIT}).")
def initialize_vector_db(docs):
# Initialize HuggingFace embeddings
embedding = HuggingFaceEmbeddings(
model_name="BAAI/bge-large-en-v1.5",
model_kwargs={'device': 'cpu'},
encode_kwargs={'normalize_embeddings': False}
)
# Shared persistent directory for long-term storage
persist_dir = "./chroma_persistent_db"
collection_name = "persistent_collection"
# Create the persistent Chroma vector store
vector_db = Chroma.from_documents(
documents=docs,
embedding=embedding,
persist_directory=persist_dir,
collection_name=collection_name
)
# Persist to disk
vector_db.persist()
return vector_db
def _split_and_load_docs(docs):
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=1000,
chunk_overlap=200,
)
chunks = text_splitter.split_documents(docs)
if "vector_db" not in st.session_state:
st.session_state.vector_db = initialize_vector_db(chunks)
else:
st.session_state.vector_db.add_documents(chunks)
st.session_state.vector_db.persist() # Save changes
# --- RAG Chain ---
def _get_context_retriever_chain(vector_db, llm):
retriever = vector_db.as_retriever()
prompt = ChatPromptTemplate.from_messages([
MessagesPlaceholder(variable_name="messages"),
("user", "{input}"),
("user", "Given the above conversation, generate a search query to find relevant information.")
])
return create_history_aware_retriever(llm, retriever, prompt)
def get_conversational_rag_chain(llm):
retriever_chain = _get_context_retriever_chain(st.session_state.vector_db, llm)
prompt = ChatPromptTemplate.from_messages([
("system",
"""You are a helpful assistant answering the user's queries using the provided context if available.\n
{context}"""),
MessagesPlaceholder(variable_name="messages"),
("user", "{input}")
])
stuff_documents_chain = create_stuff_documents_chain(llm, prompt)
return create_retrieval_chain(retriever_chain, stuff_documents_chain)
# Stream RAG LLM response
def stream_llm_rag_response(llm_stream, messages):
rag_chain = get_conversational_rag_chain(llm_stream)
# Extract latest user input and prior messages
input_text = messages[-1].content
history = messages[:-1]
# --- DEBUG: Show context retrieved ---
if st.session_state.get("debug_mode"):
retriever = st.session_state.vector_db.as_retriever()
retrieved_docs = retriever.get_relevant_documents(input_text)
st.markdown("### 🔍 Retrieved Context (Debug Mode)")
for i, doc in enumerate(retrieved_docs):
st.markdown(f"**Chunk {i+1}:**\n```\n{doc.page_content.strip()}\n```")
response_message = "*(RAG Response)*\n"
response = rag_chain.stream({
"messages": history,
"input": input_text
})
for chunk in response:
if 'answer' in chunk:
response_message += chunk['answer']
yield chunk['answer']
st.session_state.messages.append({"role": "assistant", "content": response_message})
|