Update app.py

app.py

@@ -1,4 +1,5 @@
 import os
+import tempfile
 import google.generativeai as genai
 from flask import Flask, request, jsonify, render_template
 from flask_cors import CORS
@@ -12,15 +13,15 @@ import json
 import numpy as np
 import logging
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_google_genai.llms import GoogleGenerativeAI  # Updated import statement
-from langchain_experimental.agents import create_csv_agent  # Updated import statement
+from langchain_google_genai import GoogleGenerativeAI
+from langchain_experimental.agents import create_csv_agent
 from langchain_google_genai import GoogleGenerativeAIEmbeddings
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain.chains.question_answering import load_qa_chain
 from langchain.prompts import PromptTemplate
 from langchain.docstore.document import Document
 from sklearn.metrics.pairwise import cosine_similarity
-import plotly.express as px  # For interactive charts
+import plotly.express as px
 import plotly
 from newsapi import NewsApiClient
 import certifi
@@ -45,14 +46,8 @@ newsapi.session = session
 # Initialize the model
 model = genai.GenerativeModel('gemini-pro')
 
-UPLOAD_FOLDER = 'uploads'
 ALLOWED_EXTENSIONS = {'txt', 'pdf', 'docx', 'xlsx', 'csv'}
 
-if not os.path.exists(UPLOAD_FOLDER):
-    os.makedirs(UPLOAD_FOLDER)
-
-app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
-
 # Database setup
 DATABASE = 'ai_assistant.db'
 
@@ -360,17 +355,20 @@ def upload_file():
         return jsonify({'error': 'No selected file'}), 400
     if file and allowed_file(file.filename):
         filename = secure_filename(file.filename)
-        file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
-        file.save(file_path)
         
         try:
+            # Create a temporary file
+            with tempfile.NamedTemporaryFile(delete=False) as temp_file:
+                file.save(temp_file.name)
+                temp_file_path = temp_file.name
+            
             logging.info(f"File uploaded successfully: {filename}")
-            extracted_text = process_document(file_path)
+            extracted_text = process_document(temp_file_path)
             text_chunks = get_text_chunks(extracted_text)
             analysis = analyze_document(extracted_text)
             
             db = get_db()
-            with open(file_path, 'rb') as f:
+            with open(temp_file_path, 'rb') as f:
                 file_data = f.read()
             file_data_base64 = base64.b64encode(file_data).decode('utf-8')
             cursor = db.execute('INSERT INTO files (filename, file_data, analysis) VALUES (?, ?, ?)',
@@ -380,7 +378,9 @@ def upload_file():
             
             create_vector_store(text_chunks, file_id)
             
-            os.remove(file_path)  # Remove the file after processing
+            # Remove the temporary
+            
+            os.unlink(temp_file_path)
             logging.info(f"File processing completed and saved to database with ID: {file_id}")
             
             return jsonify({'file_id': file_id, 'analysis': analysis})
@@ -400,7 +400,12 @@ def plot():
         file_data_base64 = cursor.fetchone()['file_data']
         file_data = base64.b64decode(file_data_base64)
         
-        df = pd.read_excel(pd.io.common.BytesIO(file_data))
+        with tempfile.NamedTemporaryFile(delete=False, suffix='.xlsx') as temp_file:
+            temp_file.write(file_data)
+            temp_file_path = temp_file.name
+
+        df = pd.read_excel(temp_file_path)
+        os.unlink(temp_file_path)
         
         fig = px.line(df, x=df.columns[0], y=df.columns[1:])
         graph_json = json.dumps(fig, cls=plotly.utils.PlotlyJSONEncoder)
@@ -423,9 +428,9 @@ def process_csv_query():
         file_data = base64.b64decode(file_data_base64)
         
         # Save the CSV data to a temporary file
-        temp_csv_path = f'/tmp/{file_id}.csv'
-        with open(temp_csv_path, 'wb') as temp_csv:
+        with tempfile.NamedTemporaryFile(delete=False, suffix='.csv') as temp_csv:
             temp_csv.write(file_data)
+            temp_csv_path = temp_csv.name
         
         # Create a langchain agent using the gemini-pro model
         agent = create_csv_agent(GoogleGenerativeAI(model="gemini-pro"), temp_csv_path, verbose=True)
@@ -433,6 +438,9 @@ def process_csv_query():
         # Run the query using the agent
         response = agent.run(query)
         
+        # Remove the temporary file
+        os.unlink(temp_csv_path)
+        
         return jsonify({'response': response})
     except Exception as e:
         logging.error(f'Error processing CSV query: {str(e)}', exc_info=True)
@@ -468,5 +476,4 @@ def fetch_news():
 
 if __name__ == '__main__':
     init_db()
-    port = int(os.environ.get('PORT', 7860))
-    app.run(host='0.0.0.0', port=port, debug=False)
+    app.run(debug=True, port=5000)