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

ai_data_science_team/multiagents/__init__.py

@@ -0,0 +1 @@
+from ai_data_science_team.multiagents.sql_data_analyst import SQLDataAnalyst, make_sql_data_analyst

ai_data_science_team/multiagents/sql_data_analyst.py

@@ -0,0 +1,398 @@
+
+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, Literal
+import operator
+
+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.tools.regex import remove_consecutive_duplicates, get_generic_summary
+
+
+class SQLDataAnalyst(BaseAgent):
+    """
+    SQLDataAnalyst is a multi-agent class that combines SQL database querying and data visualization capabilities.
+    
+    Parameters:
+    -----------
+    model:
+        The language model to be used for the agents.
+    sql_database_agent: SQLDatabaseAgent
+        The SQL Database Agent.
+    data_visualization_agent: DataVisualizationAgent
+        The Data Visualization Agent.
+        
+    Methods:
+    --------
+    ainvoke_agent(user_instructions, **kwargs)
+        Asynchronously invokes the SQL Data Analyst Multi-Agent with the given user instructions.
+    invoke_agent(user_instructions, **kwargs)
+        Invokes the SQL Data Analyst Multi-Agent with the given user instructions.
+    get_data_sql()
+        Returns the SQL data as a Pandas DataFrame.
+    get_plotly_graph()
+        Returns the Plotly graph as a Plotly object.
+    get_sql_query_code(markdown=False)
+        Returns the SQL query code as a string, optionally formatted as a Markdown code block.
+    get_sql_database_function(markdown=False)
+        Returns the SQL database function as a string, optionally formatted as a Markdown code block.
+    get_data_visualization_function(markdown=False)
+        Returns the data visualization function as a string, optionally formatted as a Markdown code block.
+    """
+    
+    def __init__(
+        self, 
+        model, 
+        sql_database_agent: SQLDatabaseAgent, 
+        data_visualization_agent: DataVisualizationAgent,
+        checkpointer: Checkpointer = None,
+    ):
+        self._params = {
+            "model": model,
+            "sql_database_agent": sql_database_agent,
+            "data_visualization_agent": data_visualization_agent,
+            "checkpointer": checkpointer,
+        }
+        self._compiled_graph = self._make_compiled_graph()
+        self.response = None
+    
+    def _make_compiled_graph(self):
+        """
+        Create or rebuild the compiled graph for the SQL Data Analyst Multi-Agent.
+        Running this method resets the response to None.
+        """
+        self.response = None
+        return make_sql_data_analyst(
+            model=self._params["model"],
+            sql_database_agent=self._params["sql_database_agent"]._compiled_graph,
+            data_visualization_agent=self._params["data_visualization_agent"]._compiled_graph,
+            checkpointer=self._params["checkpointer"],
+        )
+    
+    def update_params(self, **kwargs):
+        """
+        Updates the agent's parameters (e.g. model, sql_database_agent, etc.) 
+        and rebuilds the compiled graph.
+        """
+        for k, v in kwargs.items():
+            self._params[k] = v
+        self._compiled_graph = self._make_compiled_graph()
+        
+    def ainvoke_agent(self, user_instructions, max_retries:int=3, retry_count:int=0, **kwargs):
+        """
+        Asynchronosly nvokes the SQL Data Analyst Multi-Agent.
+        
+        Parameters:
+        ----------
+        user_instructions: str
+            The user's instructions for the combined SQL and (optionally) Data Visualization agents.
+        **kwargs:
+            Additional keyword arguments to pass to the compiled graph's `ainvoke` method.
+            
+        Returns:
+        -------
+        None. The response is stored in the `response` attribute.
+        
+        Example:
+        --------
+        ``` python
+        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({
+            "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, max_retries:int=3, retry_count:int=0, **kwargs):
+        """
+        Invokes the SQL Data Analyst Multi-Agent.
+        
+        Parameters:
+        ----------
+        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.
+            
+        Returns:
+        -------
+        None. The response is stored in the `response` attribute.
+        
+        Example:
+        --------
+        ``` python
+        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.
+        """
+        if self.response:
+            if self.response.get("data_sql"):
+                return pd.DataFrame(self.response.get("data_sql"))
+    
+    def get_plotly_graph(self):
+        """
+        Returns the Plotly graph as a Plotly object.
+        """
+        if self.response:
+            if self.response.get("plotly_graph"):
+                return plotly_from_dict(self.response.get("plotly_graph"))
+    
+    def get_sql_query_code(self, markdown=False):
+        """
+        Returns the SQL query code as a string.
+        
+        Parameters:
+        ----------
+        markdown: bool
+            If True, returns the code as a Markdown code block for Jupyter (IPython).
+            For streamlit, use `st.code()` instead.
+        """
+        if self.response:
+            if self.response.get("sql_query_code"):
+                if markdown:
+                    return Markdown(f"```sql\n{self.response.get('sql_query_code')}\n```")
+                return self.response.get("sql_query_code")
+    
+    def get_sql_database_function(self, markdown=False):
+        """
+        Returns the SQL database function as a string.
+        
+        Parameters:
+        ----------
+        markdown: bool
+            If True, returns the function as a Markdown code block for Jupyter (IPython).
+            For streamlit, use `st.code()` instead.
+        """
+        if self.response:
+            if self.response.get("sql_database_function"):
+                if markdown:
+                    return Markdown(f"```python\n{self.response.get('sql_database_function')}\n```")
+                return self.response.get("sql_database_function")
+    
+    def get_data_visualization_function(self, markdown=False):
+        """
+        Returns the data visualization function as a string.
+        
+        Parameters:
+        ----------
+        markdown: bool
+            If True, returns the function as a Markdown code block for Jupyter (IPython).
+            For streamlit, use `st.code()` instead.
+        """
+        if self.response:
+            if self.response.get("data_visualization_function"):
+                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)
+    
+    
+
+def make_sql_data_analyst(
+    model, 
+    sql_database_agent: CompiledStateGraph,
+    data_visualization_agent: CompiledStateGraph,
+    checkpointer: Checkpointer = None
+):
+    """
+    Creates a multi-agent system that takes in a SQL query and returns a plot or table.
+    
+    - Agent 1: SQL Database Agent made with `make_sql_database_agent()`
+    - Agent 2: Data Visualization Agent made with `make_data_visualization_agent()`
+    
+    Parameters:
+    ----------
+    model: 
+        The language model to be used for the agents.
+    sql_database_agent: CompiledStateGraph
+        The SQL Database Agent made with `make_sql_database_agent()`.
+    data_visualization_agent: CompiledStateGraph
+        The Data Visualization Agent made with `make_data_visualization_agent()`.
+    checkpointer: Checkpointer (optional)
+        The checkpointer to save the state of the multi-agent system.
+        Default: None
+        
+    Returns:
+    -------
+    CompiledStateGraph
+        The compiled multi-agent system.
+    """
+    
+    llm = model
+
+    class PrimaryState(TypedDict):
+        messages: Annotated[Sequence[BaseMessage], operator.add]    
+        user_instructions: str
+        sql_query_code: str
+        sql_database_function: str
+        data_sql: dict
+        data_raw: dict
+        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__"]]: 
+        
+        response = llm.invoke(f"Respond in 1 word ('plot' or 'table'). Is the user requesting a plot? If unknown, select 'table'. \n\n User Instructions:\n{state.get('user_instructions')}")
+        
+        if response.content == 'plot':
+            plot_required = True
+            goto="data_visualization_agent"
+        else:
+            plot_required = False
+            goto="__end__"
+        
+        return Command(
+            update={
+                'data_raw': state.get("data_sql"),
+                'plot_required': plot_required,
+            },
+            goto=goto
+        )
+
+    workflow = StateGraph(PrimaryState)
+
+    workflow.add_node("sql_database_agent", sql_database_agent)
+    workflow.add_node("route_to_visualization", route_to_visualization)
+    workflow.add_node("data_visualization_agent", data_visualization_agent)
+
+    workflow.add_edge(START, "sql_database_agent")
+    workflow.add_edge("sql_database_agent", "route_to_visualization")
+    workflow.add_edge("data_visualization_agent", END)
+
+    app = workflow.compile(checkpointer=checkpointer)
+
+    return app
+

ai_data_science_team/multiagents/supervised_data_analyst.py

@@ -0,0 +1,2 @@
+# TODO: Implement the supervised data analyst agent
+# https://langchain-ai.github.io/langgraph/tutorials/multi_agent/agent_supervisor/#create-agent-supervisor

ai_data_science_team/templates/__init__.py

@@ -3,6 +3,8 @@ 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
+    create_coding_agent_graph,
+    BaseAgent,
 )