Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[WIP] Bedrock Agents Runnable #91

Closed
wants to merge 2 commits into from
Closed

[WIP] Bedrock Agents Runnable #91

wants to merge 2 commits into from

Conversation

3coins
Copy link
Collaborator

@3coins 3coins commented Jun 28, 2024

Description

This PR introduces a new Bedrock Agents Runnable that allows using Bedrock Agents with return of control functions as tools.

Usage

Define tools

from langchain_core.tools import tool

@tool("AssetDetail::getAssetValue")
def getAssetValue(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
def getMortgageRate(asset_holder_id: str = " ", asset_value: str = " ") -> str:
    """Get the mortgage rate based on asset value"""
    return (
        f"The mortgage rate for {asset_holder_id} "
        f"with asset value of {asset_value} is 8.87%"
    )

Create the agent and use with AgentExecutor

from langchain.agents import AgentExecutor

agent = BedrockAgentsRunnable.create_agent(
    {"agent_id": "UKYYJIV1O1", "enable_trace": True}
)
tools = [getAssetValue, getMortgageRate]
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 id AVC-1234 is 8.87%"

Notes

This change follows a strategy similar to the OpenAIAssistantRunnable to create an agent. However, it's still missing the callback related code that should be plugged in before merging. This code also assumes that the agent and the corresponding actions have been setup and prepared to call before calling; this setup could be automated in the agent creation or validation function, but missing here.

@3coins 3coins mentioned this pull request Aug 17, 2024
Merged
@3coins
Copy link
Collaborator Author

3coins commented Aug 17, 2024

Closing in favor of #152

@3coins 3coins closed this Aug 17, 2024
3coins added a commit that referenced this pull request Aug 21, 2024
@kcgourishetti @3coins
Kcgourishetti/bedrock agents runnable (#152)
e86b551 
## Description
This PR introduces a new Bedrock Agents Runnable that allows using
[Bedrock
Agents](https://docs.aws.amazon.com/bedrock/latest/userguide/agents-returncontrol.html)
with return of control functions as tools. This completes the work
presented in #91.

### Usage with AgentExecutor
```python
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
```python

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)
```

---------

Co-authored-by: Piyush Jain <piyushjain@duck.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers
No reviews
Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
1 participant
@3coins