Building a graph-based implementation of create_pandas_dataframe_agent

[andrw3000]

Transitioning from LangChain to LangGraph. Very impressed with the new library for stateful implementations of agentic systems.

This discussion is to develop a mapping between libraries for the example of re-implementing the create_pandas_dataframe_agent in LangGraph. This is a ReAct agent which uses the PythonREPLTool. There are a few subtleties that I'd like to raise to the developers, so to follow the principles of the library. I'm happy to contribute the implementation as an example notebook once resolved.

This post is meant for pedagogy, to assist the transition from LangChain to LangGraph. If there is reference content that answers these points, please forward in comments. Any engagement appreciated.

Rough implementation:

• Following the creat_react_agent as a template.

• Uses a basic state:

class AgentState(TypedDict):
    """Dictionary defining a state of the Pandas Agent."""
    messages: Annotated[Sequence[BaseMessage], add_messages]
    is_last_step: IsLastStep

• Uses the following graph:

def create_pandas_agent(llm: LanguageModelLike, df: pd.DataFrame) -> CompiledGraph:

    # Bind tools to llm
    tools = [PythonAstREPLTool(locals={'df':  df}),]
    llm = llm.bind_tools(tools)

    # Get prompt and define runnable
    prompt = get_df_prompt(df, tools)
    runnable = prompt | llm
    
    # Define state
    workflow = StateGraph(AgentState)

    # Define BoTorch generation node
    workflow.add_node('generate_code', generation_node)

    # Define tool node
    workflow.add_node('python_executor', ToolNode(tools))

    # Construct graph
    workflow.set_entry_point('generate_code')
    workflow.add_conditional_edges(
        'generate_code',  # Start node
        continue_while_tool_called,  # Router function
        # Conditional edge
        {
            # If `tools`, then we call the tool node.
            'continue': 'python_executor',
            # Otherwise we finish.
            'end': END,
        },
    )
    workflow.add_edge('python_executor', 'generate_code')  # Directed edge

    # Initialize memory to persist state between graph runs
    checkpointer = MemorySaver()
    
    return workflow.compile(checkpointer=checkpointer)

Here are my queries:

1. Passing a pd.DataFrame object to the state. Does this fit the LangGraph principles, or is it intended that the state should only contain message-type objects. Currently, the tool is fixed at the creation of the nodes with python workflow.add_node('python_executor', ToolNode([PythonAstREPLTool(locals={'df': df}),])) Here I propose adding a field to the state graph state["df"] = [df,] at the input. Then I would like to be able to
   (a) Pass state["df"] to the ToolExecutor when the node is visited. This would possibly mean instantiating a PythonAstREPLTool, in which case it is unclear how to bind the tool to the llm at the start llm.bin_tools().
   (b) Allow the tool to update state["df"], say when ammending a df. A question here is, would a second df2 be added to a list state["df"] = [df, df2], and how can the tool be instructed to do so.

2. I would like to be able to replicate the following prompt: from langchain.agents.mrkl.prompt import FORMAT_INSTRUCTIONS. The takes the following form:

PREFIX = """Answer the following questions as best you can. You have access to the following tools:"""
FORMAT_INSTRUCTIONS = """Use the following format:

Question: the input question you must answer
Thought: you should always think about what to do
Action: the action to take, should be one of [{tool_names}]
Action Input: the input to the action
Observation: the result of the action
... (this Thought/Action/Action Input/Observation can repeat N times)
Thought: I now know the final answer
Final Answer: the final answer to the original input question"""
SUFFIX = """Begin!

Question: {input}
Thought:{agent_scratchpad}"""

Currently, I supply tool_names at the creation of the node where I call

prompt = prompt.partial(tool_names=tools)
runnable = prompt | llm

However, I would like to be able to parse variables to the prompt from the state, so I could update agent_scratchpad properly. How is this executed?

3. Lastly, when creating a node function, e.g.

    def generation_node(
        state: AgentState,
        config: RunnableConfig,
    ):
        messages = state["messages"]
        response = runnable.invoke(messages, config)
        if state["is_last_step"] and response.tool_calls:
            return {
                "messages": [
                    AIMessage(
                        id=response.id,
                        content="Sorry, need more steps to process this request.",
                    )
                ]
            }
        return {"messages": [response]}

Can one pass items such as the runnable in the node function arguments, like

    def generation_node(
        state: AgentState,
        config: RunnableConfig,  
        runnable
    ):

so that workflow.add_node('generate_code', generation_node) still still has access to runnable in its name space?

Once again, any comments/thoughts appreciated.

[andrw3000]

@FreddyWordingham - this will be up your street!

[hwchase17]

3/ - no, the runnable should be either in the global space or initialized inside
2/ if im understanding correctly, you can just put that parsing logic in a node:

def node(state, config):
    agent_scratchpad = parse_state(state)
    response = runnable.invoke({"agent_scratchpad": agent_scratchpad, "messages": state['messages']})
    ....

1/ is a big longer of a response. TLDR is that we're still thinking about a recomended way to handle. But regardless of what we recommend, you should be able to do it as you described. For (a) you will want to create and bind tools inside a node, eg:

def node(state, config):
      tools = [PythonREPL(df=state['df'])
      llm_with_tools = llm.bind_tools(tools)
      response = llm_with_tools.invoke....

For (b), you will want to NOT use the ToolNode (as that is a simple node that does not support this) and youll want to do the invocation your self
Again, we are working on an easier and recommended path for (1), will udpate...

[nguyenph17]

Hi @andrw3000
I appreciate it if you could share your code

[sunilkhadka]

Same. @andrw3000 . Can you share the code? I would also really appreciate it.

[sunilkhadka139]

hi @hwchase17 , anything related to the update?

[Krishn1412]

Hey @andrw3000, did you contribute this implementation?

[PvH-SPC]

Hi All, how far did you all get on this, im very interested in being able to 1. store the Python as part of the state, to allow the user to edit the python code with human in the loop, 2. have a reference to the result data as a data frame itself, and 3. have the graph able to generate charts/images with said data frame. Maybe even allow it to maintain state of a history of question and answer pairs in the state so you can do higher order logic across result sets from time. I should mention that it was pretty trival to wrap the create_pandas_dataframe_agent as a node (the solo node), in the graph, and im currently playing around with looking to use the return intermediate steps option, but this is really just an ugly hack, a complete graph based refactor is what I had in mind. Also ideally the finial implementation should probably use a better pyton repl sandbox ideally azure container apps interpreter sessions.