Merge remote-tracking branch 'origin' into wangzt-dev

Author: wzt1001

application/Dockerfile

@@ -4,16 +4,10 @@ WORKDIR /app
 
 RUN adduser --disabled-password --gecos '' appuser
 
-RUN apt-get update && apt-get install -y \
-    build-essential \
-    curl \
-    software-properties-common \
-    && rm -rf /var/lib/apt/lists/*
-
 WORKDIR /app
 
 COPY requirements.txt /app/
-RUN pip3 install -r requirements.txt
+RUN pip3 install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
 
 COPY . /app/
 

application/Dockerfile-api

@@ -3,7 +3,7 @@ FROM public.ecr.aws/docker/library/python:3.10-slim
 WORKDIR /app
 
 COPY . /app/
-RUN pip3 install -r requirements-api.txt
+RUN pip3 install -r requirements-api.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
 
 EXPOSE 8000
 

application/api/schemas.py

@@ -9,7 +9,7 @@ class Question(BaseModel):
     visualize_results_flag: bool = True
     intent_ner_recognition_flag: bool = True
     agent_cot_flag: bool = True
-    profile_name: str = "shopping-demo"
+    profile_name: str
     explain_gen_process_flag: bool = True
     gen_suggested_question_flag: bool = False
     answer_with_insights: bool = False

application/config_files/stauth_config.yaml

@@ -5,7 +5,7 @@ credentials:
       failed_login_attempts: 0 # Will be managed automatically
       logged_in: False # Will be managed automatically
       name: AWS
-      password: # Set the password following instructions in README
+      password: $2b$12$NDQv5NLaWiVlNuzQYHwAo.tv.f.TuX1nbdoUZi44/Y3xv4I4QAfjy # Set the password following instructions in README
 cookie:
   expiry_days: 30
   key: some_signature_key # Must be string

application/docker-compose.yml

@@ -1,7 +1,7 @@
 version: '3'
 services:
   opensearch-node1:
-    image: opensearchproject/opensearch:2.11.1
+    image: public.ecr.aws/opensearchproject/opensearch:2.11.1
     container_name: opensearch-node1
     environment:
       - cluster.name=opensearch-cluster
@@ -25,7 +25,7 @@ services:
     networks:
       - opensearch-net
   opensearch-node2:
-    image: opensearchproject/opensearch:2.11.1
+    image: public.ecr.aws/opensearchproject/opensearch:2.11.1
     container_name: opensearch-node2
     environment:
       - cluster.name=opensearch-cluster
@@ -46,7 +46,7 @@ services:
     networks:
       - opensearch-net
   opensearch-dashboards:
-    image: opensearchproject/opensearch-dashboards:2.11.1
+    image: public.ecr.aws/opensearchproject/opensearch-dashboards:2.11.1
     read_only: true
     container_name: opensearch-dashboards
     ports:
@@ -61,7 +61,7 @@ services:
     # 指定容器的名称
     container_name: nlq-mysql
     # 指定镜像和版本
-    image: mysql:8.0
+    image: public.ecr.aws/docker/library/mysql:8.0
     ports:
       - "3306:3306"
     restart: always

application/pages/1_🌍_Generative_BI_Playground.py

@@ -127,7 +127,8 @@ def recurrent_display(messages, i):
     return i
 
 
-def normal_text_search_streamlit(search_box, model_type, database_profile, entity_slot, opensearch_info, selected_profile,
+def normal_text_search_streamlit(search_box, model_type, database_profile, entity_slot, opensearch_info,
+                                 selected_profile,
                                  use_rag,
                                  model_provider=None):
     entity_slot_retrieve = []
@@ -241,6 +242,12 @@ def main():
     if 'current_profile' not in st.session_state:
         st.session_state['current_profile'] = ''
 
+    if 'current_model_id' not in st.session_state:
+        st.session_state['current_model_id'] = ''
+
+    if 'config_data_with_analyse' not in st.session_state:
+        st.session_state['config_data_with_analyse'] = False
+
     if 'nlq_chain' not in st.session_state:
         st.session_state['nlq_chain'] = None
 
@@ -250,7 +257,8 @@ def main():
     if "current_sql_result" not in st.session_state:
         st.session_state.current_sql_result = {}
 
-    model_ids = ['anthropic.claude-3-sonnet-20240229-v1:0', 'anthropic.claude-3-5-sonnet-20240620-v1:0', 'anthropic.claude-3-opus-20240229-v1:0',
+    model_ids = ['anthropic.claude-3-sonnet-20240229-v1:0', 'anthropic.claude-3-5-sonnet-20240620-v1:0',
+                 'anthropic.claude-3-opus-20240229-v1:0',
                  'anthropic.claude-3-haiku-20240307-v1:0', 'mistral.mixtral-8x7b-instruct-v0:1',
                  'meta.llama3-70b-instruct-v1:0']
 
@@ -263,8 +271,15 @@ def main():
     with st.sidebar:
         st.title('Setting')
         # The default option can be the first one in the profiles dictionary, if exists
-
-        selected_profile = st.selectbox("Data Profile", list(st.session_state.get('profiles', {}).keys()))
+        session_state_list = list(st.session_state.get('profiles', {}).keys())
+        if st.session_state.current_profile != "":
+            if st.session_state.current_profile in session_state_list:
+                profile_index = session_state_list.index(st.session_state.current_profile)
+                selected_profile = st.selectbox("Data Profile", session_state_list, index=profile_index)
+            else:
+                selected_profile = st.selectbox("Data Profile", session_state_list)
+        else:
+            selected_profile = st.selectbox("Data Profile", session_state_list)
         if selected_profile != st.session_state.current_profile:
             # clear session state
             st.session_state.selected_sample = ''
@@ -273,7 +288,11 @@ def main():
                 st.session_state.messages[selected_profile] = []
             st.session_state.nlq_chain = NLQChain(selected_profile)
 
-        model_type = st.selectbox("Choose your model", model_ids)
+        if st.session_state.current_model_id != "" and st.session_state.current_model_id in model_ids:
+            model_index = model_ids.index(st.session_state.current_model_id)
+            model_type = st.selectbox("Choose your model", model_ids, index=model_index)
+        else:
+            model_type = st.selectbox("Choose your model", model_ids)
 
         use_rag_flag = st.checkbox("Using RAG from Q/A Embedding", True)
         visualize_results_flag = st.checkbox("Visualize Results", True)
@@ -363,10 +382,13 @@ def main():
                         database_profile['db_url'] = db_url
                         database_profile['db_type'] = ConnectionManagement.get_db_type_by_name(conn_name)
                     prompt_map = database_profile['prompt_map']
+                    prompt_map_flag = False
                     for key in prompt_map_dict:
                         if key not in prompt_map:
                             prompt_map[key] = prompt_map_dict[key]
-                    ProfileManagement.update_table_prompt_map(selected_profile, prompt_map)
+                            prompt_map_flag = True
+                    if prompt_map_flag:
+                        ProfileManagement.update_table_prompt_map(selected_profile, prompt_map)
 
                 # Multiple rounds of dialogue, query rewriting
                 user_query_history = get_user_history(selected_profile)

application/pages/3_🪙_Data_Profile_Management.py

@@ -22,6 +22,9 @@ def main():
     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'] = ''
+
     with st.sidebar:
         st.title("Data Profile Management")
         st.selectbox("My Data Profiles", ProfileManagement.get_all_profiles(),

application/pages/4_🪙_Schema_Description_Management.py

@@ -12,9 +12,17 @@ def main():
     st.set_page_config(page_title="Schema Management", )
     make_sidebar()
 
+    if 'current_profile' not in st.session_state:
+        st.session_state['current_profile'] = ''
+
     with st.sidebar:
         st.title("Schema Management")
-        current_profile = st.selectbox("My Data Profiles", ProfileManagement.get_all_profiles(),
+        all_profiles_list = ProfileManagement.get_all_profiles()
+        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(),
                                        index=None,
                                        placeholder="Please select data profile...", key='current_profile_name')
 

application/pages/5_🪙_Prompt_Management.py

@@ -14,9 +14,17 @@ def main():
     st.set_page_config(page_title="Prompt Management")
     make_sidebar()
 
+    if 'current_profile' not in st.session_state:
+        st.session_state['current_profile'] = ''
+
     with st.sidebar:
         st.title("Prompt Management")
-        current_profile = st.selectbox("My Data Profiles", ProfileManagement.get_all_profiles(),
+        all_profiles_list = ProfileManagement.get_all_profiles()
+        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(),
                                        index=None,
                                        placeholder="Please select data profile...", key='current_profile_name')
 
@@ -29,6 +37,9 @@ def main():
                                                       format_func=lambda x: prompt_map[x].get('title'),
                                                       placeholder="Please select a prompt type")
 
+            profile_detail = ProfileManagement.get_profile_by_name(current_profile)
+            prompt_map = profile_detail.prompt_map
+
             if prompt_type_selected_table is not None:
                 single_type_prompt_map = prompt_map.get(prompt_type_selected_table)
                 system_prompt = single_type_prompt_map.get('system_prompt')

application/pages/6_📚_Index_Management.py

@@ -1,5 +1,6 @@
 import time
 
+import pandas as pd
 import streamlit as st
 from dotenv import load_dotenv
 import logging
@@ -14,6 +15,26 @@ def delete_sample(profile_name, id):
     VectorStore.delete_sample(profile_name, id)
     st.success(f'Sample {id} deleted.')
 
+def read_file(uploaded_file):
+    """
+    read upload csv file
+    :param uploaded_file:
+    :return:
+    """
+    file_type = uploaded_file.name.split('.')[-1].lower()
+    if file_type == 'csv':
+        uploaded_data = pd.read_csv(uploaded_file)
+    elif file_type in ['xls', 'xlsx']:
+        uploaded_data = pd.read_excel(uploaded_file)
+    else:
+        st.error(f"Unsupported file type: {file_type}")
+        return None
+    columns = list(uploaded_data.columns)
+    if "question" in columns and "sql" in columns:
+        return uploaded_data
+    else:
+        st.error(f"The columns need contains question and sql")
+        return None
 
 def main():
     load_dotenv()
@@ -24,13 +45,22 @@ def main():
     if 'profile_page_mode' not in st.session_state:
         st.session_state['index_mgt_mode'] = 'default'
 
+
+    if 'current_profile' not in st.session_state:
+        st.session_state['current_profile'] = ''
+
     with st.sidebar:
         st.title("Index Management")
-        current_profile = st.selectbox("My Data Profiles", ProfileManagement.get_all_profiles(),
+        all_profiles_list = ProfileManagement.get_all_profiles()
+        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(),
                                        index=None,
                                        placeholder="Please select data profile...", key='current_profile_name')
 
-    tab_view, tab_add, tab_search = st.tabs(['View Samples', 'Add New Sample', 'Sample Search'])
+    tab_view, tab_add, tab_search, batch_insert = st.tabs(['View Samples', 'Add New Sample', 'Sample Search', 'Batch Insert Samples'])
 
     if current_profile is not None:
         with tab_view:
@@ -69,6 +99,27 @@ def main():
                             st.code(sample_res)
                             st.button('Delete ' + sample['_id'], key=sample['_id'], on_click=delete_sample,
                                       args=[current_profile, sample['_id']])
+        with batch_insert:
+            if current_profile is not None:
+                st.write("This page support CSV or Excel files batch insert sql samples.")
+                st.write("**The Column Name need contain 'question' and 'sql'**")
+                uploaded_files = st.file_uploader("Choose CSV or Excel files", accept_multiple_files=True,
+                                              type=['csv', 'xls', 'xlsx'])
+                if uploaded_files:
+                    progress_bar = st.progress(0)
+                    status_text = st.empty()
+                    for i, uploaded_file in enumerate(uploaded_files):
+                        status_text.text(f"Processing file {i + 1} of {len(uploaded_files)}: {uploaded_file.name}")
+                        each_upload_data = read_file(uploaded_file)
+                        if each_upload_data is not None:
+                            for index, item in each_upload_data.iterrows():
+                                question = str(item["question"])
+                                sql = str(item["sql"])
+                                VectorStore.add_sample(current_profile, question, sql)
+                        progress_bar.progress((i + 1) / len(uploaded_files))
+
+                        st.success("{uploaded_file} uploaded successfully!".format(uploaded_file=uploaded_file.name))
+                    progress_bar.empty()
     else:
         st.info('Please select data profile in the left sidebar.')