cyyeh commited on
Commit
0208146
Β·
1 Parent(s): 50c28ba
Files changed (2) hide show
  1. src/apis.py +9 -8
  2. src/pages/4_Knowledge.py +307 -0
src/apis.py CHANGED
@@ -132,11 +132,12 @@ def generate_summary(
132
 
133
 
134
  def stream_explanation(
135
- api_key: str
 
136
  ) -> tuple[dict, str]:
137
  """Stream explanation for a query."""
138
  base_url = "https://cloud.getwren.ai/api/v1"
139
- endpoint = f"{base_url}/stream_explanation"
140
  headers = {
141
  "Authorization": f"Bearer {api_key}",
142
  "Accept": "text/event-stream",
@@ -347,7 +348,7 @@ def delete_instruction(
347
  api_key: str,
348
  project_id: str,
349
  instruction_id: str,
350
- ) -> tuple[dict, str]:
351
  """Delete an instruction."""
352
  base_url = "https://cloud.getwren.ai/api/v1"
353
  endpoint = f"{base_url}/projects/{project_id}/knowledge/instructions/{instruction_id}"
@@ -358,9 +359,9 @@ def delete_instruction(
358
  try:
359
  response = requests.delete(endpoint, headers=headers)
360
  response.raise_for_status()
361
- return response.json(), ""
362
  except requests.exceptions.RequestException as e:
363
- return {}, str(e)
364
 
365
 
366
  def get_sql_pairs(
@@ -439,7 +440,7 @@ def delete_sql_pair(
439
  api_key: str,
440
  project_id: str,
441
  sql_pair_id: str,
442
- ) -> tuple[dict, str]:
443
  """Delete an SQL pair."""
444
  base_url = "https://cloud.getwren.ai/api/v1"
445
  endpoint = f"{base_url}/projects/{project_id}/knowledge/sql_pairs/{sql_pair_id}"
@@ -450,6 +451,6 @@ def delete_sql_pair(
450
  try:
451
  response = requests.delete(endpoint, headers=headers)
452
  response.raise_for_status()
453
- return response.json(), ""
454
  except requests.exceptions.RequestException as e:
455
- return {}, str(e)
 
132
 
133
 
134
  def stream_explanation(
135
+ api_key: str,
136
+ query_id: str,
137
  ) -> tuple[dict, str]:
138
  """Stream explanation for a query."""
139
  base_url = "https://cloud.getwren.ai/api/v1"
140
+ endpoint = f"{base_url}/stream_explanation/{query_id}"
141
  headers = {
142
  "Authorization": f"Bearer {api_key}",
143
  "Accept": "text/event-stream",
 
348
  api_key: str,
349
  project_id: str,
350
  instruction_id: str,
351
+ ) -> str:
352
  """Delete an instruction."""
353
  base_url = "https://cloud.getwren.ai/api/v1"
354
  endpoint = f"{base_url}/projects/{project_id}/knowledge/instructions/{instruction_id}"
 
359
  try:
360
  response = requests.delete(endpoint, headers=headers)
361
  response.raise_for_status()
362
+ return ""
363
  except requests.exceptions.RequestException as e:
364
+ return str(e)
365
 
366
 
367
  def get_sql_pairs(
 
440
  api_key: str,
441
  project_id: str,
442
  sql_pair_id: str,
443
+ ) -> str:
444
  """Delete an SQL pair."""
445
  base_url = "https://cloud.getwren.ai/api/v1"
446
  endpoint = f"{base_url}/projects/{project_id}/knowledge/sql_pairs/{sql_pair_id}"
 
451
  try:
452
  response = requests.delete(endpoint, headers=headers)
453
  response.raise_for_status()
454
+ return ""
455
  except requests.exceptions.RequestException as e:
456
+ return str(e)
src/pages/4_Knowledge.py ADDED
@@ -0,0 +1,307 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+
3
+ from apis import (
4
+ get_instructions,
5
+ create_instruction,
6
+ update_instruction,
7
+ delete_instruction,
8
+ get_sql_pairs,
9
+ create_sql_pair,
10
+ update_sql_pair,
11
+ delete_sql_pair,
12
+ )
13
+
14
+
15
+ def main():
16
+ st.title("πŸ“š Wren AI Cloud API Demo - Knowledge Management")
17
+
18
+ if "api_key" not in st.session_state or "project_id" not in st.session_state:
19
+ st.error("Please enter your API Key and Project ID in the sidebar of Home page to get started.")
20
+ return
21
+ if not st.session_state.api_key or not st.session_state.project_id:
22
+ st.error("Please enter your API Key and Project ID in the sidebar of Home page to get started.")
23
+ return
24
+
25
+ api_key = st.session_state.api_key
26
+ project_id = st.session_state.project_id
27
+
28
+ st.markdown('Using APIs: [Knowledge Instructions](https://wrenai.readme.io/reference/knowledge-1), [Knowledge SQL Pairs](https://wrenai.readme.io/reference/knowledge-1)')
29
+
30
+ # Tabs for Instructions and SQL Pairs
31
+ tab1, tab2 = st.tabs(["πŸ’‘ Instructions", "πŸ’‘ SQL Pairs"])
32
+
33
+ with tab1:
34
+ manage_instructions(api_key, project_id)
35
+
36
+ with tab2:
37
+ manage_sql_pairs(api_key, project_id)
38
+
39
+
40
+ def manage_instructions(api_key: str, project_id: str):
41
+ st.header("Instructions Management")
42
+
43
+ # Add new instruction section
44
+ with st.expander("πŸ’‘ Add New Instruction", expanded=False):
45
+ with st.form("add_instruction_form"):
46
+ instruction_text = st.text_area(
47
+ "Instruction Text",
48
+ placeholder="Enter instruction text...",
49
+ height=100,
50
+ help="The instruction text that will guide the AI"
51
+ )
52
+ is_global = st.checkbox(
53
+ "Global Instruction",
54
+ value=True,
55
+ help="Whether this instruction applies globally"
56
+ )
57
+ questions_input = st.text_area(
58
+ "Related Questions (one per line)",
59
+ placeholder="Question 1\nQuestion 2\nQuestion 3",
60
+ height=80,
61
+ help="Optional: Questions that this instruction relates to"
62
+ )
63
+
64
+ submitted = st.form_submit_button("Create Instruction")
65
+
66
+ if submitted:
67
+ if instruction_text.strip():
68
+ questions = [q.strip() for q in questions_input.split('\n') if q.strip()]
69
+ response, error = create_instruction(
70
+ api_key,
71
+ project_id,
72
+ instruction_text,
73
+ is_global,
74
+ questions if questions else None
75
+ )
76
+
77
+ if response:
78
+ st.success("Instruction created successfully!")
79
+ st.rerun()
80
+ else:
81
+ st.error(f"Error creating instruction: {error}")
82
+ else:
83
+ st.error("Please enter instruction text")
84
+
85
+ # Display existing instructions
86
+ st.subheader("Existing Instructions")
87
+
88
+ with st.spinner("Loading instructions..."):
89
+ instructions_response, error = get_instructions(api_key, project_id)
90
+
91
+ if instructions_response:
92
+ for i, instruction in enumerate(instructions_response):
93
+ with st.expander(f"πŸ’‘ {instruction.get('instruction', 'No instruction text')[:50]}..."):
94
+ col1, col2 = st.columns([3, 1])
95
+
96
+ with col1:
97
+ st.write("**Instruction Text:**")
98
+ st.write(instruction.get("instruction", "No instruction text"))
99
+ st.write(f"**Global:** {instruction.get('isGlobal', False)}")
100
+
101
+ questions = instruction.get("questions", [])
102
+ if questions:
103
+ st.write("**Related Questions:**")
104
+ for q in questions:
105
+ st.write(f"- {q}")
106
+
107
+ with col2:
108
+ # Edit button
109
+ if st.button("πŸ”„ Edit", key=f"edit_inst_{i}"):
110
+ st.session_state[f"edit_instruction_{i}"] = True
111
+
112
+ # Delete button
113
+ if st.button("πŸ—‘οΈ Delete", key=f"delete_inst_{i}"):
114
+ error = delete_instruction(
115
+ api_key,
116
+ project_id,
117
+ instruction["id"]
118
+ )
119
+ if not error:
120
+ st.success("Instruction deleted successfully!")
121
+ st.rerun()
122
+ else:
123
+ st.error(f"Error deleting instruction: {error}")
124
+
125
+ # Edit form
126
+ if st.session_state.get(f"edit_instruction_{i}", False):
127
+ st.divider()
128
+ with st.form(f"edit_instruction_form_{i}"):
129
+ new_instruction = st.text_area(
130
+ "Edit Instruction Text",
131
+ value=instruction.get("instruction", ""),
132
+ height=100
133
+ )
134
+ new_is_global = st.checkbox(
135
+ "Global Instruction",
136
+ value=instruction.get("isGlobal", True)
137
+ )
138
+ current_questions = instruction.get("questions", [])
139
+ new_questions_input = st.text_area(
140
+ "Related Questions (one per line)",
141
+ value='\n'.join(current_questions),
142
+ height=80
143
+ )
144
+
145
+ col1, col2 = st.columns(2)
146
+ with col1:
147
+ update_submitted = st.form_submit_button("Update Instruction")
148
+ with col2:
149
+ cancel_submitted = st.form_submit_button("Cancel")
150
+
151
+ if update_submitted:
152
+ if new_instruction.strip():
153
+ new_questions = [q.strip() for q in new_questions_input.split('\n') if q.strip()]
154
+ response, error = update_instruction(
155
+ api_key,
156
+ project_id,
157
+ instruction["id"],
158
+ new_instruction,
159
+ new_is_global,
160
+ new_questions if new_questions else None
161
+ )
162
+
163
+ if response:
164
+ st.success("Instruction updated successfully!")
165
+ del st.session_state[f"edit_instruction_{i}"]
166
+ st.rerun()
167
+ else:
168
+ st.error(f"Error updating instruction: {error}")
169
+ else:
170
+ st.error("Please enter instruction text")
171
+
172
+ if cancel_submitted:
173
+ del st.session_state[f"edit_instruction_{i}"]
174
+ st.rerun()
175
+ else:
176
+ if not error:
177
+ st.info("No instructions found. Create your first instruction above!")
178
+ else:
179
+ st.error(f"Error loading instructions: {error}")
180
+
181
+
182
+ def manage_sql_pairs(api_key: str, project_id: str):
183
+ st.header("SQL Pairs Management")
184
+
185
+ # Add new SQL pair section
186
+ with st.expander("πŸ’‘ Add New SQL Pair", expanded=False):
187
+ with st.form("add_sql_pair_form"):
188
+ question = st.text_area(
189
+ "Question",
190
+ placeholder="Enter the natural language question...",
191
+ height=80,
192
+ help="The question in natural language"
193
+ )
194
+ sql = st.text_area(
195
+ "SQL Query",
196
+ placeholder="SELECT * FROM table_name WHERE condition;",
197
+ height=120,
198
+ help="The corresponding SQL query"
199
+ )
200
+
201
+ submitted = st.form_submit_button("Create SQL Pair")
202
+
203
+ if submitted:
204
+ if question.strip() and sql.strip():
205
+ response, error = create_sql_pair(
206
+ api_key,
207
+ project_id,
208
+ question,
209
+ sql
210
+ )
211
+
212
+ if response:
213
+ st.success("SQL pair created successfully!")
214
+ st.rerun()
215
+ else:
216
+ st.error(f"Error creating SQL pair: {error}")
217
+ else:
218
+ st.error("Please enter both question and SQL query")
219
+
220
+ # Display existing SQL pairs
221
+ st.subheader("Existing SQL Pairs")
222
+
223
+ with st.spinner("Loading SQL pairs..."):
224
+ sql_pairs_response, error = get_sql_pairs(api_key, project_id)
225
+
226
+ if sql_pairs_response:
227
+ for i, sql_pair in enumerate(sql_pairs_response):
228
+ with st.expander(f"πŸ’‘ {sql_pair.get('question', 'No question')[:50]}..."):
229
+ col1, col2 = st.columns([3, 1])
230
+
231
+ with col1:
232
+ st.write("**Question:**")
233
+ st.write(sql_pair.get("question", "No question"))
234
+ st.write("**SQL Query:**")
235
+ st.code(sql_pair.get("sql", "No SQL query"), language="sql")
236
+ st.write(f"**Created:** {sql_pair.get('createdAt', 'N/A')}")
237
+
238
+ with col2:
239
+ # Edit button
240
+ if st.button("πŸ”„ Edit", key=f"edit_sql_{i}"):
241
+ st.session_state[f"edit_sql_pair_{i}"] = True
242
+
243
+ # Delete button
244
+ if st.button("πŸ—‘οΈ Delete", key=f"delete_sql_{i}"):
245
+ error = delete_sql_pair(
246
+ api_key,
247
+ project_id,
248
+ sql_pair["id"]
249
+ )
250
+ if not error:
251
+ st.success("SQL pair deleted successfully!")
252
+ st.rerun()
253
+ else:
254
+ st.error(f"Error deleting SQL pair: {error}")
255
+
256
+ # Edit form
257
+ if st.session_state.get(f"edit_sql_pair_{i}", False):
258
+ st.divider()
259
+ with st.form(f"edit_sql_pair_form_{i}"):
260
+ new_question = st.text_area(
261
+ "Edit Question",
262
+ value=sql_pair.get("question", ""),
263
+ height=80
264
+ )
265
+ new_sql = st.text_area(
266
+ "Edit SQL Query",
267
+ value=sql_pair.get("sql", ""),
268
+ height=120
269
+ )
270
+
271
+ col1, col2 = st.columns(2)
272
+ with col1:
273
+ update_submitted = st.form_submit_button("Update SQL Pair")
274
+ with col2:
275
+ cancel_submitted = st.form_submit_button("Cancel")
276
+
277
+ if update_submitted:
278
+ if new_question.strip() and new_sql.strip():
279
+ response, error = update_sql_pair(
280
+ api_key,
281
+ project_id,
282
+ sql_pair["id"],
283
+ new_question,
284
+ new_sql
285
+ )
286
+
287
+ if response:
288
+ st.success("SQL pair updated successfully!")
289
+ del st.session_state[f"edit_sql_pair_{i}"]
290
+ st.rerun()
291
+ else:
292
+ st.error(f"Error updating SQL pair: {error}")
293
+ else:
294
+ st.error("Please enter both question and SQL query")
295
+
296
+ if cancel_submitted:
297
+ del st.session_state[f"edit_sql_pair_{i}"]
298
+ st.rerun()
299
+ else:
300
+ if not error:
301
+ st.info("No SQL pairs found. Create your first SQL pair above!")
302
+ else:
303
+ st.error(f"Error loading SQL pairs: {error}")
304
+
305
+
306
+ if __name__ == "__main__":
307
+ main()