File size: 13,131 Bytes
4c476be
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
import streamlit as st

from apis import (
    get_instructions,
    create_instruction,
    update_instruction,
    delete_instruction,
    get_sql_pairs,
    create_sql_pair,
    update_sql_pair,
    delete_sql_pair,
)


def main():
    st.title("πŸ“š Wren AI Cloud API Demo - Knowledge Management")

    if "api_key" not in st.session_state or "project_id" not in st.session_state:
        st.error("Please enter your API Key and Project ID in the sidebar of Home page to get started.")
        return
    if not st.session_state.api_key or not st.session_state.project_id:
        st.error("Please enter your API Key and Project ID in the sidebar of Home page to get started.")
        return
    
    api_key = st.session_state.api_key
    project_id = st.session_state.project_id

    st.markdown('Using APIs: [Knowledge Instructions](https://wrenai.readme.io/reference/knowledge-1), [Knowledge SQL Pairs](https://wrenai.readme.io/reference/knowledge-1)')

    # Tabs for Instructions and SQL Pairs
    tab1, tab2 = st.tabs(["πŸ’‘ Instructions", "πŸ’‘ SQL Pairs"])

    with tab1:
        manage_instructions(api_key, project_id)

    with tab2:
        manage_sql_pairs(api_key, project_id)


def manage_instructions(api_key: str, project_id: str):
    st.header("Instructions Management")
    
    # Add new instruction section
    with st.expander("πŸ’‘ Add New Instruction", expanded=False):
        with st.form("add_instruction_form"):
            instruction_text = st.text_area(
                "Instruction Text",
                placeholder="Enter instruction text...",
                height=100,
                help="The instruction text that will guide the AI"
            )
            is_global = st.checkbox(
                "Global Instruction",
                value=True,
                help="Whether this instruction applies globally"
            )
            questions_input = st.text_area(
                "Related Questions (one per line)",
                placeholder="Question 1\nQuestion 2\nQuestion 3",
                height=80,
                help="Optional: Questions that this instruction relates to"
            )
            
            submitted = st.form_submit_button("Create Instruction")
            
            if submitted:
                if instruction_text.strip():
                    questions = [q.strip() for q in questions_input.split('\n') if q.strip()]
                    response, error = create_instruction(
                        api_key,
                        project_id,
                        instruction_text,
                        is_global,
                        questions if questions else None
                    )
                    
                    if response:
                        st.success("Instruction created successfully!")
                        st.rerun()
                    else:
                        st.error(f"Error creating instruction: {error}")
                else:
                    st.error("Please enter instruction text")

    # Display existing instructions
    st.subheader("Existing Instructions")
    
    with st.spinner("Loading instructions..."):
        instructions_response, error = get_instructions(api_key, project_id)
    
    if instructions_response:        
        for i, instruction in enumerate(instructions_response):
            with st.expander(f"πŸ’‘ {instruction.get('instruction', 'No instruction text')[:50]}..."):
                col1, col2 = st.columns([3, 1])
                
                with col1:
                    st.write("**Instruction Text:**")
                    st.write(instruction.get("instruction", "No instruction text"))
                    st.write(f"**Global:** {instruction.get('isGlobal', False)}")
                    
                    questions = instruction.get("questions", [])
                    if questions:
                        st.write("**Related Questions:**")
                        for q in questions:
                            st.write(f"- {q}")
                                    
                with col2:
                    # Edit button
                    if st.button("πŸ”„ Edit", key=f"edit_inst_{i}"):
                        st.session_state[f"edit_instruction_{i}"] = True
                    
                    # Delete button
                    if st.button("πŸ—‘οΈ Delete", key=f"delete_inst_{i}"):
                        error = delete_instruction(
                            api_key,
                            project_id,
                            instruction["id"]
                        )
                        if not error:
                            st.success("Instruction deleted successfully!")
                            st.rerun()
                        else:
                            st.error(f"Error deleting instruction: {error}")
                
                # Edit form
                if st.session_state.get(f"edit_instruction_{i}", False):
                    st.divider()
                    with st.form(f"edit_instruction_form_{i}"):
                        new_instruction = st.text_area(
                            "Edit Instruction Text",
                            value=instruction.get("instruction", ""),
                            height=100
                        )
                        new_is_global = st.checkbox(
                            "Global Instruction",
                            value=instruction.get("isGlobal", True)
                        )
                        current_questions = instruction.get("questions", [])
                        new_questions_input = st.text_area(
                            "Related Questions (one per line)",
                            value='\n'.join(current_questions),
                            height=80
                        )
                        
                        col1, col2 = st.columns(2)
                        with col1:
                            update_submitted = st.form_submit_button("Update Instruction")
                        with col2:
                            cancel_submitted = st.form_submit_button("Cancel")
                        
                        if update_submitted:
                            if new_instruction.strip():
                                new_questions = [q.strip() for q in new_questions_input.split('\n') if q.strip()]
                                response, error = update_instruction(
                                    api_key,
                                    project_id,
                                    instruction["id"],
                                    new_instruction,
                                    new_is_global,
                                    new_questions if new_questions else None
                                )
                                
                                if response:
                                    st.success("Instruction updated successfully!")
                                    del st.session_state[f"edit_instruction_{i}"]
                                    st.rerun()
                                else:
                                    st.error(f"Error updating instruction: {error}")
                            else:
                                st.error("Please enter instruction text")
                        
                        if cancel_submitted:
                            del st.session_state[f"edit_instruction_{i}"]
                            st.rerun()
    else:
        if not error:
            st.info("No instructions found. Create your first instruction above!")
        else:
            st.error(f"Error loading instructions: {error}")


def manage_sql_pairs(api_key: str, project_id: str):
    st.header("SQL Pairs Management")
    
    # Add new SQL pair section
    with st.expander("πŸ’‘ Add New SQL Pair", expanded=False):
        with st.form("add_sql_pair_form"):
            question = st.text_area(
                "Question",
                placeholder="Enter the natural language question...",
                height=80,
                help="The question in natural language"
            )
            sql = st.text_area(
                "SQL Query",
                placeholder="SELECT * FROM table_name WHERE condition;",
                height=120,
                help="The corresponding SQL query"
            )
            
            submitted = st.form_submit_button("Create SQL Pair")
            
            if submitted:
                if question.strip() and sql.strip():
                    response, error = create_sql_pair(
                        api_key,
                        project_id,
                        question,
                        sql
                    )
                    
                    if response:
                        st.success("SQL pair created successfully!")
                        st.rerun()
                    else:
                        st.error(f"Error creating SQL pair: {error}")
                else:
                    st.error("Please enter both question and SQL query")

    # Display existing SQL pairs
    st.subheader("Existing SQL Pairs")
    
    with st.spinner("Loading SQL pairs..."):
        sql_pairs_response, error = get_sql_pairs(api_key, project_id)
    
    if sql_pairs_response:        
        for i, sql_pair in enumerate(sql_pairs_response):
            with st.expander(f"πŸ’‘ {sql_pair.get('question', 'No question')[:50]}..."):
                col1, col2 = st.columns([3, 1])
                
                with col1:
                    st.write("**Question:**")
                    st.write(sql_pair.get("question", "No question"))
                    st.write("**SQL Query:**")
                    st.code(sql_pair.get("sql", "No SQL query"), language="sql")
                    st.write(f"**Created:** {sql_pair.get('createdAt', 'N/A')}")
                
                with col2:
                    # Edit button
                    if st.button("πŸ”„ Edit", key=f"edit_sql_{i}"):
                        st.session_state[f"edit_sql_pair_{i}"] = True
                    
                    # Delete button
                    if st.button("πŸ—‘οΈ Delete", key=f"delete_sql_{i}"):
                        error = delete_sql_pair(
                            api_key,
                            project_id,
                            sql_pair["id"]
                        )
                        if not error:
                            st.success("SQL pair deleted successfully!")
                            st.rerun()
                        else:
                            st.error(f"Error deleting SQL pair: {error}")
            
                # Edit form
                if st.session_state.get(f"edit_sql_pair_{i}", False):
                    st.divider()
                    with st.form(f"edit_sql_pair_form_{i}"):
                        new_question = st.text_area(
                            "Edit Question",
                            value=sql_pair.get("question", ""),
                            height=80
                        )
                        new_sql = st.text_area(
                            "Edit SQL Query",
                            value=sql_pair.get("sql", ""),
                            height=120
                        )
                        
                        col1, col2 = st.columns(2)
                        with col1:
                            update_submitted = st.form_submit_button("Update SQL Pair")
                        with col2:
                            cancel_submitted = st.form_submit_button("Cancel")
                        
                        if update_submitted:
                            if new_question.strip() and new_sql.strip():
                                response, error = update_sql_pair(
                                    api_key,
                                    project_id,
                                    sql_pair["id"],
                                    new_question,
                                    new_sql
                                )
                                
                                if response:
                                    st.success("SQL pair updated successfully!")
                                    del st.session_state[f"edit_sql_pair_{i}"]
                                    st.rerun()
                                else:
                                    st.error(f"Error updating SQL pair: {error}")
                            else:
                                st.error("Please enter both question and SQL query")
                        
                        if cancel_submitted:
                            del st.session_state[f"edit_sql_pair_{i}"]
                            st.rerun()
    else:
        if not error:
            st.info("No SQL pairs found. Create your first SQL pair above!")
        else:
            st.error(f"Error loading SQL pairs: {error}")


if __name__ == "__main__":
    main()