diff --git a/streamlit/interface_3.py b/streamlit/interface_3.py
index a78881a..6ea16a7 100644
--- a/streamlit/interface_3.py
+++ b/streamlit/interface_3.py
@@ -445,12 +445,12 @@ if st.session_state.user_info['mail'] in user_list:
                 answer_l = list(field_dict.values())
                 page_no_l = [page_dict.get(x, "") for x in field_l]
             else:
-                field_l = field
-                answer_l = list(response_dict.values())[0]
+                field_l = [field]
+                answer_l = [list(response_dict.values())[0]]
                 try:
-                    page_no_l = list(response_dict.values())[1]
+                    page_no_l = [list(response_dict.values())[1]]
                 except:
-                    page_no_l = ''      
+                    page_no_l = ['']      
 
             field_list.extend(field_l)
             answer_list.extend(answer_l)
@@ -475,8 +475,6 @@ if st.session_state.user_info['mail'] in user_list:
             snippet_list.extend(snippet_l)
             page_no_list.extend(page_no_l)
 
-        st.write(field_list)
-        st.write(answer_list)
         df['Field Name'] = field_list
         df['Raw value'] = answer_list
         # post-processing
@@ -649,7 +647,7 @@ if st.session_state.user_info['mail'] in user_list:
     st.write(raw_response_text)
 
     try:
-        save_to_sf('load_training_results', "training_results_file_name", "results.csv", "attempt_logs_file_name", "history.csv")
+        save_to_sf('load_training_results', {"training_results_file_name": "results.csv", "attempt_logs_file_name": "history.csv"})
     except:
         st.write("running locally")
 
diff --git a/streamlit/sf_conn.py b/streamlit/sf_conn.py
index ca28bb7..12202e2 100644
--- a/streamlit/sf_conn.py
+++ b/streamlit/sf_conn.py
@@ -39,7 +39,7 @@ def get_secret():
 
 # TODO: This function needs to be changed to accept Kwargs 
 # The function name should be more generic and cofnigurable
-def save_to_sf(dag_name, conf_1, training_results_file, conf_2, attempt_logs_file):
+def save_to_sf(dag_name, dict_args):
 
     mwaa_env_name = 'doczy-dev-infra-mwaa'
     dag_name = dag_name
@@ -62,10 +62,10 @@ def save_to_sf(dag_name, conf_1, training_results_file, conf_2, attempt_logs_fil
     # attempt_logs_file = "attempt_logs_sample.csv"
     # conf = "{\"" + "training_results_file_name" + "\":\"" + {training_results_file} + "\", \"" + "attempt_logs_file_name" + "\":\"" + {attempt_logs_file} + "\"}".format(training_results_file=training_results_file, attempt_logs_file=attempt_logs_file)
 
-    # conf_1 = list(dict_args.keys())[0]
-    # conf_2 = list(dict_args.keys())[1]
-    # training_results_file = list(dict_args.values())[0]
-    # attempt_logs_file = list(dict_args.values())[1]
+    conf_1 = list(dict_args.keys())[0]
+    conf_2 = list(dict_args.keys())[1]
+    training_results_file = list(dict_args.values())[0]
+    attempt_logs_file = list(dict_args.values())[1]
     conf = f"""{{
             "{conf_1}": "{training_results_file}",
             "{conf_2}": "{attempt_logs_file}"