Update app.py

app.py

@@ -260,16 +260,15 @@ user_question = """"""
 # * Valid values for product_name include 'Tesla','iPhone' and 'Humane pin'
 
 
-def upload_file(files) -> List[str]:
+def upload_file(files) -> dict:
     model = "llama3-8b-8192"
     api_key: str = userdata.get("GROQ_API_KEY")
     
     if isinstance(files, str):
         files = [files]
     
-    stored_paths = []
+    uploaded_files = []
     stored_table_descriptions = []
-    tables = []
     
     for file in files:
         filename = Path(file.name).name
@@ -278,20 +277,15 @@ def upload_file(files) -> List[str]:
         # Copy the content of the temporary file to our destination
         shutil.copy2(file.name, path)
 
-        stored_paths.append(str(path))
-        table_description_temp = identify_column_datatypes_to_SQL_DEF(pd.read_csv(path), api_key, model)
-        desc = f"Table: {filename}\nColumns:\n{table_description_temp}"
+        uploaded_files.append(str(path))
+        table_description = identify_column_datatypes_to_SQL_DEF(pd.read_csv(path), api_key, model)
+        desc = f"Table: {filename}\nColumns:\n{table_description}"
         stored_table_descriptions.append(desc)
-        tables.append(filename)
 
-    # Update global variables
-    global table_description, tables_string, table_1, table_1_wt_xt
-    table_description = "\n".join(stored_table_descriptions)
-    tables_string = join_with_and(tables)
-    table_1 = tables[0] if tables else ""
-    table_1_wt_xt = table_1.split('.')[0] if table_1 else ""
-
-    return "\n".join(stored_table_descriptions)
+    return {
+        "files": uploaded_files,
+        "descriptions": "\n".join(stored_table_descriptions)
+    }
 
 def user_prompt_sanitization(user_prompt:str)->str:
   guide = """
@@ -314,58 +308,52 @@ def user_prompt_sanitization(user_prompt:str)->str:
   client = groq.Groq(api_key=api_key)
   return chat_with_groq(client,formatted_guide,"llama3-70b-8192",None)
 
-def queryModel(user_prompt:str,model:str = "llama3-70b-8192",api_key:str=userdata.get("GROQ_API_KEY")):
-  client = groq.Groq(api_key=api_key)
-  user_prompt = user_prompt_sanitization(user_prompt)
-  print(user_prompt)
-  full_prompt = base_prompt.format(
-      user_question=user_prompt,
-      table_description=table_description,
-      tables=tables_string,
-      table_1=table_1,
-      table_1_wt_xt=table_1_wt_xt
-  )
-  try:
-     response = chat_with_groq(client,full_prompt,model,{
-      "type":"json_object"
-     })
-  except Exception as e:
-    return [(
-        "Groq Advisor",
-        "Error: " + str(e)
-    )]
-  response = json.loads(response)
-  if "sql" in response:
-    sql_query = response["sql"]
+def queryModel(user_prompt: str, files: List[str], model: str = "llama3-70b-8192", api_key: str = userdata.get("GROQ_API_KEY")):
+    client = groq.Groq(api_key=api_key)
+    user_prompt = user_prompt_sanitization(user_prompt)
+    print(user_prompt)
+    
+    # Update this part to use the actual table names from the uploaded files
+    table_names = [Path(file).stem for file in files]
+    tables_string = join_with_and(table_names)
+    table_1 = table_names[0] if table_names else ""
+    table_1_wt_xt = table_1
+
+    full_prompt = base_prompt.format(
+        user_question=user_prompt,
+        table_description=table_description,
+        tables=tables_string,
+        table_1=table_1,
+        table_1_wt_xt=table_1_wt_xt
+    )
+    
     try:
-      results_df = execute_duckdb_query(sql_query)
+        response = chat_with_groq(client, full_prompt, model, {
+            "type": "json_object"
+        })
     except Exception as e:
-      return [(
-          "Groq Advisor",
-          "Error: " + str(e)
-      )]
-
-    fotmatted_sql_query = sqlparse.format(sql_query, reindent=True, keyword_case='upper')
-    # print(f"SQL Query: {fotmatted_sql_query}")
-    # print(results_df.to_markdown())
-    query_n_results = "SQL Query: " + fotmatted_sql_query + "\n\n" + results_df.to_markdown()
-    summarization = get_summarization(client,user_prompt,results_df,model)
-    query_n_results += "\n\n" + summarization
-
-    return [(
-        "Groq Advisor",
-        query_n_results
-    )]
-  elif "error" in response:
-    return [(
-        "Groq Advisor",
-        "Error: " + response["error"]
-    )]
-  else:
-    return [(
-        "Groq Advisor",
-        "Error: Unknown error"
-    )]
+        return [("Groq Advisor", "Error: " + str(e))]
+    
+    response = json.loads(response)
+    if "sql" in response:
+        sql_query = response["sql"]
+        try:
+            results_df = execute_duckdb_query(sql_query, files)
+        except Exception as e:
+            return [("Groq Advisor", "Error: " + str(e))]
+
+        formatted_sql_query = sqlparse.format(sql_query, reindent=True, keyword_case='upper')
+        query_n_results = "SQL Query: " + formatted_sql_query + "\n\n" + results_df.to_markdown()
+        summarization = get_summarization(client, user_prompt, results_df, model)
+        query_n_results += "\n\n" + summarization
+
+        return [("Groq Advisor", query_n_results)]
+    elif "error" in response:
+        return [("Groq Advisor", "Error: " + response["error"])]
+    else:
+        return [("Groq Advisor", "Error: Unknown error")]
+
+uploaded_files = gr.State([])
 
 with gr.Blocks() as demo:
     gr.Markdown("# CSV Database Query Interface")
@@ -375,14 +363,27 @@ with gr.Blocks() as demo:
         upload_button = gr.Button("Load CSV Files")
         upload_output = gr.Textbox(label="Upload Status", lines=5)
 
-        upload_button.click(upload_file, inputs=file_output, outputs=upload_output)
+        def handle_upload(files, uploaded_files):
+            result = upload_file(files)
+            uploaded_files.extend(result["files"])
+            return result["descriptions"], uploaded_files
+
+        upload_button.click(handle_upload, 
+                            inputs=[file_output, uploaded_files], 
+                            outputs=[upload_output, uploaded_files])
     
     with gr.Tab("Query Interface"):
         chatbot = gr.Chatbot()
         with gr.Row():
             user_input = gr.Textbox(label="Enter your question")
             submit_button = gr.Button("Submit")
-        submit_button.click(queryModel, inputs=[user_input], outputs=chatbot)
+        
+        def query_model_with_files(user_prompt, files):
+            return queryModel(user_prompt, files)
+
+        submit_button.click(query_model_with_files, 
+                            inputs=[user_input, uploaded_files], 
+                            outputs=chatbot)
 
 demo.launch()