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()