ai-data-science-team 0.0.0.9014__py3-none-any.whl → 0.0.0.9016__py3-none-any.whl

ai_data_science_team/ds_agents/eda_tools_agent.py

@@ -1,5 +1,3 @@
-
-
 from typing import Any, Optional, Annotated, Sequence, Dict
 import operator
 import pandas as pd
@@ -17,10 +15,11 @@ from ai_data_science_team.utils.regex import format_agent_name
 
 from ai_data_science_team.tools.eda import (
     explain_data,
-    describe_dataset, 
-    visualize_missing, 
-    correlation_funnel,
+    describe_dataset,
+    visualize_missing,
+    generate_correlation_funnel,
     generate_sweetviz_report,
+    generate_dtale_report,
 )
 from ai_data_science_team.utils.messages import get_tool_call_names
 
@@ -32,15 +31,17 @@ EDA_TOOLS = [
     explain_data,
     describe_dataset,
     visualize_missing,
-    correlation_funnel,
+    generate_correlation_funnel,
     generate_sweetviz_report,
+    generate_dtale_report,
 ]
 
+
 class EDAToolsAgent(BaseAgent):
     """
     An Exploratory Data Analysis Tools Agent that interacts with EDA tools to generate summary statistics,
     missing data visualizations, correlation funnels, EDA reports, etc.
-    
+
     Parameters:
     ----------
     model : langchain.llms.base.LLM
@@ -52,9 +53,9 @@ class EDAToolsAgent(BaseAgent):
     checkpointer : Checkpointer, optional
         The checkpointer for the agent.
     """
-    
+
     def __init__(
-        self, 
+        self,
         model: Any,
         create_react_agent_kwargs: Optional[Dict] = {},
         invoke_react_agent_kwargs: Optional[Dict] = {},
@@ -64,18 +65,18 @@ class EDAToolsAgent(BaseAgent):
             "model": model,
             "create_react_agent_kwargs": create_react_agent_kwargs,
             "invoke_react_agent_kwargs": invoke_react_agent_kwargs,
-            "checkpointer": checkpointer
+            "checkpointer": checkpointer,
         }
         self._compiled_graph = self._make_compiled_graph()
         self.response = None
-        
+
     def _make_compiled_graph(self):
         """
         Creates the compiled state graph for the EDA agent.
         """
         self.response = None
         return make_eda_tools_agent(**self._params)
-    
+
     def update_params(self, **kwargs):
         """
         Updates the agent's parameters and rebuilds the compiled graph.
@@ -83,16 +84,13 @@ class EDAToolsAgent(BaseAgent):
         for k, v in kwargs.items():
             self._params[k] = v
         self._compiled_graph = self._make_compiled_graph()
-        
+
     async def ainvoke_agent(
-        self, 
-        user_instructions: str = None, 
-        data_raw: pd.DataFrame = None, 
-        **kwargs
+        self, user_instructions: str = None, data_raw: pd.DataFrame = None, **kwargs
     ):
         """
         Asynchronously runs the agent with user instructions and data.
-        
+
         Parameters:
         ----------
         user_instructions : str, optional
@@ -105,20 +103,17 @@ class EDAToolsAgent(BaseAgent):
                 "user_instructions": user_instructions,
                 "data_raw": data_raw.to_dict() if data_raw is not None else None,
             },
-            **kwargs
+            **kwargs,
         )
         self.response = response
         return None
-    
+
     def invoke_agent(
-        self, 
-        user_instructions: str = None, 
-        data_raw: pd.DataFrame = None, 
-        **kwargs
+        self, user_instructions: str = None, data_raw: pd.DataFrame = None, **kwargs
     ):
         """
         Synchronously runs the agent with user instructions and data.
-        
+
         Parameters:
         ----------
         user_instructions : str, optional
@@ -131,24 +126,26 @@ class EDAToolsAgent(BaseAgent):
                 "user_instructions": user_instructions,
                 "data_raw": data_raw.to_dict() if data_raw is not None else None,
             },
-            **kwargs
+            **kwargs,
         )
         self.response = response
         return None
-    
+
     def get_internal_messages(self, markdown: bool = False):
         """
         Returns internal messages from the agent response.
         """
         pretty_print = "\n\n".join(
-            [f"### {msg.type.upper()}\n\nID: {msg.id}\n\nContent:\n\n{msg.content}" 
-             for msg in self.response["internal_messages"]]
+            [
+                f"### {msg.type.upper()}\n\nID: {msg.id}\n\nContent:\n\n{msg.content}"
+                for msg in self.response["internal_messages"]
+            ]
         )
         if markdown:
             return Markdown(pretty_print)
         else:
             return self.response["internal_messages"]
-    
+
     def get_artifacts(self, as_dataframe: bool = False):
         """
         Returns the EDA artifacts from the agent response.
@@ -157,7 +154,7 @@ class EDAToolsAgent(BaseAgent):
             return pd.DataFrame(self.response["eda_artifacts"])
         else:
             return self.response["eda_artifacts"]
-    
+
     def get_ai_message(self, markdown: bool = False):
         """
         Returns the AI message from the agent response.
@@ -166,13 +163,14 @@ class EDAToolsAgent(BaseAgent):
             return Markdown(self.response["messages"][0].content)
         else:
             return self.response["messages"][0].content
-        
+
     def get_tool_calls(self):
         """
         Returns the tool calls made by the agent.
         """
         return self.response["tool_calls"]
 
+
 def make_eda_tools_agent(
     model: Any,
     create_react_agent_kwargs: Optional[Dict] = {},
@@ -181,7 +179,7 @@ def make_eda_tools_agent(
 ):
     """
     Creates an Exploratory Data Analyst Agent that can interact with EDA tools.
-    
+
     Parameters:
     ----------
     model : Any
@@ -192,13 +190,13 @@ def make_eda_tools_agent(
         Additional kwargs for agent invocation.
     checkpointer : Checkpointer, optional
         The checkpointer for the agent.
-    
+
     Returns:
     -------
     app : langgraph.graph.CompiledStateGraph
         The compiled state graph for the EDA agent.
     """
-    
+
     class GraphState(AgentState):
         internal_messages: Annotated[Sequence[BaseMessage], operator.add]
         user_instructions: str
@@ -209,11 +207,9 @@ def make_eda_tools_agent(
     def exploratory_agent(state):
         print(format_agent_name(AGENT_NAME))
         print("    * RUN REACT TOOL-CALLING AGENT FOR EDA")
-        
-        tool_node = ToolNode(
-            tools=EDA_TOOLS
-        )
-        
+
+        tool_node = ToolNode(tools=EDA_TOOLS)
+
         eda_agent = create_react_agent(
             model,
             tools=tool_node,
@@ -221,7 +217,7 @@ def make_eda_tools_agent(
             **create_react_agent_kwargs,
             checkpointer=checkpointer,
         )
-        
+
         response = eda_agent.invoke(
             {
                 "messages": [("user", state["user_instructions"])],
@@ -229,13 +225,13 @@ def make_eda_tools_agent(
             },
             invoke_react_agent_kwargs,
         )
-        
+
         print("    * POST-PROCESSING EDA RESULTS")
-        
-        internal_messages = response['messages']
+
+        internal_messages = response["messages"]
         if not internal_messages:
             return {"internal_messages": [], "eda_artifacts": None}
-        
+
         last_ai_message = AIMessage(internal_messages[-1].content, role=AGENT_NAME)
         last_tool_artifact = None
         if len(internal_messages) > 1:
@@ -244,24 +240,24 @@ def make_eda_tools_agent(
                 last_tool_artifact = last_message.artifact
             elif isinstance(last_message, dict) and "artifact" in last_message:
                 last_tool_artifact = last_message["artifact"]
-                
+
         tool_calls = get_tool_call_names(internal_messages)
-        
+
         return {
             "messages": [last_ai_message],
             "internal_messages": internal_messages,
             "eda_artifacts": last_tool_artifact,
             "tool_calls": tool_calls,
         }
-    
+
     workflow = StateGraph(GraphState)
     workflow.add_node("exploratory_agent", exploratory_agent)
     workflow.add_edge(START, "exploratory_agent")
     workflow.add_edge("exploratory_agent", END)
-    
+
     app = workflow.compile(
         checkpointer=checkpointer,
         name=AGENT_NAME,
     )
-    
+
     return app

ai_data_science_team/multiagents/pandas_data_analyst.py

@@ -132,10 +132,10 @@ class PandasDataAnalyst(BaseAgent):
         """Returns a summary of the workflow."""
         if self.response and self.response.get("messages"):
             agents = [msg.role for msg in self.response["messages"]]
-            agent_labels = [f"- **Agent {i+1}:** {role}" for i, role in enumerate(agents)]
+            agent_labels = [f"- **Agent {i+1}:** {role}\n" for i, role in enumerate(agents)]
             header = f"# Pandas Data Analyst Workflow Summary\n\nThis workflow contains {len(agents)} agents:\n\n" + "\n".join(agent_labels)
             reports = [get_generic_summary(json.loads(msg.content)) for msg in self.response["messages"]]
-            summary = "\n" +header + "\n\n".join(reports)
+            summary = "\n\n" + header + "\n\n".join(reports)
             return Markdown(summary) if markdown else summary
 
     @staticmethod
@@ -177,15 +177,15 @@ def make_pandas_data_analyst(
     
     routing_preprocessor_prompt = PromptTemplate(
         template="""
-        You are an expert in routing decisions for a Pandas Data Manipulation Wrangling Agent, a Charting Visualization Agent, and a Pandas Table Agent. Your job is to:
+        You are an expert in routing decisions for a Pandas Data Manipulation Wrangling Agent, a Charting Visualization Agent, and a Pandas Table Agent. Your job is to tell the agents which actions to perform and determine the correct routing for the incoming user question:
         
-        1. Determine what the correct format for a Users Question should be for use with a Pandas Data Wrangling Agent based on the incoming user question. Anything related to data wrangling and manipulation should be passed along.
+        1. Determine what the correct format for a Users Question should be for use with a Pandas Data Wrangling Agent based on the incoming user question. Anything related to data wrangling and manipulation should be passed along. Anything related to data analysis can be handled by the Pandas Agent. Anything that uses Pandas can be passed along. Tables can be returned from this agent. Don't pass along anything about plotting or visualization.
         2. Determine whether or not a chart should be generated or a table should be returned based on the users question.
         3. If a chart is requested, determine the correct format of a Users Question should be used with a Data Visualization Agent. Anything related to plotting and visualization should be passed along.
         
         Use the following criteria on how to route the the initial user question:
         
-        From the incoming user question, remove any details about the format of the final response as either a Chart or Table and return only the important part of the incoming user question that is relevant for the SQL generator agent. This will be the 'user_instructions_data_wrangling'. If 'None' is found, return the original user question.
+        From the incoming user question, remove any details about the format of the final response as either a Chart or Table and return only the important part of the incoming user question that is relevant for the Pandas Data Wrangling and Transformation agent. This will be the 'user_instructions_data_wrangling'. If 'None' is found, return the original user question.
         
         Next, determine if the user would like a data visualization ('chart') or a 'table' returned with the results of the Data Wrangling Agent. If unknown, not specified or 'None' is found, then select 'table'.  
         

ai_data_science_team/multiagents/sql_data_analyst.py

@@ -301,24 +301,13 @@ class SQLDataAnalyst(BaseAgent):
         markdown: bool
             If True, returns the summary as a Markdown-formatted string.
         """
-        if self.response and self.get_response()['messages']:
-            
-            agents = [self.get_response()['messages'][i].role for i in range(len(self.get_response()['messages']))]
-            
-            agent_labels = []
-            for i in range(len(agents)):
-                agent_labels.append(f"- **Agent {i+1}:** {agents[i]}")
-            
-            # Construct header
-            header = f"# SQL Data Analyst Workflow Summary Report\n\nThis agentic workflow contains {len(agents)} agents:\n\n" + "\n".join(agent_labels)
-            
-            reports = []
-            for msg in self.get_response()['messages']:
-                reports.append(get_generic_summary(json.loads(msg.content)))
-                
-            if markdown:
-                return Markdown(header + "\n\n".join(reports))
-            return "\n\n".join(reports)
+        if self.response and self.response.get("messages"):
+            agents = [msg.role for msg in self.response["messages"]]
+            agent_labels = [f"- **Agent {i+1}:** {role}\n" for i, role in enumerate(agents)]
+            header = f"# SQL Data Analyst Workflow Summary\n\nThis workflow contains {len(agents)} agents:\n\n" + "\n".join(agent_labels)
+            reports = [get_generic_summary(json.loads(msg.content)) for msg in self.response["messages"]]
+            summary = "\n\n" + header + "\n\n".join(reports)
+            return Markdown(summary) if markdown else summary