ai-data-science-team 0.0.0.9007__py3-none-any.whl → 0.0.0.9008__py3-none-any.whl

ai_data_science_team/_version.py

@@ -1 +1 @@
-__version__ = "0.0.0.9007"
+__version__ = "0.0.0.9008"

ai_data_science_team/agents/__init__.py

@@ -1,6 +1,6 @@
 from ai_data_science_team.agents.data_cleaning_agent import make_data_cleaning_agent, DataCleaningAgent
-from ai_data_science_team.agents.feature_engineering_agent import make_feature_engineering_agent
-from ai_data_science_team.agents.data_wrangling_agent import make_data_wrangling_agent
-from ai_data_science_team.agents.sql_database_agent import make_sql_database_agent
-from ai_data_science_team.agents.data_visualization_agent import make_data_visualization_agent
+from ai_data_science_team.agents.feature_engineering_agent import make_feature_engineering_agent, FeatureEngineeringAgent
+from ai_data_science_team.agents.data_wrangling_agent import make_data_wrangling_agent, DataWranglingAgent
+from ai_data_science_team.agents.sql_database_agent import make_sql_database_agent, SQLDatabaseAgent
+from ai_data_science_team.agents.data_visualization_agent import make_data_visualization_agent, DataVisualizationAgent
 

ai_data_science_team/agents/data_cleaning_agent.py

@@ -13,21 +13,22 @@ from langchain_core.messages import BaseMessage
 from langgraph.types import Command
 from langgraph.checkpoint.memory import MemorySaver
 
-from langgraph.graph.state import CompiledStateGraph
-
 import os
 import io
 import pandas as pd
 
+from IPython.display import Markdown
+
 from ai_data_science_team.templates import(
     node_func_execute_agent_code_on_data, 
     node_func_human_review,
     node_func_fix_agent_code, 
     node_func_explain_agent_code, 
-    create_coding_agent_graph
+    create_coding_agent_graph,
+    BaseAgent,
 )
 from ai_data_science_team.tools.parsers import PythonOutputParser
-from ai_data_science_team.tools.regex import relocate_imports_inside_function, add_comments_to_top, format_agent_name
+from ai_data_science_team.tools.regex import relocate_imports_inside_function, add_comments_to_top, format_agent_name, format_recommended_steps
 from ai_data_science_team.tools.metadata import get_dataframe_summary
 from ai_data_science_team.tools.logging import log_ai_function
 
@@ -38,7 +39,109 @@ LOG_PATH = os.path.join(os.getcwd(), "logs/")
 
 
 # Class
-class DataCleaningAgent(CompiledStateGraph):
+class DataCleaningAgent(BaseAgent):
+    """
+    Creates a data cleaning agent that can process datasets based on user-defined instructions or default cleaning steps. 
+    The agent generates a Python function to clean the dataset, performs the cleaning, and logs the process, including code 
+    and errors. It is designed to facilitate reproducible and customizable data cleaning workflows.
+
+    The agent performs the following default cleaning steps unless instructed otherwise:
+
+    - Removing columns with more than 40% missing values.
+    - Imputing missing values with the mean for numeric columns.
+    - Imputing missing values with the mode for categorical columns.
+    - Converting columns to appropriate data types.
+    - Removing duplicate rows.
+    - Removing rows with missing values.
+    - Removing rows with extreme outliers (values 3x the interquartile range).
+
+    User instructions can modify, add, or remove any of these steps to tailor the cleaning process.
+
+    Parameters
+    ----------
+    model : langchain.llms.base.LLM
+        The language model used to generate the data cleaning function.
+    n_samples : int, optional
+        Number of samples used when summarizing the dataset. Defaults to 30. Reducing this number can help 
+        avoid exceeding the model's token limits.
+    log : bool, optional
+        Whether to log the generated code and errors. Defaults to False.
+    log_path : str, optional
+        Directory path for storing log files. Defaults to None.
+    file_name : str, optional
+        Name of the file for saving the generated response. Defaults to "data_cleaner.py".
+    function_name : str, optional
+        Name of the generated data cleaning function. Defaults to "data_cleaner".
+    overwrite : bool, optional
+        Whether to overwrite the log file if it exists. If False, a unique file name is created. Defaults to True.
+    human_in_the_loop : bool, optional
+        Enables user review of data cleaning instructions. Defaults to False.
+    bypass_recommended_steps : bool, optional
+        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.
+
+    Methods
+    -------
+    update_params(**kwargs)
+        Updates the agent's parameters and rebuilds the compiled state graph.
+    ainvoke_agent(user_instructions: str, data_raw: pd.DataFrame, max_retries=3, retry_count=0)
+        Cleans the provided dataset asynchronously based on user instructions.
+    invoke_agent(user_instructions: str, data_raw: pd.DataFrame, max_retries=3, retry_count=0)
+        Cleans the provided dataset synchronously based on user instructions.
+    explain_cleaning_steps()
+        Returns an explanation of the cleaning steps performed by the agent.
+    get_log_summary()
+        Retrieves a summary of logged operations if logging is enabled.
+    get_state_keys()
+        Returns a list of keys from the state graph response.
+    get_state_properties()
+        Returns detailed properties of the state graph response.
+    get_data_cleaned()
+        Retrieves the cleaned dataset as a pandas DataFrame.
+    get_data_raw()
+        Retrieves the raw dataset as a pandas DataFrame.
+    get_data_cleaner_function()
+        Retrieves the generated Python function used for cleaning the data.
+    get_recommended_cleaning_steps()
+        Retrieves the agent's recommended cleaning steps.
+    get_response()
+        Returns the response from the agent as a dictionary.
+    show()
+        Displays the agent's mermaid diagram.
+
+    Examples
+    --------
+    ```python
+    import pandas as pd
+    from langchain_openai import ChatOpenAI
+    from ai_data_science_team.agents import DataCleaningAgent
+
+    llm = ChatOpenAI(model="gpt-4o-mini")
+
+    data_cleaning_agent = DataCleaningAgent(
+        model=llm, n_samples=50, log=True, log_path="logs", human_in_the_loop=True
+    )
+
+    df = pd.read_csv("https://raw.githubusercontent.com/business-science/ai-data-science-team/refs/heads/master/data/churn_data.csv")
+
+    data_cleaning_agent.invoke_agent(
+        user_instructions="Don't remove outliers when cleaning the data.",
+        data_raw=df,
+        max_retries=3,
+        retry_count=0
+    )
+
+    cleaned_data = data_cleaning_agent.get_data_cleaned()
+    
+    response = data_cleaning_agent.response
+    ```
+    
+    Returns
+    --------
+    DataCleaningAgent : langchain.graphs.CompiledStateGraph 
+        A data cleaning agent implemented as a compiled state graph. 
+    """
     
     def __init__(
         self, 
@@ -47,6 +150,7 @@ class DataCleaningAgent(CompiledStateGraph):
         log=False, 
         log_path=None, 
         file_name="data_cleaner.py", 
+        function_name="data_cleaner",
         overwrite=True, 
         human_in_the_loop=False, 
         bypass_recommended_steps=False, 
@@ -58,6 +162,7 @@ class DataCleaningAgent(CompiledStateGraph):
             "log": log,
             "log_path": log_path,
             "file_name": file_name,
+            "function_name": function_name,
             "overwrite": overwrite,
             "human_in_the_loop": human_in_the_loop,
             "bypass_recommended_steps": bypass_recommended_steps,
@@ -67,65 +172,70 @@ class DataCleaningAgent(CompiledStateGraph):
         self.response = None
 
     def _make_compiled_graph(self):
-        self.response = None
-        return make_data_cleaning_agent(**self._params)
-
-    def update_params(self, **kwargs):
         """
-        Update one or more parameters at once, then rebuild the compiled graph.
-        e.g. agent.update_params(model=new_llm, n_samples=100)
+        Create the compiled graph for the data cleaning agent. Running this method will reset the response to None.
         """
-        self._params.update(kwargs)
-        self._compiled_graph = self._make_compiled_graph()
+        self.response=None
+        return make_data_cleaning_agent(**self._params)
 
-    def __getattr__(self, name: str):
-        """
-        Delegate attribute access to `_compiled_graph` if `name` is not
-        found in this instance. This 'inherits' methods from the compiled graph.
-        """
-        return getattr(self._compiled_graph, name)
     
-    def ainvoke(self, user_instructions: str, data_raw: pd.DataFrame, max_retries=3, retry_count=0):
+    def ainvoke_agent(self, data_raw: pd.DataFrame, user_instructions: str=None, max_retries:int=3, retry_count:int=0, **kwargs):
         """
-        Cleans the provided dataset based on user instructions.
+        Asynchronously invokes the agent. The response is stored in the response attribute.
 
         Parameters:
-            user_instructions (str): Instructions for data cleaning.
-            data_raw (pd.DataFrame): The raw dataset to be cleaned.
-            max_retries (int): Maximum retry attempts for cleaning.
-            retry_count (int): Current retry attempt.
+        ----------
+            data_raw (pd.DataFrame): 
+                The raw dataset to be cleaned.
+            user_instructions (str): 
+                Instructions for data cleaning agent.
+            max_retries (int): 
+                Maximum retry attempts for cleaning.
+            retry_count (int): 
+                Current retry attempt.
+            **kwargs
+                Additional keyword arguments to pass to ainvoke().
 
         Returns:
+        --------
             None. The response is stored in the response attribute.
         """
-        response = self.ainvoke({
+        response = self._compiled_graph.ainvoke({
             "user_instructions": user_instructions,
             "data_raw": data_raw.to_dict(),
             "max_retries": max_retries,
             "retry_count": retry_count,
-        })
+        }, **kwargs)
         self.response = response
         return None
     
-    def invoke(self, user_instructions: str, data_raw: pd.DataFrame, max_retries=3, retry_count=0):
+    def invoke_agent(self, data_raw: pd.DataFrame, user_instructions: str=None, max_retries:int=3, retry_count:int=0, **kwargs):
         """
-        Cleans the provided dataset based on user instructions.
+        Invokes the agent. The response is stored in the response attribute.
 
         Parameters:
-            user_instructions (str): Instructions for data cleaning.
-            data_raw (pd.DataFrame): The raw dataset to be cleaned.
-            max_retries (int): Maximum retry attempts for cleaning.
-            retry_count (int): Current retry attempt.
+        ----------
+            data_raw (pd.DataFrame): 
+                The raw dataset to be cleaned.
+            user_instructions (str): 
+                Instructions for data cleaning agent.
+            max_retries (int): 
+                Maximum retry attempts for cleaning.
+            retry_count (int): 
+                Current retry attempt.
+            **kwargs
+                Additional keyword arguments to pass to invoke().
 
         Returns:
+        --------
             None. The response is stored in the response attribute.
         """
-        response = self.invoke({
+        response = self._compiled_graph.invoke({
             "user_instructions": user_instructions,
             "data_raw": data_raw.to_dict(),
             "max_retries": max_retries,
             "retry_count": retry_count,
-        })
+        },**kwargs)
         self.response = response
         return None
 
@@ -139,30 +249,21 @@ class DataCleaningAgent(CompiledStateGraph):
         messages = self.response.get("messages", [])
         return messages
 
-    def get_log_summary(self):
+    def get_log_summary(self, markdown=False):
         """
         Logs a summary of the agent's operations, if logging is enabled.
         """
         if self.response:
-            if self.log: 
+            if self.response.get('data_cleaner_function_path'):
                 log_details = f"Log Path: {self.response.get('data_cleaner_function_path')}"
-                return log_details
-    
-    def get_state_keys(self):
-        """
-        Returns a list of keys that the state graph returns in a response.
-        """
-        return list(self.get_output_jsonschema()['properties'].keys())
-    
-    def get_state_properties(self):
-        """
-        Returns a list of keys that the state graph returns in a response.
-        """
-        return self.get_output_jsonschema()['properties']
+                if markdown:
+                    return Markdown(log_details) 
+                else:
+                    return log_details
     
     def get_data_cleaned(self):
         """
-        Retrieves the cleaned data stored after running invoke or clean_data methods.
+        Retrieves the cleaned data stored after running invoke_agent or clean_data methods.
         """
         if self.response:
             return pd.DataFrame(self.response.get("data_cleaned"))
@@ -174,15 +275,25 @@ class DataCleaningAgent(CompiledStateGraph):
         if self.response:
             return pd.DataFrame(self.response.get("data_raw"))
     
-    def get_data_cleaner_function(self):
+    def get_data_cleaner_function(self, markdown=False):
         """
         Retrieves the agent's pipeline function.
         """
         if self.response:
-            return self.response.get("data_cleaner_function")
-        
-    
-    
+            if markdown:
+                return Markdown(f"```python\n{self.response.get('data_cleaner_function')}\n```")
+            else:
+                return self.response.get("data_cleaner_function")
+            
+    def get_recommended_cleaning_steps(self, markdown=False):
+        """
+        Retrieves the agent's recommended cleaning steps
+        """
+        if self.response:
+            if markdown:
+                return Markdown(self.response.get('recommended_steps'))
+            else:
+                return self.response.get('recommended_steps')
 
 
 
@@ -194,6 +305,7 @@ def make_data_cleaning_agent(
     log=False, 
     log_path=None, 
     file_name="data_cleaner.py",
+    function_name="data_cleaner",
     overwrite = True, 
     human_in_the_loop=False, 
     bypass_recommended_steps=False, 
@@ -235,6 +347,8 @@ def make_data_cleaning_agent(
         "logs/".
     file_name : str, optional
         The name of the file to save the response to. Defaults to "data_cleaner.py".
+    function_name : str, optional
+        The name of the function that will be generated to clean the data. Defaults to "data_cleaner".
     overwrite : bool, optional
         Whether or not to overwrite the log file if it already exists. If False, a unique file name will be created. 
         Defaults to True.
@@ -275,6 +389,11 @@ def make_data_cleaning_agent(
     """
     llm = model
     
+    # Human in th loop requires recommended steps
+    if bypass_recommended_steps and human_in_the_loop:
+        bypass_recommended_steps = False
+        print("Bypass recommended steps set to False to enable human in the loop.")
+    
     # Setup Log Directory
     if log:
         if log_path is None:
@@ -292,6 +411,7 @@ def make_data_cleaning_agent(
         all_datasets_summary: str
         data_cleaner_function: str
         data_cleaner_function_path: str
+        data_cleaner_file_name: str
         data_cleaner_function_name: str
         data_cleaner_error: str
         max_retries: int
@@ -366,7 +486,7 @@ def make_data_cleaning_agent(
         }) 
         
         return {
-            "recommended_steps": "\n\n# Recommended Data Cleaning Steps:\n" + recommended_steps.content.strip(),
+            "recommended_steps": format_recommended_steps(recommended_steps.content.strip(), heading="# Recommended Data Cleaning Steps:"),
             "all_datasets_summary": all_datasets_summary_str
         }
     
@@ -386,42 +506,44 @@ def make_data_cleaning_agent(
         else:
             all_datasets_summary_str = state.get("all_datasets_summary")
         
+        
         data_cleaning_prompt = PromptTemplate(
             template="""
-            You are a Data Cleaning Agent. Your job is to create a data_cleaner() function that can be run on the data provided using the following recommended steps.
-            
+            You are a Data Cleaning Agent. Your job is to create a {function_name}() function that can be run on the data provided using the following recommended steps.
+
             Recommended Steps:
             {recommended_steps}
-            
+
             You can use Pandas, Numpy, and Scikit Learn libraries to clean the data.
-            
+
             Below are summaries of all datasets provided. Use this information about the data to help determine how to clean the data:
 
             {all_datasets_summary}
-            
-            Return Python code in ```python ``` format with a single function definition, data_cleaner(data_raw), that includes all imports inside the function. 
-            
+
+            Return Python code in ```python``` format with a single function definition, {function_name}(data_raw), that includes all imports inside the function.
+
             Return code to provide the data cleaning function:
-            
-            def data_cleaner(data_raw):
+
+            def {function_name}(data_raw):
                 import pandas as pd
                 import numpy as np
                 ...
                 return data_cleaned
-            
+
             Best Practices and Error Preventions:
-            
+
             Always ensure that when assigning the output of fit_transform() from SimpleImputer to a Pandas DataFrame column, you call .ravel() or flatten the array, because fit_transform() returns a 2D array while a DataFrame column is 1D.
             
             """,
-            input_variables=["recommended_steps", "all_datasets_summary"]
+            input_variables=["recommended_steps", "all_datasets_summary", "function_name"]
         )
 
         data_cleaning_agent = data_cleaning_prompt | llm | PythonOutputParser()
         
         response = data_cleaning_agent.invoke({
             "recommended_steps": state.get("recommended_steps"),
-            "all_datasets_summary": all_datasets_summary_str
+            "all_datasets_summary": all_datasets_summary_str,
+            "function_name": function_name
         })
         
         response = relocate_imports_inside_function(response)
@@ -439,19 +561,37 @@ def make_data_cleaning_agent(
         return {
             "data_cleaner_function" : response,
             "data_cleaner_function_path": file_path,
-            "data_cleaner_function_name": file_name_2,
+            "data_cleaner_file_name": file_name_2,
+            "data_cleaner_function_name": function_name,
             "all_datasets_summary": all_datasets_summary_str
         }
+        
+    # Human Review
+        
+    prompt_text_human_review = "Are the following data cleaning instructions correct? (Answer 'yes' or provide modifications)\n{steps}"
     
-    def human_review(state: GraphState) -> Command[Literal["recommend_cleaning_steps", "create_data_cleaner_code"]]:
-        return node_func_human_review(
-            state=state,
-            prompt_text="Is the following data cleaning instructions correct? (Answer 'yes' or provide modifications)\n{steps}",
-            yes_goto="create_data_cleaner_code",
-            no_goto="recommend_cleaning_steps",
-            user_instructions_key="user_instructions",
-            recommended_steps_key="recommended_steps"            
-        )
+    if not bypass_explain_code:
+        def human_review(state: GraphState) -> Command[Literal["recommend_cleaning_steps", "explain_data_cleaner_code"]]:
+            return node_func_human_review(
+                state=state,
+                prompt_text=prompt_text_human_review,
+                yes_goto= 'explain_data_cleaner_code',
+                no_goto="recommend_cleaning_steps",
+                user_instructions_key="user_instructions",
+                recommended_steps_key="recommended_steps",
+                code_snippet_key="data_cleaner_function",
+            )
+    else:
+        def human_review(state: GraphState) -> Command[Literal["recommend_cleaning_steps", "__end__"]]:
+            return node_func_human_review(
+                state=state,
+                prompt_text=prompt_text_human_review,
+                yes_goto= '__end__',
+                no_goto="recommend_cleaning_steps",
+                user_instructions_key="user_instructions",
+                recommended_steps_key="recommended_steps",
+                code_snippet_key="data_cleaner_function", 
+            )
     
     def execute_data_cleaner_code(state):
         return node_func_execute_agent_code_on_data(
@@ -460,7 +600,7 @@ def make_data_cleaning_agent(
             result_key="data_cleaned",
             error_key="data_cleaner_error",
             code_snippet_key="data_cleaner_function",
-            agent_function_name="data_cleaner",
+            agent_function_name=state.get("data_cleaner_function_name"),
             pre_processing=lambda data: pd.DataFrame.from_dict(data),
             post_processing=lambda df: df.to_dict() if isinstance(df, pd.DataFrame) else df,
             error_message_prefix="An error occurred during data cleaning: "
@@ -468,11 +608,11 @@ def make_data_cleaning_agent(
         
     def fix_data_cleaner_code(state: GraphState):
         data_cleaner_prompt = """
-        You are a Data Cleaning Agent. Your job is to create a data_cleaner() function that can be run on the data provided. The function is currently broken and needs to be fixed.
+        You are a Data Cleaning Agent. Your job is to create a {function_name}() function that can be run on the data provided. The function is currently broken and needs to be fixed.
         
-        Make sure to only return the function definition for data_cleaner().
+        Make sure to only return the function definition for {function_name}().
         
-        Return Python code in ```python``` format with a single function definition, data_cleaner(data_raw), that includes all imports inside the function.
+        Return Python code in ```python``` format with a single function definition, {function_name}(data_raw), that includes all imports inside the function.
         
         This is the broken code (please fix): 
         {code_snippet}
@@ -490,6 +630,7 @@ def make_data_cleaning_agent(
             agent_name=AGENT_NAME,
             log=log,
             file_path=state.get("data_cleaner_function_path"),
+            function_name=state.get("data_cleaner_function_name"),
         )
     
     def explain_data_cleaner_code(state: GraphState):