AWS Bedrock Agents — Token Bills 15x Higher Than Logged

# io.thecodeforge — DevOps tutorial

import boto3
import json
import os
from botocore.config import Config
from botocore.exceptions import ClientError, EndpointResolutionError

# Config-driven model ID — never hardcode this in your service layer.
# Pull from environment or parameter store so upgrades don't require redeploys.
MODEL_ID = os.environ.get("BEDROCK_MODEL_ID", "anthropic.claude-3-sonnet-20240229-v1:0")
AWS_REGION = os.environ.get("AWS_REGION", "us-east-1")

# Set explicit timeouts. Bedrock calls on large prompts can run 30-60s.
# Without this, boto3 defaults to 60s connect / no read timeout — silent hangs kill your web workers.
boto_config = Config(
    region_name=AWS_REGION,
    connect_timeout=5,        # fail fast if the endpoint is unreachable
    read_timeout=120,         # long enough for large completions, not infinite
    retries={
        "max_attempts": 3,
        "mode": "adaptive"    # exponential backoff with jitter — don't use 'legacy' mode in prod
    }
)

bedrock_runtime = boto3.client("bedrock-runtime", config=boto_config)


def invoke_document_summariser(raw_document: str, max_tokens: int = 1024) -> dict:
    """
    Production pattern: document summarisation for a content pipeline.
    Returns structured output including token usage so the caller can track cost.
    """

    # Claude models use the Messages API format — not the legacy text-completion format.
    # Mixing them up gives you a cryptic ValidationException, not a helpful error.
    request_body = {
        "anthropic_version": "bedrock-2023-05-31",  # required field for Anthropic models on Bedrock
        "max_tokens": max_tokens,
        "messages": [
            {
                "role": "user",
                "content": (
                    "Summarise the following document in 3 bullet points. "
                    "Focus on decisions made, not background context.\n\n"
                    f"{raw_document}"
                )
            }
        ],
        "temperature": 0.2,   # low temp for summarisation — you want deterministic, not creative
    }

    try:
        response = bedrock_runtime.invoke_model(
            modelId=MODEL_ID,
            contentType="application/json",
            accept="application/json",
            body=json.dumps(request_body)
        )

        response_body = json.loads(response["body"].read())

        # Always capture usage — this is your cost telemetry.
        # Log it to CloudWatch metrics or your billing system. Don't discard it.
        input_tokens = response_body["usage"]["input_tokens"]
        output_tokens = response_body["usage"]["output_tokens"]

        summary_text = response_body["content"][0]["text"]

        return {
            "summary": summary_text,
            "model_id": MODEL_ID,
            "input_tokens": input_tokens,
            "output_tokens": output_tokens,
            # Rough cost estimate for Claude 3 Sonnet at time of writing:
            # $0.003/1K input tokens, $0.015/1K output tokens
            "estimated_cost_usd": round(
                (input_tokens / 1000 * 0.003) + (output_tokens / 1000 * 0.015), 6
            )
        }

    except ClientError as e:
        error_code = e.response["Error"]["Code"]
        error_message = e.response["Error"]["Message"]

        # ThrottlingException hits when you exceed your account's TPS quota.
        # Default is 5 TPS for Claude 3 Sonnet in most regions — shockingly low for production.
        if error_code == "ThrottlingException":
            raise RuntimeError(
                f"Bedrock quota exceeded for model {MODEL_ID}. "
                "Request a limit increase via Service Quotas before going live."
            ) from e

        # ValidationException usually means malformed request body — check your model's spec.
        if error_code == "ValidationException":
            raise ValueError(f"Invalid request format for {MODEL_ID}: {error_message}") from e

        raise RuntimeError(f"Bedrock API error [{error_code}]: {error_message}") from e


if __name__ == "__main__":
    sample_doc = """
    Engineering Review — Q3 Platform Migration
    Decision: Move API gateway to AWS API Gateway v2 (HTTP APIs).
    Rationale: 60% cost reduction vs REST APIs for our traffic pattern.
    Rejected alternative: Kong on EKS — operational overhead too high for current team size.
    Timeline: Cutover scheduled for October 15th. Rollback plan approved.
    Owner: Platform team. Risk: Medium. Stakeholder sign-off: CTO, VP Engineering.
    """

    result = invoke_document_summariser(sample_doc)
    print(f"Summary:\n{result['summary']}")
    print(f"\nTokens — Input: {result['input_tokens']} | Output: {result['output_tokens']}")
    print(f"Estimated cost: ${result['estimated_cost_usd']}")

# io.thecodeforge — DevOps tutorial

import boto3
import json
import os
from botocore.config import Config
from typing import Generator

MODEL_ID = os.environ.get("BEDROCK_MODEL_ID", "anthropic.claude-3-sonnet-20240229-v1:0")

boto_config = Config(
    region_name=os.environ.get("AWS_REGION", "us-east-1"),
    connect_timeout=5,
    read_timeout=300,   # streaming completions can run long — 120s isn't always enough
    retries={"max_attempts": 1, "mode": "standard"}  # don't retry mid-stream — retry at the caller level
)

bedrock_runtime = boto3.client("bedrock-runtime", config=boto_config)


def stream_customer_support_response(
    customer_query: str,
    account_context: dict
) -> Generator[str, None, None]:
    """
    Production pattern: real-time customer support response generation.
    Yields text chunks as a generator so the caller (e.g. FastAPI StreamingResponse)
    can flush each chunk to the HTTP client immediately.

    account_context: dict with keys like 'plan', 'open_tickets', 'last_login'
    """

    # Build a system prompt from account context so the model responds with
    # customer-specific information rather than generic advice.
    system_prompt = (
        f"You are a support agent for TheCodeForge platform. "
        f"The customer is on the {account_context.get('plan', 'free')} plan. "
        f"They have {account_context.get('open_tickets', 0)} open support tickets. "
        "Be concise, direct, and actionable. Do not apologise excessively."
    )

    request_body = {
        "anthropic_version": "bedrock-2023-05-31",
        "max_tokens": 512,
        "system": system_prompt,
        "messages": [
            {"role": "user", "content": customer_query}
        ],
        "temperature": 0.3,
    }

    try:
        streaming_response = bedrock_runtime.invoke_model_with_response_stream(
            modelId=MODEL_ID,
            contentType="application/json",
            accept="application/json",
            body=json.dumps(request_body)
        )

        # event_stream is an iterator of EventStream objects — each one is a chunk
        event_stream = streaming_response["body"]

        for event in event_stream:
            chunk = event.get("chunk")
            if not chunk:
                # Non-chunk events exist (metadata, message_start, etc.) — skip them gracefully
                continue

            chunk_data = json.loads(chunk["bytes"].decode("utf-8"))

            # Claude streaming emits different event types — only 'content_block_delta' carries text
            if chunk_data.get("type") == "content_block_delta":
                delta = chunk_data.get("delta", {})
                if delta.get("type") == "text_delta":
                    text_piece = delta.get("text", "")
                    if text_piece:
                        yield text_piece  # flush this chunk to the caller immediately

            # message_stop event signals clean completion — log it for observability
            elif chunk_data.get("type") == "message_stop":
                # Amazon metrics come in the stop event — useful for billing dashboards
                amazon_metrics = chunk_data.get("amazon-bedrock-invocationMetrics", {})
                input_tokens = amazon_metrics.get("inputTokenCount", 0)
                output_tokens = amazon_metrics.get("outputTokenCount", 0)
                # In production: emit these as CloudWatch custom metrics here
                print(f"[Stream complete] input={input_tokens} output={output_tokens} tokens")

    except Exception as e:
        # Critical: yield an error marker so the client knows the stream died mid-response
        # Don't silently stop — the client will think the truncated response is complete
        yield f"\n[ERROR: Response generation interrupted — {type(e).__name__}]"
        raise


# --- Simulated FastAPI usage (shows how the generator plugs into a real web framework) ---
# In production this would be in your router module:
#
# from fastapi import FastAPI
# from fastapi.responses import StreamingResponse
#
# app = FastAPI()
#
# @app.post("/support/stream")
# async def stream_support(query: SupportQuery):
#     account_ctx = fetch_account_context(query.account_id)  # your DB call
#     return StreamingResponse(
#         stream_customer_support_response(query.text, account_ctx),
#         media_type="text/plain"
#     )


if __name__ == "__main__":
    query = "I deployed to production and my API calls are returning 429s. What do I do?"
    context = {"plan": "pro", "open_tickets": 1, "last_login": "2024-03-15"}

    print("Streaming response:\n")
    for chunk in stream_customer_support_response(query, context):
        print(chunk, end="", flush=True)  # flush=True is essential — don't buffer
    print("\n")

# io.thecodeforge — DevOps tutorial
# This Lambda is an Action Group handler for a Bedrock Agent.
# The agent calls this function when it needs to look up order status.
# Deploy this as a Lambda, then wire it to your Agent via the Bedrock console or CDK.

import json
import boto3
import os
from datetime import datetime

# In production: pull from environment, not hardcoded table names
ORDERS_TABLE = os.environ.get("ORDERS_TABLE_NAME", "platform-orders-prod")

dynamodb = boto3.resource("dynamodb")
orders_table = dynamodb.Table(ORDERS_TABLE)


def lambda_handler(event: dict, context) -> dict:
    """
    Bedrock Agent Action Group handler.
    The agent sends a specific event structure — you must return a specific structure back.
    Deviate from the response format and the agent silently fails or hallucinates an answer.
    """

    # Bedrock Agents wraps function calls in this structure
    agent_action = event.get event.get("apiPath", "")        # matches your OpenAPI schema path
    http_method = event.get("httpMethod", "")  # GET, POST, etc. — from your schema
    parameters = event.get("parameters", [])   # list of {name, type, value} dicts

    print(f"[Agent Action] group={agent_action} path={api_path} method={http_method}")

    # Route to the appropriate handler based on the API path
    if api_path == "/orders/{orderId}" and http_method == "GET":
        order_id = next(
            (p["value"] for p in parameters if("actionGroup", "")
    api_path = p["name"] == "orderId"), None
        )
        result = fetch_order_status(order_id)
    elif api_path == "/orders/{orderId}/cancel" and http_method == "POST":
        order_id = next(
            (p["value"] for p in parameters if p["name"] == "orderId"), None
        )
        result = cancel_order(order_id)
    else:
        result = {"error": f"Unknown action path: {api_path}"}

    # Bedrock Agents REQUIRES this exact response envelope.
    # Missing 'messageVersion', 'response', or 'actionGroup' fields = silent agent failure.
    return {
        "messageVersion": "1.0",
        "response": {
            "actionGroup": agent_action,
            "apiPath": api_path,
            "httpMethod": http_method,
            "httpStatusCode": 200 if "error" not in result else 400,
            "responseBody": {
                "application/json": {
                    "body": json.dumps(result)
                }
            }
        }
    }


def fetch_order_status(order_id: str) -> dict:
    """Look up a real order from DynamoDB and return structured status."""
    if not order_id:
        return {"error": "orderId is required"}

    try:
        response = orders_table.get_item(Key={"orderId": order_id})
        item = response.get("Item")

        if not item:
            # Be specific — the agent will relay this message verbatim to the user
            return {"error": f"Order {order_id} not found. It may not exist or may be archived."}

        return {
            "orderId": item["orderId"],
            "status": item["status"],                  # e.g. PROCESSING, SHIPPED, DELIVERED
            "estimatedDelivery": item.get("estimatedDelivery", "unknown"),
            "carrier": item.get("carrier", "not yet assigned"),
            "trackingNumber": item.get("trackingNumber", "not yet assigned"),
            "lastUpdated": item.get("lastUpdated", "")
        }

    except Exception as e:
        # Don't expose raw exception messages to the agent — it may relay them to the user
        print(f"[ERROR] DynamoDB lookup failed for order {order_id}: {e}")
        return {"error": "Order lookup temporarily unavailable. Please try again shortly."}


def cancel_order(order_id: str) -> dict:
    """Cancel an order if it's still in PROCESSING state."""
    if not order_id:
        return {"error": "orderId is required"}

    try:
        # Conditional update — only cancel if status is PROCESSING
        # This prevents the agent from cancelling already-shipped orders
        orders_table.update_item(
            Key={"orderId": order_id},
            UpdateExpression="SET #s = :cancelled, lastUpdated = :now",
            ConditionExpression="#s = :processing",
            ExpressionAttributeNames={"#s": "status"},   # 'status' is a reserved word in DynamoDB
            ExpressionAttributeValues={
                ":cancelled": "CANCELLED",
                ":processing": "PROCESSING",
                ":now": datetime.utcnow().isoformat()
            }
        )
        return {"orderId": order_id, "status": "CANCELLED", "message": "Order successfully cancelled."}

    except dynamodb.meta.client.exceptions.ConditionalCheckFailedException:
        # Order exists but isn't in PROCESSING — give the agent a specific reason
        return {
            "error": f"Order {order_id} cannot be cancelled — it has already been shipped or delivered."
        }
    except Exception as e:
        print(f"[ERROR] Cancel failed for order {order_id}: {e}")
        return {"error": "Cancellation temporarily unavailable."}

# io.thecodeforge — DevOps tutorial
# RAG (Retrieval Augmented Generation) pattern using Bedrock Knowledge Bases.
# Use case: internal engineering handbook bot that answers policy questions.
# The Knowledge Base is pre-populated with your company docs via the Bedrock console or CDK.

import boto3
import os
import json
from botocore.config import Config

KNOWLEDGE_BASE_ID = os.environ["BEDROCK_KB_ID"]   # e.g. "ABCD1234EF" — from Bedrock console
MODEL_ARN = (
    "arn:aws:bedrock:us-east-1::foundation-model/"
    "anthropic.claude-3-sonnet-20240229-v1:0"
)  # RetrieveAndGenerate requires the full ARN, not just the model ID

boto_config = Config(
    region_name=os.environ.get("AWS_REGION", "us-east-1"),
    connect_timeout=5,
   _attempts": 2, "mode": "adaptive"}
)

# Note: Knowledge Bases uses the 'bedrock-agent-runtime' client — NOT 'bedrock-runtime'.
# Using the wrong client gives you a NoRegionError or AttributeError with no useful message.
bedrock_agent_runtime = boto3.client("bedrock-agent-runtime", config=boto_config)


def query_engineering_handbook(question: str, max_retrieved_chunks: int = 5) -> dict:
    """
    Query the engineering handbook Knowledge Base using RetrieveAndGenerate.
    This is the fully managed RAG path — Bedrock handles retrieval + generation in one call.

    For transparency/debugging: also returns the source citations so you can verify
    the model isn't hallucinating answers that aren't in the docs.
    """

    try:
        response = bedrock_agent_runtime.retrieve_and_generate(
            input={"text": question},
            retrieveAndGenerateConfiguration={
                "type": "KNOWLEDGE_BASE",
                "knowledgeBaseConfiguration": {
                    "knowledgeBaseId": KNOWLEDGE_BASE_ID,
                    "modelArn": MODEL_ARN,
                    "retrievalConfiguration": {
                        "vectorSearchConfiguration": {
                            # Number of document chunks to retrieve before generation.
                            # Higher = more context but more tokens = higher cost and latency.
                            # 5 is a good starting point; tune based on your doc structure.
                            "numberOfResults": max_retrieved_chunks
                        }
                    },
                    "generationConfiguration": {
                        "promptTemplate": {
                            # Override the default prompt to enforce your preferred answer style.
                            # The $search_results$ placeholder is where retrieved chunks are injected.
                            "textPromptTemplate": (
                                "You are an assistant for TheCodeForge engineering team. "
                                "Answer based ONLY on the following retrieved context. "
                                "If the answer isn't in the context, say 'Not found in handbook.' "
                                "Do not invent policies or procedures.\n\n"
                                "Context:\n$search_results$\n\n"
                                f"Question: {question}"
                            )
                        }
                    }
                }
            }
        )

        answer = response["output"]["text"]

        # Extract citations — each citation maps to a specific chunk in your S3 docs.
        # In production: read_timeout=60,
    retries={"max surface these to the user so they can verify the source.
        citations = []
        for citation in response.get("citations", []):
            for reference in citation.get("retrievedReferences", []):
                location = reference.get("location", {}).get("s3Location", {})
                citations.append({
                    "source_uri": location.get("uri", "unknown"),
                    "excerpt": reference.get("content", {}).get("text", "")[:200]  # truncate for display
                })

        return {
            "answer": answer,
            "citations": citations,
            "citation_count": len(citations)
        }

    except bedrock_agent_runtime.exceptions.ResourceNotFoundException:
        raise ValueError(
            f"Knowledge Base {KNOWLEDGE_BASE_ID} not found. "
            "Check the ID and ensure the KB is in 'Active' status — "
            "embedding ingestion must complete before queries work."
        )
    except Exception as e:
        raise RuntimeError(f"Knowledge Base query failed: {e}") from e


if __name__ == "__main__":
    result = query_engineering_handbook(
        "What's our policy on hotfixing directly to the main branch?"
    )

    print(f"Answer:\n{result['answer']}\n")
    print(f"Sources ({result['citation_count']} retrieved):")
    for i, citation in enumerate(result["citations"], 1):
        print(f"  [{i}] {citation['source_uri']}")
        print(f"      Excerpt: {citation['excerpt']}...")