Update main.py

main.py

@@ -50,10 +50,11 @@ def execute_python(code: str, config: RunnableConfig):
         code: Valid Python code with correct indentation and syntax including necessary imports.
 
     Available Libraries:
-    # Use Plotly for creating visualizations
+    # Use plotly as the default charting library
 
-    plotly
+    matplotlib
     pandas
+    plotly
     groq
     yfinance
     numpy
@@ -96,33 +97,28 @@ def execute_python(code: str, config: RunnableConfig):
 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}}"),
+    ("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"
-    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"]
+    # 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,
+        token_counter=len,
+        max_tokens=16000,
+        strategy="last",
+        start_on="human",
+        include_system=True,
+        allow_partial=False,
+    )
 
 # Create the agent with the Python execution tool
 agent = create_react_agent(
@@ -147,56 +143,52 @@ async def chat(input_data: ChatInput):
     }
     
     input_message = HumanMessage(content=input_data.message)
-    try:
-        async def generate():
-            async for event in agent.astream_events(
-                {"messages": [input_message]}, 
-                config,
-                version="v2"
-            ):
-                kind = event["event"]
-                
-                if kind == "on_chat_model_stream":
-                    content = event["data"]["chunk"].content
-                    if content:
-                        yield f"{json.dumps({'type': 'token', 'content': content})}\n"
-    
-                if kind == "on_tool_start":
-                    print(f"Debug - Tool being called: {event['name']}")  # Add this debug line
-                    tool_input = event['data'].get('input', '')
-                    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")  
-                    yield f"{json.dumps({'type': 'tool_end', 'tool': event['name'], 'output': tool_output})}\n"
-        return EventSourceResponse(
-            generate(),
-            media_type="text/event-stream"
-        )
-    except Exception as e:
-        print(f"Error during event streaming: {str(e)}")
-        yield f"{json.dumps({'type': 'error', 'content': str(e)})}\n"
+    async def generate():
+        async for event in agent.astream_events(
+            {"messages": [input_message]}, 
+            config,
+            version="v2"
+        ):
+            kind = event["event"]
+            
+            if kind == "on_chat_model_stream":
+                content = event["data"]["chunk"].content
+                if content:
+                    yield f"{json.dumps({'type': 'token', 'content': content})}\n"
+
+            elif kind == "on_tool_start":
+                tool_input = event['data'].get('input', '')
+                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")  
+                yield f"{json.dumps({'type': 'tool_end', 'tool': event['name'], 'output': tool_output})}\n"
+    return EventSourceResponse(
+        generate(),
+        media_type="text/event-stream"
+    )
 
 @app.get("/health")
 async def health_check():