ai-data-science-team 0.0.0.9013__py3-none-any.whl → 0.0.0.9015__py3-none-any.whl

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/_version.py

@@ -1 +1 @@
-__version__ = "0.0.0.9013"
+__version__ = "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/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