import json, os, urllib.request

# Requires S1M_API_KEY and TENSORX_API_KEY; Python 3.10+ (standard library only).
REQUEST = "Write a Python function that returns the square of a number."
QUESTIONS = {'category': {'type': 'choice', 'criteria': {'coding': '', 'agentic': '', 'math': '', 'knowledge': '', 'long_context': '', 'tool_use': '', 'design': '', 'summarisation': '', 'general': ''}}, 'difficulty': {'type': 'score', 'criteria': ['trivial', 'easy', 'moderate', 'hard', 'frontier']}, 'stakes': {'type': 'score', 'criteria': ['negligible', 'low', 'medium', 'high']}, 'needs_tools': {'type': 'noul'}, 'needs_vision': {'type': 'noul'}, 'needs_long_context': {'type': 'noul'}, 'follow_up': {'type': 'noul'}}
FAST = "deepseek/deepseek-v3.2"
STRONG = "deepseek/deepseek-v4.1-flash"
ROUTES = {"coding": {"easy": FAST, "hard": STRONG},
          "general": {"easy": FAST, "hard": STRONG}}
FALLBACK = STRONG

def post(url, key, body):
    req = urllib.request.Request(url, data=json.dumps(body).encode(), headers={
        "Authorization": "Bearer " + key, "Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=60) as response:
        return json.load(response)

result = post("https://api.system1models.ai/v1/systemone", os.environ["S1M_API_KEY"],
              {"model": "s1-llm-auto-router", "state": {"request": REQUEST,
               "context": "(new conversation)"}, "questions": QUESTIONS})
a = result["answers"]
# score is an expected LEVEL (0..4 / 0..3), not a probability.
difficulty = "hard" if a["difficulty"]["score"] >= 2 else "easy"
model = ROUTES.get(a["category"]["choice"], {}).get(difficulty, FALLBACK)
reason = "category/difficulty table"
if a["category"]["confidence"] < 0.6 or a["stakes"]["score"] >= 2:
    model, reason = FALLBACK, "low confidence or high stakes"
# This demo handles a new text turn. Never silently lose required capabilities/context.
unsupported = [f for f in ("needs_tools", "needs_vision", "needs_long_context", "follow_up")
               if a[f]["noul"] >= 0.5]
decision = {"request_id": result["id"], "model": model, "reason": reason,
            "signals": a, "unsupported": unsupported}
print(json.dumps(decision))  # All distributions and flag probabilities; no prompts or keys.
if unsupported:
    raise SystemExit("Configure capability support / conversation history before forwarding")
completion = post("https://api.tensorx.ai/v1/chat/completions", os.environ["TENSORX_API_KEY"],
                  {"model": model, "messages": [{"role": "user", "content": REQUEST}],
                   "max_tokens": 80, "temperature": 0})
print(completion["choices"][0]["message"]["content"])