Update main.py

main.py

@@ -98,13 +98,20 @@ def execute_python(code: str, config: RunnableConfig):
     except Exception as e:
         return (f"Error executing code: {str(e)}", None)
 
+save_yf_prompt = """
+# Save downloaded data
+data = yf.download(['MSFT', 'AAPL', 'GOOG'], period='6mo')
+data.to_pickle('stock_data.pkl')
+"""
+# Load it back later
+data = pd.read_pickle('stock_data.pkl')
 
 memory = MemorySaver()
 model = ChatOpenAI(model="gpt-4o", streaming=True)
 prompt = ChatPromptTemplate.from_messages([
     ("system", f"You are a Data Visualization assistant.You have access to an jupyter notebook with access to internet for python code execution.\
     Your taks is to assist users with your data analysis and visualization expertise. Use Plotly for creating visualizations and charts. Generated artifacts\
-    are automatically rendered in the UI. Variables do not persist across tool calls, hence save any data to current directory you want to use in the next tool call to a file.(use file_name and do not add file path, as you have only permission to edit current folder) (df.to_csv('file_name.csv', index=False) Ex: save yahoofinance data as csv for reusing in next tool call) Today's date is \
+    are automatically rendered in the UI. Variables do not persist across tool calls, hence save any data to current directory you want to use in the next tool call to a file.(use file_name and do not add file path, as you have only permission to edit current folder) {save_yf_prompt} Today's date is \
     {datetime.now().strftime('%Y-%m-%d')}. Format you responses beutifully using markdown formating such as tables, lists etc. The current folder contains the following files: {{collection_files}} {yf_docs}"),
     ("placeholder", "{messages}"),
 ])