Spaces:
Running
Running
File size: 4,055 Bytes
b1b6964 212478c b1b6964 4afa186 b1b6964 212478c 8a06a9e bd6665e 66ad911 212478c b1b6964 6e1201a 3f6512f 3571b6c b1b6964 212478c 3f6512f 212478c 3571b6c 212478c 3571b6c b1b6964 212478c b1b6964 212478c 3571b6c 6e1201a 212478c 3f6512f 66ad911 212478c 3f6512f 66ad911 212478c 66ad911 3571b6c b1b6964 212478c b1b6964 212478c a8e2b6e 3571b6c b1b6964 212478c 3571b6c |
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 |
import streamlit as st
import os
import base64
from langchain.memory import ConversationBufferWindowMemory
from langchain_community.chat_message_histories import StreamlitChatMessageHistory
from utils.ingestion import DocumentProcessor
from utils.llm import LLMProcessor
from utils.qa import QAEngine
# Configure Streamlit page
st.set_page_config(page_title="AI-Powered Document QA", layout="wide")
# Function to encode image in Base64 for avatars
def encode_image(image_path):
with open(image_path, "rb") as file:
return base64.b64encode(file.read()).decode()
# Load avatar images
user_avatar = encode_image("./icons/user.jpg") # Change path if needed
ai_avatar = encode_image("./icons/ai.jpg")
# Initialize document processing & AI components
document_processor = DocumentProcessor()
llm_processor = LLMProcessor()
qa_engine = QAEngine()
# Ensure temp directory exists
os.makedirs("temp", exist_ok=True)
# Sidebar for file upload
st.sidebar.header("Upload a Document")
uploaded_file = st.sidebar.file_uploader("Choose a file", type=["pdf", "docx", "html", "pptx", "txt"])
# Initialize chat memory
memory_storage = StreamlitChatMessageHistory(key="chat_messages")
memory = ConversationBufferWindowMemory(
memory_key="chat_history", human_prefix="User", chat_memory=memory_storage, k=5
)
# Document upload & processing
if uploaded_file and "document_uploaded" not in st.session_state:
file_path = os.path.join("temp", uploaded_file.name)
with open(file_path, "wb") as f:
f.write(uploaded_file.read())
st.sidebar.success("File uploaded successfully!")
with st.spinner("Processing document..."):
document_processor.process_document(file_path)
st.sidebar.success("Document processed successfully!")
st.session_state["document_uploaded"] = True
# Chat UI Header
st.markdown("<h2 style='text-align: center;'>AI Chat Assistant</h2>", unsafe_allow_html=True)
st.markdown("---")
# Display chat history with avatars
for message in memory_storage.messages:
role = "user" if message.type == "human" else "assistant"
avatar = user_avatar if role == "user" else ai_avatar # Assign appropriate avatar
st.markdown(
f"""
<div style="display: flex; align-items: center; margin-bottom: 10px;">
<img src="data:image/jpeg;base64,{avatar}" width="40" height="40" style="border-radius: 50%; margin-right: 10px;">
<div style="background-color: #f1f1f1; padding: 10px; border-radius: 10px; max-width: 80%;">{message.content}</div>
</div>
""",
unsafe_allow_html=True
)
# User input at the bottom
user_input = st.chat_input("Ask me anything...")
if user_input:
memory_storage.add_user_message(user_input)
with st.markdown(
f"""
<div style="display: flex; align-items: center; margin-bottom: 10px;">
<img src="data:image/jpeg;base64,{user_avatar}" width="40" height="40" style="border-radius: 50%; margin-right: 10px;">
<div style="background-color: #d9f7be; padding: 10px; border-radius: 10px; max-width: 80%;">{user_input}</div>
</div>
""",
unsafe_allow_html=True
):
pass
with st.spinner("Generating response..."):
if st.session_state.get("document_uploaded", False):
answer = qa_engine.query(user_input)
else:
answer = llm_processor.generate_answer("", user_input)
st.warning("No document uploaded. This response is generated from general AI knowledge and may not be document-specific.")
memory_storage.add_ai_message(answer)
with st.markdown(
f"""
<div style="display: flex; align-items: center; margin-bottom: 10px;">
<img src="data:image/jpeg;base64,{ai_avatar}" width="40" height="40" style="border-radius: 50%; margin-right: 10px;">
<div style="background-color: #e6f7ff; padding: 10px; border-radius: 10px; max-width: 80%;">{answer.content}</div>
</div>
""",
unsafe_allow_html=True
):
pass
|