LLM Routing with Zero Code: Content-Based Model Selection on Bedrock with Step Functions

Sending every request to your biggest model is the easiest way to burn a Bedrock budget — a "what's the capital of Japan?" question doesn't need the same model as "write FizzBuzz in Python". The usual fix is a router Lambda, but that means code to write, deploy, and maintain.

In this hands-on, we'll build content-based model routing with zero application code: an AWS Step Functions state machine classifies each question with a small model, then a Choice state routes it to the right answering model — Claude Haiku for code, Amazon Nova Lite for creative writing, Nova Micro for everything else. The whole router is one JSON document.

Prefer video? This entire hands-on is also on YouTube:

What we'll build

question
   │
   ▼
[ClassifyQuestion]  ← Nova Micro, temperature 0
   │  "simple" / "code" / "creative"
   ▼
[RouteByCategory]  (Choice state)
   ├─ *code*     ──▶ [AnswerWithHaiku]     ← Claude Haiku
   ├─ *creative* ──▶ [AnswerWithNovaLite]  ← Nova Lite
   └─ default    ──▶ [AnswerWithNovaMicro] ← Nova Micro

Step Functions has an optimized integration for Bedrock (arn:aws:states:::bedrock:invokeModel), so states call models directly — no Lambda in the path. Routing logic lives in a Choice state, retries and execution history come built in, and the visual workflow in the console doubles as documentation.

Note on model IDs: Bedrock models are updated frequently. The IDs below were current when this was written — check aws bedrock list-inference-profiles (or the console) and use the latest versions.

Step 1: IAM role

Create a role for Step Functions (e.g. StepFunctionsBedrockHandsOnRole, trusted by states.amazonaws.com) and attach this inline policy (IAM → Roles → the role → Add permissions → Create inline policy → JSON):

{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Effect": "Allow",
      "Action": "bedrock:InvokeModel",
      "Resource": "*"
    }
  ]
}

(For production, scope Resource to the specific model and inference-profile ARNs.)

Step 2: The state machine

Create a state machine in the Step Functions console, switch to the Code view, and paste:

{
  "Comment": "Bedrock model routing handson",
  "StartAt": "ClassifyQuestion",
  "States": {
    "ClassifyQuestion": {
      "Type": "Task",
      "Resource": "arn:aws:states:::bedrock:invokeModel",
      "Parameters": {
        "ModelId": "amazon.nova-micro-v1:0",
        "Body": {
          "messages": [
            {
              "role": "user",
              "content": [
                {
                  "text.$": "States.Format('Classify the following question as exactly one word: simple, code, or creative. Do not output anything else. Question: {}', $.question)"
                }
              ]
            }
          ],
          "inferenceConfig": {
            "maxTokens": 10,
            "temperature": 0
          }
        }
      },
      "ResultSelector": {
        "category.$": "$.Body.output.message.content[0].text"
      },
      "ResultPath": "$.classification",
      "Next": "RouteByCategory"
    },
    "RouteByCategory": {
      "Type": "Choice",
      "Choices": [
        {
          "Variable": "$.classification.category",
          "StringMatches": "*code*",
          "Next": "AnswerWithHaiku"
        },
        {
          "Variable": "$.classification.category",
          "StringMatches": "*creative*",
          "Next": "AnswerWithNovaLite"
        }
      ],
      "Default": "AnswerWithNovaMicro"
    },
    "AnswerWithNovaMicro": {
      "Type": "Task",
      "Resource": "arn:aws:states:::bedrock:invokeModel",
      "Parameters": {
        "ModelId": "amazon.nova-micro-v1:0",
        "Body": {
          "messages": [
            {
              "role": "user",
              "content": [
                { "text.$": "$.question" }
              ]
            }
          ],
          "inferenceConfig": { "maxTokens": 300 }
        }
      },
      "End": true
    },
    "AnswerWithHaiku": {
      "Type": "Task",
      "Resource": "arn:aws:states:::bedrock:invokeModel",
      "Parameters": {
        "ModelId": "us.anthropic.claude-haiku-4-5-20251001-v1:0",
        "Body": {
          "anthropic_version": "bedrock-2023-05-31",
          "max_tokens": 500,
          "messages": [
            {
              "role": "user",
              "content": [
                {
                  "type": "text",
                  "text.$": "$.question"
                }
              ]
            }
          ]
        }
      },
      "End": true
    },
    "AnswerWithNovaLite": {
      "Type": "Task",
      "Resource": "arn:aws:states:::bedrock:invokeModel",
      "Parameters": {
        "ModelId": "amazon.nova-lite-v1:0",
        "Body": {
          "messages": [
            {
              "role": "user",
              "content": [
                { "text.$": "$.question" }
              ]
            }
          ],
          "inferenceConfig": { "maxTokens": 500 }
        }
      },
      "End": true
    }
  }
}

Four details worth reading twice:

• States.Format builds the prompt inside the state machine. The classification instruction wraps the user's question with an intrinsic function — prompt templating without a single line of code.
• ResultSelector + ResultPath keep the original input. The classifier's raw output is trimmed down to $.classification.category, while the original $.question stays in the state — so the answering states can still read it. Forgetting ResultPath here would overwrite the question with the classifier's response.
• StringMatches: "*code*" is deliberate slack. LLM classifiers don't always return exactly one clean token — wildcards tolerate whitespace, casing artifacts, or a stray period around the label.
• Each provider gets its native request body. The direct invokeModel integration doesn't normalize formats the way the Converse API does: Nova takes messages + inferenceConfig, while Claude needs anthropic_version and max_tokens. If you swap in a different provider's model, swap the body shape too.

Step 3: Test the routes

State machine → Start execution, one case at a time:

{
  "question": "What is the capital of Japan?"
}

→ classified simple, answered by Nova Micro (default route).

{
  "question": "Write FizzBuzz in Python"
}

→ classified code, answered by Claude Haiku.

{
  "question": "Write a short story with a cat as the main character"
}

→ classified creative, answered by Nova Lite.

Open the execution's graph view and you can see the path light up through the classifier, the Choice state, and the selected model — routing you can literally look at.

When you're done, delete the state machine and the IAM role to clean up (both are free while idle, but tidy is tidy).

Wrapping up

A Choice state plus Bedrock's optimized integration turns model routing from an application-code problem into an infrastructure definition — versionable, visual, and with retries and execution history included. If your Bedrock bill is dominated by simple questions hitting an expensive model, this pattern is one JSON document away.

About the author

Maruchin Tech — 12x AWS Certified | Cloud & AI for manufacturing and supply chain (AWS / Google Cloud / Azure) | Udemy instructor (100K+ students)

🎥 Video version of this hands-on:
https://youtu.be/clZCoIgFMbw

📚 Full course — AWS Certified Generative AI Developer Professional (AIP-C01) Exam Prep:
https://www.udemy.com/course/aws-certified-generative-ai-developer-professional-exam-prep/

👨‍🏫 All my courses:
(Eng) https://www.udemy.com/user/maruchin-tech-2/
(Jpn) https://www.udemy.com/user/shan-wang-wan-jun-2/

🎫 Monthly discount coupons:
https://www.youtube.com/@MaruchinTech-cloud/posts