Spaces:
Running
Running
| 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() |