121 lines
4.7 KiB
Python
121 lines
4.7 KiB
Python
import operator
|
|
from typing import TypedDict, Annotated
|
|
|
|
from langchain_aws import ChatBedrockConverse
|
|
from langchain_core.messages import AIMessage, ToolMessage
|
|
from langgraph.graph import StateGraph, END
|
|
|
|
from .config import REGION, AWS_ACCOUNT
|
|
|
|
|
|
class AgentState(TypedDict):
|
|
messages: Annotated[list, operator.add]
|
|
current_step: str
|
|
|
|
|
|
def create_bedrock_llm(model_id: str, region: str = REGION, tools: list = None):
|
|
"""
|
|
Create a ChatBedrock instance using a model ID.
|
|
|
|
Args:
|
|
model_id: Bedrock model ID (e.g., anthropic.claude-haiku-4-5-20251001-v1:0)
|
|
region: AWS region (default: REGION env var)
|
|
tools: List of LangChain tools to bind to the model
|
|
|
|
Returns:
|
|
ChatBedrock instance configured with the model
|
|
"""
|
|
MODEL_ARNS = {
|
|
"anthropic.claude-haiku-4-5-20251001-v1:0": f"arn:aws:bedrock:{REGION}:{AWS_ACCOUNT}:inference-profile/us.anthropic.claude-haiku-4-5-20251001-v1:0",
|
|
"anthropic.claude-sonnet-4-5-20250929-v1:0": f"arn:aws:bedrock:{REGION}:{AWS_ACCOUNT}:inference-profile/global.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
|
"meta.llama4-maverick-17b-instruct-v1:0": f"arn:aws:bedrock:{REGION}:{AWS_ACCOUNT}:inference-profile/us.meta.llama4-maverick-17b-instruct-v1:0",
|
|
"meta.llama4-scout-17b-instruct-v1:0": f"arn:aws:bedrock:{REGION}:{AWS_ACCOUNT}:inference-profile/us.meta.llama4-scout-17b-instruct-v1:0",
|
|
"amazon.nova-lite-v1:0": f"arn:aws:bedrock:{REGION}:{AWS_ACCOUNT}:inference-profile/us.amazon.nova-lite-v1:0",
|
|
"amazon.nova-pro-v1:0": f"arn:aws:bedrock:{REGION}:{AWS_ACCOUNT}:inference-profile/us.amazon.nova-pro-v1:0",
|
|
"amazon.nova-2-lite-v1:0": f"arn:aws:bedrock:{REGION}:{AWS_ACCOUNT}:inference-profile/global.amazon.nova-2-lite-v1:0",
|
|
}
|
|
PROVIDER = {
|
|
"anthropic.claude-haiku-4-5-20251001-v1:0": "anthropic",
|
|
"anthropic.claude-sonnet-4-5-20250929-v1:0": "anthropic",
|
|
"meta.llama4-maverick-17b-instruct-v1:0": "meta",
|
|
"meta.llama4-scout-17b-instruct-v1:0": "meta",
|
|
"amazon.nova-lite-v1:0": "amazon",
|
|
"amazon.nova-pro-v1:0": "amazon",
|
|
"amazon.nova-2-lite-v1:0": "amazon",
|
|
}
|
|
prefix = {
|
|
"anthropic.claude-haiku-4-5-20251001-v1:0": "us",
|
|
"anthropic.claude-sonnet-4-5-20250929-v1:0": "global",
|
|
"meta.llama4-maverick-17b-instruct-v1:0": "us",
|
|
"meta.llama4-scout-17b-instruct-v1:0": "us",
|
|
"amazon.nova-lite-v1:0": "us",
|
|
"amazon.nova-pro-v1:0": "us",
|
|
"amazon.nova-2-lite-v1:0": "global",
|
|
}
|
|
llm = ChatBedrockConverse(
|
|
model_id=prefix[model_id] + "." + model_id,
|
|
region_name=region,
|
|
provider=PROVIDER[model_id],
|
|
max_tokens=2048,
|
|
temperature=0.7,
|
|
)
|
|
return llm.bind_tools(tools or [])
|
|
|
|
|
|
def call_model(state: AgentState, llm) -> AgentState:
|
|
"""Call the LLM with tools."""
|
|
response = llm.invoke(state["messages"])
|
|
state["current_step"] = "model_called"
|
|
return {"messages": [response]}
|
|
|
|
|
|
def call_tools(state: AgentState, tools_map: dict) -> AgentState:
|
|
"""Execute any tool calls from the LLM response."""
|
|
last_message = state["messages"][-1]
|
|
|
|
if hasattr(last_message, "tool_calls") and last_message.tool_calls:
|
|
tool_messages = []
|
|
for tool_call in last_message.tool_calls:
|
|
result = tools_map[tool_call["name"]].invoke(tool_call["args"])
|
|
tool_messages.append(ToolMessage(content=str(result), tool_call_id=tool_call["id"]))
|
|
|
|
state["current_step"] = "tools_executed"
|
|
return {"messages": tool_messages}
|
|
else:
|
|
state["current_step"] = "no_tools"
|
|
return {"messages": []}
|
|
|
|
|
|
def should_continue(state: AgentState) -> str:
|
|
"""Determine if we should continue to tools or end."""
|
|
last_message = state["messages"][-1]
|
|
if hasattr(last_message, "tool_calls") and last_message.tool_calls:
|
|
return "tools"
|
|
return "end"
|
|
|
|
|
|
def create_agent(inference_profile_arn: str, region: str = REGION, tools: list = None):
|
|
"""
|
|
Create a LangGraph agent that uses Bedrock inference profile with tools.
|
|
|
|
Args:
|
|
inference_profile_arn: ARN of the cross-region inference profile
|
|
region: AWS region
|
|
tools: List of LangChain tools to bind to the model
|
|
|
|
Returns:
|
|
Compiled LangGraph workflow
|
|
"""
|
|
tools = tools or []
|
|
llm = create_bedrock_llm(inference_profile_arn, region, tools)
|
|
tools_map = {t.name: t for t in tools}
|
|
|
|
workflow = StateGraph(AgentState)
|
|
workflow.add_node("model", lambda state: call_model(state, llm))
|
|
workflow.add_node("tools", lambda state: call_tools(state, tools_map))
|
|
workflow.set_entry_point("model")
|
|
workflow.add_conditional_edges("model", should_continue, {"tools": "tools", "end": END})
|
|
workflow.add_edge("tools", "model")
|
|
|
|
return workflow.compile()
|