sqllite

.gitattributes

@@ -32,4 +32,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.xz filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+database/SQLAgent_DEMO_DB_V1.db filter=lfs diff=lfs merge=lfs -text

database/SQLAgent_DEMO_DB_V1.db

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:53cb4c4d3e88df7721e91ba4a4f5828e773168cbe4dc934810b57e4d179d264e
+size 7061504

database/db_tables.json

@@ -0,0 +1,71 @@
+[
+    {
+        "name": "Patient",
+        "description": "The table stores the healthcare encounter information about patients. Each row has an unique patient information. The table contains the key information by distilling and flattening the FHIR encounter schema.",
+        "fields": { "identifier_value": [ "patient identifier that uniquely identifies patient and links a patient from this to other tables", "varchar"],
+            "identifier_use": [ "if the identifier is used for any specific purpose", "varchar" ],
+            "identifier_type": ["type of identifier, ususally means the source, MR' stands for medical record", "varchar" ],
+            "identifier_start_date": ["date on since when the identifier was valid", "date"],
+            "identifier_assigner": ["Identification value assignment authority", "varchar"],
+            "active": ["if he patient is active or not", "boolean"],
+            "official_name_family": ["family name of the patient", "varchar"],
+            "official_name_given": ["given name of the patient", "varchar"],
+            "usual_name_given": ["Short form of the given name", "varchar"],
+            "gender": ["patient's gender, male or female", "varchar"],
+            "birth_date": ["date of birth of the patient", "date"],
+            "Age": ["patient age", "integer"],
+            "home_address_line": ["patient's home address street", "varchar"],
+            "home_address_city": ["patient's home address city", "varchar"],
+            "home_address_district": ["patient's home county", "varchar"],
+            "home_address_state": ["patient's home state", "varchar"],
+            "home_address_postalCode": ["patient's home address zip code", "varchar"],
+            "home_address_period_start": ["start date of the patient's home address", "date"]
+        }
+    },
+    {
+        "name": "Encounter",
+        "description": "Table that stores all encounters of each patient with the healthcare providers. Every row indicate a single encounter.",
+        "fields": { "id": [ "encounter id that identifies an encounter uniquely",  "varchar"],
+            "status": [ "encounter status, can be one of 'planned', ''completed', 'discharged', 'in-progress' ", "varchar" ],
+            "class": [ "indicates location setting of the encounter, valid values are: 'IMP' as inpatient, 'EMER' as emergency, 'AMB' as ambulatory, 'HH' as home health ", "varchar" ],
+            "priority": [ "indicates priority of the encounter, valid values are: 'UR' as urgent, 'A' as As soon as, 'S' as stat, 'R' as routine ", "varchar" ],
+            "subject_id": [ "indicates id of the patient associated with the encounter, should match with identifier_value of the Patient table", "varchar" ],
+            "service_provider_id": [ "contains the id of the care delivery organization where the patient had the encounter", "varchar" ],
+            "participant_actor_id": [ "contains the id of the provider associated with the care delivery organization who rendered the encounter", "varchar" ],
+            "diagnosis_condition_id": [ "contains list of diagnosis codes relevant to the patient of the encounter", "varchar" ],
+            "location_id": [ "location where the encounter happend or is happening or will be happening", "varchar" ],
+            "discharge_disposition": [ "how the patient was discharged at the end of the encounter", "varchar" ],
+            "diagnosis_condition_text": [ "clinical description of the diagnosis codes", "varchar" ],
+            "condition_class": [ "condition of the patient classified into specific broad classe., may contain multiple coditions. All lower case.", "varchar" ]
+        }
+    },
+    {
+        "name": "EpisodeOfCare",
+        "description": "contains continuous period of engagement by a care manager and/or a care management organization with the patient. Every row indicates a unique episide of care for a patient. One patient may have multiple episodes of care ",
+        "fields": { "identifier_value": [ "unique identifier of the episode", "varchar" ],
+            "type": [ "type of episode, can be disease management, post acute care or specialist referral", "varchar" ],
+            "diagnosis_condition_id": [ "ICD-10 diagnosis code assiciated with the episode of care", "varchar" ],
+            "subject_id": [ "id of the patient associated with episode, should have a corresponding 'identifier_value' in the Patient table", "varchar" ],
+            "managing_organization_id": [ "contains the id of the organization managing the episode", "varchar" ],
+            "care_manager_id": [ "contains the id of the care manager managing the episode", "varchar" ],
+            "care_team_id": [ "contains the id of the care team managing the episode. Care manager is part of the care team", "varchar" ]
+        }
+    },
+    {
+        "name": "RiskScore",
+        "description": "Contains the health risk scores of each of the patients. Only the latest risk score is stored. Every row has risk score of an unique patient",
+        "fields": { "patient_id": [ "identifier that uniquely identifies a patient. Matches with at least one identifier_value of Patient table.",  "varchar"],
+            "risk_score": [ "decimal number between 0 and 1 indicating the risk score", "decimal number" ]
+        }
+    },
+    {
+        "name": "patient_sdoh_scores",
+        "description": "table stores the various social determinants of quality scores about a patient obtained through assessment. Each row indicate score about one patient and about one type of assessment",
+        "fields": { "Patient_Id": [ "unique identifier of the patient.  Matches with at least one identifier_value of Patient table.",  "varchar"],
+            "Assessment_Id": [ "name of the assessment", "varchar" ],
+            "Answer": [ "The actual answer provided in the assessment", "integer" ],
+            "Assessment_Type": [ "type of the assessment, can be 'Financial', 'Home', 'Food' and 'Physical'", "varchar" ],
+            "score": [ "Derived standardized score based on the answer provided", "decimal number" ]
+        }
+    }
+]

database/gravity_sdoh.sqbpro

@@ -0,0 +1,8 @@
+<?xml version="1.0" encoding="UTF-8"?><sqlb_project><db path="/Users/niloy/Documents/GitHub/SQL Agent Demo/sdoh/gravity_sdoh_observations.db" readonly="0" foreign_keys="1" case_sensitive_like="0" temp_store="0" wal_autocheckpoint="1000" synchronous="2"/><attached/><window><main_tabs open="structure browser pragmas query" current="0"/></window><tab_structure><column_width id="0" width="300"/><column_width id="1" width="0"/><column_width id="2" width="100"/><column_width id="3" width="10670"/><column_width id="4" width="0"/><expanded_item id="0" parent="1"/><expanded_item id="1" parent="1"/><expanded_item id="2" parent="1"/><expanded_item id="3" parent="1"/></tab_structure><tab_browse><current_table name="4,7:mainPatient"/><default_encoding codec=""/><browse_table_settings/></tab_browse><tab_sql><sql name="SQL 1">CREATE TABLE RiskScore (
+    patient_id TEXT,       -- Reference to the patient
+    risk_score DECIMAL(4, 2),  -- Risk score, a decimal value between 0.0 and 1.0
+    risk_score_date TEXT,  -- Date of the risk score
+
+    PRIMARY KEY (patient_id, risk_score_date)  -- Composite primary key: patient ID and risk score date
+);
+</sql><current_tab id="0"/></tab_sql></sqlb_project>

pages/solution.py

@@ -1,5 +1,5 @@
 import json
-# import sqlite3
+import sqlite3
 # import pyodbc
 import mysql.connector
 import boto3
@@ -26,8 +26,6 @@ from loguru import logger
 from st_aggrid import AgGrid, GridOptionsBuilder
 from datetime import datetime
 
-APP_TITLE = '**Social <br>Determinant<br>of Health**'
-
 # Initialize token storage
 token_file = "token_usage.json"
 if not os.path.exists(token_file):
@@ -64,6 +62,8 @@ def show_messages(message):
     success_msg.empty()
 
 # Locations of various files
+APP_TITLE = '**Social <br>Determinant<br>of Health**'
+
 sql_dir = 'generated_sql/'
 method_dir = 'generated_method/'
 insight_lib = 'insight_library/'
@@ -72,6 +72,8 @@ report_path = 'Reports/'
 connection_string = "DefaultEndpointsProtocol=https;AccountName=phsstorageacc;AccountKey=cEvoESH5CknyeZtbe8eCFuebwr7lRFi1EyO8smA35i5EuoSOfnzRXX/4337Y743B05tQsGPoQbsr+AStNRWeBg==;EndpointSuffix=core.windows.net"
 container_name = "insights-lab"
 persona_list = ["Population Analyst", "SDoH Specialist"]
+DB_List=["Patient SDOH"]
+
 
 def getBlobContent(dir_path):
     try:
@@ -447,22 +449,28 @@ def get_existing_insight(base_code, user_persona):
     insights_directory = f"insight_library/{user_persona}/{st.session_state.userId}/"
     try:
         blobs = container_client.list_blobs(name_starts_with=insights_directory)
-        for blob in blobs:
+        for index, blob in enumerate(blobs):
+            # Skip the first item
+            if index == 0:
+                continue
             blob_name = blob.name  # Extract the blob names
-            print(blob_name)
             file_name_with_extension = blob_name.split('/')[-1]
             file_name = file_name_with_extension.split('.')[0]
+            
             blob_client = container_client.get_blob_client(blob_name)
             blob_content = blob_client.download_blob().readall()
-            print(blob_content)
+            
             insight_data = json.loads(blob_content)
             if insight_data['base_code'] == base_code:
-                logger.info("Existing insight found for base code: {}", base_code)
+                logger.info("Existing insight found for base code: %s", base_code)
                 return insight_data, file_name
-        logger.info("No existing insight found for base code: {}", base_code)
+        logger.info("No existing insight found for base code: %s", base_code)
+        return None
+    except json.JSONDecodeError as e:
+        logger.error("Error while retrieving insight: %s", e)
         return None
     except Exception as e:
-        logger.error("Error while retrieving insight: {}", e)
+        logger.error("Error while retrieving insight: %s", e)
         return None
 
 def update_insight(insight_data, user_persona, file_number):
@@ -474,10 +482,10 @@ def update_insight(insight_data, user_persona, file_number):
         file_path = f"{user_directory}/{file_number}.json"
         file_content = json.dumps(insight_data, indent=4)
         container_client.upload_blob(file_path, data=file_content, overwrite=True)
-        logger.info("Insight updated successfully: {}", file_number)
+        logger.info("Insight updated successfully: %s", file_number)
         return True
     except Exception as e:
-        logger.error("Error while updating insight: {}", e)
+        logger.error("Error while updating insight: %s", e)
         return False
 
 def save_insight(next_file_number, user_persona, insight_desc, base_prompt, base_code,selected_db, insight_prompt, insight_code, chart_prompt, chart_code):
@@ -590,44 +598,14 @@ def generate_sql(query, table_descriptions, table_details,selected_db):
 #             logger.error("Error connecting to MySQL: {}", err)
 #             return None
         
-def execute_sql(query, selected_db):
-    update_config(selected_db)
-    engine = create_sqlalchemy_engine()
-    if engine:
-        connection = engine.connect()
-        logger.info(f"Connected to the database {selected_db}.")
-        try:
-            df = pd.read_sql_query(query, connection)
-            logger.info("Query executed successfully.")
-            return df
-        except Exception as e:
-            logger.error(f"Query execution failed: {e}")
-            return pd.DataFrame()
-        finally:
-            connection.close()
-    else:
-        logger.error("Failed to create a SQLAlchemy engine.")
-        return None
-
-# def execute_sql(query, selected_db, offset, limit=100):
+# def execute_sql(query, selected_db):
 #     update_config(selected_db)
 #     engine = create_sqlalchemy_engine()
 #     if engine:
 #         connection = engine.connect()
 #         logger.info(f"Connected to the database {selected_db}.")
 #         try:
-#             # Modify the query to use ROW_NUMBER() for pagination
-#             paginated_query = f"""
-#             WITH CTE AS (
-#                 SELECT *,
-#                        ROW_NUMBER() OVER (ORDER BY (SELECT NULL)) AS RowNum
-#                 FROM ({query.rstrip(';')}) AS subquery
-#             )
-#             SELECT *
-#             FROM CTE
-#             WHERE RowNum BETWEEN {offset + 1} AND {offset + limit};
-#             """
-#             df = pd.read_sql_query(paginated_query, connection)
+#             df = pd.read_sql_query(query, connection)
 #             logger.info("Query executed successfully.")
 #             return df
 #         except Exception as e:
@@ -638,10 +616,23 @@ def execute_sql(query, selected_db):
 #     else:
 #         logger.error("Failed to create a SQLAlchemy engine.")
 #         return None
-    
-# def fetch_data(query, selected_db, offset, limit):
-#     df = execute_sql(query, selected_db, offset, limit)   
-#     return drop_duplicate_columns(df)
+
+def execute_sql(query,selected_db):
+    df = None
+    try:
+        conn = sqlite3.connect(selected_db)
+        curr = conn.cursor()
+        curr.execute(query)
+
+        results = curr.fetchall()
+        columns = [desc[0] for desc in curr.description]
+        df = pd.DataFrame(results, columns=columns).copy()
+        logger.info("Query executed successfully.")
+    except sqlite3.Error as e:
+        logger.error(f"Error while querying the DB : {e}")
+    finally:
+        conn.close()
+    return df
     
 def handle_retrieve_request(prompt):
     sql_generated = generate_sql(prompt, st.session_state['table_master'], st.session_state['table_details'], st.session_state['selected_db'])
@@ -830,12 +821,6 @@ def answer_guide_question(question, dframe, df_structure, selected_db):
         logger.debug("Code execution error state: {}", st.session_state['code_execution_error'])
     return result_df, last_method_num + 1, analysis_code
 
-# def get_metadata(table):
-#     table_details = st.session_state['table_details'][table]
-#     matadata = [[field, details[0], details[1]] for field, details in table_details.items()]
-#     metadata_df = pd.DataFrame(matadata, columns=['Field Name', 'Field Description', 'Field Type'])
-#     return metadata_df
-
 def generate_graph(query, df, df_structure,generate_graph):
     if query is None or df is None or df_structure is None:
         logger.error("generate_graph received None values for query, df, or df_structure")
@@ -974,79 +959,62 @@ def get_table_details(engine,selected_db):
     return tables_master_dict, tables_details_dict
 
 # Function to fetch database names from SQL Server
-def get_database_names():
-    query = """
-    SELECT name 
-    FROM sys.databases
-    WHERE name NOT IN ('master', 'tempdb', 'model', 'msdb');
-    """
-    connection_string = (
-        f"DRIVER={SQL_SERVER_CONFIG['driver']};"
-        f"SERVER={SQL_SERVER_CONFIG['server']};"
-        f"UID={SQL_SERVER_CONFIG['username']};"  # Use SQL Server authentication username
-        f"PWD={SQL_SERVER_CONFIG['password']}"   # Use SQL Server authentication password
-    )
-    engine = create_engine(f"mssql+pyodbc:///?odbc_connect={connection_string}")
-    try:
-        with engine.connect() as conn:
-            result = conn.execute(query)
-            databases = [row['name'] for row in result]
-        logger.info("Database names fetched successfully.")
-        return databases
-    except Exception as e:
-        logger.error("Error fetching database names: {}", e)
-        return []
-
-def get_metadata(selected_table):
-    try:
-        metadata_df = pd.DataFrame(st.session_state['table_details'][selected_table])
-        logger.info("Metadata fetched for table: {}", selected_table)
-        return metadata_df
-    except Exception as e:
-        logger.error("Error fetching metadata for table {}: {}", selected_table, e)
-        return pd.DataFrame()
-
-# def load_data(sql_generated, selected_db):
-#     # Fetch data in chunks of 100 rows
-#     if 'offset' not in st.session_state:
-#         st.session_state['offset'] = 0
-
-#     if 'data' not in st.session_state:
-#         st.session_state['data'] = pd.DataFrame()  # Initialize as an empty DataFrame
-
-#     new_data = fetch_data(sql_generated, selected_db, st.session_state['offset'], 100)
-#     if not new_data.empty:
-#         if st.session_state['offset'] == 0:
-#             st.session_state['data'] = new_data
-#         else:
-#             st.session_state['data'] = pd.concat([st.session_state['data'], new_data], ignore_index=True)
-
-#     grid_options = get_ag_grid_options(st.session_state['data'])
-#     AgGrid(st.session_state['data'], gridOptions=grid_options, key=f'query_grid_{st.session_state["offset"]}', lazyloading=True)
-#     i=0
-#     if not new_data.empty :
-#         button_clicked = False
-#         while not button_clicked:
-#             i+=1
-#             if st.button('Load more', key=i):
-#                 button_clicked = True
-#                 st.write('Button clicked!')
-#                 st.session_state['offset'] += 100
-#                 load_data(sql_generated, selected_db)
-#             else:
-#                 st.write('Waiting for button click...')
-#             time.sleep(1)
-#         # if st.button("Load more"):
-#         #     logger.info(st.session_state['offset'])
-#         #     logger.info("hi............................................................")
-#         #     st.session_state['offset'] += 100
-#         #     load_data(sql_generated, selected_db)
-#         # else:
-#         #     logger.info(st.session_state['offset'])
-#         #     logger.info("hi buttoon............................................................")
-#     else:
-#         logger.info(st.session_state['offset'])
-#         logger.info("hi next data............................................................")
+# def get_database_names():
+#     query = """
+#     SELECT name 
+#     FROM sys.databases
+#     WHERE name NOT IN ('master', 'tempdb', 'model', 'msdb');
+#     """
+#     connection_string = (
+#         f"DRIVER={SQL_SERVER_CONFIG['driver']};"
+#         f"SERVER={SQL_SERVER_CONFIG['server']};"
+#         f"UID={SQL_SERVER_CONFIG['username']};"  # Use SQL Server authentication username
+#         f"PWD={SQL_SERVER_CONFIG['password']}"   # Use SQL Server authentication password
+#     )
+#     engine = create_engine(f"mssql+pyodbc:///?odbc_connect={connection_string}")
+#     try:
+#         with engine.connect() as conn:
+#             result = conn.execute(query)
+#             databases = [row['name'] for row in result]
+#         logger.info("Database names fetched successfully.")
+#         return databases
+#     except Exception as e:
+#         logger.error("Error fetching database names: {}", e)
+#         return []
+
+# def get_metadata(selected_table):
+#     try:
+#         metadata_df = pd.DataFrame(st.session_state['table_details'][selected_table])
+#         logger.info("Metadata fetched for table: {}", selected_table)
+#         return metadata_df
+#     except Exception as e:
+#         logger.error("Error fetching metadata for table {}: {}", selected_table, e)
+#         return pd.DataFrame()
+
+def get_metadata(table):
+    table_details = st.session_state['table_details'][table]
+    matadata = [[field, details[0], details[1]] for field, details in table_details.items()]
+    metadata_df = pd.DataFrame(matadata, columns=['Field Name', 'Field Description', 'Field Type'])
+    return metadata_df
+
+def get_meta():
+    print("---------------step1 -------------------------")
+    if 'table_master' not in st.session_state:
+        # load db metadata file
+        print("---------------step2 -------------------------")
+        db_js = json.load(open('./database/db_tables.json'))
+        tables_master_dict = {}
+        tables_details_dict = {}
+        for j in db_js:
+            tables_master_dict[j['name']] = j['description']
+            tables_details_dict[j['name']] = j['fields']
+        print(tables_details_dict)
+        print(tables_master_dict)
+        st.session_state['table_master'] = tables_master_dict
+        st.session_state['table_details'] = tables_details_dict
+    return
+
+get_meta()
 
 def compose_dataset():
     if "messages" not in st.session_state:
@@ -1064,24 +1032,30 @@ def compose_dataset():
     with col_cc:
         st.markdown(APP_TITLE, unsafe_allow_html=True)
 
-    databases = get_database_names()
-    selected_db = st.selectbox('Select Database:', [''] + databases)
+    # databases = get_database_names()
+    # selected_db = st.selectbox('Select Database:', [''] + databases)
+    selected = st.selectbox('Select Database:', DB_List)
 
+    if selected == "Patient SDOH":
+        selected_db = './gravity_sdoh_observations.db'
+        st.session_state['selected_db'] = selected_db
+    
     if selected_db:
         if 'selected_db' in st.session_state and st.session_state['selected_db'] != selected_db:
             # Clear session state data related to the previous database
             st.session_state['messages'] = []
             st.session_state['selected_table'] = None
             logger.debug('Session state cleared due to database change.')
+            st.session_state['selected_db'] = selected_db
 
-        update_config(selected_db)
-        engine = create_sqlalchemy_engine()
+        # update_config(selected_db)
+        # engine = create_sqlalchemy_engine()
 
-        if 'table_master' not in st.session_state or st.session_state.get('selected_db') != selected_db:
-            tables_master_dict, tables_details_dict = get_table_details(engine, selected_db)
-            st.session_state['table_master'] = tables_master_dict
-            st.session_state['table_details'] = tables_details_dict
-            st.session_state['selected_db'] = selected_db
+        # if 'table_master' not in st.session_state or st.session_state.get('selected_db') != selected_db:
+        #     tables_master_dict, tables_details_dict = get_table_details(engine, selected_db)
+        #     st.session_state['table_master'] = tables_master_dict
+        #     st.session_state['table_details'] = tables_details_dict
+        #     st.session_state['selected_db'] = selected_db
 
         tables = list(st.session_state['table_master'].keys())
         selected_table = st.selectbox('Tables available:', [''] + tables)
@@ -1095,7 +1069,7 @@ def compose_dataset():
                     st.session_state.messages.append({"role": "assistant", "type": "text", "content": table_desc})
                     st.session_state.messages.append({"role": "assistant", "type": "dataframe", "content": table_metadata_df})
                     logger.debug('Table metadata and description added to session state messages.')
-                    # st.session_state.messages.append({"role": "assistant", "type": "text", "content":  ""})
+                    st.session_state.messages.append({"role": "assistant", "type": "text", "content":  ""})
                     # display_paginated_dataframe(table_metadata_df, "table_metadata")
                 except Exception as e:
                     st.error("Please try again")
@@ -1166,16 +1140,7 @@ def compose_dataset():
                         st.write(f"Query saved in the library with id {st.session_state['retrieval_query_no']}.")
                         logger.info("Query saved in the library with id {}.", st.session_state['retrieval_query_no'])
 
-if 'graph_obj' not in st.session_state:
-    st.session_state['graph_obj'] = None
-if 'graph_prompt' not in st.session_state:
-    st.session_state['graph_prompt'] = ''
-if 'data_obj' not in st.session_state:
-    st.session_state['data_obj'] = None
-if 'data_prompt' not in st.session_state:
-    st.session_state['data_prompt'] = ''
-if 'code_execution_error' not in st.session_state:
-    st.session_state['code_execution_error'] = (None, None)
+
 
 def design_insight():
     col_aa, col_bb, col_cc = st.columns([1, 4, 1], gap="small", vertical_alignment="center")
@@ -1186,6 +1151,17 @@ def design_insight():
         st.markdown('**Select a dataset that you generated and ask for different types of tabular insight or graphical charts.**')
     with col_cc:
         st.markdown(APP_TITLE, unsafe_allow_html=True)
+    
+    if 'graph_obj' not in st.session_state:
+        st.session_state['graph_obj'] = None
+    if 'graph_prompt' not in st.session_state:
+        st.session_state['graph_prompt'] = ''
+    if 'data_obj' not in st.session_state:
+        st.session_state['data_obj'] = None
+    if 'data_prompt' not in st.session_state:
+        st.session_state['data_prompt'] = ''
+    if 'code_execution_error' not in st.session_state:
+        st.session_state['code_execution_error'] = (None, None)
 
     get_saved_query_blob_list()
     selected_query = st.selectbox('Select a saved query', [""] + list(st.session_state['query_display_dict'].keys()))
@@ -1272,7 +1248,7 @@ def design_insight():
                 if st.session_state['data_obj'] is not None:
                     display_paginated_dataframe(st.session_state['data_obj'], "ag_grid_insight")
                     st.session_state['data_prompt'] = data_prompt
-
+                    
             with st.container():
                 st.subheader('Generate Graph')
                 graph_prompt_value = st.session_state.get('graph_prompt', '')
@@ -1296,7 +1272,10 @@ def design_insight():
                                 logger.error("Error in generating graph: %s", e)
                                 st.write("Error in generating graph, please try again")
                         else:
-                            st.plotly_chart(st.session_state['graph_obj'], use_container_width=True)
+                            try:
+                                st.plotly_chart(st.session_state['graph_obj'], use_container_width=True)
+                            except Exception as e:
+                                st.write("Error in displaying graph, please try again")
                     st.session_state['graph_prompt'] = graph_prompt
                 else:
                     if st.session_state['graph_obj'] is not None:
@@ -1304,7 +1283,7 @@ def design_insight():
                             st.plotly_chart(st.session_state['graph_obj'], use_container_width=True)
                         except Exception as e:
                             st.write("Error in displaying graph, please try again")
-                            logger.error("Error in displaying graph: %s", e)
+                            logger.error("Error in displaying graph: %s", e)                      
             with st.container():
                 if 'graph_obj' in st.session_state or 'data_obj' in st.session_state:
                     user_persona = st.selectbox('Select a persona to save the result of your exploration', persona_list)
@@ -1322,15 +1301,17 @@ def design_insight():
                         try:
                             result = get_existing_insight(base_code, user_persona)
                             if result:
-                                existing_insight, file_number = result
-                                existing_insight['prompt'][f'prompt_{len(existing_insight["prompt"]) + 1}'] = {
-                                    'insight_prompt': insight_prompt,
-                                    'insight_code': insight_code
-                                }
-                                existing_insight['chart'][f'chart_{len(existing_insight["chart"]) + 1}'] = {
-                                    'chart_prompt': chart_prompt,
-                                    'chart_code': chart_code
-                                }
+                                existing_insight, file_number = result 
+                                if insight_prompt and insight_code is not None:
+                                    existing_insight['prompt'][f'prompt_{len(existing_insight["prompt"]) + 1}'] = {
+                                        'insight_prompt': insight_prompt,
+                                        'insight_code': insight_code
+                                    }
+                                if chart_prompt and chart_code is not None:
+                                    existing_insight['chart'][f'chart_{len(existing_insight["chart"]) + 1}'] = {
+                                        'chart_prompt': chart_prompt,
+                                        'chart_code': chart_code
+                                    }
                                 try:
                                     update_insight(existing_insight, user_persona, file_number)
                                     st.text('Insight updated with new Graph and/or Data.')

requirements.txt

@@ -6,6 +6,7 @@ altair==5.4.1
 reportlab==4.2.4
 streamlit_navigation_bar==3.3.0
 altair_saver==0.5.0
+sqlite3
 plotly
 boto3
 azure.storage.blob