Model-Based Routing

cat <<'EOF' > kong-recipe.yaml
_defaults:
  kongctl:
    namespace: model-based-routing-recipe

ai_gateway_model_providers:
  - ref: model-based-routing-openai-provider
    ai_gateway: !lookup {id: !env AI_GATEWAY_ID}
    name: model-based-routing-openai-provider
    display_name: "model-based-routing (OpenAI)"
    type: openai
    config:
      auth:
        type: basic
        headers:
          - name: Authorization
            value: !secret {source: !env OPENAI_TOKEN}

  - ref: model-based-routing-bedrock-provider
    ai_gateway: !lookup {id: !env AI_GATEWAY_ID}
    name: model-based-routing-bedrock-provider
    display_name: "model-based-routing (AWS Bedrock)"
    type: bedrock
    config:
      auth:
        type: aws
        access_key_id: !env AWS_ACCESS_KEY_ID
        secret_access_key: !secret {source: !env AWS_SECRET_ACCESS_KEY}

ai_gateway_auth_strategies:
  - ref: model-based-routing-auth
    ai_gateway: !lookup {id: !env AI_GATEWAY_ID}
    name: model-based-routing-auth
    display_name: "model-based-routing key auth"
    type: key-auth
    config:
      key_names:
        - apikey
      hide_credentials: true

ai_gateway_consumers:
  - ref: model-based-routing-consumer
    ai_gateway: !lookup {id: !env AI_GATEWAY_ID}
    name: model-based-routing-consumer
    display_name: "model-based-routing demo consumer"
    type: api-key
    credentials:
      - ref: model-based-routing-credential
        name: model-based-routing-credential
        display_name: "model-based-routing demo API key"
        type: api-key
        api_key: !secret {source: !env CONSUMER_API_KEY}

ai_gateway_policies:
  - ref: model-based-routing-decorator
    ai_gateway: !lookup {id: !env AI_GATEWAY_ID}
    name: model-based-routing-decorator
    display_name: "model-based-routing classifier prompt"
    type: ai-prompt-decorator
    config:
      llm_format: openai
      prompts:
        prepend:
          - role: system
            content: >
              You are a model router. Analyze the user's prompt and recommend
              the most appropriate model tier.

              Return ONLY ONE of these values:
              - "fast" for simple tasks (greetings, basic questions, straightforward requests)
              - "smart" for complex tasks (reasoning, analysis, coding, creative writing)

              Respond with just the single word "fast" or "smart", nothing else.

  - ref: model-based-routing-orchestrator
    ai_gateway: !lookup {id: !env AI_GATEWAY_ID}
    name: model-based-routing-orchestrator
    display_name: "model-based-routing tier orchestration"
    type: datakit
    config:
      nodes:
        - name: EXTRACT_PROMPT
          type: jq
          input: request.body
          output: prompt_data
          jq: '{messages: .messages}'

        - name: EXTRACT_AUTH
          type: jq
          input: request.headers
          output: auth_header
          jq: '{apikey: (.apikey // .Apikey // .APIKey)}'

        - name: CALL_MODEL_SELECTOR
          type: call
          url: http://localhost:8000/model-selection
          method: POST
          inputs:
            body: EXTRACT_PROMPT.prompt_data
            headers: EXTRACT_AUTH.auth_header

        - name: EXTRACT_TIER
          type: jq
          input: CALL_MODEL_SELECTOR.body
          output: tier
          jq: '.choices[0].message.content | gsub("^\\s+|\\s+$"; "")'

        - name: IS_SMART
          type: jq
          input: EXTRACT_TIER.tier
          output: value
          jq: '. == "smart"'

        - name: ROUTE_DECISION
          type: branch
          input: IS_SMART.value
          then:
            - CALL_SMART
            - RESPOND_SMART
          else:
            - CALL_FAST
            - RESPOND_FAST

        - name: CALL_SMART
          type: call
          url: http://localhost:8000/model-based-routing-smart
          method: POST
          inputs:
            body: EXTRACT_PROMPT.prompt_data
            headers: EXTRACT_AUTH.auth_header

        - name: RESPOND_SMART
          type: exit
          inputs:
            body: CALL_SMART.body
            headers: CALL_SMART.headers

        - name: CALL_FAST
          type: call
          url: http://localhost:8000/model-based-routing-fast
          method: POST
          inputs:
            body: EXTRACT_PROMPT.prompt_data
            headers: EXTRACT_AUTH.auth_header

        - name: RESPOND_FAST
          type: exit
          inputs:
            body: CALL_FAST.body
            headers: CALL_FAST.headers

ai_gateway_models:
  - ref: model-based-routing-selector
    ai_gateway: !lookup {id: !env AI_GATEWAY_ID}
    name: model-based-routing-selector
    display_name: "model-based-routing (selector)"
    type: model
    formats:
      - type: openai
    capabilities:
      - generate
    access:
      auth_strategies:
        - model-based-routing-auth
    policies:
      - model-based-routing-decorator
    config:
      max_request_body_size: 5242880
      response_streaming: deny
      logging:
        payloads: false
      route:
        paths:
          - /model-selection
        protocols:
          - http
          - https
        methods:
          - POST
          - OPTIONS
        strip_path: true
    targets:
      - name: !env OPENAI_SELECTOR_SLM
        provider: model-based-routing-openai-provider
        config:
          type: openai

  - ref: model-based-routing-fast
    ai_gateway: !lookup {id: !env AI_GATEWAY_ID}
    name: model-based-routing-fast
    display_name: "model-based-routing (fast tier)"
    type: model
    formats:
      - type: openai
    capabilities:
      - generate
    access:
      auth_strategies:
        - model-based-routing-auth
    config:
      max_request_body_size: 10485760
      response_streaming: deny
      logging:
        payloads: false
      route:
        paths:
          - /model-based-routing-fast
        protocols:
          - http
          - https
        methods:
          - POST
          - OPTIONS
        strip_path: true
    targets:
      - name: !env OPENAI_FAST_MODEL
        provider: model-based-routing-openai-provider
        config:
          type: openai

  - ref: model-based-routing-smart
    ai_gateway: !lookup {id: !env AI_GATEWAY_ID}
    name: model-based-routing-smart
    display_name: "model-based-routing (smart tier)"
    type: model
    formats:
      - type: openai
    capabilities:
      - generate
    access:
      auth_strategies:
        - model-based-routing-auth
    config:
      max_request_body_size: 10485760
      response_streaming: deny
      logging:
        payloads: false
      route:
        paths:
          - /model-based-routing-smart
        protocols:
          - http
          - https
        methods:
          - POST
          - OPTIONS
        strip_path: true
    targets:
      - name: !env BEDROCK_SMART_MODEL
        provider: model-based-routing-bedrock-provider
        config:
          region: !env AWS_REGION
          type: bedrock

  - ref: model-based-routing-chat
    ai_gateway: !lookup {id: !env AI_GATEWAY_ID}
    name: model-based-routing-chat
    display_name: "model-based-routing (chat)"
    type: model
    formats:
      - type: openai
    capabilities:
      - generate
    access:
      auth_strategies:
        - model-based-routing-auth
    policies:
      - model-based-routing-orchestrator
    config:
      max_request_body_size: 10485760
      response_streaming: deny
      logging:
        payloads: false
      route:
        paths:
          - /chat
        protocols:
          - http
          - https
        methods:
          - POST
          - OPTIONS
        strip_path: true
    targets:
      - name: !env OPENAI_FAST_MODEL
        provider: model-based-routing-openai-provider
        config:
          type: openai
EOF
kongctl apply -f kong-recipe.yaml -o text --auto-approve --pat "${KONNECT_TOKEN}"

rm -f kong-recipe.yaml

cat <<'EOF' > demo.py
"""Model-based routing demo. See README for context.

Sends a simple prompt and a complex prompt to Kong's /chat endpoint. Kong classifies each
prompt's complexity internally (via an OpenAI o3-mini call) and dynamically routes the
request to either the OpenAI fast tier or the AWS Bedrock (Claude) smart tier. The demo
prints the X-Kong-LLM-Model header on each response so you can confirm which upstream model
actually served the request. A final call with an invalid API key shows Kong rejecting the
request with 401 before any classification or upstream call happens.

Run with:

    python demo.py

Requires the Kong configuration from the README to be applied and PROXY_URL (defaults to
http://localhost:8000) reachable.
"""

import os
import sys
import time

from openai import APIStatusError, OpenAI

PROXY_URL = os.getenv("PROXY_URL", "http://localhost:8000")
API_KEY = "demo-api-key"

SIMPLE_PROMPT = "Hi there! What's 2 + 2?"
COMPLEX_PROMPT = (
    "Perform in depth analysis of the ancient Roman world during the time "
    "of Pliny the Elder in two paragraphs."
)

# ANSI color codes. Disabled when stdout isn't a TTY or NO_COLOR is set.
_USE_COLOR = sys.stdout.isatty() and "NO_COLOR" not in os.environ
def _c(code: str, s: str) -> str:
    return f"\033[{code}m{s}\033[0m" if _USE_COLOR else s
BOLD  = lambda s: _c("1", s)
DIM   = lambda s: _c("2", s)
GREEN = lambda s: _c("32", s)
CYAN  = lambda s: _c("36", s)
RED   = lambda s: _c("31", s)


def make_client(api_key: str) -> OpenAI:
    """Construct an OpenAI client that sends the given API key in the apikey header."""
    return OpenAI(
        base_url=f"{PROXY_URL}/chat",
        api_key="unused",  # required by the SDK; Kong reads the apikey header instead
        default_headers={"apikey": api_key},
    )


def call(client: OpenAI, label: str, prompt: str) -> None:
    """Send one chat request and print which upstream model Kong routed it to."""
    print(f"\n{BOLD('[REQUEST]')} {label}: {prompt!r}")
    start_ms = round(time.time() * 1000)
    try:
        # The body's "model" field is a placeholder. Kong's internal classifier
        # overwrites it with "fast" or "smart" before the request reaches the
        # upstream provider, so the value sent here never reaches OpenAI or Bedrock.
        raw = client.chat.completions.with_raw_response.create(
            model="fast",
            messages=[{"role": "user", "content": prompt}],
        )
    except APIStatusError as e:
        elapsed_ms = round(time.time() * 1000) - start_ms
        print(f"{RED(BOLD('[BLOCKED]'))} {RED(BOLD(str(e.status_code)))} {e.message}  ({elapsed_ms}ms)")
        return

    elapsed_ms = round(time.time() * 1000) - start_ms
    completion = raw.parse()
    upstream_model = raw.headers.get("x-kong-llm-model", ".")
    upstream_latency = raw.headers.get("x-kong-upstream-latency", ".")
    proxy_latency = raw.headers.get("x-kong-proxy-latency", ".")
    answer = completion.choices[0].message.content

    print(f"[RESPONSE] {DIM(answer)}")
    print(f"{GREEN(BOLD('[ROUTED TO]'))} upstream model={CYAN(BOLD(upstream_model))}")
    print(f"[LATENCY] {DIM(f'upstream={upstream_latency}ms  proxy={proxy_latency}ms  total={elapsed_ms}ms')}")


def section(title: str) -> None:
    bar = "=" * 70
    print(f"\n{bar}\n{BOLD(title)}\n{bar}")


def main() -> None:
    section("1. Simple prompt. Kong's classifier should recommend the fast tier")
    client = make_client(API_KEY)
    call(client, "simple", SIMPLE_PROMPT)

    section("2. Complex prompt. Kong's classifier should recommend the smart tier")
    call(client, "complex", COMPLEX_PROMPT)

    section("3. Invalid API key. Kong rejects before any classification or upstream call")
    bad_client = make_client("not-a-real-key")
    call(bad_client, "simple", SIMPLE_PROMPT)

    section("Done.")


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        sys.exit(130)
EOF