Update app.py

app.py

@@ -12,6 +12,9 @@ from dotenv import load_dotenv
 load_dotenv()
 userdata = os.environ
 
+DATA_DIR = Path(os.getcwd()) / "data"
+DATA_DIR.mkdir(parents=True, exist_ok=True)
+
 
 def chat_with_groq(client:groq.Groq,
                    prompt:str,
@@ -46,29 +49,19 @@ def chat_with_groq(client:groq.Groq,
                   #  logger.info(f"Completion: {completion}")
                    return completion.choices[0].message.content
 
-def execute_duckdb_query(query:str)->pd.DataFrame:
-    """
-    Execute a DuckDB query and return the result as a pandas DataFrame.
-
-    Args:
-        query (str): The DuckDB query to execute.
-
-    Returns:
-        pd.DataFrame: The result of the query as a pandas DataFrame.
-    """
-    original_cwd = os.getcwd()
-    print(f"PATH:{original_cwd}")
-    os.chdir('data')
-    print(f"PATH:{os.getcwd()}")
-
+def execute_duckdb_query(query: str) -> pd.DataFrame:
     try:
         conn = duckdb.connect(database=":memory:", read_only=False)
+        
+        # Load all CSV files from the data directory
+        for csv_file in DATA_DIR.glob("*.csv"):
+            table_name = csv_file.stem
+            conn.execute(f"CREATE TABLE {table_name} AS SELECT * FROM read_csv_auto('{csv_file}')")
+        
         query_result = conn.execute(query).fetch_df().reset_index()
-        os.chdir(original_cwd)
         return query_result
     except Exception as e:
-        print(f"Error: {e}")
-        os.chdir(original_cwd)
+        print(f"Error executing query: {e}")
         raise e
 def get_summarization(client:groq.Groq,
                       use_question:str,
@@ -258,10 +251,6 @@ base_prompt =  """
     * Ensure that the entire output is returned on only one single line
     * Keep your query as simple and straightforward as possible; do not use subqueries
     """
-table_description = """"""
-tables_string = """"""
-table_1 = """"""
-table_1_wt_xt = """"""
 user_question = """"""
 
 # And some rules for querying the dataset:
@@ -272,38 +261,37 @@ user_question = """"""
 
 
 def upload_file(files) -> List[str]:
-    # will have to change to the private system is initiializes
     model = "llama3-8b-8192"
-    api_key:str=userdata.get("GROQ_API_KEY")
-    data_dir = Path("data")
-    data_dir.mkdir(parents=True, exist_ok=True)
-    if type(files) == str:
+    api_key: str = userdata.get("GROQ_API_KEY")
+    
+    if isinstance(files, str):
         files = [files]
+    
     stored_paths = []
     stored_table_descriptions = []
     tables = []
+    
     for file in files:
         filename = Path(file.name).name
-        path = data_dir / filename
+        path = DATA_DIR / filename
 
         # Copy the content of the temporary file to our destination
-        with open(file.name, "rb") as source, open(path, "wb") as destination:
-            destination.write(source.read())
+        shutil.copy2(file.name, path)
 
-        stored_paths.append(str(path.absolute()))
-        table_description = identify_column_datatypes_to_SQL_DEF(pd.read_csv(path),api_key,model)
-        desc = "Table: " + filename + "\n Columns:\n" + table_description
+        stored_paths.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)
-    # constructing a string
+
+    # 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)
-    final = "\n".join(stored_table_descriptions)
-    table_1_wt_xt = tables[0].split('.')[0]
-    table_description = final
-    tables_string = tables_string
-    table_1 = tables[0]
-    table_1_wt_xt = table_1_wt_xt
-    return final
+    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)
 
 def user_prompt_sanitization(user_prompt:str)->str:
   guide = """
@@ -388,6 +376,7 @@ with gr.Blocks() as demo:
         upload_output = gr.Textbox(label="Upload Status", lines=5)
 
         upload_button.click(upload_file, inputs=file_output, outputs=upload_output)
+    
     with gr.Tab("Query Interface"):
         chatbot = gr.Chatbot()
         with gr.Row():
@@ -395,8 +384,6 @@ with gr.Blocks() as demo:
             submit_button = gr.Button("Submit")
         submit_button.click(queryModel, inputs=[user_input], outputs=chatbot)
 
-
-
-demo.launch(share=True)
+demo.launch()