482 lines
17 KiB
Python
482 lines
17 KiB
Python
"""
|
|
LangGraph Agent using AWS Bedrock Cross-Region Inference Profile with Tools
|
|
|
|
This script demonstrates how to create a LangGraph agent that uses
|
|
an AWS Bedrock inference profile with custom tools (add and multiply).
|
|
"""
|
|
import boto3
|
|
from typing import TypedDict, Annotated
|
|
from langgraph.graph import StateGraph, END
|
|
from langchain_aws import ChatBedrockConverse
|
|
from langchain_core.tools import tool
|
|
from langchain_core.messages import HumanMessage, AIMessage, ToolMessage, SystemMessage
|
|
import operator
|
|
import json
|
|
import time
|
|
from langfuse import Langfuse
|
|
from langfuse.langchain import CallbackHandler
|
|
from botocore.exceptions import ClientError
|
|
import os
|
|
WORKGROUP = "iceberg-workgroup"
|
|
DATABASE = "dnx_warehouse"
|
|
TABLE = "poc_dnx_monthly_summary"
|
|
REGION = "us-east-1"
|
|
|
|
# DynamoDB client
|
|
dynamodb = boto3.resource("dynamodb", region_name=REGION)
|
|
|
|
|
|
@tool
|
|
def get_monthly_report(id: str, variable: str) -> str:
|
|
"""
|
|
Get a specific variable's data from DynamoDB for a specific id.
|
|
|
|
Args:
|
|
id: The id of the data
|
|
variable: The variable/column name to retrieve from the table
|
|
|
|
Returns:
|
|
The content of the specified variable for the given id
|
|
"""
|
|
print(f"\n🔧 [TOOL CALLED] get_monthly_report for month: {id}, variable: {variable}")
|
|
|
|
try:
|
|
table = dynamodb.Table(TABLE)
|
|
response = table.get_item(Key={"id": id})
|
|
|
|
if "Item" not in response:
|
|
return f"No report found for month: {id}"
|
|
|
|
item = response["Item"]
|
|
content = item.get(variable, "")
|
|
|
|
if not content:
|
|
return f"Variable '{variable}' not found for month: {id}"
|
|
|
|
result = f"<{id}>\n{content}\n</{id}>"
|
|
return result
|
|
|
|
except ClientError as e:
|
|
error_message = e.response["Error"]["Message"]
|
|
return f"Error fetching report: {error_message}"
|
|
|
|
|
|
@tool
|
|
def get_consolidated_keys(id: str) -> str:
|
|
"""
|
|
Get the list of consolidated keys (variables) available in the table for a specific month.
|
|
|
|
Args:
|
|
id: The id of the data
|
|
|
|
Returns:
|
|
The list of available variables/keys for the specified data
|
|
"""
|
|
print(f"\n🔧 [TOOL CALLED] get_consolidated_keys for id: {id}")
|
|
|
|
try:
|
|
table = dynamodb.Table(TABLE)
|
|
response = table.get_item(Key={"id": id})
|
|
|
|
if "Item" not in response:
|
|
return f"No data found for month: {id}"
|
|
|
|
item = response["Item"]
|
|
chaves_consolidadas = item.get("chaves_consolidadas", "")
|
|
|
|
if not chaves_consolidadas:
|
|
return f"No consolidated keys found for id: {id}"
|
|
|
|
return chaves_consolidadas
|
|
|
|
except ClientError as e:
|
|
error_message = e.response["Error"]["Message"]
|
|
return f"Error fetching consolidated keys: {error_message}"
|
|
|
|
|
|
def get_contexto() -> dict:
|
|
"""
|
|
Get contexto, filter, and items_disponiveis from DynamoDB where id=DASHBOARD+'_contexto'.
|
|
|
|
Returns:
|
|
Dict with 'contexto', 'filter', and 'items_disponiveis' keys
|
|
"""
|
|
try:
|
|
table = dynamodb.Table(TABLE)
|
|
response = table.get_item(Key={"id": DASHBOARD + "_contexto"})
|
|
|
|
if "Item" not in response:
|
|
return {"contexto": "", "filter": "", "items_disponiveis": {}}
|
|
|
|
item = response["Item"]
|
|
return {
|
|
"contexto": item.get("contexto", ""),
|
|
"filter": item.get("filter_key", ""),
|
|
"items_disponiveis": item.get("itens_disponiveis", {}),
|
|
}
|
|
|
|
except ClientError as e:
|
|
error_message = e.response["Error"]["Message"]
|
|
return {"contexto": f"Error: {error_message}", "filter": "", "items_disponiveis": {}}
|
|
|
|
|
|
def get_secret():
|
|
|
|
secret_name = "assistente-db-secrets-manager"
|
|
region_name = "us-east-1"
|
|
|
|
# Create a Secrets Manager client
|
|
session = boto3.session.Session()
|
|
client = session.client(
|
|
service_name='secretsmanager',
|
|
region_name=region_name
|
|
)
|
|
|
|
try:
|
|
get_secret_value_response = client.get_secret_value(
|
|
SecretId=secret_name
|
|
)
|
|
except ClientError as e:
|
|
# For a list of exceptions thrown, see
|
|
# https://docs.aws.amazon.com/secretsmanager/latest/apireference/API_GetSecretValue.html
|
|
raise e
|
|
|
|
secret = get_secret_value_response['SecretString']
|
|
return secret
|
|
secrets=json.loads(get_secret())
|
|
langfuse = Langfuse(
|
|
public_key=secrets['LANGFUSE-PUBLIC-KEY'],
|
|
secret_key=secrets['LANGFUSE-SECRET-KEY'],
|
|
host=os.environ["LANGFUSE_HOST"]
|
|
)
|
|
session = boto3.Session()
|
|
athena = session.client("athena", region_name="us-east-1")
|
|
|
|
# ==============================================
|
|
# QUERY
|
|
# ==============================================
|
|
|
|
def exec_athena_query(query):
|
|
print("Executando query no Athena...")
|
|
response = athena.start_query_execution(
|
|
QueryString=query,
|
|
QueryExecutionContext={"Database": DATABASE},
|
|
WorkGroup=WORKGROUP
|
|
)
|
|
|
|
query_execution_id = response["QueryExecutionId"]
|
|
print(f"QueryExecutionId: {query_execution_id}")
|
|
|
|
# ==============================================
|
|
# AGUARDAR RESULTADO
|
|
# ==============================================
|
|
|
|
while True:
|
|
result = athena.get_query_execution(QueryExecutionId=query_execution_id)
|
|
state = result["QueryExecution"]["Status"]["State"]
|
|
|
|
if state in ["SUCCEEDED", "FAILED", "CANCELLED"]:
|
|
print("Estado final:", state)
|
|
break
|
|
|
|
print("Aguardando execução...")
|
|
time.sleep(1)
|
|
if state == "SUCCEEDED":
|
|
output = athena.get_query_results(QueryExecutionId=query_execution_id)
|
|
print(f"\n🔧 [TOOL CALLED] consult answer")
|
|
return output["ResultSet"]["Rows"]
|
|
else:
|
|
print("Erro ao executar a query.")
|
|
# Define tools
|
|
|
|
|
|
|
|
# Define@tool the agent state
|
|
class AgentState(TypedDict):
|
|
messages: Annotated[list, operator.add]
|
|
current_step: str
|
|
|
|
|
|
# Initialize Bedrock client with inference profile
|
|
def create_bedrock_llm(model_id: str, region: str = "us-east-1"):
|
|
"""
|
|
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: us-east-1)
|
|
|
|
Returns:
|
|
ChatBedrock instance configured with the model
|
|
"""
|
|
# Determine provider and model_kwargs based on model ID
|
|
MODEL_ARNS = {
|
|
"anthropic.claude-haiku-4-5-20251001-v1:0": "arn:aws:bedrock:us-east-1:305427701314:inference-profile/us.anthropic.claude-haiku-4-5-20251001-v1:0",
|
|
"anthropic.claude-sonnet-4-5-20250929-v1:0": "arn:aws:bedrock:us-east-1:305427701314:inference-profile/global.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
|
"meta.llama4-maverick-17b-instruct-v1:0": "arn:aws:bedrock:us-east-1:305427701314:inference-profile/us.meta.llama4-maverick-17b-instruct-v1:0",
|
|
"meta.llama4-scout-17b-instruct-v1:0": "arn:aws:bedrock:us-east-1:305427701314:inference-profile/us.meta.llama4-scout-17b-instruct-v1:0",
|
|
"amazon.nova-lite-v1:0": "arn:aws:bedrock:us-east-1:305427701314:inference-profile/us.amazon.nova-lite-v1:0",
|
|
"amazon.nova-pro-v1:0": "arn:aws:bedrock:us-east-1:305427701314:inference-profile/us.amazon.nova-pro-v1:0",
|
|
"amazon.nova-2-lite-v1:0": "arn:aws:bedrock:us-east-1:305427701314: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
|
|
)
|
|
|
|
# Bind tools to the LLM
|
|
tools = [get_monthly_report, get_consolidated_keys]
|
|
llm_with_tools = llm.bind_tools(tools)
|
|
|
|
return llm_with_tools
|
|
|
|
|
|
# Define agent nodes
|
|
def call_model(state: AgentState, llm) -> AgentState:
|
|
"""Call the LLM with tools."""
|
|
print(f"[MODEL] Calling Bedrock inference profile...")
|
|
|
|
messages = state["messages"]
|
|
response = llm.invoke(messages)
|
|
state["current_step"] = "model_called"
|
|
return {"messages": [response]}
|
|
|
|
|
|
def call_tools(state: AgentState) -> AgentState:
|
|
"""Execute any tool calls from the LLM response."""
|
|
print(f"[TOOLS] Checking for tool calls...")
|
|
|
|
messages = state["messages"]
|
|
last_message = messages[-1]
|
|
|
|
# Check if there are tool calls
|
|
if hasattr(last_message, 'tool_calls') and last_message.tool_calls:
|
|
print(f"[TOOLS] Found {len(last_message.tool_calls)} tool call(s)")
|
|
|
|
tool_messages = []
|
|
tools_map = {
|
|
"get_monthly_report": get_monthly_report,
|
|
"get_consolidated_keys": get_consolidated_keys
|
|
}
|
|
|
|
# Execute each tool call
|
|
for tool_call in last_message.tool_calls:
|
|
tool_name = tool_call["name"]
|
|
tool_args = tool_call["args"]
|
|
|
|
print(f"[TOOLS] Executing: {tool_name}")
|
|
|
|
# Call the appropriate tool
|
|
tool_func = tools_map[tool_name]
|
|
result = tool_func.invoke(tool_args)
|
|
|
|
# Create tool message
|
|
tool_message = ToolMessage(
|
|
content=str(result),
|
|
tool_call_id=tool_call["id"]
|
|
)
|
|
tool_messages.append(tool_message)
|
|
|
|
state["current_step"] = "tools_executed"
|
|
return {"messages": tool_messages}
|
|
else:
|
|
print(f"[TOOLS] No tool calls found")
|
|
state["current_step"] = "no_tools"
|
|
return {"messages": []}
|
|
|
|
|
|
def should_continue(state: AgentState) -> str:
|
|
"""Determine if we should continue to tools or end."""
|
|
messages = state["messages"]
|
|
last_message = messages[-1]
|
|
|
|
# If there are tool calls, continue to tools node
|
|
if hasattr(last_message, 'tool_calls') and last_message.tool_calls:
|
|
print("[ROUTER] Routing to tools...")
|
|
return "tools"
|
|
|
|
# Otherwise, end
|
|
print("[ROUTER] No more tool calls, ending...")
|
|
return "end"
|
|
|
|
|
|
# Build the LangGraph agent
|
|
def create_agent(inference_profile_arn: str, region: str = "us-east-1"):
|
|
"""
|
|
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
|
|
|
|
Returns:
|
|
Compiled LangGraph workflow
|
|
"""
|
|
# Initialize the LLM with tools
|
|
llm = create_bedrock_llm(inference_profile_arn, region)
|
|
|
|
# Create the graph
|
|
workflow = StateGraph(AgentState)
|
|
|
|
# Add nodes
|
|
workflow.add_node("model", lambda state: call_model(state, llm))
|
|
workflow.add_node("tools", call_tools)
|
|
|
|
# Define the workflow
|
|
workflow.set_entry_point("model")
|
|
|
|
# Add conditional edges
|
|
workflow.add_conditional_edges(
|
|
"model",
|
|
should_continue,
|
|
{
|
|
"tools": "tools",
|
|
"end": END
|
|
}
|
|
)
|
|
|
|
# After tools, go back to model
|
|
workflow.add_edge("tools", "model")
|
|
|
|
# Compile the graph
|
|
app = workflow.compile()
|
|
|
|
return app
|
|
|
|
|
|
def main(user_query,history,model,base):
|
|
"""Main execution function."""
|
|
global DASHBOARD
|
|
DASHBOARD = base
|
|
|
|
# Configuration - Update with your actual inference profile ARN
|
|
|
|
INFERENCE_PROFILE_ARN = model
|
|
REGION = "us-east-1"
|
|
# System prompt for the agent
|
|
|
|
contexto_data = get_contexto()
|
|
if contexto_data["filter"]=="period":
|
|
CONSULT_RULES="""To use the tools you must give the id of the correspondant data, which can be associated to a given month and year in the following format year_month, which:
|
|
-Year is the year in 4 digits (2025,2024,2023,2022,2021,...)
|
|
-Month is th two digit representation: 01,02,03,04,05,06,07,08,09,10,11,12
|
|
The format of the dict is: {id1:year_month1,id2:year_month2...}
|
|
Choose the correct id based on the following dict:
|
|
"""
|
|
elif contexto_data["filter"]=="event":
|
|
CONSULT_RULES="""To use the tools you must give the id of the correspondant data, which can be associated to a event, which is in the format "Name - City DD/MM/YYYY", where the last is a date in the format day/month/year. Theformat of elements in dict is {id1:event_description1,id2:event_description2...}"""
|
|
else:
|
|
CONSULT_RULES="""Wrong filter value, you must terminate the workflow and ask the user to contact the technical team"""
|
|
SYSTEM_PROMPT=""" You are a analitical agent in Brazilian Portuguese, with acess to monthly reports about a specific company, specified in the context. You have access to tools that lets you consult present variables in table, you always have access to "context", which keeps inside answers to different questions, that you may consult as you desire.
|
|
Do not access other variables besides the ones reported by the tool and "context".
|
|
You currently have access to data in a period specified in the context, so only answer questions inside the time window.
|
|
<context>
|
|
"""+contexto_data["contexto"]+"""
|
|
</context>
|
|
"""+CONSULT_RULES+"""
|
|
<correlation>
|
|
"""+str(contexto_data["items_disponiveis"])+"""
|
|
</correlation>
|
|
Here is the chat history:"""+history+"""
|
|
Inside the "NPS" in data is some useful values to calculate the NPS, which includes "distribuicao".
|
|
Inside of it are grades and the amount of people who given that grade.
|
|
Grades from 0 to 6 are detractors.
|
|
Grades from 7 to 8 are neutral.
|
|
Grades from 9 to 10 are promoters.
|
|
Calculate the percentage of them when prompted about NPS and then calculate the nps using the following formula: NPS = %promoter - %detractor, never use the medium of the notes.
|
|
|
|
You have access to the tools:
|
|
-get_consolidated_keys: Given a id returns the column names inside of a entity of a given table element.
|
|
- get_monthly_report: given a id and a variable name, either one listed in the previous tool output or "context", returns its value. Using "context" gives you a summarization of many answers of questions asked to the customers.
|
|
Answer, in Brazilian Portuguese, to the user the best you can with the given information, if you don't know the answer or how to answer say so, only answer from what you know.
|
|
Always consult the most recent information when a date is not given, like questions "Quanto é meu nps?" """
|
|
|
|
print("=" * 60)
|
|
print("LangGraph Agent with AWS Bedrock Inference Profile + Tools")
|
|
print("=" * 60)
|
|
print(f"\nUsing inference profile: {INFERENCE_PROFILE_ARN}")
|
|
print(f"Region: {REGION}\n")
|
|
print("Available Tools:")
|
|
print(" - add_numbers(a, b): Add two numbers")
|
|
print(" - multiply_numbers(a, b): Multiply two numbers")
|
|
print("\nSystem Prompt: Configured ✓")
|
|
print("=" * 60)
|
|
|
|
# Create the agent with a unique session_id to group all steps
|
|
langfuse_handler = CallbackHandler()
|
|
agent = create_agent(INFERENCE_PROFILE_ARN, REGION)
|
|
|
|
# Initialize state with system prompt
|
|
initial_state = {
|
|
"messages": [
|
|
SystemMessage(content=SYSTEM_PROMPT),
|
|
HumanMessage(content=user_query)
|
|
],
|
|
"current_step": "init"
|
|
}
|
|
|
|
print(f"\nUser Query: {user_query}\n")
|
|
print("-" * 60)
|
|
|
|
# Run the agent with callbacks at graph level
|
|
config = {"callbacks": [langfuse_handler], "tags": [DASHBOARD]}
|
|
final_state = agent.invoke(initial_state, config=config)
|
|
|
|
# Display results
|
|
print("-" * 60)
|
|
print("\n[FINAL RESULT]")
|
|
print("\nConversation History:")
|
|
for i, msg in enumerate(final_state["messages"], 1):
|
|
if isinstance(msg, SystemMessage):
|
|
print(f"\n{i}. System: [System prompt configured]")
|
|
elif isinstance(msg, HumanMessage):
|
|
print(f"\n{i}. User: {msg.content}")
|
|
elif isinstance(msg, AIMessage):
|
|
if hasattr(msg, 'tool_calls') and msg.tool_calls:
|
|
print(f"\n{i}. AI: [Calling tools...]")
|
|
else:
|
|
print(f"\n{i}. AI: {msg.content}")
|
|
elif isinstance(msg, ToolMessage):
|
|
print(f"\n{i}. Tool Result: {msg.content}")
|
|
|
|
print("\n" + "=" * 60)
|
|
print(f"Agent completed successfully. Final step: {final_state['current_step']}")
|
|
|
|
# Aggregate token usage from all AIMessage objects
|
|
total_input_tokens = 0
|
|
total_output_tokens = 0
|
|
for msg in final_state["messages"]:
|
|
if isinstance(msg, AIMessage) and hasattr(msg, 'usage_metadata') and msg.usage_metadata:
|
|
total_input_tokens += msg.usage_metadata.get("input_tokens", 0)
|
|
total_output_tokens += msg.usage_metadata.get("output_tokens", 0)
|
|
|
|
langfuse.flush()
|
|
return {
|
|
"response": final_state['messages'][-1].content,
|
|
"input_tokens": total_input_tokens,
|
|
"output_tokens": total_output_tokens,
|
|
"total_tokens": total_input_tokens + total_output_tokens,
|
|
}
|
|
if __name__=="__main__":
|
|
main("Liste o nps mês a mês desde maio 2025 até dezembro 2025","","anthropic.claude-sonnet-4-5-20250929-v1:0") |