Update app.py

app.py

@@ -12,9 +12,6 @@ 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,
@@ -49,19 +46,29 @@ 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:
+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()}")
+
     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 executing query: {e}")
+        print(f"Error: {e}")
+        os.chdir(original_cwd)
         raise e
 def get_summarization(client:groq.Groq,
                       use_question:str,
@@ -251,6 +258,10 @@ 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:
@@ -260,32 +271,39 @@ user_question = """"""
 # * Valid values for product_name include 'Tesla','iPhone' and 'Humane pin'
 
 
-def upload_file(files) -> dict:
+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")
-    
-    if isinstance(files, str):
+    api_key:str=userdata.get("GROQ_API_KEY")
+    data_dir = Path("data")
+    data_dir.mkdir(parents=True, exist_ok=True)
+    if type(files) == str:
         files = [files]
-    
-    uploaded_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
-        shutil.copy2(file.name, path)
+        with open(file.name, "rb") as source, open(path, "wb") as destination:
+            destination.write(source.read())
 
-        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_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_table_descriptions.append(desc)
-
-    return {
-        "files": uploaded_files,
-        "descriptions": "\n".join(stored_table_descriptions)
-    }
+        tables.append(filename)
+    # constructing a string
+    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
 
 def user_prompt_sanitization(user_prompt:str)->str:
   guide = """
@@ -308,52 +326,56 @@ 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, 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
-    )
-    
+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"]
     try:
-        response = chat_with_groq(client, full_prompt, model, {
-            "type": "json_object"
-        })
+      results_df = execute_duckdb_query(sql_query)
     except Exception as e:
-        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([])
+      return [(
+          "Groq Advisor",
+          "Error: " + str(e)
+      )]
+
+    fotmatted_sql_query = sqlparse.format(sql_query, reindent=True, keyword_case='upper')
+    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"
+    )]
 
 with gr.Blocks() as demo:
     gr.Markdown("# CSV Database Query Interface")
@@ -363,28 +385,16 @@ with gr.Blocks() as demo:
         upload_button = gr.Button("Load CSV Files")
         upload_output = gr.Textbox(label="Upload Status", lines=5)
 
-        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])
-    
+        upload_button.click(upload_file, inputs=file_output, outputs=upload_output)
     with gr.Tab("Query Interface"):
         chatbot = gr.Chatbot()
         with gr.Row():
             user_input = gr.Textbox(label="Enter your question")
             submit_button = gr.Button("Submit")
-        
-        def query_model_with_files(user_prompt, files):
-            return queryModel(user_prompt, files)
+        submit_button.click(queryModel, inputs=[user_input], outputs=chatbot)
+
 
-        submit_button.click(query_model_with_files, 
-                            inputs=[user_input, uploaded_files], 
-                            outputs=chatbot)
 
-demo.launch()
+demo.launch(share=True)