|
import streamlit as st |
|
from pathlib import Path |
|
from data_preprocessing import process_docs |
|
from rag import create_rag_chain |
|
import time |
|
import pandas as pd |
|
import os |
|
from datetime import datetime |
|
|
|
|
|
FEEDBACK_FILE = "feedback.csv" |
|
if not os.path.exists(FEEDBACK_FILE): |
|
pd.DataFrame(columns=["timestamp", "query", "response", "rating"]).to_csv(FEEDBACK_FILE, index=False) |
|
|
|
def save_feedback(query, response, rating): |
|
feedback = { |
|
"timestamp": datetime.now().isoformat(), |
|
"query": query, |
|
"response": response, |
|
"rating": rating |
|
} |
|
pd.DataFrame([feedback]).to_csv(FEEDBACK_FILE, mode="a", header=False, index=False) |
|
|
|
def response_generator(prompt, chain): |
|
response = chain.invoke(prompt) |
|
for word in response.split(): |
|
yield word + " " |
|
time.sleep(0.05) |
|
|
|
|
|
save_directory = "docs" |
|
save_path = "docs/file.pdf" |
|
Path(save_directory).mkdir(parents=True, exist_ok=True) |
|
|
|
st.title("π InsureAgent") |
|
|
|
with st.sidebar: |
|
uploaded_file = st.file_uploader("Upload a document", type=("pdf")) |
|
if uploaded_file is not None: |
|
with open(save_path, "wb") as f: |
|
f.write(uploaded_file.getbuffer()) |
|
st.success(f"File saved successfully: {save_path}") |
|
|
|
|
|
show_feedback = st.checkbox("Show feedback data") |
|
|
|
|
|
retriever = process_docs(save_path) |
|
chain, chain_with_sources = create_rag_chain(retriever) |
|
|
|
|
|
if "messages" not in st.session_state: |
|
st.session_state.messages = [] |
|
|
|
|
|
for idx, message in enumerate(st.session_state.messages): |
|
with st.chat_message(message["role"]): |
|
st.markdown(message["content"]) |
|
|
|
|
|
if message["role"] == "assistant": |
|
if "rating" not in message: |
|
col1, col2 = st.columns(2) |
|
with col1: |
|
if st.button("π Good", key=f"good_{idx}"): |
|
message["rating"] = "good" |
|
query = st.session_state.messages[idx-1]["content"] |
|
save_feedback(query, message["content"], "good") |
|
st.rerun() |
|
with col2: |
|
if st.button("π Bad", key=f"bad_{idx}"): |
|
message["rating"] = "bad" |
|
query = st.session_state.messages[idx-1]["content"] |
|
save_feedback(query, message["content"], "bad") |
|
st.rerun() |
|
else: |
|
st.write(f"Rated: {message['rating'].capitalize()}") |
|
|
|
|
|
if show_feedback: |
|
st.sidebar.subheader("User Feedback") |
|
try: |
|
feedback_df = pd.read_csv(FEEDBACK_FILE) |
|
st.sidebar.dataframe(feedback_df) |
|
|
|
|
|
csv = feedback_df.to_csv(index=False).encode('utf-8') |
|
st.sidebar.download_button( |
|
label="Download feedback as CSV", |
|
data=csv, |
|
file_name="feedback_data.csv", |
|
mime="text/csv" |
|
) |
|
except FileNotFoundError: |
|
st.sidebar.warning("No feedback data yet") |
|
|
|
|
|
if prompt := st.chat_input("Ask about your insurance document:"): |
|
|
|
st.session_state.messages.append({"role": "user", "content": prompt}) |
|
|
|
|
|
with st.chat_message("user"): |
|
st.markdown(prompt) |
|
|
|
|
|
with st.chat_message("assistant"): |
|
response = st.write_stream(response_generator(prompt, chain)) |
|
|
|
|
|
st.session_state.messages.append({"role": "assistant", "content": response}) |