Custom LLM Routing

cat <<'EOF' > kong-recipe.yaml
_defaults:
  kongctl:
    namespace: custom-llm-routing-recipe
ai_gateway_model_providers:
- ref: custom-llm-routing-provider
  ai_gateway: !lookup {id: !env 'AI_GATEWAY_ID'}
  name: custom-llm-routing-provider
  display_name: custom-llm-routing
  type: openai
  config:
    auth:
      type: basic
      headers:
      - name: Authorization
        value: !secret {source: !env 'DECK_OPENAI_TOKEN'}
ai_gateway_auth_strategies:
- ref: custom-llm-routing-auth
  ai_gateway: !lookup {id: !env 'AI_GATEWAY_ID'}
  name: custom-llm-routing-auth
  display_name: custom-llm-routing key auth
  type: key-auth
  config:
    key_names:
    - apikey
    hide_credentials: true
ai_gateway_consumers:
- ref: custom-llm-routing-consumer
  ai_gateway: !lookup {id: !env 'AI_GATEWAY_ID'}
  name: custom-llm-routing-consumer
  display_name: custom-llm-routing demo consumer
  type: api-key
  credentials:
  - ref: custom-llm-routing-credential
    name: custom-llm-routing-credential
    display_name: custom-llm-routing demo API key
    type: api-key
    api_key: !secret {source: !env 'DECK_CONSUMER_API_KEY'}
ai_gateway_policies:
- ref: custom-llm-routing-orchestrator
  ai_gateway: !lookup {id: !env 'AI_GATEWAY_ID'}
  name: custom-llm-routing-orchestrator
  display_name: custom-llm-routing token size orchestrator
  type: datakit
  config:
    nodes:
    - name: EXTRACT_PROMPT
      type: jq
      input: request.body
      jq: '{messages: .messages}'
    - name: EXTRACT_AUTH
      type: jq
      input: request.headers
      jq: '{apikey: (.apikey // .Apikey // .APIKEY)}'
    - name: ESTIMATE_TIER
      type: jq
      input: request.body
      jq: "(
  [ (.messages // [])[]
    | .content
    | if type == \"string\"\
        \ then length
      elif type == \"array\" then
        ([.[] | if type\
        \ == \"object\" then ((.text // \"\") | length) elif type == \"string\" then\
        \ length else 0 end] | add // 0)
      else 0 end
  ] | add // 0
) as $chars
\
        | ([ ($chars / 4 | ceil), 1 ] | max) as $tokens
| {
    tokens: $tokens,
\
        \    tier: (if $tokens >= 80 then \"large\" else \"small\" end)
  }
"
    - name: BUILD_DIAG_HEADERS
      type: jq
      input: ESTIMATE_TIER
      jq: "{
  \"X-LLM-Tier\": .tier,
  \"X-Prompt-Est-Tokens\": (.tokens | tostring)
\
        }
"
    - name: IS_LARGE
      type: jq
      input: ESTIMATE_TIER
      jq: .tier == "large"
    - name: ROUTE_DECISION
      type: branch
      input: IS_LARGE
      then:
      - CALL_LARGE
      - RESPOND_LARGE
      else:
      - CALL_SMALL
      - RESPOND_SMALL
    - name: CALL_SMALL
      type: call
      url: http://localhost:8000/custom-llm-routing-small
      method: POST
      inputs:
        body: EXTRACT_PROMPT
        headers: EXTRACT_AUTH
    - name: CALL_LARGE
      type: call
      url: http://localhost:8000/custom-llm-routing-large
      method: POST
      inputs:
        body: EXTRACT_PROMPT
        headers: EXTRACT_AUTH
    - name: RESPOND_SMALL
      type: exit
      inputs:
        body: CALL_SMALL.body
        headers: BUILD_DIAG_HEADERS
    - name: RESPOND_LARGE
      type: exit
      inputs:
        body: CALL_LARGE.body
        headers: BUILD_DIAG_HEADERS
ai_gateway_models:
- ref: custom-llm-routing-small
  ai_gateway: !lookup {id: !env 'AI_GATEWAY_ID'}
  name: custom-llm-routing-small
  display_name: custom-llm-routing (small internal)
  type: model
  formats:
  - type: openai
  capabilities:
  - generate
  access:
    auth_strategies:
    - custom-llm-routing-auth
  config:
    max_request_body_size: 10485760
    response_streaming: deny
    logging:
      payloads: true
    route:
      paths:
      - /custom-llm-routing-small
      protocols:
      - http
      - https
      methods:
      - POST
      - OPTIONS
      strip_path: true
  targets:
  - name: !env 'DECK_CHAT_MODEL_1'
    provider: custom-llm-routing-provider
    config:
      type: openai
- ref: custom-llm-routing-large
  ai_gateway: !lookup {id: !env 'AI_GATEWAY_ID'}
  name: custom-llm-routing-large
  display_name: custom-llm-routing (large internal)
  type: model
  formats:
  - type: openai
  capabilities:
  - generate
  access:
    auth_strategies:
    - custom-llm-routing-auth
  config:
    max_request_body_size: 10485760
    response_streaming: deny
    logging:
      payloads: true
    route:
      paths:
      - /custom-llm-routing-large
      protocols:
      - http
      - https
      methods:
      - POST
      - OPTIONS
      strip_path: true
  targets:
  - name: !env 'DECK_CHAT_MODEL_2'
    provider: custom-llm-routing-provider
    config:
      type: openai
- ref: custom-llm-routing-chat
  ai_gateway: !lookup {id: !env 'AI_GATEWAY_ID'}
  name: custom-llm-routing-chat
  display_name: custom-llm-routing (public)
  type: model
  formats:
  - type: openai
  capabilities:
  - generate
  access:
    auth_strategies:
    - custom-llm-routing-auth
  policies:
  - custom-llm-routing-orchestrator
  config:
    max_request_body_size: 10485760
    response_streaming: deny
    logging:
      payloads: true
    route:
      paths:
      - /custom-llm-routing
      protocols:
      - http
      - https
      methods:
      - POST
      - OPTIONS
      strip_path: true
  targets:
  - name: !env 'DECK_CHAT_MODEL_1'
    provider: custom-llm-routing-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
"""Custom LLM routing demo.

Sends a short and a long prompt through Kong. The pre-function Policy estimates
input size, rewrites the request model to small or large, and AI Models route
to the matching OpenAI target. The client model field is ignored.

Expected output:
  - Short prompt -> X-LLM-Tier=small and X-Kong-LLM-Model ending in the small model
  - Long prompt  -> X-LLM-Tier=large and X-Kong-LLM-Model ending in the large model

Run:
  export PROXY_URL=http://localhost:8000
  python demo.py
"""

from __future__ import annotations

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"
THRESHOLD = int(os.getenv("DECK_TOKEN_THRESHOLD", "80"))

_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


def BOLD(s: str) -> str:
    return _c("1", s)


def DIM(s: str) -> str:
    return _c("2", s)


def GREEN(s: str) -> str:
    return _c("32", s)


def CYAN(s: str) -> str:
    return _c("36", s)


def RED(s: str) -> str:
    return _c("31", s)


def YELLOW(s: str) -> str:
    return _c("33", s)


def make_client() -> OpenAI:
    return OpenAI(
        base_url=f"{PROXY_URL}/custom-llm-routing",
        api_key="unused",
        default_headers={"apikey": API_KEY},
    )


def est_tokens(content: str) -> int:
    # Matches the Pre-function heuristic: ceil(utf8-bytes / 4).
    return max(1, (len(content.encode("utf-8")) + 3) // 4)


def call(client: OpenAI, label: str, content: str) -> None:
    tokens = est_tokens(content)
    expected = "large" if tokens >= THRESHOLD else "small"
    print(f"\n{BOLD('[REQUEST]')} {label}")
    print(
        f"  {DIM(f'est_tokens={tokens} threshold={THRESHOLD} expected_tier={expected}')}"
    )
    print(f"  {DIM(f'client model field=ignored-by-router prompt_preview={content[:48]!r}')}")

    start_ms = round(time.time() * 1000)
    try:
        raw = client.chat.completions.with_raw_response.create(
            model="ignored-by-router",
            messages=[{"role": "user", "content": content}],
        )
    except APIStatusError as e:
        elapsed_ms = round(time.time() * 1000) - start_ms
        print(
            f"{RED(BOLD('[ERROR]'))} {RED(BOLD(str(e.status_code)))} "
            f"{e.message}  ({elapsed_ms}ms)"
        )
        raise SystemExit(1) from e

    elapsed_ms = round(time.time() * 1000) - start_ms
    completion = raw.parse()
    tier = raw.headers.get("x-llm-tier", ".")
    est = raw.headers.get("x-prompt-est-tokens", ".")
    upstream_model = raw.headers.get("x-kong-llm-model", ".")
    answer = (completion.choices[0].message.content or "")[:120]

    print(f"[RESPONSE] {DIM(answer)}")
    print(
        f"{GREEN(BOLD('[ROUTED TO]'))} tier={GREEN(BOLD(tier))} "
        f"est_tokens={YELLOW(BOLD(est))} "
        f"upstream={CYAN(BOLD(upstream_model))}"
    )
    print(f"[LATENCY] {DIM(f'total={elapsed_ms}ms')}")

    if tier != expected:
        print(
            f"{RED(BOLD('[MISMATCH]'))} expected tier={expected!r}, got={tier!r}"
        )
        raise SystemExit(1)


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


def main() -> None:
    section("1. Short prompt routes to small")
    client = make_client()
    call(client, "short", "hi")

    section("2. Long prompt routes to large")
    long_prompt = "x" * 400
    call(client, "long", long_prompt)

    section("Done.")


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