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

ai_data_science_team/ds_agents/eda_tools_agent.py

@@ -1,12 +1,8 @@
 
 
-from typing import Any, Optional, Annotated, Sequence, List, Dict, Tuple
+from typing import Any, Optional, Annotated, Sequence, Dict
 import operator
 import pandas as pd
-import os
-from io import StringIO, BytesIO
-import base64
-import matplotlib.pyplot as plt
 
 from IPython.display import Markdown
 
@@ -14,22 +10,26 @@ from langchain_core.messages import BaseMessage, AIMessage
 from langgraph.prebuilt import create_react_agent, ToolNode
 from langgraph.prebuilt.chat_agent_executor import AgentState
 from langgraph.graph import START, END, StateGraph
+from langgraph.types import Checkpointer
 
 from ai_data_science_team.templates import BaseAgent
 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,
     generate_sweetviz_report,
 )
+from ai_data_science_team.utils.messages import get_tool_call_names
 
 
 AGENT_NAME = "exploratory_data_analyst_agent"
 
 # Updated tool list for EDA
 EDA_TOOLS = [
+    explain_data,
     describe_dataset,
     visualize_missing,
     correlation_funnel,
@@ -49,6 +49,8 @@ class EDAToolsAgent(BaseAgent):
         Additional kwargs for create_react_agent.
     invoke_react_agent_kwargs : dict
         Additional kwargs for agent invocation.
+    checkpointer : Checkpointer, optional
+        The checkpointer for the agent.
     """
     
     def __init__(
@@ -56,11 +58,13 @@ class EDAToolsAgent(BaseAgent):
         model: Any,
         create_react_agent_kwargs: Optional[Dict] = {},
         invoke_react_agent_kwargs: Optional[Dict] = {},
+        checkpointer: Optional[Checkpointer] = None,
     ):
         self._params = {
             "model": model,
             "create_react_agent_kwargs": create_react_agent_kwargs,
             "invoke_react_agent_kwargs": invoke_react_agent_kwargs,
+            "checkpointer": checkpointer
         }
         self._compiled_graph = self._make_compiled_graph()
         self.response = None
@@ -162,11 +166,18 @@ 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] = {},
     invoke_react_agent_kwargs: Optional[Dict] = {},
+    checkpointer: Optional[Checkpointer] = None,
 ):
     """
     Creates an Exploratory Data Analyst Agent that can interact with EDA tools.
@@ -179,6 +190,8 @@ def make_eda_tools_agent(
         Additional kwargs for create_react_agent.
     invoke_react_agent_kwargs : dict
         Additional kwargs for agent invocation.
+    checkpointer : Checkpointer, optional
+        The checkpointer for the agent.
     
     Returns:
     -------
@@ -191,6 +204,7 @@ def make_eda_tools_agent(
         user_instructions: str
         data_raw: dict
         eda_artifacts: dict
+        tool_calls: list
 
     def exploratory_agent(state):
         print(format_agent_name(AGENT_NAME))
@@ -205,6 +219,7 @@ def make_eda_tools_agent(
             tools=tool_node,
             state_schema=GraphState,
             **create_react_agent_kwargs,
+            checkpointer=checkpointer,
         )
         
         response = eda_agent.invoke(
@@ -229,11 +244,14 @@ 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)
@@ -241,5 +259,9 @@ def make_eda_tools_agent(
     workflow.add_edge(START, "exploratory_agent")
     workflow.add_edge("exploratory_agent", END)
     
-    app = workflow.compile()
+    app = workflow.compile(
+        checkpointer=checkpointer,
+        name=AGENT_NAME,
+    )
+    
     return app

ai_data_science_team/ml_agents/h2o_ml_agent.py

@@ -5,7 +5,7 @@
 
 import os
 import json
-from typing import TypedDict, Annotated, Sequence, Literal
+from typing import TypedDict, Annotated, Sequence, Literal, Optional
 import operator
 
 import pandas as pd
@@ -14,7 +14,7 @@ from IPython.display import Markdown
 from langchain.prompts import PromptTemplate
 from langchain_core.messages import BaseMessage
 
-from langgraph.types import Command
+from langgraph.types import Command, Checkpointer
 from langgraph.checkpoint.memory import MemorySaver
 
 from ai_data_science_team.templates import(
@@ -79,6 +79,8 @@ class H2OMLAgent(BaseAgent):
         Name of the MLflow experiment (created if doesn't exist).
     mlflow_run_name : str, default None
         A custom name for the MLflow run.
+    checkpointer : langgraph.checkpoint.memory.MemorySaver, optional
+        A checkpointer object for saving the agent's state. Defaults to None.
     
     
     Methods
@@ -176,6 +178,7 @@ class H2OMLAgent(BaseAgent):
         mlflow_tracking_uri=None,
         mlflow_experiment_name="H2O AutoML",
         mlflow_run_name=None,
+        checkpointer: Optional[Checkpointer]=None,
     ):
         self._params = {
             "model": model,
@@ -193,6 +196,7 @@ class H2OMLAgent(BaseAgent):
             "mlflow_tracking_uri": mlflow_tracking_uri,
             "mlflow_experiment_name": mlflow_experiment_name,
             "mlflow_run_name": mlflow_run_name,
+            "checkpointer": checkpointer,
         }
         self._compiled_graph = self._make_compiled_graph()
         self.response = None
@@ -350,6 +354,7 @@ def make_h2o_ml_agent(
     mlflow_tracking_uri=None,
     mlflow_experiment_name="H2O AutoML",
     mlflow_run_name=None,
+    checkpointer=None,
 ):
     """
     Creates a machine learning agent that uses H2O for AutoML. 
@@ -384,6 +389,12 @@ def make_h2o_ml_agent(
             "    pip install h2o\n\n"
             "Visit https://docs.h2o.ai/h2o/latest-stable/h2o-docs/downloading.html for details."
         ) from e
+        
+    if human_in_the_loop:
+        if checkpointer is None:
+            print("Human in the loop is enabled. A checkpointer is required. Setting to MemorySaver().")
+            checkpointer = MemorySaver()
+        
 
     # Define GraphState
     class GraphState(TypedDict):
@@ -844,9 +855,10 @@ def make_h2o_ml_agent(
         retry_count_key="retry_count",
         human_in_the_loop=human_in_the_loop,
         human_review_node_name="human_review",  
-        checkpointer=MemorySaver(),
+        checkpointer=checkpointer,
         bypass_recommended_steps=bypass_recommended_steps,
         bypass_explain_code=bypass_explain_code,
+        agent_name=AGENT_NAME,
     )
 
     return app

ai_data_science_team/ml_agents/mlflow_tools_agent.py

@@ -10,6 +10,7 @@ from langchain_core.messages import BaseMessage, AIMessage
 
 from langgraph.prebuilt import create_react_agent, ToolNode
 from langgraph.prebuilt.chat_agent_executor import AgentState
+from langgraph.types import Checkpointer
 from langgraph.graph import START, END, StateGraph
 
 from ai_data_science_team.templates import BaseAgent
@@ -27,6 +28,7 @@ from ai_data_science_team.tools.mlflow import (
     mlflow_search_registered_models,
     mlflow_get_model_version_details,
 )
+from ai_data_science_team.utils.messages import get_tool_call_names
 
 AGENT_NAME = "mlflow_tools_agent"
 
@@ -67,6 +69,8 @@ class MLflowToolsAgent(BaseAgent):
         Additional keyword arguments to pass to the create_react_agent function.
     invoke_react_agent_kwargs : dict
         Additional keyword arguments to pass to the invoke method of the react agent.
+    checkpointer : langchain.checkpointing.Checkpointer, optional
+        A checkpointer to use for saving and loading the agent's state. Defaults to None.
     
     Methods:
     --------
@@ -118,6 +122,7 @@ class MLflowToolsAgent(BaseAgent):
         mlflow_registry_uri: Optional[str]=None,
         create_react_agent_kwargs: Optional[Dict]={},
         invoke_react_agent_kwargs: Optional[Dict]={},
+        checkpointer: Optional[Checkpointer]=None,
     ):
         self._params = {
             "model": model,
@@ -125,6 +130,7 @@ class MLflowToolsAgent(BaseAgent):
             "mlflow_registry_uri": mlflow_registry_uri,
             "create_react_agent_kwargs": create_react_agent_kwargs,
             "invoke_react_agent_kwargs": invoke_react_agent_kwargs,
+            "checkpointer": checkpointer,            
         }
         self._compiled_graph = self._make_compiled_graph()
         self.response = None
@@ -228,6 +234,12 @@ class MLflowToolsAgent(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"]
             
     
     
@@ -238,6 +250,7 @@ def make_mlflow_tools_agent(
     mlflow_registry_uri: str=None,
     create_react_agent_kwargs: Optional[Dict]={},
     invoke_react_agent_kwargs: Optional[Dict]={},
+    checkpointer: Optional[Checkpointer]=None,
 ):
     """
     MLflow Tool Calling Agent
@@ -254,6 +267,8 @@ def make_mlflow_tools_agent(
         Additional keyword arguments to pass to the agent's create_react_agent method.
     invoke_react_agent_kwargs : dict, optional
         Additional keyword arguments to pass to the agent's invoke method.
+    checkpointer : langchain.checkpointing.Checkpointer, optional
+        A checkpointer to use for saving and loading the agent's state. Defaults to None.
         
     Returns
     -------
@@ -296,6 +311,7 @@ def make_mlflow_tools_agent(
             model, 
             tools=tool_node, 
             state_schema=GraphState,
+            checkpointer=checkpointer,
             **create_react_agent_kwargs,
         )
         
@@ -330,10 +346,13 @@ def make_mlflow_tools_agent(
             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,
             "mlflow_artifacts": last_tool_artifact,
+            "tool_calls": tool_calls,
         }
 
     
@@ -344,7 +363,10 @@ def make_mlflow_tools_agent(
     workflow.add_edge(START, "mlflow_tools_agent")
     workflow.add_edge("mlflow_tools_agent", END)
     
-    app = workflow.compile()
+    app = workflow.compile(
+        checkpointer=checkpointer,
+        name=AGENT_NAME,
+    )
 
     return app
     

ai_data_science_team/multiagents/__init__.py

@@ -1 +1,2 @@
-from ai_data_science_team.multiagents.sql_data_analyst import SQLDataAnalyst, make_sql_data_analyst
+from ai_data_science_team.multiagents.sql_data_analyst import SQLDataAnalyst, make_sql_data_analyst
+from ai_data_science_team.multiagents.pandas_data_analyst import PandasDataAnalyst, make_pandas_data_analyst

ai_data_science_team/multiagents/pandas_data_analyst.py

@@ -0,0 +1,305 @@
+from langchain_core.messages import BaseMessage
+from langchain.prompts import PromptTemplate
+from langchain_core.output_parsers import JsonOutputParser
+from langgraph.types import Checkpointer
+from langgraph.graph import START, END, StateGraph
+from langgraph.graph.state import CompiledStateGraph
+
+from typing import TypedDict, Annotated, Sequence, Union
+import operator
+
+import pandas as pd
+import json
+from IPython.display import Markdown
+
+from ai_data_science_team.templates import BaseAgent
+from ai_data_science_team.agents import DataWranglingAgent, DataVisualizationAgent
+from ai_data_science_team.utils.plotly import plotly_from_dict
+from ai_data_science_team.utils.regex import remove_consecutive_duplicates, get_generic_summary
+
+AGENT_NAME = "pandas_data_analyst"
+
+class PandasDataAnalyst(BaseAgent):
+    """
+    PandasDataAnalyst is a multi-agent class that combines data wrangling and visualization capabilities.
+
+    Parameters:
+    -----------
+    model:
+        The language model to be used for the agents.
+    data_wrangling_agent: DataWranglingAgent
+        The Data Wrangling Agent for transforming raw data.
+    data_visualization_agent: DataVisualizationAgent
+        The Data Visualization Agent for generating plots.
+    checkpointer: Checkpointer (optional)
+        The checkpointer to save the state of the multi-agent system.
+
+    Methods:
+    --------
+    ainvoke_agent(user_instructions, data_raw, **kwargs)
+        Asynchronously invokes the multi-agent with user instructions and raw data.
+    invoke_agent(user_instructions, data_raw, **kwargs)
+        Synchronously invokes the multi-agent with user instructions and raw data.
+    get_data_wrangled()
+        Returns the wrangled data as a Pandas DataFrame.
+    get_plotly_graph()
+        Returns the Plotly graph as a Plotly object.
+    get_data_wrangler_function(markdown=False)
+        Returns the data wrangling function as a string, optionally in Markdown.
+    get_data_visualization_function(markdown=False)
+        Returns the data visualization function as a string, optionally in Markdown.
+    """
+
+    def __init__(
+        self,
+        model,
+        data_wrangling_agent: DataWranglingAgent,
+        data_visualization_agent: DataVisualizationAgent,
+        checkpointer: Checkpointer = None,
+    ):
+        self._params = {
+            "model": model,
+            "data_wrangling_agent": data_wrangling_agent,
+            "data_visualization_agent": data_visualization_agent,
+            "checkpointer": checkpointer,
+        }
+        self._compiled_graph = self._make_compiled_graph()
+        self.response = None
+
+    def _make_compiled_graph(self):
+        """Create or rebuild the compiled graph. Resets response to None."""
+        self.response = None
+        return make_pandas_data_analyst(
+            model=self._params["model"],
+            data_wrangling_agent=self._params["data_wrangling_agent"]._compiled_graph,
+            data_visualization_agent=self._params["data_visualization_agent"]._compiled_graph,
+            checkpointer=self._params["checkpointer"],
+        )
+
+    def update_params(self, **kwargs):
+        """Updates parameters and rebuilds the compiled graph."""
+        for k, v in kwargs.items():
+            self._params[k] = v
+        self._compiled_graph = self._make_compiled_graph()
+
+    async def ainvoke_agent(self, user_instructions, data_raw: Union[pd.DataFrame, dict, list], max_retries: int = 3, retry_count: int = 0, **kwargs):
+        """Asynchronously invokes the multi-agent."""
+        response = await self._compiled_graph.ainvoke({
+            "user_instructions": user_instructions,
+            "data_raw": self._convert_data_input(data_raw),
+            "max_retries": max_retries,
+            "retry_count": retry_count,
+        }, **kwargs)
+        if response.get("messages"):
+            response["messages"] = remove_consecutive_duplicates(response["messages"])
+        self.response = response
+
+    def invoke_agent(self, user_instructions, data_raw: Union[pd.DataFrame, dict, list], max_retries: int = 3, retry_count: int = 0, **kwargs):
+        """Synchronously invokes the multi-agent."""
+        response = self._compiled_graph.invoke({
+            "user_instructions": user_instructions,
+            "data_raw": self._convert_data_input(data_raw),
+            "max_retries": max_retries,
+            "retry_count": retry_count,
+        }, **kwargs)
+        if response.get("messages"):
+            response["messages"] = remove_consecutive_duplicates(response["messages"])
+        self.response = response
+
+    def get_data_wrangled(self):
+        """Returns the wrangled data as a Pandas DataFrame."""
+        if self.response and self.response.get("data_wrangled"):
+            return pd.DataFrame(self.response.get("data_wrangled"))
+
+    def get_plotly_graph(self):
+        """Returns the Plotly graph as a Plotly object."""
+        if self.response and self.response.get("plotly_graph"):
+            return plotly_from_dict(self.response.get("plotly_graph"))
+
+    def get_data_wrangler_function(self, markdown=False):
+        """Returns the data wrangling function as a string."""
+        if self.response and self.response.get("data_wrangler_function"):
+            code = self.response.get("data_wrangler_function")
+            return Markdown(f"```python\n{code}\n```") if markdown else code
+
+    def get_data_visualization_function(self, markdown=False):
+        """Returns the data visualization function as a string."""
+        if self.response and self.response.get("data_visualization_function"):
+            code = self.response.get("data_visualization_function")
+            return Markdown(f"```python\n{code}\n```") if markdown else code
+
+    def get_workflow_summary(self, markdown=False):
+        """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)]
+            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)
+            return Markdown(summary) if markdown else summary
+
+    @staticmethod
+    def _convert_data_input(data_raw: Union[pd.DataFrame, dict, list]) -> Union[dict, list]:
+        """Converts input data to the expected format (dict or list of dicts)."""
+        if isinstance(data_raw, pd.DataFrame):
+            return data_raw.to_dict()
+        if isinstance(data_raw, dict):
+            return data_raw
+        if isinstance(data_raw, list):
+            return [item.to_dict() if isinstance(item, pd.DataFrame) else item for item in data_raw]
+        raise ValueError("data_raw must be a DataFrame, dict, or list of DataFrames/dicts")
+
+def make_pandas_data_analyst(
+    model,
+    data_wrangling_agent: CompiledStateGraph,
+    data_visualization_agent: CompiledStateGraph,
+    checkpointer: Checkpointer = None
+):
+    """
+    Creates a multi-agent system that wrangles data and optionally visualizes it.
+
+    Parameters:
+    -----------
+    model: The language model to be used.
+    data_wrangling_agent: CompiledStateGraph
+        The Data Wrangling Agent.
+    data_visualization_agent: CompiledStateGraph
+        The Data Visualization Agent.
+    checkpointer: Checkpointer (optional)
+        The checkpointer to save the state.
+
+    Returns:
+    --------
+    CompiledStateGraph: The compiled multi-agent system.
+    """
+    
+    llm = model
+    
+    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:
+        
+        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.
+        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.
+        
+        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'.  
+        
+        If a 'chart' is requested, return the 'user_instructions_data_visualization'. If 'None' is found, return None.
+        
+        Return JSON with 'user_instructions_data_wrangling', 'user_instructions_data_visualization' and 'routing_preprocessor_decision'.
+        
+        INITIAL_USER_QUESTION: {user_instructions}
+        """,
+        input_variables=["user_instructions"]
+    )
+
+    routing_preprocessor = routing_preprocessor_prompt | llm | JsonOutputParser()
+
+    class PrimaryState(TypedDict):
+        messages: Annotated[Sequence[BaseMessage], operator.add]
+        user_instructions: str
+        user_instructions_data_wrangling: str
+        user_instructions_data_visualization: str
+        routing_preprocessor_decision: str
+        data_raw: Union[dict, list]
+        data_wrangled: dict
+        data_wrangler_function: str
+        data_visualization_function: str
+        plotly_graph: dict
+        plotly_error: str
+        max_retries: int
+        retry_count: int
+        
+        
+    def preprocess_routing(state: PrimaryState):
+        print("---PANDAS DATA ANALYST---")
+        print("*************************")
+        print("---PREPROCESS ROUTER---")
+        question = state.get("user_instructions")
+        
+        # Chart Routing and SQL Prep
+        response = routing_preprocessor.invoke({"user_instructions": question})
+        
+        return {
+            "user_instructions_data_wrangling": response.get('user_instructions_data_wrangling'),
+            "user_instructions_data_visualization": response.get('user_instructions_data_visualization'),
+            "routing_preprocessor_decision": response.get('routing_preprocessor_decision'),
+        }
+    
+    def router_chart_or_table(state: PrimaryState):
+        print("---ROUTER: CHART OR TABLE---")
+        return "chart" if state.get('routing_preprocessor_decision') == "chart" else "table"
+    
+    
+    def invoke_data_wrangling_agent(state: PrimaryState):
+        
+        response = data_wrangling_agent.invoke({
+            "user_instructions": state.get("user_instructions_data_wrangling"),
+            "data_raw": state.get("data_raw"),
+            "max_retries": state.get("max_retries"),
+            "retry_count": state.get("retry_count"),
+        })
+
+        return {
+            "messages": response.get("messages"),
+            "data_wrangled": response.get("data_wrangled"),
+            "data_wrangler_function": response.get("data_wrangler_function"),
+            "plotly_error": response.get("data_visualization_error"),
+            
+        }
+        
+    def invoke_data_visualization_agent(state: PrimaryState):
+        
+        response = data_visualization_agent.invoke({
+            "user_instructions": state.get("user_instructions_data_visualization"),
+            "data_raw": state.get("data_wrangled") if state.get("data_wrangled") else state.get("data_raw"),
+            "max_retries": state.get("max_retries"),
+            "retry_count": state.get("retry_count"),
+        })
+        
+        return {
+            "messages": response.get("messages"),
+            "data_visualization_function": response.get("data_visualization_function"),
+            "plotly_graph": response.get("plotly_graph"),
+            "plotly_error": response.get("data_visualization_error"),
+        }
+
+    def route_printer(state: PrimaryState):
+        print("---ROUTE PRINTER---")
+        print(f"    Route: {state.get('routing_preprocessor_decision')}")
+        print("---END---")
+        return {}
+    
+    workflow = StateGraph(PrimaryState)
+    
+    workflow.add_node("routing_preprocessor", preprocess_routing)
+    workflow.add_node("data_wrangling_agent", invoke_data_wrangling_agent)
+    workflow.add_node("data_visualization_agent", invoke_data_visualization_agent)
+    workflow.add_node("route_printer", route_printer)
+
+    workflow.add_edge(START, "routing_preprocessor")
+    workflow.add_edge("routing_preprocessor", "data_wrangling_agent")
+    
+    workflow.add_conditional_edges(
+        "data_wrangling_agent", 
+        router_chart_or_table,
+        {
+            "chart": "data_visualization_agent",
+            "table": "route_printer"
+        }
+    )
+    
+    workflow.add_edge("data_visualization_agent", "route_printer")
+    workflow.add_edge("route_printer", END)
+
+    app = workflow.compile(
+        checkpointer=checkpointer, 
+        name=AGENT_NAME
+    )
+    
+    return app