rename

src/pages/4_Knowledge_Mangement.py

@@ -0,0 +1,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()