ai-data-science-team 0.0.0.9013__tar.gz → 0.0.0.9015__tar.gz

{ai_data_science_team-0.0.0.9013/ai_data_science_team.egg-info → ai_data_science_team-0.0.0.9015}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: ai-data-science-team
-Version: 0.0.0.9013
+Version: 0.0.0.9015
 Summary: Build and run an AI-powered data science team.
 Home-page: https://github.com/business-science/ai-data-science-team
 Author: Matt Dancho
@@ -18,7 +18,7 @@ Requires-Dist: langchain
 Requires-Dist: langchain_community
 Requires-Dist: langchain_openai
 Requires-Dist: langchain_experimental
-Requires-Dist: langgraph>=0.2.57
+Requires-Dist: langgraph>=0.2.74
 Requires-Dist: openai
 Requires-Dist: pandas
 Requires-Dist: sqlalchemy
@@ -136,9 +136,9 @@ This project is a work in progress. New data science agents will be released soo
 
 ### NEW: Multi-Agents
 
-This is the internals of the SQL Data Analyst Agent that connects to SQL databases to pull data into the data science environment. It creates pipelines to automate data extraction, performs Joins, Aggregations, and other SQL Query operations. And it includes a Data Visualization Agent that creates visualizations to help you understand your data.:
+**🔥 Pandas Data Analyst Agent:** Combines the ability to wrangle, transform, and analyze data with an optional data visualization agent that can create interactive plots.
 
-![Business Intelligence SQL Agent](/img/multi_agent_sql_data_visualization.jpg)
+![Business Intelligence SQL Agent](/img/multi_agent_pandas_data_analyst.jpg)
 
 ### Data Science Apps
 
@@ -182,7 +182,8 @@ This is a top secret project I'm working on. It's a multi-agent data science app
 
 #### Multi-Agents
 
-1. **SQL Data Analyst Agent:** Connects to SQL databases to pull data into the data science environment. Creates pipelines to automate data extraction. Performs Joins, Aggregations, and other SQL Query operations. Includes a Data Visualization Agent that creates visualizations to help you understand your data. [See Example](https://github.com/business-science/ai-data-science-team/blob/master/examples/multiagents/sql_data_analyst.ipynb)
+1. **🔥🔥 Pandas Data Analyst Agent:** Combines the ability to wrangle, transform, and analyze data with an optional data visualization agent that can create interactive plots. [See Example](https://github.com/business-science/ai-data-science-team/blob/master/examples/multiagents/pandas_data_analyst.ipynb)
+2. **🔥🔥 SQL Data Analyst Agent:** Connects to SQL databases to pull data into the data science environment. Creates pipelines to automate data extraction. Performs Joins, Aggregations, and other SQL Query operations. Includes a Data Visualization Agent that creates visualizations to help you understand your data. [See Example](https://github.com/business-science/ai-data-science-team/blob/master/examples/multiagents/sql_data_analyst.ipynb)
 
 ### Agents Coming Soon
 

{ai_data_science_team-0.0.0.9013 → ai_data_science_team-0.0.0.9015}/README.md

@@ -82,9 +82,9 @@ This project is a work in progress. New data science agents will be released soo
 
 ### NEW: Multi-Agents
 
-This is the internals of the SQL Data Analyst Agent that connects to SQL databases to pull data into the data science environment. It creates pipelines to automate data extraction, performs Joins, Aggregations, and other SQL Query operations. And it includes a Data Visualization Agent that creates visualizations to help you understand your data.:
+**🔥 Pandas Data Analyst Agent:** Combines the ability to wrangle, transform, and analyze data with an optional data visualization agent that can create interactive plots.
 
-![Business Intelligence SQL Agent](/img/multi_agent_sql_data_visualization.jpg)
+![Business Intelligence SQL Agent](/img/multi_agent_pandas_data_analyst.jpg)
 
 ### Data Science Apps
 
@@ -128,7 +128,8 @@ This is a top secret project I'm working on. It's a multi-agent data science app
 
 #### Multi-Agents
 
-1. **SQL Data Analyst Agent:** Connects to SQL databases to pull data into the data science environment. Creates pipelines to automate data extraction. Performs Joins, Aggregations, and other SQL Query operations. Includes a Data Visualization Agent that creates visualizations to help you understand your data. [See Example](https://github.com/business-science/ai-data-science-team/blob/master/examples/multiagents/sql_data_analyst.ipynb)
+1. **🔥🔥 Pandas Data Analyst Agent:** Combines the ability to wrangle, transform, and analyze data with an optional data visualization agent that can create interactive plots. [See Example](https://github.com/business-science/ai-data-science-team/blob/master/examples/multiagents/pandas_data_analyst.ipynb)
+2. **🔥🔥 SQL Data Analyst Agent:** Connects to SQL databases to pull data into the data science environment. Creates pipelines to automate data extraction. Performs Joins, Aggregations, and other SQL Query operations. Includes a Data Visualization Agent that creates visualizations to help you understand your data. [See Example](https://github.com/business-science/ai-data-science-team/blob/master/examples/multiagents/sql_data_analyst.ipynb)
 
 ### Agents Coming Soon
 

ai_data_science_team-0.0.0.9015/ai_data_science_team/__init__.py

@@ -0,0 +1,22 @@
+from ai_data_science_team.agents import (
+    DataCleaningAgent,
+    DataLoaderToolsAgent,
+    DataVisualizationAgent,
+    SQLDatabaseAgent,
+    DataWranglingAgent,
+    FeatureEngineeringAgent,
+)
+
+from ai_data_science_team.ds_agents import (
+    EDAToolsAgent,
+)
+
+from ai_data_science_team.ml_agents import (
+    H2OMLAgent,
+    MLflowToolsAgent,
+)
+
+from ai_data_science_team.multiagents import (
+    SQLDataAnalyst, 
+    PandasDataAnalyst, 
+)

ai_data_science_team-0.0.0.9015/ai_data_science_team/_version.py

@@ -0,0 +1 @@
+__version__ = "0.0.0.9015"

{ai_data_science_team-0.0.0.9013 → ai_data_science_team-0.0.0.9015}/ai_data_science_team/agents/data_cleaning_agent.py

@@ -12,6 +12,7 @@ from langchain_core.messages import BaseMessage
 
 from langgraph.types import Command
 from langgraph.checkpoint.memory import MemorySaver
+from langgraph.types import Checkpointer
 
 import os
 import json
@@ -85,6 +86,8 @@ class DataCleaningAgent(BaseAgent):
         If True, skips the default recommended cleaning steps. Defaults to False.
     bypass_explain_code : bool, optional
         If True, skips the step that provides code explanations. Defaults to False.
+    checkpointer : langgraph.types.Checkpointer, optional
+        Checkpointer to save and load the agent's state. Defaults to None.
 
     Methods
     -------
@@ -159,7 +162,8 @@ class DataCleaningAgent(BaseAgent):
         overwrite=True, 
         human_in_the_loop=False, 
         bypass_recommended_steps=False, 
-        bypass_explain_code=False
+        bypass_explain_code=False,
+        checkpointer: Checkpointer = None
     ):
         self._params = {
             "model": model,
@@ -172,6 +176,7 @@ class DataCleaningAgent(BaseAgent):
             "human_in_the_loop": human_in_the_loop,
             "bypass_recommended_steps": bypass_recommended_steps,
             "bypass_explain_code": bypass_explain_code,
+            "checkpointer": checkpointer
         }
         self._compiled_graph = self._make_compiled_graph()
         self.response = None
@@ -320,7 +325,8 @@ def make_data_cleaning_agent(
     overwrite = True, 
     human_in_the_loop=False, 
     bypass_recommended_steps=False, 
-    bypass_explain_code=False
+    bypass_explain_code=False,
+    checkpointer: Checkpointer = None
 ):
     """
     Creates a data cleaning agent that can be run on a dataset. The agent can be used to clean a dataset in a variety of
@@ -369,6 +375,8 @@ def make_data_cleaning_agent(
         Bypass the recommendation step, by default False
     bypass_explain_code : bool, optional
         Bypass the code explanation step, by default False.
+    checkpointer : langgraph.types.Checkpointer, optional
+        Checkpointer to save and load the agent's state. Defaults to None.
         
     Examples
     -------
@@ -400,6 +408,11 @@ def make_data_cleaning_agent(
     """
     llm = model
     
+    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()
+    
     # Human in th loop requires recommended steps
     if bypass_recommended_steps and human_in_the_loop:
         bypass_recommended_steps = False
@@ -680,9 +693,10 @@ def make_data_cleaning_agent(
         error_key="data_cleaner_error",
         human_in_the_loop=human_in_the_loop,
         human_review_node_name="human_review",
-        checkpointer=MemorySaver() if human_in_the_loop else None,
+        checkpointer=checkpointer,
         bypass_recommended_steps=bypass_recommended_steps,
         bypass_explain_code=bypass_explain_code,
+        agent_name=AGENT_NAME,
     )
 
     return app

{ai_data_science_team-0.0.0.9013 → ai_data_science_team-0.0.0.9015}/ai_data_science_team/agents/data_loader_tools_agent.py

@@ -13,6 +13,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
@@ -50,6 +51,8 @@ class DataLoaderToolsAgent(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 : langgraph.types.Checkpointer
+        A checkpointer to use for saving and loading the agent's state.
         
     Methods:
     --------
@@ -73,11 +76,13 @@ class DataLoaderToolsAgent(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
@@ -188,6 +193,7 @@ def make_data_loader_tools_agent(
     model: Any,
     create_react_agent_kwargs: Optional[Dict]={},
     invoke_react_agent_kwargs: Optional[Dict]={},
+    checkpointer: Optional[Checkpointer]=None,
 ):
     """
     Creates a Data Loader Agent that can interact with data loading tools.
@@ -200,6 +206,8 @@ def make_data_loader_tools_agent(
         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 : langgraph.types.Checkpointer
+        A checkpointer to use for saving and loading the agent's state.
     
     Returns:
     --------
@@ -228,6 +236,7 @@ def make_data_loader_tools_agent(
             model, 
             tools=tool_node, 
             state_schema=GraphState,
+            checkpointer=checkpointer,
             **create_react_agent_kwargs,
         )
         
@@ -277,7 +286,10 @@ def make_data_loader_tools_agent(
     workflow.add_edge(START, "data_loader_agent")
     workflow.add_edge("data_loader_agent", END)
     
-    app = workflow.compile()
+    app = workflow.compile(
+        checkpointer=checkpointer,
+        name=AGENT_NAME,
+    )
 
     return app