AI-inovyo-assistende-db/code/app/backend/BDAgent.py

"""
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
from backend.utils import dynamodb_read_table as drt
WORKGROUP = "iceberg-workgroup"
DATABASE = "dnx_warehouse"
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/global.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": "global",
            "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 = [consult_answers,count_table_rows]
    tools=[]
    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"]
    langfuse_handler = CallbackHandler()
    config = {"configurable": {"thread_id": "abc123"},"callbacks": [langfuse_handler]}
    response = llm.invoke(messages,config=config)
    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 = {
        }
        
        # 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):
    """Main execution function."""

    # Configuration - Update with your actual inference profile ARN
    
    INFERENCE_PROFILE_ARN = model
    REGION = "us-east-1"
    # System prompt for the agent

    SYSTEM_PROMPT=""" You are a analitical agent, with acess to monthly reports about Bacio di latte
    <context>
    A Bacio di Latte é uma rede de gelaterias artesanais fundada em São Paulo, Brasil, em 2011, pelos irmãos milaneses Edoardo e Luigi Tonolli, que trouxeram a tradição do gelato italiano com ingredientes de alta qualidade, resultando em um sorvete cremoso e fresco, produzido diariamente, sem gordura hidrogenada ou trans, e que se tornou popular não só no Brasil, mas também nos EUA, representando uma experiência autêntica de gelato. 
    <\context>
    <reports>
    """+drt.read_table_as_xml("poc_dnx_monthly_summary","us-east-1")+""""
    <\reports>
    Here is the chat history:"""+history+"""
    Aswer 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."""

    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
    agent = create_agent(INFERENCE_PROFILE_ARN, REGION)

    # Example query that requires tools

    # 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
    final_state = agent.invoke(initial_state)

    # 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']}")
    langfuse.flush()
    return final_state['messages'][-1].content
if __name__=="__main__":
    main("oi","ancar_nps_tradicional","","")