ai-data-science-team 0.0.0.9009__py3-none-any.whl → 0.0.0.9011__py3-none-any.whl

ai_data_science_team/_version.py

@@ -1 +1 @@
-__version__ = "0.0.0.9009"
+__version__ = "0.0.0.9011"

ai_data_science_team/agents/__init__.py

@@ -3,3 +3,4 @@ from ai_data_science_team.agents.feature_engineering_agent import make_feature_e
 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
+from ai_data_science_team.agents.data_loader_tools_agent import make_data_loader_tools_agent, DataLoaderToolsAgent

ai_data_science_team/agents/data_cleaning_agent.py

@@ -27,16 +27,16 @@ from ai_data_science_team.templates import(
     create_coding_agent_graph,
     BaseAgent,
 )
-from ai_data_science_team.tools.parsers import PythonOutputParser
-from ai_data_science_team.tools.regex import (
+from ai_data_science_team.parsers.parsers import PythonOutputParser
+from ai_data_science_team.utils.regex import (
     relocate_imports_inside_function, 
     add_comments_to_top, 
     format_agent_name, 
     format_recommended_steps, 
     get_generic_summary,
 )
-from ai_data_science_team.tools.metadata import get_dataframe_summary
-from ai_data_science_team.tools.logging import log_ai_function
+from ai_data_science_team.tools.dataframe import get_dataframe_summary
+from ai_data_science_team.utils.logging import log_ai_function
 
 # Setup
 AGENT_NAME = "data_cleaning_agent"
@@ -183,7 +183,7 @@ class DataCleaningAgent(BaseAgent):
         self.response=None
         return make_data_cleaning_agent(**self._params)
 
-    def ainvoke_agent(self, data_raw: pd.DataFrame, user_instructions: str=None, max_retries:int=3, retry_count:int=0, **kwargs):
+    async def ainvoke_agent(self, data_raw: pd.DataFrame, user_instructions: str=None, max_retries:int=3, retry_count:int=0, **kwargs):
         """
         Asynchronously invokes the agent. The response is stored in the response attribute.
 
@@ -204,7 +204,7 @@ class DataCleaningAgent(BaseAgent):
         --------
             None. The response is stored in the response attribute.
         """
-        response = self._compiled_graph.ainvoke({
+        response = await self._compiled_graph.ainvoke({
             "user_instructions": user_instructions,
             "data_raw": data_raw.to_dict(),
             "max_retries": max_retries,

ai_data_science_team/agents/data_loader_tools_agent.py

@@ -0,0 +1,272 @@
+
+
+
+from typing import Any, Optional, Annotated, Sequence, List, Dict
+import operator
+
+import pandas as pd
+import os
+
+from IPython.display import Markdown
+
+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 ai_data_science_team.templates import BaseAgent
+from ai_data_science_team.utils.regex import format_agent_name
+from ai_data_science_team.tools.data_loader import (
+    load_directory,
+    load_file,
+    list_directory_contents,
+    list_directory_recursive,
+    get_file_info,
+    search_files_by_pattern,
+)
+
+AGENT_NAME = "data_loader_tools_agent"
+
+tools = [
+    load_directory,
+    load_file,
+    list_directory_contents,
+    list_directory_recursive,
+    get_file_info,
+    search_files_by_pattern,
+]
+
+class DataLoaderToolsAgent(BaseAgent):
+    """
+    A Data Loader Agent that can interact with data loading tools and search for files in your file system.
+    
+    Parameters:
+    ----------
+    model : langchain.llms.base.LLM
+        The language model used to generate the tool calling agent.
+    react_agent_kwargs : dict
+        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.
+        
+    Methods:
+    --------
+    update_params(**kwargs)
+        Updates the agent's parameters and rebuilds the compiled graph.
+    ainvoke_agent(user_instructions: str=None, **kwargs)
+        Runs the agent with the given user instructions asynchronously.
+    invoke_agent(user_instructions: str=None, **kwargs)
+        Runs the agent with the given user instructions.
+    get_internal_messages(markdown: bool=False)
+        Returns the internal messages from the agent's response.
+    get_artifacts(as_dataframe: bool=False)
+        Returns the MLflow artifacts from the agent's response.
+    get_ai_message(markdown: bool=False)
+        Returns the AI message from the agent's response.
+    
+    """
+    
+    def __init__(
+        self, 
+        model: Any,
+        create_react_agent_kwargs: Optional[Dict]={},
+        invoke_react_agent_kwargs: Optional[Dict]={},
+    ):
+        self._params = {
+            "model": model,
+            "create_react_agent_kwargs": create_react_agent_kwargs,
+            "invoke_react_agent_kwargs": invoke_react_agent_kwargs,
+        }
+        self._compiled_graph = self._make_compiled_graph()
+        self.response = None
+        
+    def _make_compiled_graph(self):
+        """
+        Creates the compiled graph for the agent.
+        """
+        self.response = None
+        return make_data_loader_tools_agent(**self._params)
+    
+    
+    def update_params(self, **kwargs):
+        """
+        Updates the agent's 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: str=None, 
+        **kwargs
+    ):
+        """
+        Runs the agent with the given user instructions.
+        
+        Parameters:
+        ----------
+        user_instructions : str, optional
+            The user instructions to pass to the agent.
+        kwargs : dict, optional
+            Additional keyword arguments to pass to the agents ainvoke method.
+        
+        """
+        response = await self._compiled_graph.ainvoke(
+            {
+                "user_instructions": user_instructions,
+            }, 
+            **kwargs
+        )
+        self.response = response
+        return None
+    
+    def invoke_agent(
+        self, 
+        user_instructions: str=None, 
+        **kwargs
+    ):
+        """
+        Runs the agent with the given user instructions.
+        
+        Parameters:
+        ----------
+        user_instructions : str, optional
+            The user instructions to pass to the agent.
+        kwargs : dict, optional
+            Additional keyword arguments to pass to the agents invoke method.
+        
+        """
+        response = self._compiled_graph.invoke(
+            {
+                "user_instructions": user_instructions,
+            },
+            **kwargs
+        )
+        self.response = response
+        return None
+    
+    def get_internal_messages(self, markdown: bool=False):
+        """
+        Returns the internal messages from the agent's 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"]])       
+        if markdown:
+            return Markdown(pretty_print)
+        else:
+            return self.response["internal_messages"]
+    
+    def get_artifacts(self, as_dataframe: bool=False):
+        """
+        Returns the MLflow artifacts from the agent's response.
+        """
+        if as_dataframe:
+            return pd.DataFrame(self.response["data_loader_artifacts"])
+        else:
+            return self.response["data_loader_artifacts"]
+    
+    def get_ai_message(self, markdown: bool=False):
+        """
+        Returns the AI message from the agent's response.
+        """
+        if markdown:
+            return Markdown(self.response["messages"][0].content)
+        else:
+            return self.response["messages"][0].content
+
+    
+
+def make_data_loader_tools_agent(
+    model: Any,
+    create_react_agent_kwargs: Optional[Dict]={},
+    invoke_react_agent_kwargs: Optional[Dict]={},
+):
+    """
+    Creates a Data Loader Agent that can interact with data loading tools.
+    
+    Parameters:
+    ----------
+    model : langchain.llms.base.LLM
+        The language model used to generate the tool calling agent.
+    react_agent_kwargs : dict
+        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.
+    
+    Returns:
+    --------
+    app : langchain.graphs.CompiledStateGraph
+        An agent that can interact with data loading tools.
+    """
+    
+    class GraphState(AgentState):
+        internal_messages: Annotated[Sequence[BaseMessage], operator.add]
+        user_instructions: str
+        data_loader_artifacts: dict
+        
+    def data_loader_agent(state):
+        
+        print(format_agent_name(AGENT_NAME))
+        print("    ")
+        
+        print("    * RUN REACT TOOL-CALLING AGENT")
+        
+        tool_node = ToolNode(
+            tools=tools
+        )
+        
+        data_loader_agent = create_react_agent(
+            model, 
+            tools=tool_node, 
+            state_schema=GraphState,
+            **create_react_agent_kwargs,
+        )
+        
+        response = data_loader_agent.invoke(
+            {
+                "messages": [("user", state["user_instructions"])],
+            },
+            invoke_react_agent_kwargs,
+        )
+        
+        print("    * POST-PROCESS RESULTS")
+        
+        internal_messages = response['messages']
+
+        # Ensure there is at least one AI message
+        if not internal_messages:
+            return {
+                "internal_messages": [],
+                "mlflow_artifacts": None,
+            }
+
+        # Get the last AI message
+        last_ai_message = AIMessage(internal_messages[-1].content, role = AGENT_NAME)
+
+        # Get the last tool artifact safely
+        last_tool_artifact = None
+        if len(internal_messages) > 1:
+            last_message = internal_messages[-2]  # Get second-to-last message
+            if hasattr(last_message, "artifact"):  # Check if it has an "artifact"
+                last_tool_artifact = last_message.artifact
+            elif isinstance(last_message, dict) and "artifact" in last_message:
+                last_tool_artifact = last_message["artifact"]
+
+        return {
+            "messages": [last_ai_message], 
+            "internal_messages": internal_messages,
+            "data_loader_artifacts": last_tool_artifact,
+        }
+        
+    workflow = StateGraph(GraphState)
+    
+    workflow.add_node("data_loader_agent", data_loader_agent)
+    
+    workflow.add_edge(START, "data_loader_agent")
+    workflow.add_edge("data_loader_agent", END)
+    
+    app = workflow.compile()
+
+    return app
+

ai_data_science_team/agents/data_visualization_agent.py

@@ -10,7 +10,6 @@ from typing import TypedDict, Annotated, Sequence, Literal
 import operator
 
 from langchain.prompts import PromptTemplate
-from langchain_core.output_parsers import StrOutputParser
 from langchain_core.messages import BaseMessage
 
 from langgraph.types import Command
@@ -30,16 +29,16 @@ from ai_data_science_team.templates import(
     create_coding_agent_graph,
     BaseAgent,
 )
-from ai_data_science_team.tools.parsers import PythonOutputParser
-from ai_data_science_team.tools.regex import (
+from ai_data_science_team.parsers.parsers import PythonOutputParser
+from ai_data_science_team.utils.regex import (
     relocate_imports_inside_function, 
     add_comments_to_top, 
     format_agent_name, 
     format_recommended_steps, 
     get_generic_summary,
 )
-from ai_data_science_team.tools.metadata import get_dataframe_summary
-from ai_data_science_team.tools.logging import log_ai_function
+from ai_data_science_team.tools.dataframe import get_dataframe_summary
+from ai_data_science_team.utils.logging import log_ai_function
 from ai_data_science_team.utils.plotly import plotly_from_dict
 
 # Setup
@@ -197,7 +196,7 @@ class DataVisualizationAgent(BaseAgent):
         # Rebuild the compiled graph
         self._compiled_graph = self._make_compiled_graph()
 
-    def ainvoke_agent(self, data_raw: pd.DataFrame, user_instructions: str=None, max_retries:int=3, retry_count:int=0, **kwargs):
+    async def ainvoke_agent(self, data_raw: pd.DataFrame, user_instructions: str=None, max_retries:int=3, retry_count:int=0, **kwargs):
         """
         Asynchronously invokes the agent to generate a visualization. 
         The response is stored in the 'response' attribute.
@@ -219,7 +218,7 @@ class DataVisualizationAgent(BaseAgent):
         -------
         None
         """
-        response = self._compiled_graph.ainvoke({
+        response = await self._compiled_graph.ainvoke({
             "user_instructions": user_instructions,
             "data_raw": data_raw.to_dict(),
             "max_retries": max_retries,

ai_data_science_team/agents/data_wrangling_agent.py

@@ -24,16 +24,16 @@ from ai_data_science_team.templates import(
     create_coding_agent_graph,
     BaseAgent,
 )
-from ai_data_science_team.tools.parsers import PythonOutputParser
-from ai_data_science_team.tools.regex import (
+from ai_data_science_team.parsers.parsers import PythonOutputParser
+from ai_data_science_team.utils.regex import (
     relocate_imports_inside_function, 
     add_comments_to_top, 
     format_agent_name, 
     format_recommended_steps, 
     get_generic_summary,
 )
-from ai_data_science_team.tools.metadata import get_dataframe_summary
-from ai_data_science_team.tools.logging import log_ai_function
+from ai_data_science_team.tools.dataframe import get_dataframe_summary
+from ai_data_science_team.utils.logging import log_ai_function
 
 # Setup Logging Path
 AGENT_NAME = "data_wrangling_agent"
@@ -213,7 +213,7 @@ class DataWranglingAgent(BaseAgent):
             self._params[k] = v
         self._compiled_graph = self._make_compiled_graph()
 
-    def ainvoke_agent(
+    async def ainvoke_agent(
         self,
         data_raw: Union[pd.DataFrame, dict, list],
         user_instructions: str=None,
@@ -245,7 +245,7 @@ class DataWranglingAgent(BaseAgent):
         None
         """
         data_input = self._convert_data_input(data_raw)
-        response = self._compiled_graph.ainvoke({
+        response = await self._compiled_graph.ainvoke({
             "user_instructions": user_instructions,
             "data_raw": data_input,
             "max_retries": max_retries,

ai_data_science_team/agents/feature_engineering_agent.py

@@ -27,16 +27,16 @@ from ai_data_science_team.templates import(
     create_coding_agent_graph,
     BaseAgent,
 )
-from ai_data_science_team.tools.parsers import PythonOutputParser
-from ai_data_science_team.tools.regex import (
+from ai_data_science_team.parsers.parsers import PythonOutputParser
+from ai_data_science_team.utils.regex import (
     relocate_imports_inside_function, 
     add_comments_to_top, 
     format_agent_name, 
     format_recommended_steps, 
     get_generic_summary,
 )
-from ai_data_science_team.tools.metadata import get_dataframe_summary
-from ai_data_science_team.tools.logging import log_ai_function
+from ai_data_science_team.tools.dataframe import get_dataframe_summary
+from ai_data_science_team.utils.logging import log_ai_function
 
 # Setup
 AGENT_NAME = "feature_engineering_agent"
@@ -203,7 +203,7 @@ class FeatureEngineeringAgent(BaseAgent):
             self._params[k] = v
         self._compiled_graph = self._make_compiled_graph()
 
-    def ainvoke_agent(
+    async def ainvoke_agent(
         self, 
         data_raw: pd.DataFrame, 
         user_instructions: str=None, 
@@ -235,7 +235,7 @@ class FeatureEngineeringAgent(BaseAgent):
         -------
         None
         """
-        response = self._compiled_graph.ainvoke({
+        response = await self._compiled_graph.ainvoke({
             "user_instructions": user_instructions,
             "data_raw": data_raw.to_dict(),
             "target_variable": target_variable,

ai_data_science_team/agents/sql_database_agent.py

@@ -25,15 +25,15 @@ from ai_data_science_team.templates import(
     create_coding_agent_graph,
     BaseAgent,
 )
-from ai_data_science_team.tools.parsers import SQLOutputParser  
-from ai_data_science_team.tools.regex import (
+from ai_data_science_team.parsers.parsers import SQLOutputParser  
+from ai_data_science_team.utils.regex import (
     add_comments_to_top, 
     format_agent_name, 
     format_recommended_steps, 
     get_generic_summary,
 )
-from ai_data_science_team.tools.metadata import get_database_metadata
-from ai_data_science_team.tools.logging import log_ai_function
+from ai_data_science_team.tools.sql import get_database_metadata
+from ai_data_science_team.utils.logging import log_ai_function
 
 # Setup
 AGENT_NAME = "sql_database_agent"
@@ -193,7 +193,7 @@ class SQLDatabaseAgent(BaseAgent):
             self._params[k] = v
         self._compiled_graph = self._make_compiled_graph()
 
-    def ainvoke_agent(self, user_instructions: str=None, max_retries=3, retry_count=0, **kwargs):
+    async def ainvoke_agent(self, user_instructions: str=None, max_retries=3, retry_count=0, **kwargs):
         """
         Asynchronously runs the SQL Database Agent based on user instructions.
 
@@ -212,7 +212,7 @@ class SQLDatabaseAgent(BaseAgent):
         -------
         None
         """
-        response = self._compiled_graph.ainvoke({
+        response = await self._compiled_graph.ainvoke({
             "user_instructions": user_instructions,
             "max_retries": max_retries,
             "retry_count": retry_count

ai_data_science_team/ml_agents/__init__.py

@@ -1 +1,2 @@
 from ai_data_science_team.ml_agents.h2o_ml_agent import make_h2o_ml_agent, H2OMLAgent
+from ai_data_science_team.ml_agents.mlflow_tools_agent import make_mlflow_tools_agent, MLflowToolsAgent