diff --git "a/application/pages/1_\360\237\214\215_Natural_Language_Querying.py" "b/application/pages/1_\360\237\214\215_Natural_Language_Querying.py" index 688fc7b..f25652a 100644 --- "a/application/pages/1_\360\237\214\215_Natural_Language_Querying.py" +++ "b/application/pages/1_\360\237\214\215_Natural_Language_Querying.py" @@ -361,8 +361,8 @@ def main(): st.session_state.messages[selected_profile].append( {"role": "assistant", "content": current_nlq_chain.get_generated_sql_explain()}) - st.markdown('The generated SQL statement is:') - st.code(current_nlq_chain.get_generated_sql(), language="sql") + with st.expander("The generated SQL"): + st.code(current_nlq_chain.get_generated_sql(), language="sql") st.markdown('Generation process explanations:') st.markdown(current_nlq_chain.get_generated_sql_explain()) diff --git a/application/utils/llm.py b/application/utils/llm.py index 30cebe2..87ed638 100644 --- a/application/utils/llm.py +++ b/application/utils/llm.py @@ -61,7 +61,7 @@ def invoke_model_claude3(model_id, system_prompt, messages, max_tokens, with_res return response_body -@logger.catch + def get_sagemaker_client(): global sagemaker_client if not sagemaker_client: @@ -191,7 +191,7 @@ def claude3_to_sql(ddl, hints, search_box, sql_examples=None, ner_example=None, return final_response -@logger.catch + def sagemaker_to_explain(endpoint_name: str, sql: str, with_response_stream=False): body = json.dumps({"query": generate_sagemaker_explain_prompt(sql), "stream": with_response_stream,}) @@ -204,7 +204,7 @@ def sagemaker_to_explain(endpoint_name: str, sql: str, with_response_stream=Fals return response -@logger.catch + def sagemaker_to_sql(ddl, hints, search_box, endpoint_name, sql_examples=None, ner_example=None, dialect='mysql', model_provider=None, with_response_stream=False): body = json.dumps({"prompt": generate_sagemaker_sql_prompt(ddl, hints, search_box, sql_examples, ner_example, diff --git a/application/utils/prompts/generate_prompt.py b/application/utils/prompts/generate_prompt.py index 27c5a37..2f51abb 100644 --- a/application/utils/prompts/generate_prompt.py +++ b/application/utils/prompts/generate_prompt.py @@ -142,7 +142,7 @@ def generate_llm_prompt(ddl, hints, search_box, sql_examples=None, ner_example=N example_sql_prompt += "A: ```sql\n" + item['_source']['sql'] + "```\n" if ner_example: - for item in sql_examples: + for item in ner_example: example_ner_prompt += "ner: " + item['_source']['entity'] + "\n" example_ner_prompt += "ner info:" + item['_source']['comment'] + "\n"