Jev as an LLM Router: Fast Typed Decisions for Model Selection

"""LLM router demo: Jev classifies the query, code routes it to a Bedrock model.

A user query comes in. Jev (TypeSafe's System One model) classifies it with a
couple of typed questions, and ordinary code maps the result to a model tier on
Amazon Bedrock. The cheap tier handles simple work, the mid tier handles routine
work, and the frontier tier handles complex reasoning. When Jev is not confident
about the route, the query falls back to the frontier tier.

The route types come from the R.E.S.E.T. framework in AI Route Advisor:
https://github.com/enitinb/ai-engineering-lab/tree/main/ai-route-advisor-skill
There, the route class is chosen at design time. Here
the same route types are applied at runtime and mapped to a concrete Bedrock stack.

Docs:
- TypeSafe primitives : https://docs.typesafe.ai/primitives
- TypeSafe confidence : https://docs.typesafe.ai/confidence
- Bedrock Converse API : https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference.html
"""

import boto3
from typesafe_sdk import Choice, Score, TypeSafeClient

# ---- Config ----
TYPESAFE_API_KEY = "<<TYPESAFE_API_KEY>>"
AWS_PROFILE = None # None = default credential chain
AWS_REGION = "us-east-1"

# Route type (R.E.S.E.T. taxonomy) -> Bedrock model.
ROUTE_TO_MODEL = {
"low_cost_tier": "us.openai.gpt-6-luna", # repeatable work at scale
"mid_tier": "us.openai.gpt-6-sol", # demanding development work
"frontier_or_reasoning_tier": "us.openai.gpt-6-astra", # most capable, hardest work
}
FRONTIER_ROUTE = "frontier_or_reasoning_tier" # fallback route
CONFIDENCE_FLOOR = 0.6 # route confidence below this falls back to frontier

typesafe = TypeSafeClient(api_key=TYPESAFE_API_KEY)
bedrock = boto3.Session(profile_name=AWS_PROFILE, region_name=AWS_REGION).client(
"bedrock-runtime"
)

# ---- 1. Jev: classify the query ----
def classify(query: str) -> dict:
r = typesafe.system_one(
state=query,
questions={
"route_type": Choice(
instructions="Which route type best fits this request",
criteria={
"low_cost_tier": "Simple, high-volume work: extraction, formatting, tagging, short summaries",
"mid_tier": "Routine work: drafting, longer summarization, structured analysis, or writing straightforward code",
"frontier_or_reasoning_tier": "Complex, multi-step reasoning: architecture decisions, strategy, or hard problem solving",
},
),
"complexity": Score(
instructions="How much reasoning does this request require",
criteria=[
"Trivial, single step",
"Moderate, a few steps",
"Hard, deep multi-step reasoning",
],
),
},
)
return {
"route_type": r.answers["route_type"].choice,
"route_confidence": r.answers["route_type"].confidence,
"complexity": r.answers["complexity"].score,
}

# ---- 2. Code: pick the model ----
def pick_route(c: dict) -> str:
# Low confidence on the route -> send to the frontier tier.
if c["route_confidence"] < CONFIDENCE_FLOOR:
return FRONTIER_ROUTE
# High complexity -> upgrade to the frontier tier regardless of route.
if c["complexity"] >= 2.0:
return FRONTIER_ROUTE
return c["route_type"]

# ---- 3. Bedrock: answer with the chosen model ----
def answer(query: str, model_id: str) -> str:
resp = bedrock.converse(
modelId=model_id,
messages=[{"role": "user", "content": [{"text": query}]}],
inferenceConfig={"maxTokens": 1500},
)
# These models can return reasoning blocks alongside the answer, so pick the
# text block rather than assuming it is first.
blocks = resp["output"]["message"]["content"]
texts = [b["text"] for b in blocks if "text" in b]
if texts:
return "\n".join(texts)
reason = resp.get("stopReason", "unknown")
return f"(no text returned; stopReason={reason})"

# ---- Demo: one query per route ----
QUERIES = [
"Extract the names, dates, and dollar amounts from this note: "
"Invoice #204 dated March 3 for $1,250 to Acme, due April 2.",
"Summarize this in two sentences: Our Q3 launch shipped three features, "
"grew signups 18 percent, and cut onboarding time from 12 minutes to 4.",
"Write a Python function that reads a CSV file and returns the rows where "
"the 'amount' column is greater than 1000, with basic error handling.",
"Compare a monolith and microservices for a 5-person startup that expects "
"rapid growth. Recommend one and explain the tradeoffs.",
"Thoughts?",
]

def main() -> None:
for query in QUERIES:
c = classify(query)
route = pick_route(c)
model_id = ROUTE_TO_MODEL[route]

print("=" * 70)
print(f"QUERY : {query}")
print(f"JEV : route={c['route_type']} (conf {c['route_confidence']:.2f}) "
f"complexity={c['complexity']:.2f}")
print(f"ROUTE : {route} -> {model_id}")
print("-" * 70)
print(answer(query, model_id))
print()

if __name__ == "__main__":
main()