Source code for langchain_aws.agents.base

from __future__ import annotations

import json
import logging
import time
import uuid
from collections import defaultdict
from typing import Any, Dict, List, Optional, Tuple, TypedDict, Union

import boto3
from botocore.client import Config
from botocore.exceptions import UnknownServiceError
from langchain_core.agents import AgentAction, AgentFinish
from langchain_core.callbacks import CallbackManager
from langchain_core.load import dumpd
from langchain_core.messages import AIMessage
from langchain_core.pydantic_v1 import root_validator
from langchain_core.runnables import RunnableConfig, RunnableSerializable, ensure_config
from langchain_core.tools import BaseTool

_DEFAULT_ACTION_GROUP_NAME = "DEFAULT_AG_"
_TEST_AGENT_ALIAS_ID = "TSTALIASID"


[docs]def get_boto_session( credentials_profile_name: Optional[str] = None, region_name: Optional[str] = None, endpoint_url: Optional[str] = None, ) -> Any: """ Construct the boto3 session """ if credentials_profile_name: session = boto3.Session(profile_name=credentials_profile_name) else: # use default credentials session = boto3.Session() client_params = { "config": Config( connect_timeout=120, read_timeout=120, retries={"max_attempts": 0} ) } if region_name: client_params["region_name"] = region_name if endpoint_url: client_params["endpoint_url"] = endpoint_url return client_params, session
[docs]def parse_agent_response(response: Any) -> OutputType: """ Parses the raw response from Bedrock Agent Args: response: The raw response from Bedrock Agent Returns Either a BedrockAgentAction or a BedrockAgentFinish """ response_text = "" event_stream = response["completion"] session_id = response["sessionId"] for event in event_stream: if "returnControl" in event: response_text = json.dumps(event) break if "chunk" in event: response_text = event["chunk"]["bytes"].decode("utf-8") agent_finish = BedrockAgentFinish( return_values={"output": response_text}, log=response_text, session_id=session_id, ) if not response_text: return agent_finish if "returnControl" not in response_text: return agent_finish return_control = json.loads(response_text).get("returnControl") if not return_control: return agent_finish invocation_inputs = return_control.get("invocationInputs") if not invocation_inputs: return agent_finish try: invocation_input = invocation_inputs[0].get("functionInvocationInput", {}) action_group = invocation_input.get("actionGroup", "") function = invocation_input.get("function", "") parameters = invocation_input.get("parameters", []) parameters_json = {} for parameter in parameters: parameters_json[parameter.get("name")] = parameter.get("value", "") tool = f"{action_group}::{function}" if _DEFAULT_ACTION_GROUP_NAME in action_group: tool = f"{function}" return [ BedrockAgentAction( tool=tool, tool_input=parameters_json, log=response_text, session_id=session_id, ) ] except Exception as ex: raise Exception("Parse exception encountered {}".format(repr(ex)))
def _create_bedrock_agent( bedrock_client: Any, agent_name: str, agent_resource_role_arn: str, instruction: str, foundation_model: str, client_token: Optional[str] = None, customer_encryption_key_arn: Optional[str] = None, description: Optional[str] = None, guardrail_configuration: Optional[GuardrailConfiguration] = None, idle_session_ttl_in_seconds: Optional[int] = None, ) -> Union[str, None]: """ Creates the bedrock agent """ create_agent_request: dict = { "agentName": agent_name, "agentResourceRoleArn": agent_resource_role_arn, "foundationModel": foundation_model, "instruction": instruction, } if description: create_agent_request["description"] = description if client_token: create_agent_request["clientToken"] = client_token if customer_encryption_key_arn: create_agent_request["customerEncryptionKeyArn"] = customer_encryption_key_arn if guardrail_configuration is not None: create_agent_request["guardrailConfiguration"] = { "guardrailIdentifier": guardrail_configuration["guardrail_identifier"], "guardrailVersion": guardrail_configuration["guardrail_version"] or "DRAFT", } if idle_session_ttl_in_seconds: create_agent_request["idleSessionTTLInSeconds"] = idle_session_ttl_in_seconds create_agent_response = bedrock_client.create_agent(**create_agent_request) request_id = create_agent_response.get("ResponseMetadata", {}).get("RequestId", "") logging.info(f"Create bedrock agent call successful with request id: {request_id}") agent_id = create_agent_response["agent"]["agentId"] create_agent_start_time = time.time() while time.time() - create_agent_start_time < 10: agent_creation_status = ( bedrock_client.get_agent(agentId=agent_id) .get("agent", {}) .get("agentStatus", {}) ) if agent_creation_status == "NOT_PREPARED": return agent_id else: time.sleep(2) logging.error(f"Failed to create bedrock agent {agent_id}") raise Exception(f"Failed to create bedrock agent {agent_id}") def _get_action_group_and_function_names(tool: BaseTool) -> Tuple[str, str]: """ Convert the LangChain 'Tool' into Bedrock Action Group name and Function name """ action_group_name = _DEFAULT_ACTION_GROUP_NAME function_name = tool.name tool_name_split = tool.name.split("::") if len(tool_name_split) > 1: action_group_name = tool_name_split[0] function_name = tool_name_split[1] return action_group_name, function_name def _create_bedrock_action_groups( bedrock_client: Any, agent_id: str, tools: List[BaseTool] ) -> None: """Create the bedrock action groups for the agent""" tools_by_action_group = defaultdict(list) for tool in tools: action_group_name, function_name = _get_action_group_and_function_names(tool) tools_by_action_group[action_group_name].append(tool) for action_group_name, functions in tools_by_action_group.items(): bedrock_client.create_agent_action_group( actionGroupName=action_group_name, actionGroupState="ENABLED", actionGroupExecutor={"customControl": "RETURN_CONTROL"}, functionSchema={ "functions": [_tool_to_function(function) for function in functions] }, agentId=agent_id, agentVersion="DRAFT", ) def _tool_to_function(tool: BaseTool) -> dict: """ Convert LangChain tool to a Bedrock function schema """ _, function_name = _get_action_group_and_function_names(tool) function_parameters = {} for arg_name, arg_details in tool.args.items(): function_parameters[arg_name] = { "description": arg_details.get( "description", arg_details.get("title", arg_name) ), "type": arg_details.get("type", "string"), "required": not bool(arg_details.get("default", None)), } return { "description": tool.description, "name": function_name, "parameters": function_parameters, } def _prepare_agent(bedrock_client: Any, agent_id: str) -> None: """ Prepare the agent for invocations """ bedrock_client.prepare_agent(agentId=agent_id) prepare_agent_start_time = time.time() while time.time() - prepare_agent_start_time < 10: agent_status = bedrock_client.get_agent(agentId=agent_id) if agent_status.get("agent", {}).get("agentStatus", "") == "PREPARED": return else: time.sleep(2) raise Exception(f"Timed out while preparing the agent with id {agent_id}") def _get_bedrock_agent(bedrock_client: Any, agent_name: str) -> Any: """ Get the agent by name """ next_token = None while True: if next_token: list_agents_response = bedrock_client.list_agents( maxResults=1000, nextToken=next_token ) else: list_agents_response = bedrock_client.list_agents(maxResults=1000) agent_summaries = list_agents_response.get("agentSummaries", []) next_token = list_agents_response.get("nextToken") agent_summary = next( (x for x in agent_summaries if x.get("agentName") == agent_name), None ) if agent_summary: return agent_summary if next_token is None: return None
[docs]class BedrockAgentFinish(AgentFinish): """AgentFinish with session id information. Parameters: session_id: Session id """ session_id: str @classmethod def is_lc_serializable(cls) -> bool: """Check if the class is serializable by LangChain. Returns: False """ return False
[docs]class BedrockAgentAction(AgentAction): """AgentAction with session id information. Parameters: session_id: session id """ session_id: str @classmethod def is_lc_serializable(cls) -> bool: """Check if the class is serializable by LangChain. Returns: False """ return False
OutputType = Union[List[BedrockAgentAction], BedrockAgentFinish]
[docs]class GuardrailConfiguration(TypedDict): guardrail_identifier: str guardrail_version: str
[docs]class BedrockAgentsRunnable(RunnableSerializable[Dict, OutputType]): """ Invoke a Bedrock Agent """ agent_id: Optional[str] """Bedrock Agent Id""" agent_alias_id: Optional[str] = _TEST_AGENT_ALIAS_ID """Bedrock Agent Alias Id""" client: Any """Boto3 client""" region_name: Optional[str] = None """Region""" credentials_profile_name: Optional[str] = None """Credentials to use to invoke the agent""" endpoint_url: Optional[str] = None """Endpoint URL""" @root_validator(skip_on_failure=True) def validate_agent(cls, values: dict) -> dict: if values.get("client") is not None: return values try: client_params, session = get_boto_session( credentials_profile_name=values["credentials_profile_name"], region_name=values["region_name"], endpoint_url=values["endpoint_url"], ) values["client"] = session.client("bedrock-agent-runtime", **client_params) return values except ImportError: raise ModuleNotFoundError( "Could not import boto3 python package. " "Please install it with `pip install boto3`." ) except UnknownServiceError as e: raise ModuleNotFoundError( "Ensure that you have installed the latest boto3 package " "that contains the API for `bedrock-runtime-agent`." ) from e except Exception as e: raise ValueError( "Could not load credentials to authenticate with AWS client. " "Please check that credentials in the specified " "profile name are valid." ) from e
[docs] @classmethod def create_agent( cls, agent_name: str, agent_resource_role_arn: str, foundation_model: str, instruction: str, tools: List[BaseTool] = [], *, client_token: Optional[str] = None, customer_encryption_key_arn: Optional[str] = None, description: Optional[str] = None, guardrail_configuration: Optional[GuardrailConfiguration] = None, idle_session_ttl_in_seconds: Optional[int] = None, credentials_profile_name: Optional[str] = None, region_name: Optional[str] = None, bedrock_endpoint_url: Optional[str] = None, runtime_endpoint_url: Optional[str] = None, **kwargs: Any, ) -> BedrockAgentsRunnable: """ Creates a Bedrock Agent Runnable that can be used with an AgentExecutor or with LangGraph. This also sets up the Bedrock agent, actions and action groups infrastructure if they don't exist, ensures the agent is in PREPARED state so that it is ready to be called. Args: agent_name: Name of the agent agent_resource_role_arn: The Amazon Resource Name (ARN) of the IAM role with permissions to invoke API operations on the agent. foundation_model: The foundation model to be used for orchestration by the agent you create instruction: Instructions that tell the agent what it should do and how it should interact with users tools: List of tools. Accepts LangChain's BaseTool format client_token: A unique, case-sensitive identifier to ensure that the API request completes no more than one time. If this token matches a previous request, Amazon Bedrock ignores the request, but does not return an error customer_encryption_key_arn: The Amazon Resource Name (ARN) of the KMS key with which to encrypt the agent description: A description of the agent guardrail_configuration: The unique Guardrail configuration assigned to the agent when it is created. idle_session_ttl_in_seconds: The number of seconds for which Amazon Bedrock keeps information about a user's conversation with the agent. A user interaction remains active for the amount of time specified. If no conversation occurs during this time, the session expires and Amazon Bedrock deletes any data provided before the timeout credentials_profile_name: The profile name to use if different from default region_name: Region for the Bedrock agent bedrock_endpoint_url: Endpoint URL for bedrock agent runtime_endpoint_url: Endpoint URL for bedrock agent runtime **kwargs: Additional arguments Returns: BedrockAgentsRunnable configured to invoke the Bedrock agent """ client_params, session = get_boto_session( credentials_profile_name=credentials_profile_name, region_name=region_name, endpoint_url=bedrock_endpoint_url, ) bedrock_client = session.client("bedrock-agent", **client_params) bedrock_agent = _get_bedrock_agent( bedrock_client=bedrock_client, agent_name=agent_name ) if bedrock_agent: agent_id = bedrock_agent["agentId"] agent_status = bedrock_agent["agentStatus"] if agent_status != "PREPARED": _prepare_agent(bedrock_client, agent_id) else: try: agent_id = _create_bedrock_agent( bedrock_client=bedrock_client, agent_name=agent_name, agent_resource_role_arn=agent_resource_role_arn, instruction=instruction, foundation_model=foundation_model, client_token=client_token, customer_encryption_key_arn=customer_encryption_key_arn, description=description, guardrail_configuration=guardrail_configuration, idle_session_ttl_in_seconds=idle_session_ttl_in_seconds, ) _create_bedrock_action_groups(bedrock_client, agent_id, tools) _prepare_agent(bedrock_client, agent_id) except Exception as exception: logging.error(f"Error in create agent call: {exception}") raise exception return cls( agent_id=agent_id, region_name=region_name, credentials_profile_name=credentials_profile_name, endpoint_url=runtime_endpoint_url, **kwargs, )
[docs] def invoke( self, input: Dict, config: Optional[RunnableConfig] = None ) -> OutputType: """ Invoke the Bedrock agent. Args: input: The LangChain Runnable input dictionary that can include: input: The input text to the agent memory_id: The memory id to use for an agent with memory enabled session_id: The session id to use. If not provided, a new session will be started end_session: Boolean indicating whether to end a session or not intermediate_steps: The intermediate steps that are used to provide RoC invocation details enable_trace: Boolean flag to enable trace when invoke bedrock agent Returns: Union[List[BedrockAgentAction], BedrockAgentFinish] """ config = ensure_config(config) callback_manager = CallbackManager.configure( inheritable_callbacks=config.get("callbacks"), inheritable_tags=config.get("tags"), inheritable_metadata=config.get("metadata"), ) run_manager = callback_manager.on_chain_start( dumpd(self), input, name=config.get("run_name") ) try: agent_input = { "agentId": self.agent_id, "agentAliasId": self.agent_alias_id, "enableTrace": input.get("enable_trace", False), "endSession": bool(input.get("end_session", False)), } if input.get("memory_id"): agent_input["memoryId"] = input.get("memory_id") if input.get("intermediate_steps"): session_id, session_state = self._parse_intermediate_steps( input.get("intermediate_steps") # type: ignore[arg-type] ) if session_id is not None: agent_input["sessionId"] = session_id if session_state is not None: agent_input["sessionState"] = session_state else: agent_input["inputText"] = input.get("input", "") agent_input["sessionId"] = input.get("session_id", str(uuid.uuid4())) output = self.client.invoke_agent(**agent_input) except Exception as e: run_manager.on_chain_error(e) raise e try: response = parse_agent_response(output) except Exception as e: run_manager.on_chain_error(e) raise e else: run_manager.on_chain_end(response) return response
def _parse_intermediate_steps( self, intermediate_steps: List[Tuple[BedrockAgentAction, str]] ) -> Tuple[Union[str, None], Union[Dict[str, Any], None]]: last_step = max(0, len(intermediate_steps) - 1) action = intermediate_steps[last_step][0] tool_invoked = action.tool messages = action.messages session_id = action.session_id if tool_invoked: action_group_name = _DEFAULT_ACTION_GROUP_NAME function_name = tool_invoked tool_name_split = tool_invoked.split("::") if len(tool_name_split) > 1: action_group_name = tool_name_split[0] function_name = tool_name_split[1] if messages: last_message = max(0, len(messages) - 1) message = messages[last_message] if type(message) is AIMessage: response = intermediate_steps[last_step][1] session_state = { "invocationId": json.loads(message.content) # type: ignore[arg-type] .get("returnControl", {}) .get("invocationId", ""), "returnControlInvocationResults": [ { "functionResult": { "actionGroup": action_group_name, "function": function_name, "responseBody": {"TEXT": {"body": response}}, } } ], } return session_id, session_state return None, None