sivan22 commited on
Commit
2c517f0
·
verified ·
1 Parent(s): 360f294

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +318 -318
app.py CHANGED
@@ -1,318 +1,318 @@
1
- import streamlit as st
2
- import os
3
- from typing import Optional, List
4
- from dotenv import load_dotenv
5
- import gdown
6
- import llm_providers
7
- import tantivy_search
8
- import agent
9
- import json
10
- import zipfile
11
-
12
-
13
- INDEX_PATH = "./index"
14
-
15
- # Load environment variables
16
- load_dotenv()
17
-
18
- class SearchAgentUI:
19
- index_path = INDEX_PATH
20
- gdrive_index_id = os.getenv("GDRIVE_INDEX_ID", "1lpbBCPimwcNfC0VZOlQueA4SHNGIp5_t")
21
-
22
-
23
- @st.cache_resource
24
- def get_agent(_self):
25
- index_path = INDEX_PATH
26
- return agent.Agent(index_path)
27
-
28
- def download_index_from_gdrive(self) -> bool:
29
- try:
30
- zip_path = "index.zip"
31
- url = f"https://drive.google.com/uc?id={self.gdrive_index_id}"
32
- gdown.download(url, zip_path, quiet=False)
33
- with zipfile.ZipFile(zip_path, 'r') as zip_ref:
34
- zip_ref.extractall(".")
35
- os.remove(zip_path)
36
- return True
37
-
38
- except Exception as e:
39
- st.error(f"Failed to download index: {str(e)}")
40
- return False
41
-
42
-
43
- @st.cache_resource
44
- def initialize_system(_self,api_keys:dict[str,str]) -> tuple[bool, str, List[str]]:
45
-
46
- try:
47
- # download index
48
- if not os.path.exists(_self.index_path):
49
- st.warning("Index folder not found. Attempting to download from Google Drive...")
50
- if not _self.download_index_from_gdrive():
51
- return False, "שגיאה: לא ניתן להוריד את האינדקס", []
52
- st.success("Index downloaded successfully!")
53
- _self.llm_providers = llm_providers.LLMProvider(api_keys)
54
- available_providers = _self.llm_providers.get_available_providers()
55
- if not available_providers:
56
- return False, "שגיאה: לא נמצאו ספקי AI זמינים. אנא הזן מפתח API אחד לפחות.", []
57
- return True, "המערכת מוכנה לחי শবפש", available_providers
58
-
59
- except Exception as ex:
60
- return False, f"שגיאה באתחול המערכת: {str(ex)}", []
61
-
62
- def update_messages(self, messages):
63
- st.session_state.messages = messages
64
-
65
- def main(self):
66
- st.set_page_config(
67
- page_title="איתוריא",
68
- layout="wide",
69
- initial_sidebar_state="expanded"
70
- )
71
-
72
-
73
- # Enhanced styling with better visual hierarchy and modern design
74
- st.markdown("""
75
- <style>
76
- /* Global RTL Support */
77
- .stApp {
78
- direction: rtl;
79
- background-color: #f8f9fa;
80
- }
81
-
82
- /* Input Fields RTL */
83
- .stTextInput > div > div > input,
84
- .stSelectbox > div > div > div,
85
- .stNumberInput > div > div > input {
86
- direction: rtl;
87
- border-radius: 8px !important;
88
- border: 2px solid #e2e8f0 !important;
89
- padding: 0.75rem !important;
90
- transition: all 0.3s ease;
91
- }
92
-
93
- .stTextInput > div > div > input:focus,
94
- .stSelectbox > div > div > div:focus {
95
- border-color: #4299e1 !important;
96
- box-shadow: 0 0 0 1px #4299e1 !important;
97
- }
98
-
99
- /* Message Containers */
100
- .chat-container {
101
- background: white;
102
- border-radius: 12px;
103
- padding: 1.5rem;
104
- margin: 1rem 0;
105
- box-shadow: 0 2px 4px rgba(0,0,0,0.1);
106
- }
107
-
108
- /* Tool Calls Styling */
109
- .tool-call {
110
- background: #f0f7ff;
111
- border-radius: 8px;
112
- padding: 1rem;
113
- margin: 0.5rem 0;
114
- border-right: 4px solid #3182ce;
115
- }
116
-
117
- /* Search Results */
118
- .search-step {
119
- background: white;
120
- border-radius: 10px;
121
- padding: 1.25rem;
122
- margin: 1rem 0;
123
- box-shadow: 0 2px 4px rgba(0,0,0,0.05);
124
- border: 1px solid #e2e8f0;
125
- }
126
-
127
- .document-group {
128
- background: #f7fafc;
129
- border-radius: 8px;
130
- padding: 1rem;
131
- margin: 0.75rem 0;
132
- border: 1px solid #e2e8f0;
133
- }
134
-
135
- .document-item {
136
- background: white;
137
- border-radius: 6px;
138
- padding: 1rem;
139
- margin: 0.5rem 0;
140
- border: 1px solid #edf2f7;
141
- }
142
-
143
- /* Sidebar Styling */
144
- [data-testid="stSidebar"] {
145
- direction: rtl;
146
- background-color: #f8fafc;
147
- padding: 2rem 1rem;
148
- }
149
-
150
- .sidebar-content {
151
- padding: 1rem;
152
- }
153
-
154
- /* Chat Messages */
155
- .stChatMessage {
156
- direction: rtl;
157
- background: white !important;
158
- border-radius: 12px !important;
159
- padding: 1rem !important;
160
- margin: 0.75rem 0 !important;
161
- box-shadow: 0 2px 4px rgba(0,0,0,0.05) !important;
162
- }
163
-
164
- /* Buttons */
165
- .stButton > button {
166
- border-radius: 8px !important;
167
- padding: 0.5rem 1.5rem !important;
168
- background-color: #3182ce !important;
169
- color: white !important;
170
- border: none !important;
171
- transition: all 0.3s ease !important;
172
- }
173
-
174
- .stButton > button:hover {
175
- background-color: #2c5282 !important;
176
- transform: translateY(-1px);
177
- }
178
-
179
- /* Code Blocks */
180
- .stCodeBlock {
181
- direction: ltr;
182
- text-align: left;
183
- border-radius: 8px !important;
184
- background: #2d3748 !important;
185
- }
186
-
187
- /* Links */
188
- a {
189
- color: #3182ce;
190
- text-decoration: none;
191
- transition: color 0.2s ease;
192
- }
193
-
194
- a:hover {
195
- color: #2c5282;
196
- text-decoration: underline;
197
- }
198
-
199
- /* Error Messages */
200
- .stAlert {
201
- border-radius: 8px !important;
202
- border: none !important;
203
- }
204
- </style>
205
- """, unsafe_allow_html=True)
206
-
207
- # Initialize session state for message deduplication
208
- if "messages" not in st.session_state:
209
- st.session_state.messages = []
210
-
211
- st.session_state.api_keys = {
212
- 'google': "",
213
- 'openai': "",
214
- 'anthropic': ""
215
- }
216
-
217
- # Sidebar settings
218
- with st.sidebar:
219
- st.title("הגדרות")
220
-
221
- st.subheader("הגדרת מפתחות API")
222
-
223
- # API Key inputs with improved styling
224
- for provider, label in [
225
- ('google', 'Google API Key'),
226
- ('openai', 'OpenAI API Key'),
227
- ('anthropic', 'Anthropic API Key')
228
- ]:
229
- key = st.text_input(
230
- label,
231
- value=st.session_state.api_keys[provider],
232
- type="password",
233
- key=f"{provider}_key",
234
- help=f"הזן את מפתח ה-API של {label}"
235
- )
236
- st.session_state.api_keys[provider] = key
237
-
238
- # Provider-specific links
239
- links = {
240
- 'google': 'https://aistudio.google.com/app/apikey',
241
- 'openai': 'https://platform.openai.com/account/api-keys',
242
- 'anthropic': 'https://console.anthropic.com/'
243
- }
244
- st.html(f'<small> ניתן להשיג מפתח <a href="{links[provider]}">כאן</a> </small>')
245
-
246
- st.markdown("---")
247
-
248
- # Initialize system
249
- success, status_msg, available_providers = self.initialize_system(st.session_state.api_keys)
250
-
251
- if not success:
252
- st.error(status_msg)
253
- return
254
-
255
- agent = self.get_agent()
256
-
257
- # Provider selection in sidebar
258
- with st.sidebar:
259
-
260
- if 'provider' not in st.session_state or st.session_state.provider not in available_providers:
261
- st.session_state.provider = available_providers[0]
262
-
263
- provider = st.selectbox(
264
- "ספק בינה מלאכותית",
265
- options=available_providers,
266
- key='provider',
267
- help="בחר את מודל הAI לשימוש (רק מודלים עם מפתח API זמין יוצגו)"
268
- )
269
- if agent:
270
- agent.set_llm(provider)
271
-
272
-
273
-
274
- # Main chat interface
275
-
276
- query = st.chat_input("הזן שאלה", key="chat_input")
277
- if query:
278
- stream = agent.chat(query)
279
- for chunk in stream:
280
- st.session_state.messages = chunk["messages"]
281
- if st.button("צ'אט חדש"):
282
- st.session_state.messages = []
283
- agent.clear_chat()
284
-
285
- for message in st.session_state.messages:
286
- if message.type == "tool":
287
- if message.name == "search":
288
- results =json.loads(message.content) if message.content else []
289
- with st.expander(f"🔍 תוצאות חיפוש: {len(results)}"):
290
- for result in results:
291
- st.write(result['reference'])
292
- st.info(result['text'])
293
- elif message.name == "get_text":
294
- st.expander(f"📝 טקסט: {message.content}")
295
-
296
- elif message.type == "ai" :
297
- if message.content != "":
298
-
299
- with st.chat_message(message.type):
300
- if isinstance(message.content, list):
301
- for item in message.content:
302
- if ('text' in item):
303
- st.write(item['text'])
304
-
305
- else:
306
- st.write(message.content)
307
-
308
- for tool_call in message.tool_calls:
309
- with st.expander(f"🛠️ שימוש בכלי: {tool_call["name"]}"):
310
- st.json(tool_call["args"])
311
- else:
312
- with st.chat_message(message.type):
313
- st.write(message.content)
314
-
315
-
316
- if __name__ == "__main__":
317
- app = SearchAgentUI()
318
- app.main()
 
1
+ import streamlit as st
2
+ import os
3
+ from typing import Optional, List
4
+ from dotenv import load_dotenv
5
+ import gdown
6
+ import llm_providers
7
+ import tantivy_search
8
+ import agent
9
+ import json
10
+ import zipfile
11
+
12
+
13
+ INDEX_PATH = "./index"
14
+
15
+ # Load environment variables
16
+ load_dotenv()
17
+
18
+ class SearchAgentUI:
19
+ index_path = INDEX_PATH
20
+ gdrive_index_id = os.getenv("GDRIVE_INDEX_ID", "1lpbBCPimwcNfC0VZOlQueA4SHNGIp5_t")
21
+
22
+
23
+ @st.cache_resource
24
+ def get_agent(_self):
25
+ index_path = INDEX_PATH
26
+ return agent.Agent(index_path)
27
+
28
+ def download_index_from_gdrive(self) -> bool:
29
+ try:
30
+ zip_path = "index.zip"
31
+ url = f"https://drive.google.com/uc?id={self.gdrive_index_id}"
32
+ gdown.download(url, zip_path, quiet=False)
33
+ with zipfile.ZipFile(zip_path, 'r') as zip_ref:
34
+ zip_ref.extractall(".")
35
+ os.remove(zip_path)
36
+ return True
37
+
38
+ except Exception as e:
39
+ st.error(f"Failed to download index: {str(e)}")
40
+ return False
41
+
42
+
43
+ @st.cache_resource
44
+ def initialize_system(_self,api_keys:dict[str,str]) -> tuple[bool, str, List[str]]:
45
+
46
+ try:
47
+ # download index
48
+ if not os.path.exists(_self.index_path):
49
+ st.warning("Index folder not found. Attempting to download from Google Drive...")
50
+ if not _self.download_index_from_gdrive():
51
+ return False, "שגיאה: לא ניתן להוריד את האינדקס", []
52
+ st.success("Index downloaded successfully!")
53
+ _self.llm_providers = llm_providers.LLMProvider(api_keys)
54
+ available_providers = _self.llm_providers.get_available_providers()
55
+ if not available_providers:
56
+ return False, "שגיאה: לא נמצאו ספקי AI זמינים. אנא הזן מפתח API אחד לפחות.", []
57
+ return True, "המערכת מוכנה לחי শবפש", available_providers
58
+
59
+ except Exception as ex:
60
+ return False, f"שגיאה באתחול המערכת: {str(ex)}", []
61
+
62
+ def update_messages(self, messages):
63
+ st.session_state.messages = messages
64
+
65
+ def main(self):
66
+ st.set_page_config(
67
+ page_title="איתוריא",
68
+ layout="wide",
69
+ initial_sidebar_state="expanded"
70
+ )
71
+
72
+
73
+ # Enhanced styling with better visual hierarchy and modern design
74
+ st.markdown("""
75
+ <style>
76
+ /* Global RTL Support */
77
+ .stApp {
78
+ direction: rtl;
79
+ background-color: #f8f9fa;
80
+ }
81
+
82
+ /* Input Fields RTL */
83
+ .stTextInput > div > div > input,
84
+ .stSelectbox > div > div > div,
85
+ .stNumberInput > div > div > input {
86
+ direction: rtl;
87
+ border-radius: 8px !important;
88
+ border: 2px solid #e2e8f0 !important;
89
+ padding: 0.75rem !important;
90
+ transition: all 0.3s ease;
91
+ }
92
+
93
+ .stTextInput > div > div > input:focus,
94
+ .stSelectbox > div > div > div:focus {
95
+ border-color: #4299e1 !important;
96
+ box-shadow: 0 0 0 1px #4299e1 !important;
97
+ }
98
+
99
+ /* Message Containers */
100
+ .chat-container {
101
+ background: white;
102
+ border-radius: 12px;
103
+ padding: 1.5rem;
104
+ margin: 1rem 0;
105
+ box-shadow: 0 2px 4px rgba(0,0,0,0.1);
106
+ }
107
+
108
+ /* Tool Calls Styling */
109
+ .tool-call {
110
+ background: #f0f7ff;
111
+ border-radius: 8px;
112
+ padding: 1rem;
113
+ margin: 0.5rem 0;
114
+ border-right: 4px solid #3182ce;
115
+ }
116
+
117
+ /* Search Results */
118
+ .search-step {
119
+ background: white;
120
+ border-radius: 10px;
121
+ padding: 1.25rem;
122
+ margin: 1rem 0;
123
+ box-shadow: 0 2px 4px rgba(0,0,0,0.05);
124
+ border: 1px solid #e2e8f0;
125
+ }
126
+
127
+ .document-group {
128
+ background: #f7fafc;
129
+ border-radius: 8px;
130
+ padding: 1rem;
131
+ margin: 0.75rem 0;
132
+ border: 1px solid #e2e8f0;
133
+ }
134
+
135
+ .document-item {
136
+ background: white;
137
+ border-radius: 6px;
138
+ padding: 1rem;
139
+ margin: 0.5rem 0;
140
+ border: 1px solid #edf2f7;
141
+ }
142
+
143
+ /* Sidebar Styling */
144
+ [data-testid="stSidebar"] {
145
+ direction: rtl;
146
+ background-color: #f8fafc;
147
+ padding: 2rem 1rem;
148
+ }
149
+
150
+ .sidebar-content {
151
+ padding: 1rem;
152
+ }
153
+
154
+ /* Chat Messages */
155
+ .stChatMessage {
156
+ direction: rtl;
157
+ background: white !important;
158
+ border-radius: 12px !important;
159
+ padding: 1rem !important;
160
+ margin: 0.75rem 0 !important;
161
+ box-shadow: 0 2px 4px rgba(0,0,0,0.05) !important;
162
+ }
163
+
164
+ /* Buttons */
165
+ .stButton > button {
166
+ border-radius: 8px !important;
167
+ padding: 0.5rem 1.5rem !important;
168
+ background-color: #3182ce !important;
169
+ color: white !important;
170
+ border: none !important;
171
+ transition: all 0.3s ease !important;
172
+ }
173
+
174
+ .stButton > button:hover {
175
+ background-color: #2c5282 !important;
176
+ transform: translateY(-1px);
177
+ }
178
+
179
+ /* Code Blocks */
180
+ .stCodeBlock {
181
+ direction: ltr;
182
+ text-align: left;
183
+ border-radius: 8px !important;
184
+ background: #2d3748 !important;
185
+ }
186
+
187
+ /* Links */
188
+ a {
189
+ color: #3182ce;
190
+ text-decoration: none;
191
+ transition: color 0.2s ease;
192
+ }
193
+
194
+ a:hover {
195
+ color: #2c5282;
196
+ text-decoration: underline;
197
+ }
198
+
199
+ /* Error Messages */
200
+ .stAlert {
201
+ border-radius: 8px !important;
202
+ border: none !important;
203
+ }
204
+ </style>
205
+ """, unsafe_allow_html=True)
206
+
207
+ # Initialize session state for message deduplication
208
+ if "messages" not in st.session_state:
209
+ st.session_state.messages = []
210
+
211
+ st.session_state.api_keys = {
212
+ 'google': "",
213
+ 'openai': "",
214
+ 'anthropic': ""
215
+ }
216
+
217
+ # Sidebar settings
218
+ with st.sidebar:
219
+ st.title("הגדרות")
220
+
221
+ st.subheader("הגדרת מפתחות API")
222
+
223
+ # API Key inputs with improved styling
224
+ for provider, label in [
225
+ ('google', 'Google API Key'),
226
+ ('openai', 'OpenAI API Key'),
227
+ ('anthropic', 'Anthropic API Key')
228
+ ]:
229
+ key = st.text_input(
230
+ label,
231
+ value=st.session_state.api_keys[provider],
232
+ type="password",
233
+ key=f"{provider}_key",
234
+ help=f"הזן את מפתח ה-API של {label}"
235
+ )
236
+ st.session_state.api_keys[provider] = key
237
+
238
+ # Provider-specific links
239
+ links = {
240
+ 'google': 'https://aistudio.google.com/app/apikey',
241
+ 'openai': 'https://platform.openai.com/account/api-keys',
242
+ 'anthropic': 'https://console.anthropic.com/'
243
+ }
244
+ st.html(f'<small> ניתן להשיג מפתח <a href="{links[provider]}">כאן</a> </small>')
245
+
246
+ st.markdown("---")
247
+
248
+ # Initialize system
249
+ success, status_msg, available_providers = self.initialize_system(st.session_state.api_keys)
250
+
251
+ if not success:
252
+ st.error(status_msg)
253
+ return
254
+
255
+ agent = self.get_agent()
256
+
257
+ # Provider selection in sidebar
258
+ with st.sidebar:
259
+
260
+ if 'provider' not in st.session_state or st.session_state.provider not in available_providers:
261
+ st.session_state.provider = available_providers[0]
262
+
263
+ provider = st.selectbox(
264
+ "ספק בינה מלאכותית",
265
+ options=available_providers,
266
+ key='provider',
267
+ help="בחר את מודל הAI לשימוש (רק מודלים עם מפתח API זמין יוצגו)"
268
+ )
269
+ if agent:
270
+ agent.set_llm(provider)
271
+
272
+
273
+
274
+ # Main chat interface
275
+
276
+ query = st.chat_input("הזן שאלה", key="chat_input")
277
+ if query:
278
+ stream = agent.chat(query)
279
+ for chunk in stream:
280
+ st.session_state.messages = chunk["messages"]
281
+ if st.button("צ'אט חדש"):
282
+ st.session_state.messages = []
283
+ agent.clear_chat()
284
+
285
+ for message in st.session_state.messages:
286
+ if message.type == "tool":
287
+ if message.name == "search":
288
+ results =json.loads(message.content) if message.content else []
289
+ with st.expander(f"🔍 תוצאות חיפוש: {len(results)}"):
290
+ for result in results:
291
+ st.write(result['reference'])
292
+ st.info(result['text'])
293
+ elif message.name == "get_text":
294
+ st.expander(f"📝 טקסט: {message.content}")
295
+
296
+ elif message.type == "ai" :
297
+ if message.content != "":
298
+
299
+ with st.chat_message(message.type):
300
+ if isinstance(message.content, list):
301
+ for item in message.content:
302
+ if ('text' in item):
303
+ st.write(item['text'])
304
+
305
+ else:
306
+ st.write(message.content)
307
+
308
+ for tool_call in message.tool_calls:
309
+ with st.expander(f"🛠️ שימוש בכלי: {tool_call['name']}"):
310
+ st.json(tool_call["args"])
311
+ else:
312
+ with st.chat_message(message.type):
313
+ st.write(message.content)
314
+
315
+
316
+ if __name__ == "__main__":
317
+ app = SearchAgentUI()
318
+ app.main()