-
Notifications
You must be signed in to change notification settings - Fork 0
/
app.py
136 lines (113 loc) · 5.13 KB
/
app.py
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
import os
import traceback
import streamlit as st
from streamlit.runtime.uploaded_file_manager import UploadedFile
from benchmark.helix_ft import invoke
from dotenv import load_dotenv
from db.ASMEKnowledgeStore import ASMEKnowledgeStore
import firebase_admin
from firebase_admin import credentials, firestore
# Check if there's any Firebase app already initialized, if not, initialize one.
if not firebase_admin._apps:
cred = credentials.Certificate('ASME_key.json')
firebase_admin.initialize_app(cred)
db = firestore.client()
# ... rest of your Streamlit app code ...
load_dotenv()
def add_docs(vectordb: ASMEKnowledgeStore, files: list[UploadedFile]):
paths = []
for file in files:
tmp_location = os.path.join('/tmp', file.name)
with open(tmp_location, 'wb') as f:
f.write(file.read())
paths.append(tmp_location)
vectordb.add_docs(paths)
def get_response_and_reference(vectordb: ASMEKnowledgeStore, question):
try:
search_results = vectordb.similarity_search(question, k=3)
context = ""
for sr in search_results:
context += sr[0] + "\n"
# context = search_results[0][0] # This is the page_content of the top document
print(context)
try:
model_response = invoke({"question": question, "context": context})
resp = model_response['choices'][0]['message']['content'].strip()
except Exception as e:
return f"Failed to get model response: {e}"
ref = search_results[0][1] # This is the metadata of the top document
return resp, ref
except Exception as e:
print(f"Failed to get response and reference: {e}. Traceback: {traceback.format_exc()}")
return "None", "No reference available"
knowledge_store = ASMEKnowledgeStore(index_name="asme-bot-knowledge")
# Setting up the Streamlit UI
st.title("Piping Code Assistant")
uploaded_files: list[UploadedFile] = st.file_uploader("Upload PDF documents:", accept_multiple_files=True, type=['pdf'])
if uploaded_files:
add_docs(knowledge_store, uploaded_files)
user_question = st.text_input("Enter your question:", "")
# Check if there's any Firebase app already initialized, if not, initialize one.
if not firebase_admin._apps:
cred = credentials.Certificate('ASME_key.json')
firebase_admin.initialize_app(cred)
db = firestore.client()
# Define a function to reset the feedback form state
def reset_feedback_form_state():
st.session_state['feedback_text'] = ""
st.session_state['rating'] = None
st.session_state['feedback_submitted'] = False
# Initialize the session state variables
if 'feedback_text' not in st.session_state:
st.session_state['feedback_text'] = ""
if 'rating' not in st.session_state:
st.session_state['rating'] = None
if 'feedback_submitted' not in st.session_state:
st.session_state['feedback_submitted'] = False
# ... your existing code ...
def create_feedback_form():
with st.form(key='feedback_form'):
st.markdown("### Your feedback:")
st.markdown("<span style='color:grey;'>(Please drag the red dot towards right to adjust the rating)</span>",unsafe_allow_html=True)
# Use a select slider for the smiley ratings
smiley_ratings = {
'Very Unsatisfied ☹️': 1,
'Unsatisfied 🙁': 2,
'Neutral 😐': 3,
'Satisfied 🙂': 4,
'Very Satisfied 😄': 5
}
# Display the slider only if feedback is not submitted
if not st.session_state['feedback_submitted']:
rating = st.select_slider("Rating:", options=list(smiley_ratings.keys()), format_func=lambda x: x)
st.session_state['rating'] = smiley_ratings[rating]
# Text input for additional feedback
feedback_text = st.text_area("Feedback:", placeholder="Your feedback is valuable to us!",
value=st.session_state['feedback_text'], max_chars=200)
else:
rating = st.empty()
feedback_text = st.empty()
# Submit button for the form
submitted = st.form_submit_button("Submit Feedback")
if submitted and not st.session_state['feedback_submitted']:
feedback_data = {
"feedback": feedback_text,
"rating": st.session_state['rating'],
}
# Write data to Firestore
db.collection('feedback').add(feedback_data)
st.success("Thank you for your feedback!")
st.write("Your Feedback:", feedback_text)
st.write("Your Rating:", st.session_state['rating'])
# Mark feedback as submitted
st.session_state['feedback_submitted'] = True
# Clearing the input fields
reset_feedback_form_state()
# ... rest of your code ...
if user_question:
if not st.session_state['feedback_submitted']:
with st.spinner('Getting response...'):
response, reference = get_response_and_reference(knowledge_store, user_question)
st.text_area("Response:", value=response, height=300)
st.text_area("Reference:", value=reference, height=150)
create_feedback_form()