Update main.py

main.py

@@ -19,9 +19,6 @@ 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()
 
@@ -43,15 +40,18 @@ class CodeExecutionResult:
 API_URL = "https://pvanand-code-execution-files-v5.hf.space"
 
 @tool
-def execute_python(code: str):
+def execute_python(code: str) -> str:
     """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: Valid Python code with correct indentation and syntax including necessary imports.
-
+        code: The Python code to execute
     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
     plotly
@@ -96,22 +96,10 @@ def execute_python(code: 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. 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(
-        formatted_prompt,
+        state["messages"],
         token_counter=len,
         max_tokens=16000,
         strategy="last",
@@ -162,28 +150,7 @@ async def chat(input_data: ChatInput):
                 yield f"{json.dumps({'type': 'tool_start', 'tool': event['name'], 'input': tool_input})}\n"
             
             elif kind == "on_tool_end":
-                tool_output = event['data'].get('output', '').content
-                #print(type(tool_output))
-                #print(dir(tool_output))
-                #print the keys
-                pattern = r'data: (.*?)\ndata:'
-                match = re.search(pattern, tool_output)
-                print(tool_output)
-
-                if match:
-                    tool_output_json = match.group(1).strip()
-                    try:
-                        tool_output = json.loads(tool_output_json)
-                        if "artifacts" in tool_output:
-                            for artifact in tool_output["artifacts"]:
-                                artifact_content = requests.get(f"{API_URL}/artifact/{artifact['artifact_id']}").content
-                                print(artifact_content)
-                                tool_output["artifacts"][artifact["artifact_id"]] = artifact_content
-                    except Exception as e:
-                        print(e)
-                        print("Error parsing tool output as json: ", tool_output)
-                else:
-                    print("No match found in tool output")  
+                tool_output = event['data'].get('output', '').content 
                 yield f"{json.dumps({'type': 'tool_end', 'tool': event['name'], 'output': tool_output})}\n"
     return EventSourceResponse(
         generate(),