File size: 3,173 Bytes
b1b6964
 
8a06a9e
212478c
 
 
b1b6964
 
4afa186
b1b6964
212478c
8a06a9e
bd6665e
d08fec1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
212478c
d08fec1
 
 
212478c
 
b1b6964
 
 
 
6e1201a
 
 
3f6512f
212478c
b1b6964
 
212478c
 
 
 
 
 
3f6512f
212478c
6e1201a
212478c
b1b6964
 
 
212478c
b1b6964
212478c
b1b6964
6e1201a
212478c
3f6512f
 
212478c
 
3f6512f
 
212478c
 
 
 
 
b1b6964
212478c
 
b1b6964
212478c
 
 
a8e2b6e
212478c
 
b1b6964
212478c
 
 
 
 
 
 
 
 
 
 
 
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
import streamlit as st
import os
import json
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")

# # Background Image
# def add_bg_from_local(image_file):
#     with open(image_file, "rb") as image_file:
#         encoded_string = base64.b64encode(image_file.read())
#     st.markdown(
#         f"""
#         <style>
#             .stApp {{
#                 background-image: url(data:image/png;base64,{encoded_string.decode()});
#                 background-size: cover;
#             }}
#         </style>
#         """,
#         unsafe_allow_html=True,
#     )

# # Path to background image
# image_bg = "./image/background.jpeg"  # Change this path accordingly
# add_bg_from_local(image_bg)

# 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 PDF")
uploaded_file = st.sidebar.file_uploader("Choose a PDF file", type=["pdf"])

# 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:
    pdf_path = os.path.join("temp", uploaded_file.name)

    with open(pdf_path, "wb") as f:
        f.write(uploaded_file.read())

    st.sidebar.success("File uploaded successfully!")

    with st.spinner("Processing document..."):
        document_processor.process_document(pdf_path)

    st.sidebar.success("Document processed successfully!")
    st.session_state["document_uploaded"] = True

# Chat interface layout
st.markdown("<h2 style='text-align: center;'>AI Chat Assistant</h2>", unsafe_allow_html=True)
st.markdown("---")

# Display chat history
for message in memory_storage.messages:
    role = "user" if message.type == "human" else "assistant"
    with st.chat_message(role):
        st.markdown(message.content)

# User input at the bottom
user_input = st.chat_input("Ask me anything...")

if user_input:
    # Store user message in memory
    memory_storage.add_user_message(user_input)

    with st.chat_message("user"):
        st.markdown(user_input)

    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.")

    # Store AI response in memory
    memory_storage.add_ai_message(answer)

    with st.chat_message("assistant"):
        st.markdown(answer)