Update app.py
Browse files
app.py
CHANGED
@@ -9,6 +9,9 @@ from sentence_transformers import SentenceTransformer
|
|
9 |
import google.generativeai as genai
|
10 |
import uuid
|
11 |
|
|
|
|
|
|
|
12 |
# Initialize Gemini API
|
13 |
genai.configure(api_key="AIzaSyAxUd2tS-qj9C7frYuHRsv92tziXHgIvLo")
|
14 |
|
@@ -21,6 +24,8 @@ if 'scraped' not in st.session_state:
|
|
21 |
st.session_state.scraped = False
|
22 |
if 'collection_name' not in st.session_state:
|
23 |
st.session_state.collection_name = "default_collection"
|
|
|
|
|
24 |
|
25 |
# Initialize embedding model
|
26 |
embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
|
@@ -87,12 +92,37 @@ def ask_question(query, collection_name):
|
|
87 |
response = model.generate_content(full_prompt)
|
88 |
return response.text
|
89 |
|
90 |
-
#
|
91 |
-
st.
|
92 |
|
93 |
-
#
|
94 |
-
with
|
95 |
-
st.
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
96 |
|
97 |
url = st.text_input("Enter the URL to scrape:")
|
98 |
|
@@ -102,63 +132,64 @@ with st.container():
|
|
102 |
result = scrape_text(url)
|
103 |
st.success(result)
|
104 |
|
105 |
-
#
|
106 |
-
|
107 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
108 |
st.subheader("Step 2: Ask Questions About the Scraped Content")
|
109 |
-
st.write("Ask a question about the content you've scraped:")
|
110 |
|
111 |
-
#
|
112 |
-
|
113 |
-
|
114 |
-
|
115 |
# Display chat history
|
116 |
for message in st.session_state.chat_history:
|
117 |
-
with
|
118 |
st.write(message["content"])
|
119 |
|
120 |
-
|
|
|
|
|
121 |
user_query = st.chat_input("Ask your question here")
|
122 |
|
123 |
if user_query:
|
124 |
# Add user question to chat history
|
125 |
st.session_state.chat_history.append({"role": "user", "content": user_query})
|
126 |
|
127 |
-
#
|
128 |
-
with st.
|
129 |
-
st.
|
130 |
-
|
131 |
-
# Get and display answer
|
132 |
-
with st.chat_message("assistant"):
|
133 |
-
with st.spinner("Searching database..."):
|
134 |
-
answer = ask_question(user_query, st.session_state.collection_name)
|
135 |
-
st.write(answer)
|
136 |
-
|
137 |
# Add answer to chat history
|
138 |
st.session_state.chat_history.append({"role": "assistant", "content": answer})
|
139 |
-
|
140 |
-
# Selection of existing collections
|
141 |
-
with st.sidebar:
|
142 |
-
st.header("Database Management")
|
143 |
-
|
144 |
-
# List available collections
|
145 |
-
try:
|
146 |
-
# Fix for ChromaDB v0.6.0 - list_collections() now returns only names
|
147 |
-
collection_names = chroma_client.list_collections()
|
148 |
-
|
149 |
-
if collection_names:
|
150 |
-
st.write("Available data collections:")
|
151 |
-
selected_collection = st.selectbox("Select a collection to query:", collection_names)
|
152 |
|
153 |
-
|
154 |
-
|
155 |
-
|
156 |
-
|
157 |
-
st.rerun()
|
158 |
-
except Exception as e:
|
159 |
-
st.error(f"Error: {str(e)}")
|
160 |
-
|
161 |
-
# Add a button to clear the session and start over
|
162 |
-
if st.button("Clear Chat History"):
|
163 |
-
st.session_state.chat_history = []
|
164 |
-
st.rerun()
|
|
|
9 |
import google.generativeai as genai
|
10 |
import uuid
|
11 |
|
12 |
+
# Page configuration
|
13 |
+
st.set_page_config(layout="wide")
|
14 |
+
|
15 |
# Initialize Gemini API
|
16 |
genai.configure(api_key="AIzaSyAxUd2tS-qj9C7frYuHRsv92tziXHgIvLo")
|
17 |
|
|
|
24 |
st.session_state.scraped = False
|
25 |
if 'collection_name' not in st.session_state:
|
26 |
st.session_state.collection_name = "default_collection"
|
27 |
+
if 'chat_history' not in st.session_state:
|
28 |
+
st.session_state.chat_history = []
|
29 |
|
30 |
# Initialize embedding model
|
31 |
embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
|
|
|
92 |
response = model.generate_content(full_prompt)
|
93 |
return response.text
|
94 |
|
95 |
+
# Create two columns: sidebar for database and main content
|
96 |
+
col1, main_col = st.columns([1, 3])
|
97 |
|
98 |
+
# Database management sidebar
|
99 |
+
with col1:
|
100 |
+
st.header("Database Management")
|
101 |
+
|
102 |
+
# List available collections
|
103 |
+
try:
|
104 |
+
# Fix for ChromaDB v0.6.0 - list_collections() now returns only names
|
105 |
+
collection_names = chroma_client.list_collections()
|
106 |
+
|
107 |
+
if collection_names:
|
108 |
+
st.write("Available data collections:")
|
109 |
+
selected_collection = st.selectbox("Select a collection to query:", collection_names)
|
110 |
+
|
111 |
+
if selected_collection and st.button("Load Selected Collection"):
|
112 |
+
st.session_state.collection_name = selected_collection
|
113 |
+
st.session_state.scraped = True
|
114 |
+
st.success(f"Loaded collection: {selected_collection}")
|
115 |
+
st.rerun()
|
116 |
+
except Exception as e:
|
117 |
+
st.error(f"Error: {str(e)}")
|
118 |
+
|
119 |
+
# Add a button to clear the session and start over
|
120 |
+
if st.button("Clear Chat History"):
|
121 |
+
st.session_state.chat_history = []
|
122 |
+
st.rerun()
|
123 |
+
|
124 |
+
# Scraping section
|
125 |
+
st.header("Step 1: Scrape a Website")
|
126 |
|
127 |
url = st.text_input("Enter the URL to scrape:")
|
128 |
|
|
|
132 |
result = scrape_text(url)
|
133 |
st.success(result)
|
134 |
|
135 |
+
# Main content area
|
136 |
+
with main_col:
|
137 |
+
st.title("Web Scraper & Q&A Chatbot")
|
138 |
+
|
139 |
+
# Use a container with custom CSS for the scrollable chat area
|
140 |
+
chat_container = st.container()
|
141 |
+
|
142 |
+
# Apply custom CSS for the chat container
|
143 |
+
st.markdown("""
|
144 |
+
<style>
|
145 |
+
.chat-container {
|
146 |
+
height: 500px;
|
147 |
+
overflow-y: auto;
|
148 |
+
border: 1px solid #ddd;
|
149 |
+
border-radius: 5px;
|
150 |
+
padding: 15px;
|
151 |
+
margin-bottom: 10px;
|
152 |
+
background-color: #f9f9f9;
|
153 |
+
}
|
154 |
+
.stChatInputContainer {
|
155 |
+
position: sticky;
|
156 |
+
bottom: 0;
|
157 |
+
background-color: white;
|
158 |
+
padding-top: 10px;
|
159 |
+
z-index: 100;
|
160 |
+
}
|
161 |
+
</style>
|
162 |
+
""", unsafe_allow_html=True)
|
163 |
+
|
164 |
+
# Q&A section - only appears after scraping is complete
|
165 |
+
if st.session_state.scraped:
|
166 |
st.subheader("Step 2: Ask Questions About the Scraped Content")
|
|
|
167 |
|
168 |
+
# Use a div with our custom class for the scrollable area
|
169 |
+
st.markdown('<div class="chat-container">', unsafe_allow_html=True)
|
170 |
+
|
|
|
171 |
# Display chat history
|
172 |
for message in st.session_state.chat_history:
|
173 |
+
with chat_container.chat_message(message["role"]):
|
174 |
st.write(message["content"])
|
175 |
|
176 |
+
st.markdown('</div>', unsafe_allow_html=True)
|
177 |
+
|
178 |
+
# Input for new question - always at the bottom
|
179 |
user_query = st.chat_input("Ask your question here")
|
180 |
|
181 |
if user_query:
|
182 |
# Add user question to chat history
|
183 |
st.session_state.chat_history.append({"role": "user", "content": user_query})
|
184 |
|
185 |
+
# Get answer
|
186 |
+
with st.spinner("Searching database..."):
|
187 |
+
answer = ask_question(user_query, st.session_state.collection_name)
|
188 |
+
|
|
|
|
|
|
|
|
|
|
|
|
|
189 |
# Add answer to chat history
|
190 |
st.session_state.chat_history.append({"role": "assistant", "content": answer})
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
191 |
|
192 |
+
# Rerun to update the UI with new messages
|
193 |
+
st.rerun()
|
194 |
+
else:
|
195 |
+
st.info("Please scrape a website or load a collection to start chatting.")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|