fix DataSourceFactory

application/nlq/data_access/database.py

@@ -21,14 +21,25 @@ class RelationDatabase():
 
     @classmethod
     def get_db_url(cls, db_type, user, password, host, port, db_name):
-        db_url = db.engine.URL.create(
-            drivername=cls.db_mapping[db_type],
-            username=user,
-            password=password,
-            host=host,
-            port=port,
-            database=db_name
-        )
+        if db_type == "hive":
+            db_url = db.engine.URL.create(
+                drivername=cls.db_mapping[db_type],
+                username=user,
+                password=password,
+                host=host,
+                port=port,
+                database=db_name,
+                query={'auth': 'LDAP'}
+            )
+        else:
+            db_url = db.engine.URL.create(
+                drivername=cls.db_mapping[db_type],
+                username=user,
+                password=password,
+                host=host,
+                port=port,
+                database=db_name
+            )
         return db_url
 
     @classmethod
@@ -53,7 +64,7 @@ def get_all_schema_names_by_connection(cls, connection: ConnectConfigEntity):
         if db_type == 'postgresql':
             schemas = [schema for schema in inspector.get_schema_names() if
                        schema not in ('pg_catalog', 'information_schema', 'public')]
-        elif db_type in ('redshift', 'mysql', 'starrocks', 'clickhouse'):
+        elif db_type in ('redshift', 'mysql', 'starrocks', 'clickhouse', 'hive'):
             schemas = inspector.get_schema_names()
         else:
             raise ValueError("Unsupported database type")

application/pages/1_🌍_Generative_BI_Playground.py

@@ -3,7 +3,7 @@
 import pandas as pd
 import plotly.express as px
 from dotenv import load_dotenv
-
+import logging
 from api.service import user_feedback_downvote
 from nlq.business.connection import ConnectionManagement
 from nlq.business.profile import ProfileManagement
@@ -13,9 +13,9 @@
 from nlq.core.state_machine import QueryStateMachine
 from utils.navigation import make_sidebar
 from utils.env_var import opensearch_info
-from utils.logging import getLogger
 
-logger = getLogger()
+logger = logging.getLogger(__name__)
+
 
 def sample_question_clicked(sample):
     """Update the selected_sample variable with the text of the clicked button"""
@@ -148,17 +148,28 @@ def main():
     # Title and Description
     st.subheader('Generative BI Playground')
 
-    st.write('Current Username: ' + st.session_state['auth_username'])
-
+    demo_profile_suffix = '(demo)'
     # Initialize or set up state variables
+
+    if "update_profile" not in st.session_state:
+        st.session_state.update_profile = False
+
+    if "profiles_list" not in st.session_state:
+        st.session_state["profiles_list"] = []
+
     if 'profiles' not in st.session_state:
         # get all user defined profiles with info (db_url, conn_name, tables_info, hints, search_samples)
         all_profiles = ProfileManagement.get_all_profiles_with_info()
         # all_profiles.update(demo_profile)
         st.session_state['profiles'] = all_profiles
+        st.session_state["profiles_list"] = list(all_profiles.keys())
     else:
-        all_profiles = ProfileManagement.get_all_profiles_with_info()
-        st.session_state['profiles'] = all_profiles
+        if st.session_state.update_profile:
+            logger.info("session_state update_profile get_all_profiles_with_info")
+            all_profiles = ProfileManagement.get_all_profiles_with_info()
+            st.session_state["profiles_list"] = list(all_profiles.keys())
+            st.session_state['profiles'] = all_profiles
+            st.session_state.update_profile = False
 
     if "vision_change" not in st.session_state:
         st.session_state["vision_change"] = False
@@ -321,7 +332,7 @@ def main():
                     search_box=search_box,
                     query_rewrite="",
                     session_id="",
-                    user_id=st.session_state['auth_username'],
+                    user_id="",
                     selected_profile=selected_profile,
                     database_profile=database_profile,
                     model_type=model_type,
@@ -392,21 +403,24 @@ def main():
                             state_machine.handle_sql_generation()
                             sql = state_machine.get_answer().sql_search_result.sql
                             st.code(sql, language="sql")
-                            st.session_state.messages[selected_profile].append(
+                            if not visualize_results_flag:
+                                st.session_state.messages[selected_profile].append(
                                 {"role": "assistant", "content": sql, "type": "sql"})
                             feedback = st.columns(2)
                             feedback[0].button('👍 Upvote (save as embedding for retrieval)', type='secondary',
+                                               key="upvote",
                                                use_container_width=True,
                                                on_click=upvote_clicked,
                                                args=[search_box,
                                                      sql])
                             feedback[1].button('👎 Downvote', type='secondary', use_container_width=True,
+                                               key="downvote",
                                                on_click=downvote_clicked,
                                                args=[search_box, sql])
                             status_text.update(
                                 label=f"Generating SQL Done",
                                 state="complete", expanded=True)
-                        if state_machine.context.gen_suggested_question_flag:
+                        if state_machine.context.explain_gen_process_flag:
                             with st.status("Generating explanations...") as status_text:
                                 st.markdown(state_machine.get_answer().sql_search_result.sql_gen_process)
                                 status_text.update(
@@ -421,7 +435,39 @@ def main():
                         status_text.update(label=f"Intent Recognition Completed: This is a **{intent}** question",
                                            state="complete", expanded=False)
                     elif state_machine.get_state() == QueryState.EXECUTE_QUERY:
-                        state_machine.handle_execute_query()
+                        with st.status("Execute SQL...") as status_text:
+                            state_machine.handle_execute_query()
+                        status_text.update(label=f"Execute SQL Done",
+                                           state="complete", expanded=False)
+                        sql = state_machine.get_answer().sql_search_result.sql
+                        if state_machine.use_auto_correction_flag:
+                            with st.expander("The SQL Error Info"):
+                                st.markdown(state_machine.first_sql_execute_info["error_info"])
+                            with st.status("Generating SQL Again ... ") as status_text:
+                                st.code(sql, language="sql")
+                                st.session_state.messages[selected_profile].append(
+                                    {"role": "assistant", "content": sql, "type": "sql"})
+                                feedback = st.columns(2)
+                                feedback[0].button('👍 Upvote (save as embedding for retrieval)', type='secondary',
+                                                   key="upvote_again",
+                                                   use_container_width=True,
+                                                   on_click=upvote_clicked,
+                                                   args=[search_box,
+                                                         sql])
+                                feedback[1].button('👎 Downvote', type='secondary', use_container_width=True,
+                                                   key="downcote_again",
+                                                   on_click=downvote_clicked,
+                                                   args=[search_box, sql])
+                                status_text.update(
+                                    label=f"Generating SQL Done",
+                                    state="complete", expanded=True)
+                        else:
+                            st.session_state.messages[selected_profile].append(
+                                {"role": "assistant", "content": sql, "type": "sql"})
+                        if state_machine.get_answer().sql_search_result.sql_data is not None:
+                            st.session_state.messages[selected_profile].append(
+                                {"role": "assistant", "content": state_machine.get_answer().sql_search_result.sql_data, "type": "pandas"})
+
                     elif state_machine.get_state() == QueryState.ANALYZE_DATA:
                         with st.spinner('Generating data summarize...'):
                             state_machine.handle_analyze_data()
@@ -482,6 +528,12 @@ def main():
                             st.session_state.current_sql_result = \
                                 state_machine.intent_search_result["sql_execute_result"]["data"]
                             do_visualize_results()
+                elif state_machine.get_state() == QueryState.ERROR:
+                    with st.status("The Error Info Please Check") as status_text:
+                        st.write(state_machine.error_log)
+                    status_text.update(label=f"The Error Info Please Check",
+                                       state="error", expanded=False)
+
                 if processing_context.gen_suggested_question_flag:
                     if state_machine.search_intent_flag or state_machine.agent_intent_flag:
                         st.markdown('You might want to further ask:')
@@ -502,7 +554,8 @@ def main():
                                                 on_click=sample_question_clicked,
                                                 args=[gen_sq_list[2]])
         else:
-            do_visualize_results()
+            if visualize_results_flag:
+                do_visualize_results()
 
 
 if __name__ == '__main__':

application/pages/3_🪙_Data_Profile_Management.py

@@ -51,6 +51,7 @@ def get_table_name_by_config(_conn_config, schema_names, default_values):
 
 def show_delete_profile(profile_name):
     if st.button('Delete Profile'):
+        st.session_state.update_profile = True
         ProfileManagement.delete_profile(profile_name)
         st.success(f"{profile_name} deleted successfully!")
         st.session_state.profile_page_mode = 'default'
@@ -64,12 +65,30 @@ def main():
     st.set_page_config(page_title="Data Profile Management", )
     make_sidebar()
 
+    if "update_profile" not in st.session_state:
+        st.session_state.update_profile = False
+
     if 'profile_page_mode' not in st.session_state:
         st.session_state['profile_page_mode'] = 'default'
 
     if 'current_profile' not in st.session_state:
         st.session_state['current_profile'] = ''
 
+    if "profiles_list" not in st.session_state:
+        st.session_state["profiles_list"] = []
+
+    if 'profiles' not in st.session_state:
+        all_profiles = ProfileManagement.get_all_profiles_with_info()
+        st.session_state['profiles'] = all_profiles
+        st.session_state["profiles_list"] = list(all_profiles.keys())
+
+    if st.session_state.update_profile:
+        logger.info("session_state update_profile get_all_profiles_with_info")
+        all_profiles = ProfileManagement.get_all_profiles_with_info()
+        st.session_state["profiles_list"] = list(all_profiles.keys())
+        st.session_state['profiles'] = all_profiles
+        st.session_state.update_profile = False
+
     with st.sidebar:
         st.title("Data Profile Management")
         st.selectbox("My Data Profiles", get_all_profiles(),
@@ -94,6 +113,7 @@ def main():
             comments = st.text_input("Comments")
 
             if st.button('Create Profile', type='primary'):
+                st.session_state.update_profile = True
                 if not selected_tables:
                     st.error('Please select at least one table.')
                     return
@@ -156,6 +176,7 @@ def main():
         column_value: $login_user.username""", disabled=not st_enable_rls, height=240)
 
         if st.button('Update Profile', type='primary'):
+            st.session_state.update_profile = True
             if not selected_tables:
                 st.error('Please select at least one table.')
                 return
@@ -173,6 +194,7 @@ def main():
                 st.cache_data.clear()
 
         if st.button('Fetch table definition'):
+            st.session_state.update_profile = True
             if not selected_tables:
                 st.error('Please select at least one table.')
             with st.spinner('fetching...'):

application/pages/4_🪙_Schema_Description_Management.py

@@ -15,14 +15,32 @@ def main():
     if 'current_profile' not in st.session_state:
         st.session_state['current_profile'] = ''
 
+    if "update_profile" not in st.session_state:
+        st.session_state.update_profile = False
+
+    if "profiles_list" not in st.session_state:
+        st.session_state["profiles_list"] = []
+
+    if 'profiles' not in st.session_state:
+        all_profiles = ProfileManagement.get_all_profiles_with_info()
+        st.session_state['profiles'] = all_profiles
+        st.session_state["profiles_list"] = list(all_profiles.keys())
+
+    if st.session_state.update_profile:
+        logger.info("session_state update_profile get_all_profiles_with_info")
+        all_profiles = ProfileManagement.get_all_profiles_with_info()
+        st.session_state["profiles_list"] = list(all_profiles.keys())
+        st.session_state['profiles'] = all_profiles
+        st.session_state.update_profile = False
+
     with st.sidebar:
         st.title("Schema Management")
-        all_profiles_list = ProfileManagement.get_all_profiles()
+        all_profiles_list = st.session_state["profiles_list"]
         if st.session_state.current_profile != "" and st.session_state.current_profile in all_profiles_list:
             profile_index = all_profiles_list.index(st.session_state.current_profile)
             current_profile = st.selectbox("My Data Profiles", all_profiles_list, index=profile_index)
         else:
-            current_profile = st.selectbox("My Data Profiles", ProfileManagement.get_all_profiles(),
+            current_profile = st.selectbox("My Data Profiles", all_profiles_list,
                                        index=None,
                                        placeholder="Please select data profile...", key='current_profile_name')
 
@@ -56,6 +74,7 @@ def main():
 );
     ''')
                 if st.button('Save', type='primary'):
+                    st.session_state.update_profile = True
                     origin_tables_info = profile_detail.tables_info
                     origin_table_info = origin_tables_info[selected_table]
                     origin_table_info['tbl_a'] = tbl_annotation

application/pages/5_🪙_Prompt_Management.py

@@ -17,14 +17,32 @@ def main():
     if 'current_profile' not in st.session_state:
         st.session_state['current_profile'] = ''
 
+    if "update_profile" not in st.session_state:
+        st.session_state.update_profile = False
+
+    if "profiles_list" not in st.session_state:
+        st.session_state["profiles_list"] = []
+
+    if 'profiles' not in st.session_state:
+        all_profiles = ProfileManagement.get_all_profiles_with_info()
+        st.session_state['profiles'] = all_profiles
+        st.session_state["profiles_list"] = list(all_profiles.keys())
+
+    if st.session_state.update_profile:
+        logger.info("session_state update_profile get_all_profiles_with_info")
+        all_profiles = ProfileManagement.get_all_profiles_with_info()
+        st.session_state["profiles_list"] = list(all_profiles.keys())
+        st.session_state['profiles'] = all_profiles
+        st.session_state.update_profile = False
+
     with st.sidebar:
         st.title("Prompt Management")
-        all_profiles_list = ProfileManagement.get_all_profiles()
+        all_profiles_list = st.session_state["profiles_list"]
         if st.session_state.current_profile != "" and st.session_state.current_profile in all_profiles_list:
             profile_index = all_profiles_list.index(st.session_state.current_profile)
             current_profile = st.selectbox("My Data Profiles", all_profiles_list, index=profile_index)
         else:
-            current_profile = st.selectbox("My Data Profiles", ProfileManagement.get_all_profiles(),
+            current_profile = st.selectbox("My Data Profiles", all_profiles_list,
                                        index=None,
                                        placeholder="Please select data profile...", key='current_profile_name')
 
@@ -54,6 +72,7 @@ def main():
 
                     if st.button('Save', type='primary'):
                         # check prompt syntax, missing placeholder will cause backend execution failure
+                        st.session_state.update_profile = True
                         if check_prompt_syntax(system_prompt_input, user_prompt_input,
                                                prompt_type_selected_table, model_selected_table):
                             # assign new system/user prompt by selected model

application/pages/6_📚_Index_Management.py

@@ -51,14 +51,32 @@ def main():
     if 'current_profile' not in st.session_state:
         st.session_state['current_profile'] = ''
 
+    if "update_profile" not in st.session_state:
+        st.session_state.update_profile = False
+
+    if "profiles_list" not in st.session_state:
+        st.session_state["profiles_list"] = []
+
+    if 'profiles' not in st.session_state:
+        all_profiles = ProfileManagement.get_all_profiles_with_info()
+        st.session_state['profiles'] = all_profiles
+        st.session_state["profiles_list"] = list(all_profiles.keys())
+
+    if st.session_state.update_profile:
+        logger.info("session_state update_profile get_all_profiles_with_info")
+        all_profiles = ProfileManagement.get_all_profiles_with_info()
+        st.session_state["profiles_list"] = list(all_profiles.keys())
+        st.session_state['profiles'] = all_profiles
+        st.session_state.update_profile = False
+
     with st.sidebar:
         st.title("Index Management")
-        all_profiles_list = ProfileManagement.get_all_profiles()
+        all_profiles_list = st.session_state["profiles_list"]
         if st.session_state.current_profile != "" and st.session_state.current_profile in all_profiles_list:
             profile_index = all_profiles_list.index(st.session_state.current_profile)
             current_profile = st.selectbox("My Data Profiles", all_profiles_list, index=profile_index)
         else:
-            current_profile = st.selectbox("My Data Profiles", ProfileManagement.get_all_profiles(),
+            current_profile = st.selectbox("My Data Profiles", all_profiles_list,
                                            index=None,
                                            placeholder="Please select data profile...", key='current_profile_name')