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,7 +43,7 @@ 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:
@@ -54,7 +57,7 @@ def execute_python(code: str) -> str:
         'Content-Type': 'application/json'
     }
     data = {
-        "session_token": "test12345", #config.configurable.get("thread_id", "test"),
+        "session_token": config.configurable.get("thread_id", "test"),
         "code": code
     }
     response = requests.post(
@@ -69,20 +72,36 @@ def execute_python(code: str) -> str:
         response_json = response.json()
         return f"data: {json.dumps(response_json)} \ndata:"
 
-# 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]:
-    return trim_messages(
-        state["messages"],
-        token_counter=len,
-        max_tokens=16000,
-        strategy="last",
-        start_on="human",
-        include_system=True,
-        allow_partial=False,
-    )
+    collection_files = "None"
+    try:
+        formatted_prompt = prompt.invoke({
+            "collection_files": collection_files,
+            "messages": state["messages"]
+        })
+        print(state["messages"])
+        return trim_messages(
+            formatted_prompt,
+            token_counter=len,
+            max_tokens=16000,
+            strategy="last",
+            start_on="human",
+            include_system=True,
+            allow_partial=False,
+        )
+        
+    except Exception as e:
+        print(f"Error in state modifier: {str(e)}")
+        return state["messages"]
 
 # Create the agent with the Python execution tool
 agent = create_react_agent(