Add headers to file chat function

functions/chat_functions.py

@@ -3,11 +3,12 @@ from utils import TEMP_DIR, message_dict
 from haystack.dataclasses import ChatMessage
 from haystack.components.generators.chat import OpenAIChatGenerator
 
+import sqlite3
+
 chat_generator = OpenAIChatGenerator(model="gpt-4o")
 response = None
 
 def example_question_generator(session_hash):
-    import sqlite3
     example_response = None
     example_messages = [
         ChatMessage.from_system(
@@ -64,7 +65,7 @@ def doc_db_example_question_generator(session_hash, db_collections, db_name, db_
             f"You are a helpful and knowledgeable agent who has access to an MongoDB NoSQL document database called {db_name}."
         )
     ]
-    
+
     example_messages.append(ChatMessage.from_user(text=f"""We have a MongoDB NoSQL document database with the following collections: {db_collections}.
                                                   The schema of these collections is: {db_schema}. 
                                                   We also have an AI agent with access to the same database that will be performing data analysis.
@@ -87,13 +88,22 @@ def chatbot_with_fc(message, history, session_hash):
                            "scatter_chart_generation_func":scatter_chart_generation_func, "pie_chart_generation_func":pie_chart_generation_func,
                            "histogram_generation_func":histogram_generation_func,
                            "regression_func":regression_func }
+    
+    session_path = 'file_upload'
+
+    dir_path = TEMP_DIR / str(session_hash) / str(session_path)
+    connection = sqlite3.connect(f'{dir_path}/data_source.db')
+    cur=connection.execute('select * from data_source')
+    columns = [i[0] for i in cur.description]
+    cur.close()
+    connection.close()
 
     if message_dict[session_hash]['file_upload'] != None:
         message_dict[session_hash]['file_upload'].append(ChatMessage.from_user(message))
     else:
         messages = [
             ChatMessage.from_system(
-                """You are a helpful and knowledgeable agent who has access to an SQLite database which has a table called 'data_source'. 
+                f"""You are a helpful and knowledgeable agent who has access to an SQLite database which has a table called 'data_source' that contains the following columns: {columns}. 
                 You also have access to a function, called table_generation_func, that can take a query.csv file generated from our sql query and returns an iframe that we should display in our chat window.
                 You also have access to a scatter plot function, called scatter_chart_generation_func, that can take a query.csv file generated from our sql query and uses plotly dictionaries to generate a scatter plot and returns an iframe that we should display in our chat window.
                 You also have access to a line chart function, called line_chart_generation_func, that can take a query.csv file generated from our sql query and uses plotly dictionaries to generate a line chart and returns an iframe that we should display in our chat window.