Merge pull request #58 from dasmy/dev/improve_prompt

Add some information in the data to the prompt

Author: ricklamers

frontend/src/App.tsx

@@ -155,9 +155,9 @@ function App() {
     });
   }
 
-  function completeUpload(filename: string) {
+  function completeUpload(message: string) {
     addMessage({
-      text: `File ${filename} was uploaded successfully.`,
+      text: message,
       type: "message",
       role: "system",
     });
@@ -171,7 +171,7 @@ function App() {
         "Content-Type": "application/json",
       },
       body: JSON.stringify({
-        prompt: `File ${filename} was uploaded successfully.`,
+        prompt: message,
       }),
     })
       .then(() => {})

frontend/src/components/Input.tsx

@@ -43,11 +43,13 @@ export default function Input(props: { onSendMessage: any, onStartUpload: any, o
           body: formData,
         });
 
-        props.onCompletedUpload(file.name);
-
         if (!response.ok) {
           throw new Error("Network response was not ok");
         }
+
+        const json = await response.json();
+        props.onCompletedUpload(json["message"]);
+
       } catch (error) {
         console.error("Error:", error);
       }

gpt_code_ui/webapp/main.py

@@ -6,6 +6,7 @@
 import re
 import logging
 import sys
+import pandas as pd
 
 from collections import deque
 
@@ -56,6 +57,33 @@ def allowed_file(filename):
     return True
 
 
+def inspect_file(filename: str) -> str:
+    READER_MAP = {
+        '.csv': pd.read_csv,
+        '.tsv': pd.read_csv,
+        '.xlsx': pd.read_excel,
+        '.xls': pd.read_excel,
+        '.xml': pd.read_xml,
+        '.json': pd.read_json,
+        '.hdf': pd.read_hdf,
+        '.hdf5': pd.read_hdf,
+        '.feather': pd.read_feather,
+        '.parquet': pd.read_parquet,
+        '.pkl': pd.read_pickle,
+        '.sql': pd.read_sql,
+    }
+
+    _, ext = os.path.splitext(filename)
+
+    try:
+        df = READER_MAP[ext.lower()](filename)
+        return f'The file contains the following columns: {", ".join(df.columns)}'
+    except KeyError:
+        return ''  # unsupported file type
+    except Exception:
+        return ''  # file reading failed. - Don't want to know why.
+
+
 async def get_code(user_prompt, user_openai_key=None, model="gpt-3.5-turbo"):
 
     prompt = f"""First, here is a history of what I asked you to do earlier. 
@@ -253,8 +281,10 @@ def upload_file():
     if file.filename == '':
         return jsonify({'error': 'No selected file'}), 400
     if file and allowed_file(file.filename):
-        file.save(os.path.join(app.config['UPLOAD_FOLDER'], file.filename))
-        return jsonify({'message': 'File successfully uploaded'}), 200
+        file_target = os.path.join(app.config['UPLOAD_FOLDER'], file.filename)
+        file.save(file_target)
+        file_info = inspect_file(file_target)
+        return jsonify({'message': f'File {file.filename} uploaded successfully.\n{file_info}'}), 200
     else:
         return jsonify({'error': 'File type not allowed'}), 400