Kcgourishetti/bedrock agents runnable

[kcgourishetti]

Description

This PR introduces a new Bedrock Agents Runnable that allows using Bedrock Agents with return of control functions as tools. This completes the work presented in #91.

Usage with AgentExecutor

import boto3

from langchain.agents import AgentExecutor
from langchain_core.tools import tool

from langchain_aws.agents.base import  BedrockAgentsRunnable

def delete_agent(agent_id: str) -> None:
    bedrock_client = boto3.client("bedrock-agent")
    bedrock_client.delete_agent(agentId=agent_id, skipResourceInUseCheck=True)

@tool("AssetDetail::getAssetValue")
def get_asset_value(asset_holder_id: str) -> str:
    """Get the asset value for an owner id"""
    return f"The total asset value for {asset_holder_id} is 100K"

@tool("AssetDetail::getMortgageRate")
def get_mortgage_rate(asset_holder_id: str, asset_value: str) -> str:
    """Get the mortgage rate based on asset value"""
    return (
        f"The mortgage rate for the asset holder id {asset_holder_id}"
        f"with asset value of {asset_value} is 8.87%"
    )

foundation_model = "anthropic.claude-3-sonnet-20240229-v1:0"
tools = [get_asset_value, get_mortgage_rate]
agent = None
try:
    agent = BedrockAgentsRunnable.create_agent(
        agent_name="mortgage_interest_rate_agent",
        agent_resource_role_arn="<agent-resource-role-arn>",
        foundation_model=foundation_model,
        instruction=(
            "You are an agent who helps with getting the mortgage rate based on "
            "the current asset valuation"
        ),
        tools=tools,
    )
    agent_executor = AgentExecutor(agent=agent, tools=tools)  # type: ignore[arg-type]
    output = agent_executor.invoke(
        {"input": "what is my mortgage rate for id AVC-1234"}
    )

    assert output["output"] == (
        "The mortgage rate for the asset holder id AVC-1234 "
        "with an asset value of 100K is 8.87%."
    )
except Exception as ex:
    raise ex
finally:
    if agent:
        delete_agent(agent.agent_id)

Usage with LangGraph

import boto3

from langgraph.graph import END, START, StateGraph
from langgraph.prebuilt.tool_executor import ToolExecutor

from langchain_aws.agents.base import (
    BedrockAgentAction,
    BedrockAgentFinish,
    BedrockAgentsRunnable,
)

def delete_agent(agent_id: str) -> None:
    bedrock_client = boto3.client("bedrock-agent")
    bedrock_client.delete_agent(agentId=agent_id, skipResourceInUseCheck=True)

@tool
def get_weather(location: str = "") -> str:
    """
    Get the weather of a location

    Args:
        location: location of the place
    """
    if location.lower() == "seattle":
        return f"It is raining in {location}"
    return f"It is hot and humid in {location}"

class AgentState(TypedDict):
    input: str
    output: Union[BedrockAgentAction, BedrockAgentFinish, None]
    intermediate_steps: Annotated[
        list[tuple[BedrockAgentAction, str]], operator.add
    ]

def get_weather_agent_node() -> Tuple[BedrockAgentsRunnable, str]:
    foundation_model = "anthropic.claude-3-sonnet-20240229-v1:0"
    tools = [get_weather]
    try:
        agent_resource_role_arn = _create_agent_role(
            agent_region="us-west-2", 
            foundation_model=foundation_model
        )
        agent = BedrockAgentsRunnable.create_agent(
            agent_name="weather_agent",
            agent_resource_role_arn=agent_resource_role_arn,
            foundation_model=foundation_model,
            instruction=(
                "You are an agent who helps with getting weather for a given "
                "location"
            ),
            tools=tools,
        )

        return agent, agent_resource_role_arn
    except Exception as e:
        raise e

agent_runnable, agent_resource_role_arn = get_weather_agent_node()

def run_agent(data):
    agent_outcome = agent_runnable.invoke(data)
    return {"output": agent_outcome}

tool_executor = ToolExecutor([get_weather])

# Define the function to execute tools
def execute_tools(data):
    # Get the most recent output - this is the key added in the `agent` above
    agent_action = data["output"]
    output = tool_executor.invoke(agent_action[0])
    tuple_output = agent_action[0], output
    return {"intermediate_steps": [tuple_output]}

def should_continue(data):
    output_ = data["output"]

    # If the agent outcome is a list of BedrockAgentActions,
    # then we continue to tool execution
    if (
        isinstance(output_, list)
        and len(output_) > 0
        and isinstance(output_[0], BedrockAgentAction)
    ):
        return "continue"

    # If the agent outcome is an AgentFinish, then we return `exit` string
    # This will be used when setting up the graph to define the flow
    if isinstance(output_, BedrockAgentFinish):
        return "end"

    # Unknown output from the agent, end the graph
    return "end"

try:
    # Define a new graph
    workflow = StateGraph(AgentState)

    # Define the two nodes we will cycle between
    workflow.add_node("agent", run_agent)
    workflow.add_node("action", execute_tools)

    # Set the entrypoint as `agent`
    # This means that this node is the first one called
    workflow.add_edge(START, "agent")

    # We now add a conditional edge
    workflow.add_conditional_edges(
        # First, we define the start node. We use `agent`.
        # This means these are the edges taken after the `agent` node is called.
        "agent",
        # Next, we pass in the function that will determine which node
        # will be called next.
        should_continue,
        # Finally we pass in a mapping.
        # The keys are strings, and the values are other nodes.
        # END is a special node marking that the graph should finish.
        # What will happen is we will call `should_continue`, and then the output
        # of that will be matched against the keys in this mapping.
        # The matched node will then be called.
        {
            # If `tools`, then we call the tool node.
            "continue": "action",
            # Otherwise we finish.
            "end": END,
        },
    )

    # We now add a normal edge from `tools` to `agent`.
    # This means that after `tools` is called, `agent` node is called next.
    workflow.add_edge("action", "agent")

    # Finally, we compile it!
    # This compiles it into a LangChain Runnable,
    # meaning you can use it as you would any other runnable
    app = workflow.compile()

    inputs = {"input": "what is the weather in seattle?"}
    final_state = app.invoke(inputs)

    assert isinstance(final_state.get("output", {}), BedrockAgentFinish)
    assert (
        final_state.get("output").return_values["output"]
        == "It is raining in Seattle"
    )
finally:
    if agent_runnable:
        delete_agent(agent_id=agent_runnable.agent_id)

[3coins]

@kcgourishetti
Thanks for submitting these changes. Some minor suggestions before we merge.

[3coins]

Not a priority, but time.sleep is blocking, we might want to look into asyncio.sleep.

[3coins]

@kcgourishetti
Great job on all the changes. We are almost there, some minor comments on aligning the create_job with the create_agent API, and moving notebooks to samples directory.

[3coins]

Is it possible that these 2 can be different?

[kcgourishetti]

yes, one is the endpoint for control plane responsible for CRUDL of the agents and one is the runtime responsible for invoke etc. the boto client names for each of these are also different one is bedrock-agent and other is bedrock-agent-runtime.

[3coins]

@kcgourishetti
I took care of some of the formatting, linting, typing and CI errors. One last thing I noticed is that the notebook bedrock_agent_langgraph.ipynb seems to be incomplete, it doesn't have any of the code for the LangGraph executor, you seem to have some of this code in the integration test.

A couple of other observations, not blockers on this PR:

1. Even though there is no LLM present here (agent manages the LLM calls), a lot of the LangChain components work around messages which is central to it's working. We should follow up with a PR to support messages.
2. The LangGraph integration test is using the ToolExecutor, which has been deprecated in favor of ToolNode, which works only with messages. To future proof this work, we should move to use ToolNode in LangGraph samples and tests.

[kcgourishetti]

@kcgourishetti I took care of some of the formatting, linting, typing and CI errors. One last thing I noticed is that the notebook bedrock_agent_langgraph.ipynb seems to be incomplete, it doesn't have any of the code for the LangGraph executor, you seem to have some of this code in the integration test.

A couple of other observations, not blockers on this PR:

1. Even though there is no LLM present here (agent manages the LLM calls), a lot of the LangChain components work around messages which is central to it's working. We should follow up with a PR to support messages.
2. The LangGraph integration test is using the ToolExecutor, which has been deprecated in favor of ToolNode, which works only with messages. To future proof this work, we should move to use ToolNode in LangGraph samples and tests.

• bedrock_agent_langgraph.ipynb: weird. looks like the file got overwritten. updated and committed it now.
• messages: ack. happy to hear if you have any thoughts on it.
• ToolExecutor: ah, thats a nice catch. we gotta make this work with messages then. thanks for the detailed info.