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

ai_data_science_team/ml_agents/mlflow_tools_agent.py

@@ -10,6 +10,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
@@ -68,6 +69,8 @@ class MLflowToolsAgent(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 : langchain.checkpointing.Checkpointer, optional
+        A checkpointer to use for saving and loading the agent's state. Defaults to None.
     
     Methods:
     --------
@@ -119,6 +122,7 @@ class MLflowToolsAgent(BaseAgent):
         mlflow_registry_uri: Optional[str]=None,
         create_react_agent_kwargs: Optional[Dict]={},
         invoke_react_agent_kwargs: Optional[Dict]={},
+        checkpointer: Optional[Checkpointer]=None,
     ):
         self._params = {
             "model": model,
@@ -126,6 +130,7 @@ class MLflowToolsAgent(BaseAgent):
             "mlflow_registry_uri": mlflow_registry_uri,
             "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
@@ -245,6 +250,7 @@ def make_mlflow_tools_agent(
     mlflow_registry_uri: str=None,
     create_react_agent_kwargs: Optional[Dict]={},
     invoke_react_agent_kwargs: Optional[Dict]={},
+    checkpointer: Optional[Checkpointer]=None,
 ):
     """
     MLflow Tool Calling Agent
@@ -261,6 +267,8 @@ def make_mlflow_tools_agent(
         Additional keyword arguments to pass to the agent's create_react_agent method.
     invoke_react_agent_kwargs : dict, optional
         Additional keyword arguments to pass to the agent's invoke method.
+    checkpointer : langchain.checkpointing.Checkpointer, optional
+        A checkpointer to use for saving and loading the agent's state. Defaults to None.
         
     Returns
     -------
@@ -303,6 +311,7 @@ def make_mlflow_tools_agent(
             model, 
             tools=tool_node, 
             state_schema=GraphState,
+            checkpointer=checkpointer,
             **create_react_agent_kwargs,
         )
         
@@ -354,7 +363,10 @@ def make_mlflow_tools_agent(
     workflow.add_edge(START, "mlflow_tools_agent")
     workflow.add_edge("mlflow_tools_agent", END)
     
-    app = workflow.compile()
+    app = workflow.compile(
+        checkpointer=checkpointer,
+        name=AGENT_NAME,
+    )
 
     return app
     

ai_data_science_team/multiagents/__init__.py

@@ -1 +1,2 @@
-from ai_data_science_team.multiagents.sql_data_analyst import SQLDataAnalyst, make_sql_data_analyst
+from ai_data_science_team.multiagents.sql_data_analyst import SQLDataAnalyst, make_sql_data_analyst
+from ai_data_science_team.multiagents.pandas_data_analyst import PandasDataAnalyst, make_pandas_data_analyst

ai_data_science_team/multiagents/pandas_data_analyst.py

@@ -0,0 +1,305 @@
+from langchain_core.messages import BaseMessage
+from langchain.prompts import PromptTemplate
+from langchain_core.output_parsers import JsonOutputParser
+from langgraph.types import Checkpointer
+from langgraph.graph import START, END, StateGraph
+from langgraph.graph.state import CompiledStateGraph
+
+from typing import TypedDict, Annotated, Sequence, Union
+import operator
+
+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 DataWranglingAgent, 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
+
+AGENT_NAME = "pandas_data_analyst"
+
+class PandasDataAnalyst(BaseAgent):
+    """
+    PandasDataAnalyst is a multi-agent class that combines data wrangling and visualization capabilities.
+
+    Parameters:
+    -----------
+    model:
+        The language model to be used for the agents.
+    data_wrangling_agent: DataWranglingAgent
+        The Data Wrangling Agent for transforming raw data.
+    data_visualization_agent: DataVisualizationAgent
+        The Data Visualization Agent for generating plots.
+    checkpointer: Checkpointer (optional)
+        The checkpointer to save the state of the multi-agent system.
+
+    Methods:
+    --------
+    ainvoke_agent(user_instructions, data_raw, **kwargs)
+        Asynchronously invokes the multi-agent with user instructions and raw data.
+    invoke_agent(user_instructions, data_raw, **kwargs)
+        Synchronously invokes the multi-agent with user instructions and raw data.
+    get_data_wrangled()
+        Returns the wrangled data as a Pandas DataFrame.
+    get_plotly_graph()
+        Returns the Plotly graph as a Plotly object.
+    get_data_wrangler_function(markdown=False)
+        Returns the data wrangling function as a string, optionally in Markdown.
+    get_data_visualization_function(markdown=False)
+        Returns the data visualization function as a string, optionally in Markdown.
+    """
+
+    def __init__(
+        self,
+        model,
+        data_wrangling_agent: DataWranglingAgent,
+        data_visualization_agent: DataVisualizationAgent,
+        checkpointer: Checkpointer = None,
+    ):
+        self._params = {
+            "model": model,
+            "data_wrangling_agent": data_wrangling_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. Resets response to None."""
+        self.response = None
+        return make_pandas_data_analyst(
+            model=self._params["model"],
+            data_wrangling_agent=self._params["data_wrangling_agent"]._compiled_graph,
+            data_visualization_agent=self._params["data_visualization_agent"]._compiled_graph,
+            checkpointer=self._params["checkpointer"],
+        )
+
+    def update_params(self, **kwargs):
+        """Updates 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, data_raw: Union[pd.DataFrame, dict, list], max_retries: int = 3, retry_count: int = 0, **kwargs):
+        """Asynchronously invokes the multi-agent."""
+        response = await self._compiled_graph.ainvoke({
+            "user_instructions": user_instructions,
+            "data_raw": self._convert_data_input(data_raw),
+            "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, data_raw: Union[pd.DataFrame, dict, list], max_retries: int = 3, retry_count: int = 0, **kwargs):
+        """Synchronously invokes the multi-agent."""
+        response = self._compiled_graph.invoke({
+            "user_instructions": user_instructions,
+            "data_raw": self._convert_data_input(data_raw),
+            "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_wrangled(self):
+        """Returns the wrangled data as a Pandas DataFrame."""
+        if self.response and self.response.get("data_wrangled"):
+            return pd.DataFrame(self.response.get("data_wrangled"))
+
+    def get_plotly_graph(self):
+        """Returns the Plotly graph as a Plotly object."""
+        if self.response and self.response.get("plotly_graph"):
+            return plotly_from_dict(self.response.get("plotly_graph"))
+
+    def get_data_wrangler_function(self, markdown=False):
+        """Returns the data wrangling function as a string."""
+        if self.response and self.response.get("data_wrangler_function"):
+            code = self.response.get("data_wrangler_function")
+            return Markdown(f"```python\n{code}\n```") if markdown else code
+
+    def get_data_visualization_function(self, markdown=False):
+        """Returns the data visualization function as a string."""
+        if self.response and self.response.get("data_visualization_function"):
+            code = self.response.get("data_visualization_function")
+            return Markdown(f"```python\n{code}\n```") if markdown else code
+
+    def get_workflow_summary(self, markdown=False):
+        """Returns a summary of the workflow."""
+        if self.response and self.response.get("messages"):
+            agents = [msg.role for msg in self.response["messages"]]
+            agent_labels = [f"- **Agent {i+1}:** {role}" for i, role in enumerate(agents)]
+            header = f"# Pandas Data Analyst Workflow Summary\n\nThis workflow contains {len(agents)} agents:\n\n" + "\n".join(agent_labels)
+            reports = [get_generic_summary(json.loads(msg.content)) for msg in self.response["messages"]]
+            summary = "\n" +header + "\n\n".join(reports)
+            return Markdown(summary) if markdown else summary
+
+    @staticmethod
+    def _convert_data_input(data_raw: Union[pd.DataFrame, dict, list]) -> Union[dict, list]:
+        """Converts input data to the expected format (dict or list of dicts)."""
+        if isinstance(data_raw, pd.DataFrame):
+            return data_raw.to_dict()
+        if isinstance(data_raw, dict):
+            return data_raw
+        if isinstance(data_raw, list):
+            return [item.to_dict() if isinstance(item, pd.DataFrame) else item for item in data_raw]
+        raise ValueError("data_raw must be a DataFrame, dict, or list of DataFrames/dicts")
+
+def make_pandas_data_analyst(
+    model,
+    data_wrangling_agent: CompiledStateGraph,
+    data_visualization_agent: CompiledStateGraph,
+    checkpointer: Checkpointer = None
+):
+    """
+    Creates a multi-agent system that wrangles data and optionally visualizes it.
+
+    Parameters:
+    -----------
+    model: The language model to be used.
+    data_wrangling_agent: CompiledStateGraph
+        The Data Wrangling Agent.
+    data_visualization_agent: CompiledStateGraph
+        The Data Visualization Agent.
+    checkpointer: Checkpointer (optional)
+        The checkpointer to save the state.
+
+    Returns:
+    --------
+    CompiledStateGraph: The compiled multi-agent system.
+    """
+    
+    llm = model
+    
+    routing_preprocessor_prompt = PromptTemplate(
+        template="""
+        You are an expert in routing decisions for a Pandas Data Manipulation Wrangling Agent, a Charting Visualization Agent, and a Pandas Table Agent. Your job is to:
+        
+        1. Determine what the correct format for a Users Question should be for use with a Pandas Data Wrangling Agent based on the incoming user question. Anything related to data wrangling and manipulation should be passed along.
+        2. Determine whether or not a chart should be generated or a table should be returned based on the users question.
+        3. If a chart is requested, determine the correct format of a Users Question should be used with a Data Visualization Agent. Anything related to plotting and visualization should be passed along.
+        
+        Use the following criteria on how to route the the initial user question:
+        
+        From the incoming user question, remove any details about the format of the final response as either a Chart or Table and return only the important part of the incoming user question that is relevant for the SQL generator agent. This will be the 'user_instructions_data_wrangling'. If 'None' is found, return the original user question.
+        
+        Next, determine if the user would like a data visualization ('chart') or a 'table' returned with the results of the Data Wrangling Agent. If unknown, not specified or 'None' is found, then select 'table'.  
+        
+        If a 'chart' is requested, return the 'user_instructions_data_visualization'. If 'None' is found, return None.
+        
+        Return JSON with 'user_instructions_data_wrangling', 'user_instructions_data_visualization' and 'routing_preprocessor_decision'.
+        
+        INITIAL_USER_QUESTION: {user_instructions}
+        """,
+        input_variables=["user_instructions"]
+    )
+
+    routing_preprocessor = routing_preprocessor_prompt | llm | JsonOutputParser()
+
+    class PrimaryState(TypedDict):
+        messages: Annotated[Sequence[BaseMessage], operator.add]
+        user_instructions: str
+        user_instructions_data_wrangling: str
+        user_instructions_data_visualization: str
+        routing_preprocessor_decision: str
+        data_raw: Union[dict, list]
+        data_wrangled: dict
+        data_wrangler_function: str
+        data_visualization_function: str
+        plotly_graph: dict
+        plotly_error: str
+        max_retries: int
+        retry_count: int
+        
+        
+    def preprocess_routing(state: PrimaryState):
+        print("---PANDAS DATA ANALYST---")
+        print("*************************")
+        print("---PREPROCESS ROUTER---")
+        question = state.get("user_instructions")
+        
+        # Chart Routing and SQL Prep
+        response = routing_preprocessor.invoke({"user_instructions": question})
+        
+        return {
+            "user_instructions_data_wrangling": response.get('user_instructions_data_wrangling'),
+            "user_instructions_data_visualization": response.get('user_instructions_data_visualization'),
+            "routing_preprocessor_decision": response.get('routing_preprocessor_decision'),
+        }
+    
+    def router_chart_or_table(state: PrimaryState):
+        print("---ROUTER: CHART OR TABLE---")
+        return "chart" if state.get('routing_preprocessor_decision') == "chart" else "table"
+    
+    
+    def invoke_data_wrangling_agent(state: PrimaryState):
+        
+        response = data_wrangling_agent.invoke({
+            "user_instructions": state.get("user_instructions_data_wrangling"),
+            "data_raw": state.get("data_raw"),
+            "max_retries": state.get("max_retries"),
+            "retry_count": state.get("retry_count"),
+        })
+
+        return {
+            "messages": response.get("messages"),
+            "data_wrangled": response.get("data_wrangled"),
+            "data_wrangler_function": response.get("data_wrangler_function"),
+            "plotly_error": response.get("data_visualization_error"),
+            
+        }
+        
+    def invoke_data_visualization_agent(state: PrimaryState):
+        
+        response = data_visualization_agent.invoke({
+            "user_instructions": state.get("user_instructions_data_visualization"),
+            "data_raw": state.get("data_wrangled") if state.get("data_wrangled") else state.get("data_raw"),
+            "max_retries": state.get("max_retries"),
+            "retry_count": state.get("retry_count"),
+        })
+        
+        return {
+            "messages": response.get("messages"),
+            "data_visualization_function": response.get("data_visualization_function"),
+            "plotly_graph": response.get("plotly_graph"),
+            "plotly_error": response.get("data_visualization_error"),
+        }
+
+    def route_printer(state: PrimaryState):
+        print("---ROUTE PRINTER---")
+        print(f"    Route: {state.get('routing_preprocessor_decision')}")
+        print("---END---")
+        return {}
+    
+    workflow = StateGraph(PrimaryState)
+    
+    workflow.add_node("routing_preprocessor", preprocess_routing)
+    workflow.add_node("data_wrangling_agent", invoke_data_wrangling_agent)
+    workflow.add_node("data_visualization_agent", invoke_data_visualization_agent)
+    workflow.add_node("route_printer", route_printer)
+
+    workflow.add_edge(START, "routing_preprocessor")
+    workflow.add_edge("routing_preprocessor", "data_wrangling_agent")
+    
+    workflow.add_conditional_edges(
+        "data_wrangling_agent", 
+        router_chart_or_table,
+        {
+            "chart": "data_visualization_agent",
+            "table": "route_printer"
+        }
+    )
+    
+    workflow.add_edge("data_visualization_agent", "route_printer")
+    workflow.add_edge("route_printer", END)
+
+    app = workflow.compile(
+        checkpointer=checkpointer, 
+        name=AGENT_NAME
+    )
+    
+    return app

ai_data_science_team/multiagents/sql_data_analyst.py

@@ -1,12 +1,14 @@
 
 from langchain_core.messages import BaseMessage
-from langgraph.types import Checkpointer
+
+from langchain.prompts import PromptTemplate
+from langchain_core.output_parsers import JsonOutputParser
 
 from langgraph.graph import START, END, StateGraph
 from langgraph.graph.state import CompiledStateGraph
-from langgraph.types import Command
+from langgraph.types import Checkpointer
 
-from typing import TypedDict, Annotated, Sequence, Literal
+from typing import TypedDict, Annotated, Sequence
 import operator
 
 from typing_extensions import TypedDict
@@ -20,6 +22,7 @@ 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
 
+AGENT_NAME = "sql_data_analyst"
 
 class SQLDataAnalyst(BaseAgent):
     """
@@ -33,6 +36,8 @@ class SQLDataAnalyst(BaseAgent):
         The SQL Database Agent.
     data_visualization_agent: DataVisualizationAgent
         The Data Visualization Agent.
+    checkpointer: Checkpointer (optional)
+        The checkpointer to save the state of the multi-agent system.
         
     Methods:
     --------
@@ -326,17 +331,17 @@ def make_sql_data_analyst(
     """
     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()`
+    - Agent 1: SQL Database Agent made with `SQLDatabaseAgent()`
+    - Agent 2: Data Visualization Agent made with `DataVisualizationAgent()`
     
     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()`.
+        The SQL Database Agent made with `SQLDatabaseAgent()`.
     data_visualization_agent: CompiledStateGraph
-        The Data Visualization Agent made with `make_data_visualization_agent()`.
+        The Data Visualization Agent made with `DataVisualizationAgent()`.
     checkpointer: Checkpointer (optional)
         The checkpointer to save the state of the multi-agent system.
         Default: None
@@ -348,10 +353,39 @@ def make_sql_data_analyst(
     """
     
     llm = model
+    
+    
+    routing_preprocessor_prompt = PromptTemplate(
+        template="""
+        You are an expert in routing decisions for a SQL Database Agent, a Charting Visualization Agent, and a Pandas Table Agent. Your job is to:
+        
+        1. Determine what the correct format for a Users Question should be for use with a SQL Database Agent based on the incoming user question. Anything related to database and data manipulation should be passed along.
+        2. Determine whether or not a chart should be generated or a table should be returned based on the users question.
+        3. If a chart is requested, determine the correct format of a Users Question should be used with a Data Visualization Agent. Anything related to plotting and visualization should be passed along.
+        
+        Use the following criteria on how to route the the initial user question:
+        
+        From the incoming user question, remove any details about the format of the final response as either a Chart or Table and return only the important part of the incoming user question that is relevant for the SQL generator agent. This will be the 'user_instructions_sql_database'. If 'None' is found, return the original user question.
+        
+        Next, determine if the user would like a data visualization ('chart') or a 'table' returned with the results of the Data Wrangling Agent. If unknown, not specified or 'None' is found, then select 'table'.  
+        
+        If a 'chart' is requested, return the 'user_instructions_data_visualization'. If 'None' is found, return None.
+        
+        Return JSON with 'user_instructions_sql_database', 'user_instructions_data_visualization' and 'routing_preprocessor_decision'.
+        
+        INITIAL_USER_QUESTION: {user_instructions}
+        """,
+        input_variables=["user_instructions"]
+    )
+
+    routing_preprocessor = routing_preprocessor_prompt | llm | JsonOutputParser()
 
     class PrimaryState(TypedDict):
         messages: Annotated[Sequence[BaseMessage], operator.add]    
         user_instructions: str
+        user_instructions_sql_database: str
+        user_instructions_data_visualization: str
+        routing_preprocessor_decision: str
         sql_query_code: str
         sql_database_function: str
         data_sql: dict
@@ -359,39 +393,94 @@ def make_sql_data_analyst(
         plot_required: bool
         data_visualization_function: str
         plotly_graph: dict
+        plotly_error: str
         max_retries: int
         retry_count: int
         
-    def route_to_visualization(state) -> Command[Literal["data_visualization_agent", "__end__"]]: 
+    def preprocess_routing(state: PrimaryState):
+        print("---SQL DATA ANALYST---")
+        print("*************************")
+        print("---PREPROCESS ROUTER---")
+        question = state.get("user_instructions")
         
-        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')}")
+        # Chart Routing and SQL Prep
+        response = routing_preprocessor.invoke({"user_instructions": question})
         
-        if response.content == 'plot':
-            plot_required = True
-            goto="data_visualization_agent"
-        else:
-            plot_required = False
-            goto="__end__"
+        return {
+            "user_instructions_sql_database": response.get('user_instructions_sql_database'),
+            "user_instructions_data_visualization": response.get('user_instructions_data_visualization'),
+            "routing_preprocessor_decision": response.get('routing_preprocessor_decision'),
+        }
+    
+    def router_chart_or_table(state: PrimaryState):
+        print("---ROUTER: CHART OR TABLE---")
+        return "chart" if state.get('routing_preprocessor_decision') == "chart" else "table"
+    
+    
+    def invoke_sql_database_agent(state: PrimaryState):
         
-        return Command(
-            update={
-                'data_raw': state.get("data_sql"),
-                'plot_required': plot_required,
-            },
-            goto=goto
-        )
+        response = sql_database_agent.invoke({
+            "user_instructions": state.get("user_instructions_sql_database"),
+            "max_retries": state.get("max_retries"),
+            "retry_count": state.get("retry_count"),
+        })
 
-    workflow = StateGraph(PrimaryState)
+        return {
+            "messages": response.get("messages"),
+            "data_sql": response.get("data_sql"),
+            "sql_query_code": response.get("sql_query_code"),
+            "sql_database_function": response.get("sql_database_function"),
+            
+        }
+        
+    def invoke_data_visualization_agent(state: PrimaryState):
+        
+        response = data_visualization_agent.invoke({
+            "user_instructions": state.get("user_instructions_data_visualization"),
+            "data_raw": state.get("data_sql"),
+            "max_retries": state.get("max_retries"),
+            "retry_count": state.get("retry_count"),
+        })
+        
+        return {
+            "messages": response.get("messages"),
+            "data_visualization_function": response.get("data_visualization_function"),
+            "plotly_graph": response.get("plotly_graph"),
+            "plotly_error": response.get("data_visualization_error"),
+        }
 
-    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)
+    def route_printer(state: PrimaryState):
+        print("---ROUTE PRINTER---")
+        print(f"    Route: {state.get('routing_preprocessor_decision')}")
+        print("---END---")
+        return {}
+
+    workflow = StateGraph(PrimaryState)
+    
+    workflow.add_node("routing_preprocessor", preprocess_routing)
+    workflow.add_node("sql_database_agent", invoke_sql_database_agent)
+    workflow.add_node("data_visualization_agent", invoke_data_visualization_agent)
+    workflow.add_node("route_printer", route_printer)
 
-    workflow.add_edge(START, "sql_database_agent")
-    workflow.add_edge("sql_database_agent", "route_to_visualization")
-    workflow.add_edge("data_visualization_agent", END)
+    workflow.add_edge(START, "routing_preprocessor")
+    workflow.add_edge("routing_preprocessor", "sql_database_agent")
+    
+    workflow.add_conditional_edges(
+        "sql_database_agent", 
+        router_chart_or_table,
+        {
+            "chart": "data_visualization_agent",
+            "table": "route_printer"
+        }
+    )
+    
+    workflow.add_edge("data_visualization_agent", "route_printer")
+    workflow.add_edge("route_printer", END)
 
-    app = workflow.compile(checkpointer=checkpointer)
+    app = workflow.compile(
+        checkpointer=checkpointer, 
+        name=AGENT_NAME
+    )
 
     return app
 

ai_data_science_team/templates/agent_templates.py

@@ -40,6 +40,21 @@ class BaseAgent(CompiledStateGraph):
         self._params = params
         self._compiled_graph = self._make_compiled_graph()
         self.response = None
+        self.name = self._compiled_graph.name
+        self.checkpointer = self._compiled_graph.checkpointer
+        self.store = self._compiled_graph.store
+        self.output_channels = self._compiled_graph.output_channels
+        self.nodes = self._compiled_graph.nodes
+        self.stream_mode = self._compiled_graph.stream_mode
+        self.builder = self._compiled_graph.builder
+        self.channels = self._compiled_graph.channels
+        self.input_channels = self._compiled_graph.input_channels
+        self.input_schema = self._compiled_graph.input_schema
+        self.output_schema = self._compiled_graph.output_schema
+        self.debug = self._compiled_graph.debug
+        self.interrupt_after_nodes = self._compiled_graph.interrupt_after_nodes
+        self.interrupt_before_nodes = self._compiled_graph.interrupt_before_nodes
+        self.config = self._compiled_graph.config
 
     def _make_compiled_graph(self):
         """
@@ -197,6 +212,24 @@ class BaseAgent(CompiledStateGraph):
         """
         return self.get_output_jsonschema()['properties']
     
+    def get_state(self, config, *, subgraphs = False):
+        """
+        Returns the state of the agent.
+        """
+        return self._compiled_graph.get_state(config, subgraphs=subgraphs)
+    
+    def get_state_history(self, config, *, filter = None, before = None, limit = None):
+        """
+        Returns the state history of the agent.
+        """
+        return self._compiled_graph.get_state_history(config, filter=filter, before=before, limit=limit)
+    
+    def update_state(self, config, values, as_node = None):
+        """
+        Updates the state of the agent.
+        """
+        return self._compiled_graph.update_state(config, values, as_node)
+    
     def get_response(self):
         """
         Returns the response generated by the agent.
@@ -237,6 +270,7 @@ def create_coding_agent_graph(
     checkpointer: Optional[Callable] = None,
     bypass_recommended_steps: bool = False,
     bypass_explain_code: bool = False,
+    agent_name: str = "coding_agent"
 ):
     """
     Creates a generic agent graph using the provided node functions and node names.
@@ -281,6 +315,8 @@ def create_coding_agent_graph(
         Whether to skip the recommended steps node.
     bypass_explain_code : bool, optional
         Whether to skip the final explain code node.
+    name : str, optional
+        The name of the agent graph.
 
     Returns
     -------
@@ -366,10 +402,10 @@ def create_coding_agent_graph(
         workflow.add_edge(explain_code_node_name, END)
     
     # Finally, compile
-    if human_in_the_loop:
-        app = workflow.compile(checkpointer=checkpointer)
-    else:
-        app = workflow.compile()
+    app = workflow.compile(
+        checkpointer=checkpointer,
+        name=agent_name,
+    )
     
     return app
 
@@ -574,7 +610,7 @@ def node_func_execute_agent_from_sql_connection(
     
     # Retrieve SQLAlchemy connection and code snippet from the state
     is_engine = isinstance(connection, sql.engine.base.Engine)
-    conn = connection.connect() if is_engine else connection
+    connection = connection.connect() if is_engine else connection
     agent_code = state.get(code_snippet_key)
     
     # Ensure the connection object is provided