File size: 6,393 Bytes
a2ac738 13ee6ba a2ac738 13ee6ba a2ac738 13ee6ba a2ac738 13ee6ba a2ac738 13ee6ba a2ac738 13ee6ba a2ac738 13ee6ba a2ac738 13ee6ba a2ac738 13ee6ba a2ac738 13ee6ba a2ac738 13ee6ba a2ac738 13ee6ba a2ac738 13ee6ba a2ac738 13ee6ba a2ac738 13ee6ba |
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 |
import os
import dotenv
from time import time
import streamlit as st
import logging
# Configure environment for Hugging Face Spaces
os.environ["HF_HOME"] = "/tmp/.cache/huggingface"
os.environ["TRANSFORMERS_CACHE"] = "/tmp/.cache/huggingface"
os.environ["HUGGINGFACE_HUB_CACHE"] = "/tmp/.cache/huggingface"
# Create necessary directories
os.makedirs("/tmp/.cache/huggingface", exist_ok=True)
os.makedirs("/tmp/chroma_persistent_db", exist_ok=True)
os.makedirs("/tmp/source_files", exist_ok=True)
# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
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
def clean_temp_files():
"""Clean up temporary files to prevent storage issues"""
try:
for folder in ["/tmp/source_files"]:
for filename in os.listdir(folder):
file_path = os.path.join(folder, filename)
if os.path.isfile(file_path):
os.unlink(file_path)
except Exception as e:
logger.warning(f"Error cleaning temp files: {e}")
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})
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:
try:
file_path = f"/tmp/source_files/{doc_file.name}"
with open(file_path, "wb") as file:
file.write(doc_file.getbuffer())
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)
logger.info(f"Successfully loaded document: {doc_file.name}")
except Exception as e:
st.toast(f"Error loading document {doc_file.name}: {str(e)}", icon="⚠️")
logger.error(f"Error loading document: {e}")
finally:
if os.path.exists(file_path):
os.remove(file_path)
else:
st.error(f"Max documents reached ({DB_DOCS_LIMIT}).")
if docs:
_split_and_load_docs(docs)
st.toast("Documents loaded successfully.", icon="✅")
clean_temp_files()
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)
logger.info(f"Successfully loaded URL: {url}")
except Exception as e:
st.error(f"Error loading from URL {url}: {str(e)}")
logger.error(f"Error loading 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):
embedding = HuggingFaceEmbeddings(
model_name="BAAI/bge-large-en-v1.5",
model_kwargs={'device': 'cpu'},
encode_kwargs={'normalize_embeddings': False},
cache_folder="/tmp/.cache"
)
persist_dir = "/tmp/chroma_persistent_db"
collection_name = "persistent_collection"
vector_db = Chroma.from_documents(
documents=docs,
embedding=embedding,
persist_directory=persist_dir,
collection_name=collection_name
)
vector_db.persist()
logger.info("Vector database initialized and persisted")
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()
logger.info("Added new documents to existing vector database")
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 |