|
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)') |
|
|
|
|
|
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") |
|
|
|
|
|
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") |
|
|
|
|
|
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: |
|
|
|
if st.button("π Edit", key=f"edit_inst_{i}"): |
|
st.session_state[f"edit_instruction_{i}"] = True |
|
|
|
|
|
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}") |
|
|
|
|
|
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") |
|
|
|
|
|
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") |
|
|
|
|
|
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: |
|
|
|
if st.button("π Edit", key=f"edit_sql_{i}"): |
|
st.session_state[f"edit_sql_pair_{i}"] = True |
|
|
|
|
|
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}") |
|
|
|
|
|
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() |