Update main.py

main.py

@@ -19,6 +19,9 @@ import requests
 import uvicorn
 import re
 from fastapi.staticfiles import StaticFiles
+from langchain_core.runnables import RunnableConfig
+from langchain_core.prompts import ChatPromptTemplate
+from datetime import datetime
 
 app = FastAPI()
 
@@ -40,23 +43,17 @@ class CodeExecutionResult:
 API_URL = "https://pvanand-code-execution-files-v5.hf.space"
 
 @tool
-def execute_python(code: str) -> str:
+def execute_python(code: str, config: RunnableConfig):
     """Execute Python code in an IPython interactiveshell and return the output.
        The returned artifacts (if present) are automatically rendered in the UI and visible to the user. 
     Args:
-        code: The Python code to execute
+        code: Valid Python code with correct indentation and syntax including necessary imports.
 
     Available Libraries:
-    # Use plotly as the default charting library
-    # While using yfinance to pull stock data, Always clean the multiindex columns as this might cause issues in plotting plotly charts
-    # Remove the ticker level from columns if it exists
-    yf_data = yf.download(symbol, start=start_date, end=end_date)
-    if isinstance(yf_data.columns, pd.MultiIndex):
-        yf_data.columns = yf_data.columns.get_level_values(0)
-
-    matplotlib
-    pandas
+    # Use Plotly for creating visualizations
+
     plotly
+    pandas
     groq
     yfinance
     numpy
@@ -98,10 +95,24 @@ def execute_python(code: str) -> str:
 # Configure the memory and model"
 memory = MemorySaver()
 model = ChatOpenAI(model="gpt-4o-mini", streaming=True)
+prompt = ChatPromptTemplate.from_messages([
+    ("system", f"You are a Data Visualization assistant.You have access to a jupyter client 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. Today's date is \
+    {datetime.now().strftime('%Y-%m-%d')}. The current folder contains the following files: {{collection_files}}"),
+    ("placeholder", "{messages}"),
+])
 
 def state_modifier(state) -> list[BaseMessage]:
+    collection_files = "None"
+    # Format the prompt with the current state
+    formatted_prompt = prompt.invoke({
+        "collection_files": collection_files,
+        "messages": state["messages"]
+    })
+    
+    # Trim the messages
     return trim_messages(
-        state["messages"],
+        formatted_prompt,
         token_counter=len,
         max_tokens=16000,
         strategy="last",