Monitoring Amazon Bedrock Agents with W&B Weave

pip install boto3 weave 

import boto3
import botocore.exceptions

# Initialize AWS clients
bedrock = boto3.client('bedrock', region_name='us-east-2')  # Adjust region as necessary
s3 = boto3.client('s3', region_name='us-east-2')

# Step 1: Create the Text File
def create_text_file(file_path):
    print("Creating text file...")
    data = [
        "Store: Walmart, City: New York",
        "Store: Target, City: Los Angeles",
        "Store: Costco, City: Chicago",
        "Store: Best Buy, City: Houston",
        "Store: IKEA, City: Seattle"
    ]
    with open(file_path, "w") as file:
        file.write("\n".join(data))
    print(f"Text file created: {file_path}")

# Step 2: Create S3 Bucket (if it doesn't exist)
def create_bucket(bucket_name, region='us-east-2'):
    try:
        print(f"Checking if bucket '{bucket_name}' exists...")
        s3.head_bucket(Bucket=bucket_name)
        print(f"Bucket '{bucket_name}' already exists.")
    except botocore.exceptions.ClientError:
        print(f"Bucket '{bucket_name}' does not exist. Creating it...")
        s3.create_bucket(
            Bucket=bucket_name,
            CreateBucketConfiguration={'LocationConstraint': region}
        )
        print(f"Bucket '{bucket_name}' created successfully.")

# Step 3: Upload the Text File to S3
def upload_file_to_s3(bucket_name, file_path, s3_key):
    print("Uploading file to S3...")
    s3.upload_file(file_path, bucket_name, s3_key)
    print(f"File uploaded successfully: s3://{bucket_name}/{s3_key}")
    return f"s3://{bucket_name}/{s3_key}"

# Step 4: Main Function to Orchestrate the Process
def main():
    # Update these variables
    bucket_name = "mybucket-agent-12345"  # Replace with a globally unique S3 bucket name
    file_path = "stores_and_cities.txt"  # Local file path for the text file
    s3_key = "stores/stores_and_cities.txt"  # Path to store the file in the S3 bucket

    # Step 1: Create the text file
    create_text_file(file_path)
    
    # Step 2: Create the bucket if it doesn't exist
    create_bucket(bucket_name)
    
    # Step 3: Upload the file to S3
    upload_file_to_s3(bucket_name, file_path, s3_key)
    

if __name__ == "__main__":
    main()


import json
import urllib.request
import urllib.parse
import urllib.error
import re

def lambda_handler(event, context):
    def build_response(body, status="success"):
        return {
            'response': {
                'actionGroup': event.get('actionGroup', 'unknown_actionGroup'),
                'function': event.get('function', 'unknown_function'),
                'functionResponse': {
                    'responseBody': {"TEXT": {"body": body}}
                }
            },
            'messageVersion': event.get('messageVersion', '1.0')
        }

    # Extract 'city' from parameters
    city = next((p['value'] for p in event.get('parameters', []) if p.get('name') == 'city'), None)
    if not city:
        return build_response('Error: Missing "city" parameter in request.')

    if not re.match(r'^[A-Za-z\s\-]+$', city):
        return build_response('Error: Invalid city name provided.')

    try:
        # Fetch weather data
        url = f"https://wttr.in/{urllib.parse.quote(city)}?format=j1"
        with urllib.request.urlopen(url) as response:
            if response.status == 200:
                data = json.loads(response.read())
                current = data.get('current_condition', [{}])[0]
                weather = (
                    f"Weather in {city}:\n"
                    f"Temperature: {current.get('temp_C', 'N/A')}°C\n"
                    f"Description: {current.get('weatherDesc', [{'value': 'N/A'}])[0].get('value', 'N/A')}\n"
                    f"Humidity: {current.get('humidity', 'N/A')}%"
                )
                return build_response(weather)
            else:
                return build_response(f"Error: Failed to fetch weather data. Status code: {response.status}")
    except urllib.error.HTTPError as e:
        return build_response(f"HTTP Error: {e.reason}")
    except urllib.error.URLError as e:
        return build_response(f"URL Error: {e.reason}")
    except Exception as e:
        return build_response(f"Internal Server Error: {str(e)}")


import boto3
import json
import weave

weave.init("bedrock_invoke_agent")  # Initialize weave tracking


def format_prompt_for_traces(traces, input_text):
    """
    Formats the traces into a prompt for the LLM.

    Parameters:
        traces (list): List of traces to be summarized.

    Returns:
        str: The formatted prompt.
    """
    return f"""
    <|begin_of_text|><|start_header_id|>user<|end_header_id|>
    Please analyze the following traces from an agent and provide a step-by-step summary:
    Original user input: {input_text}
    Traces:
    {json.dumps(traces, indent=4)}

    The summary should clearly outline:
    - The main steps executed
    - Any tools used
    - Results retrieved from knowledge base
    - Final outcome

    Generate a coherent and concise summary.
    <|eot_id|>
    <|start_header_id|>assistant<|end_header_id|>
    """


def summarize_traces_with_llm(traces, input_text):
    """
    Summarizes traces using an LLM.

    Parameters:
        traces (list): The list of trace dictionaries.

    Returns:
        str: A summary of the traces or an error message if summarization fails.
    """
    try:
        # Initialize the Bedrock runtime client
        client = boto3.client(service_name="bedrock-runtime", region_name="us-east-1")
        prompt = format_prompt_for_traces(traces, input_text)

        # Prepare messages for Nova Pro API
        messages = [
            {
                "role": "user",
                "content": [{"text": prompt}]
            }
        ]

        # Invoke the model
        print("Invoking Nova Pro model for summarization...")
        response = client.converse(
            modelId="us.amazon.nova-pro-v1:0",
            messages=messages
        )

        # Extract response content
        prediction = response["output"]["message"]["content"][0]["text"].strip()
        return prediction

    except Exception as e:
        print(f"Error during LLM summarization: {e}")
        return "Error summarizing traces."
    

@weave.op
def invoke_agent(agent_id, agent_alias_id, session_id, input_text, enable_trace=False):
    """
    Invokes an Amazon Bedrock agent with optional debugging.

    Parameters:
        agent_id (str): The unique identifier of the agent.
        agent_alias_id (str): The alias of the agent.
        session_id (str): The session ID for maintaining context.
        input_text (str): The prompt or query to send to the agent.
        enable_trace (bool): Whether to enable trace mode for debugging.

    Returns:
        dict: A dictionary with 'response', 'traces', and a summary.
    """
    try:
        # Initialize the Bedrock runtime client
        bedrock_client = boto3.client("bedrock-agent-runtime", region_name="us-east-2")

        print("Invoking agent...")
        print(f"Agent ID: {agent_id}, Alias ID: {agent_alias_id}, Session ID: {session_id}")
        print(f"Input Text: {input_text}, Enable Trace: {enable_trace}")

        # Call the invoke_agent API
        response = bedrock_client.invoke_agent(
            agentId=agent_id,
            agentAliasId=agent_alias_id,
            sessionId=session_id,
            inputText=input_text,
            enableTrace=enable_trace # enable tracing 
        )

        # Initialize variables for response and trace
        completion = ""
        trace_data = []

        # Iterate over the response stream
        for event in response["completion"]:
            try:
                if "chunk" in event:
                    completion += event["chunk"]["bytes"].decode()
                else:
                    trace_data.append(event)  # Collect non-chunk trace data
            except Exception as e:
                print(f"Error processing event: {event}. Exception: {str(e)}")

        # Summarize traces using the LLM
        summary = summarize_traces_with_llm(trace_data, input_text)

        # Return both response, trace data, and summary
        return {
            "response": completion,
            "traces": trace_data,
            "summary": summary
        }

    except Exception as e:
        print("Exception occurred:")

        return {
            "error": str(e),
            "response": None,
            "traces": None,
            "summary": "Error during invocation or summarization."
        }

if __name__ == "__main__":
    # Replace with your agent details
    agent_id = "RTWDM0FTEU"  # Your agent ID
    agent_alias_id = "JFMNFQSLDS"  # Your agent alias ID
    session_id = "default-session"  # A unique session ID
    input_text = "what is the weather in each of the locations of my stores?"
    enable_trace = True  # Toggle tracing

    print(f"Invoking agent with input: '{input_text}'")

    # Call the invoke_agent function
    result = invoke_agent(agent_id, agent_alias_id, session_id, input_text, enable_trace)

    # Output the results
    print("Agent Response:")
    print(result["response"])

    print("\nTraces:")
    print(json.dumps(result["traces"], indent=4))

    print("\nSummary:")
    print(result["summary"])