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

ai_data_science_team/ml_agents/mlflow_tools_agent.py

@@ -0,0 +1,327 @@
+
+from typing import Any, Optional, Annotated, Sequence
+import operator
+
+import pandas as pd
+
+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.mlflow import (
+    mlflow_search_experiments, 
+    mlflow_search_runs,
+    mlflow_create_experiment, 
+    mlflow_predict_from_run_id,
+    mlflow_launch_ui,
+    mlflow_stop_ui,
+    mlflow_list_artifacts,
+    mlflow_download_artifacts,
+    mlflow_list_registered_models,
+    mlflow_search_registered_models,
+    mlflow_get_model_version_details,
+)
+
+AGENT_NAME = "mlflow_tools_agent"
+
+# TOOL SETUP
+tools = [
+    mlflow_search_experiments, 
+    mlflow_search_runs, 
+    mlflow_create_experiment, 
+    mlflow_predict_from_run_id,
+    mlflow_launch_ui,
+    mlflow_stop_ui,
+    mlflow_list_artifacts,
+    mlflow_download_artifacts,
+    mlflow_list_registered_models,
+    mlflow_search_registered_models,
+    mlflow_get_model_version_details,
+]
+
+class MLflowToolsAgent(BaseAgent):
+    """
+    An agent that can interact with MLflow by calling tools.
+    
+    Current tools include:
+    - List Experiments
+    - Search Runs
+    - Create Experiment
+    - Predict (from a Run ID)
+    
+    Parameters:
+    ----------
+    model : langchain.llms.base.LLM
+        The language model used to generate the tool calling agent.
+    mlfow_tracking_uri : str, optional
+        The tracking URI for MLflow. Defaults to None.
+    mlflow_registry_uri : str, optional
+        The registry URI for MLflow. Defaults to None.
+    **react_agent_kwargs : dict, optional
+        Additional keyword arguments to pass to the agent's react agent.
+    
+    Methods:
+    --------
+    update_params(**kwargs):
+        Updates the agent's parameters and rebuilds the compiled graph.
+    ainvoke_agent(user_instructions: str=None, data_raw: pd.DataFrame=None, **kwargs):
+        Asynchronously runs the agent with the given user instructions.
+    invoke_agent(user_instructions: str=None, data_raw: pd.DataFrame=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_mlflow_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
+    
+    
+    
+    Examples:
+    --------
+    ```python
+    from ai_data_science_team.ml_agents import MLflowToolsAgent
+    
+    mlflow_agent = MLflowToolsAgent(llm)
+
+    mlflow_agent.invoke_agent(user_instructions="List the MLflow experiments")
+
+    mlflow_agent.get_response()
+
+    mlflow_agent.get_internal_messages(markdown=True)
+
+    mlflow_agent.get_ai_message(markdown=True)
+
+    mlflow_agent.get_mlflow_artifacts(as_dataframe=True)
+    
+    ```
+    
+    Returns
+    -------
+    MLflowToolsAgent : langchain.graphs.CompiledStateGraph 
+        An instance of the MLflow Tools Agent.
+    
+    """
+    
+    def __init__(
+        self, 
+        model: Any,
+        mlflow_tracking_uri: Optional[str]=None,
+        mlflow_registry_uri: Optional[str]=None,
+        **react_agent_kwargs,
+    ):
+        self._params = {
+            "model": model,
+            "mlflow_tracking_uri": mlflow_tracking_uri,
+            "mlflow_registry_uri": mlflow_registry_uri,
+            **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_mlflow_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, 
+        data_raw: pd.DataFrame=None, 
+        **kwargs
+    ):
+        """
+        Runs the agent with the given user instructions.
+        
+        Parameters:
+        ----------
+        user_instructions : str, optional
+            The user instructions to pass to the agent.
+        data_raw : pd.DataFrame, optional
+            The data to pass to the agent. Used for prediction and tool calls where data is required.
+        kwargs : dict, optional
+            Additional keyword arguments to pass to the agents ainvoke method.
+        
+        """
+        response = await self._compiled_graph.ainvoke(
+            {
+                "user_instructions": user_instructions,
+                "data_raw": data_raw.to_dict() if data_raw is not None else None,
+            }, 
+            **kwargs
+        )
+        self.response = response
+        return None
+    
+    def invoke_agent(
+        self, 
+        user_instructions: str=None, 
+        data_raw: pd.DataFrame=None, 
+        **kwargs
+    ):
+        """
+        Runs the agent with the given user instructions.
+        
+        Parameters:
+        ----------
+        user_instructions : str, optional
+            The user instructions to pass to the agent.
+        data_raw : pd.DataFrame, optional
+            The raw data to pass to the agent. Used for prediction and tool calls where data is required.
+        kwargs : dict, optional
+            Additional keyword arguments to pass to the agents invoke method.
+        
+        """
+        response = self._compiled_graph.invoke(
+            {
+                "user_instructions": user_instructions,
+                "data_raw": data_raw.to_dict() if data_raw is not None else None,
+            },
+            **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_mlflow_artifacts(self, as_dataframe: bool=False):
+        """
+        Returns the MLflow artifacts from the agent's response.
+        """
+        if as_dataframe:
+            return pd.DataFrame(self.response["mlflow_artifacts"])
+        else:
+            return self.response["mlflow_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_mlflow_tools_agent(
+    model: Any,
+    mlflow_tracking_uri: str=None,
+    mlflow_registry_uri: str=None,
+    **react_agent_kwargs,
+):
+    """
+    MLflow Tool Calling Agent
+    """
+    
+    try:
+        import mlflow
+    except ImportError:
+        return "MLflow is not installed. Please install it by running: !pip install mlflow"
+    
+    if mlflow_tracking_uri is not None:
+        mlflow.set_tracking_uri(mlflow_tracking_uri)
+    
+    if mlflow_registry_uri is not None:
+        mlflow.set_registry_uri(mlflow_registry_uri)
+    
+    class GraphState(AgentState):
+        internal_messages: Annotated[Sequence[BaseMessage], operator.add]
+        user_instructions: str
+        data_raw: dict
+        mlflow_artifacts: dict
+
+    
+    def mflfow_tools_agent(state):
+        """
+        Postprocesses the MLflow state, keeping only the last message
+        and extracting the last tool artifact.
+        """
+        print(format_agent_name(AGENT_NAME))
+        print("    * RUN REACT TOOL-CALLING AGENT")
+        
+        tool_node = ToolNode(
+            tools=tools
+        )
+        
+        mlflow_agent = create_react_agent(
+            model, 
+            tools=tool_node, 
+            state_schema=GraphState,
+            **react_agent_kwargs,
+        )
+        
+        response = mlflow_agent.invoke(
+            {
+                "messages": [("user", state["user_instructions"])],
+                "data_raw": state["data_raw"],
+            },
+        )
+        
+        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,
+            "mlflow_artifacts": last_tool_artifact,
+        }
+
+    
+    workflow = StateGraph(GraphState)
+    
+    workflow.add_node("mlflow_tools_agent", mflfow_tools_agent)
+    
+    workflow.add_edge(START, "mlflow_tools_agent")
+    workflow.add_edge("mlflow_tools_agent", END)
+    
+    app = workflow.compile()
+
+    return app
+    

ai_data_science_team/multiagents/sql_data_analyst.py

@@ -1,24 +1,24 @@
 
 from langchain_core.messages import BaseMessage
-from langgraph.checkpoint.memory import MemorySaver
 from langgraph.types import Checkpointer
 
 from langgraph.graph import START, END, StateGraph
 from langgraph.graph.state import CompiledStateGraph
 from langgraph.types import Command
 
-from typing import TypedDict, Annotated, Sequence
+from typing import TypedDict, Annotated, Sequence, Literal
 import operator
 
-from typing_extensions import TypedDict, Literal
+from typing_extensions import TypedDict
 
 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 SQLDatabaseAgent, 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
 
 
 class SQLDataAnalyst(BaseAgent):
@@ -90,7 +90,7 @@ class SQLDataAnalyst(BaseAgent):
             self._params[k] = v
         self._compiled_graph = self._make_compiled_graph()
         
-    def ainvoke_agent(self, user_instructions, **kwargs):
+    async def ainvoke_agent(self, user_instructions, max_retries:int=3, retry_count:int=0, **kwargs):
         """
         Asynchronosly nvokes the SQL Data Analyst Multi-Agent.
         
@@ -108,15 +108,53 @@ class SQLDataAnalyst(BaseAgent):
         Example:
         --------
         ``` python
-        # TODO
+        from langchain_openai import ChatOpenAI
+        import sqlalchemy as sql
+        from ai_data_science_team.multiagents import SQLDataAnalyst
+        from ai_data_science_team.agents import SQLDatabaseAgent, DataVisualizationAgent
+        
+        llm = ChatOpenAI(model = "gpt-4o-mini")
+        
+        sql_engine = sql.create_engine("sqlite:///data/northwind.db")
+
+        conn = sql_engine.connect()
+        
+        sql_data_analyst = SQLDataAnalyst(
+            model = llm,
+            sql_database_agent = SQLDatabaseAgent(
+                model = llm,
+                connection = conn,
+                n_samples = 1,
+            ),
+            data_visualization_agent = DataVisualizationAgent(
+                model = llm,
+                n_samples = 10,
+            )
+        )
+        
+        sql_data_analyst.ainvoke_agent(
+            user_instructions = "Make a plot of sales revenue by month by territory. Make a dropdown for the user to select the territory.",
+        )
+        
+        sql_data_analyst.get_sql_query_code()
+        
+        sql_data_analyst.get_data_sql()
+        
+        sql_data_analyst.get_plotly_graph()
         ```
         """
-        response = self._compiled_graph.ainvoke({
+        response = await self._compiled_graph.ainvoke({
             "user_instructions": user_instructions,
+            "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, **kwargs):
+    def invoke_agent(self, user_instructions, max_retries:int=3, retry_count:int=0, **kwargs):
         """
         Invokes the SQL Data Analyst Multi-Agent.
         
@@ -124,6 +162,10 @@ class SQLDataAnalyst(BaseAgent):
         ----------
         user_instructions: str
             The user's instructions for the combined SQL and (optionally) Data Visualization agents.
+        max_retries (int): 
+                Maximum retry attempts for cleaning.
+        retry_count (int): 
+            Current retry attempt.
         **kwargs:
             Additional keyword arguments to pass to the compiled graph's `invoke` method.
             
@@ -134,14 +176,53 @@ class SQLDataAnalyst(BaseAgent):
         Example:
         --------
         ``` python
-        # TODO
+        from langchain_openai import ChatOpenAI
+        import sqlalchemy as sql
+        from ai_data_science_team.multiagents import SQLDataAnalyst
+        from ai_data_science_team.agents import SQLDatabaseAgent, DataVisualizationAgent
+        
+        llm = ChatOpenAI(model = "gpt-4o-mini")
+        
+        sql_engine = sql.create_engine("sqlite:///data/northwind.db")
+
+        conn = sql_engine.connect()
+        
+        sql_data_analyst = SQLDataAnalyst(
+            model = llm,
+            sql_database_agent = SQLDatabaseAgent(
+                model = llm,
+                connection = conn,
+                n_samples = 1,
+            ),
+            data_visualization_agent = DataVisualizationAgent(
+                model = llm,
+                n_samples = 10,
+            )
+        )
+        
+        sql_data_analyst.invoke_agent(
+            user_instructions = "Make a plot of sales revenue by month by territory. Make a dropdown for the user to select the territory.",
+        )
+        
+        sql_data_analyst.get_sql_query_code()
+        
+        sql_data_analyst.get_data_sql()
+        
+        sql_data_analyst.get_plotly_graph()
         ```
         """
         response = self._compiled_graph.invoke({
             "user_instructions": user_instructions,
+            "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_sql(self):
         """
         Returns the SQL data as a Pandas DataFrame.
@@ -205,6 +286,34 @@ class SQLDataAnalyst(BaseAgent):
                 if markdown:
                     return Markdown(f"```python\n{self.response.get('data_visualization_function')}\n```")
                 return self.response.get("data_visualization_function")
+            
+    def get_workflow_summary(self, markdown=False):
+        """
+        Returns a summary of the SQL Data Analyst workflow.
+        
+        Parameters:
+        ----------
+        markdown: bool
+            If True, returns the summary as a Markdown-formatted string.
+        """
+        if self.response and self.get_response()['messages']:
+            
+            agents = [self.get_response()['messages'][i].role for i in range(len(self.get_response()['messages']))]
+            
+            agent_labels = []
+            for i in range(len(agents)):
+                agent_labels.append(f"- **Agent {i+1}:** {agents[i]}")
+            
+            # Construct header
+            header = f"# SQL Data Analyst Workflow Summary Report\n\nThis agentic workflow contains {len(agents)} agents:\n\n" + "\n".join(agent_labels)
+            
+            reports = []
+            for msg in self.get_response()['messages']:
+                reports.append(get_generic_summary(json.loads(msg.content)))
+                
+            if markdown:
+                return Markdown(header + "\n\n".join(reports))
+            return "\n\n".join(reports)
     
     
 
@@ -250,6 +359,8 @@ def make_sql_data_analyst(
         plot_required: bool
         data_visualization_function: str
         plotly_graph: dict
+        max_retries: int
+        retry_count: int
         
     def route_to_visualization(state) -> Command[Literal["data_visualization_agent", "__end__"]]: 
         

ai_data_science_team/{tools → parsers}/parsers.py

@@ -3,7 +3,6 @@
 # ***
 # Parsers
 
-from langchain_core.output_parsers import JsonOutputParser
 from langchain_core.output_parsers import BaseOutputParser
 
 import re

ai_data_science_team/templates/__init__.py

@@ -3,6 +3,7 @@ from ai_data_science_team.templates.agent_templates import(
     node_func_human_review,
     node_func_fix_agent_code, 
     node_func_explain_agent_code, 
+    node_func_report_agent_outputs,
     node_func_execute_agent_from_sql_connection,
     create_coding_agent_graph,
     BaseAgent,

ai_data_science_team/templates/agent_templates.py

@@ -8,11 +8,16 @@ from langgraph.pregel.types import StreamMode
 
 import pandas as pd
 import sqlalchemy as sql
+import json
 
-from typing import Any, Callable, Dict, Type, Optional, Union
+from typing import Any, Callable, Dict, Type, Optional, Union, List
 
-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
+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,
+    remove_consecutive_duplicates
+)
 
 from IPython.display import Image, display
 import pandas as pd
@@ -82,9 +87,13 @@ class BaseAgent(CompiledStateGraph):
             Any: The agent's response.
         """
         self.response = self._compiled_graph.invoke(input=input, config=config,**kwargs)
+        
+        if self.response.get("messages"):
+            self.response["messages"] = remove_consecutive_duplicates(self.response["messages"])
+        
         return self.response
     
-    def ainvoke(
+    async def ainvoke(
         self, 
         input: Union[dict[str, Any], Any], 
         config: Optional[RunnableConfig] = None, 
@@ -101,7 +110,11 @@ class BaseAgent(CompiledStateGraph):
         Returns:
             Any: The agent's response.
         """
-        self.response = self._compiled_graph.ainvoke(input=input, config=config,**kwargs)
+        self.response = await self._compiled_graph.ainvoke(input=input, config=config,**kwargs)
+        
+        if self.response.get("messages"):
+            self.response["messages"] = remove_consecutive_duplicates(self.response["messages"])
+        
         return self.response
     
     def stream(
@@ -129,9 +142,13 @@ class BaseAgent(CompiledStateGraph):
             Any: The agent's response.
         """
         self.response = self._compiled_graph.stream(input=input, config=config, stream_mode=stream_mode, **kwargs)
+        
+        if self.response.get("messages"):
+            self.response["messages"] = remove_consecutive_duplicates(self.response["messages"])        
+        
         return self.response
     
-    def astream(
+    async def astream(
         self,
         input: dict[str, Any] | Any,
         config: RunnableConfig | None = None,
@@ -155,7 +172,11 @@ class BaseAgent(CompiledStateGraph):
         Returns:
             Any: The agent's response.
         """
-        self.response = self._compiled_graph.astream(input=input, config=config, stream_mode=stream_mode, **kwargs)
+        self.response = await self._compiled_graph.astream(input=input, config=config, stream_mode=stream_mode, **kwargs)
+        
+        if self.response.get("messages"):
+            self.response["messages"] = remove_consecutive_duplicates(self.response["messages"])
+        
         return self.response
     
     def get_state_keys(self):
@@ -183,6 +204,9 @@ class BaseAgent(CompiledStateGraph):
         Returns:
             Any: The agent's response.
         """
+        if self.response.get("messages"):
+            self.response["messages"] = remove_consecutive_duplicates(self.response["messages"])  
+        
         return self.response
 
     def show(self, xray: int = 0):
@@ -729,3 +753,50 @@ def node_func_explain_agent_code(
         # Return an error message if there was a problem with the code
         message = AIMessage(content=error_message)
         return {result_key: [message]}
+
+
+
+def node_func_report_agent_outputs(
+    state: Dict[str, Any],
+    keys_to_include: List[str],
+    result_key: str,
+    role: str,
+    custom_title: str = "Agent Output Summary"
+) -> Dict[str, Any]:
+    """
+    Gathers relevant data directly from the state (filtered by `keys_to_include`) 
+    and returns them as a structured message in `state[result_key]`.
+
+    No LLM is used.
+
+    Parameters
+    ----------
+    state : Dict[str, Any]
+        The current state dictionary holding all agent variables.
+    keys_to_include : List[str]
+        The list of keys in `state` to include in the output.
+    result_key : str
+        The key in `state` under which we'll store the final structured message.
+    role : str
+        The role that will be used in the final AIMessage (e.g., "DataCleaningAgent").
+    custom_title : str, optional
+        A title or heading for your report. Defaults to "Agent Output Summary".
+    """
+    print("    * REPORT AGENT OUTPUTS")
+
+    final_report = {"report_title": custom_title}
+
+    for key in keys_to_include:
+        final_report[key] = state.get(key, f"<{key}_not_found_in_state>")
+
+    # Wrap it in a list of messages (like the current "messages" pattern).
+    # You can serialize this dictionary as JSON or just cast it to string.
+    return {
+        result_key: [
+            AIMessage(
+                content=json.dumps(final_report, indent=2), 
+                role=role
+            )
+        ]
+    }
+