Text to SQL Agent

upgrade_output = !pip install --upgrade pip
install_boto3_output = !pip install boto3
##Ensure your boto3 and botocore libraries are up to date
upgrade_output_botocore_boto3= !pip install --upgrade boto3 botocore 

from dependencies.config import *

!python ./dependencies/build_infrastructure.py

from dependencies.config import *

list_agent=bedrock_agent_client.list_agents()['agentSummaries']
#print(list_agent)
#print(agent_name)
<h2>Search f</h2>
agent_id = next((agent['agentId'] for agent in list_agent if agent['agentName'] == agent_name), None)

print(agent_id)

response = bedrock_agent_client.list_agent_aliases(
    agentId=agent_id,
)
response['agentAliasSummaries']
agentAliasId=next((agent['agentAliasId'] for agent in response['agentAliasSummaries'] if agent['agentAliasName'] == agent_alias_name), None)
agent_alias_id=agentAliasId
print(agent_alias_id)

<h2>### Invoke Agent</h2>
<h2>Now that we've created the agent, let's use the `bedrock-agent-runtime` client to invoke this agent and perform some tasks.</h2>
query_to_agent = """What was John Denny's salary in 1986?"""

<h2>create a random id for session initiator id</h2>
session_id:str = str(uuid.uuid1())
enable_trace:bool = True
end_session:bool = False


<h2>invoke the agent API</h2>
agentResponse = bedrock_agent_runtime_client.invoke_agent(
    inputText=query_to_agent,
    agentId=agent_id,
    agentAliasId=agent_alias_id, 
    sessionId=session_id,
    enableTrace=enable_trace, 
    endSession= end_session
)
logger.info(pprint.pprint(agentResponse))

%%time
event_stream = agentResponse['completion']
try:
    for event in event_stream:        
        if 'chunk' in event:
            data = event['chunk']['bytes']
            logger.info(f"Final answer ->\n{data.decode('utf8')}")
            agent_answer = data.decode('utf8')
            end_event_received = True
            # End event indicates that the request finished successfully
        elif 'trace' in event:
            logger.info(json.dumps(event['trace'], indent=2))
        else:
            raise Exception("unexpected event.", event)
except Exception as e:
    raise Exception("unexpected event.", e)

<h2>And here is the response if you just want to see agent's reply</h2>
print(agent_answer)

#Create function to invoke agent
def invoke_agent(query):
    ## create a random id for session initiator id
    session_id:str = str(uuid.uuid1())
    enable_trace:bool = False
    end_session:bool = False


    # invoke the agent API
    agentResponse = bedrock_agent_runtime_client.invoke_agent(
        inputText=query,
        agentId=agent_id,
        agentAliasId=agent_alias_id, 
        sessionId=session_id,
        enableTrace=enable_trace, 
        endSession= end_session
    )
    event_stream = agentResponse['completion']
    print("Fetching answer...")
    try:
        for event in event_stream:        
            if 'chunk' in event:
                data = event['chunk']['bytes']
                logger.info(f"Final answer ->\n{data.decode('utf8')}")
                agent_answer = data.decode('utf8')
                end_event_received = True
                # End event indicates that the request finished successfully
            elif 'trace' in event:
                logger.info(json.dumps(event['trace'], indent=2))
            else:
                raise Exception("unexpected event.", event)
    except Exception as e:
        raise Exception("unexpected event.", e)

invoke_agent("What year was Nolan Ryan inducted into the Hall of Fame?")

invoke_agent("What year was Nolan Ryan inducted into the Hall of Fame?")

invoke_agent("In what year did Hank Aaron hit the most home runs?")

#This query requires a join of two tables to be able to answer the question
invoke_agent("What player has received the most All-Star Game selections?")

#This query should say there is no data available for this year!!
invoke_agent("What was Babe Ruth's salary in 1930?")

invoke_agent("Who is the richest player in baseball history? ")

invoke_agent("What was John Denny's salary in 1986? ")

<h2>This will stop the notebook's execution and wait for the user to press Enter</h2>
input("Press Enter to continue...")

!python ./dependencies/clean.py