# /// script
# requires-python = ">=3.10"
# dependencies = ["openai>=3", "python-dotenv>=1"]
# ///
"""hello_three.py: one question, three models.

Web3 Geeks AI Engineering, Week 0.

Sends the same question to a Gemini model on Google's free tier, a free model
on OpenRouter, and a small model running on your own laptop with Ollama. Prints
each answer with its token counts and the time it took.

Run it from the folder that holds your .env file:

    uv run hello_three.py
"""

import os
import sys
import time

from dotenv import load_dotenv
from openai import OpenAI

load_dotenv()

QUESTION = "In one sentence, what is a token for a language model?"

# All three providers accept the same OpenAI-style API. Switching model is a
# change of base URL and model name, nothing else in the code.
MODELS = [
    {
        "label": "Gemini, free tier",
        "base_url": "https://generativelanguage.googleapis.com/v1beta/openai/",
        "key_name": "GEMINI_API_KEY",
        "model": "gemini-3.5-flash-lite",
    },
    {
        "label": "OpenRouter, free model",
        "base_url": "https://openrouter.ai/api/v1",
        "key_name": "OPENROUTER_API_KEY",
        "model": "google/gemma-4-26b-a4b-it:free",
    },
    {
        "label": "Ollama, on your laptop",
        "base_url": "http://localhost:11434/v1",
        "key_name": None,  # Ollama runs locally and needs no key
        "model": "llama3.2:3b",
    },
]


def ask(cfg):
    """Send QUESTION to one model. Returns (answer, tokens_in, tokens_out, seconds)."""
    if cfg["key_name"]:
        key = os.getenv(cfg["key_name"])
        if not key:
            raise RuntimeError(f"{cfg['key_name']} is missing from your .env file")
    else:
        key = "ollama"  # the client needs some value; Ollama ignores it

    client = OpenAI(base_url=cfg["base_url"], api_key=key, timeout=120, max_retries=0)
    start = time.perf_counter()
    reply = client.chat.completions.create(
        model=cfg["model"],
        messages=[{"role": "user", "content": QUESTION}],
    )
    seconds = time.perf_counter() - start

    answer = (reply.choices[0].message.content or "").strip()
    usage = reply.usage
    tokens_in = usage.prompt_tokens if usage else None
    tokens_out = usage.completion_tokens if usage else None
    return answer, tokens_in, tokens_out, seconds


def main():
    print(f"Question: {QUESTION}\n")
    failed = 0
    for cfg in MODELS:
        print(f"{cfg['label']} ({cfg['model']})")
        try:
            answer, tokens_in, tokens_out, seconds = ask(cfg)
        except Exception as err:  # one failing provider shouldn't stop the others
            failed += 1
            print(f"  FAILED: {type(err).__name__}: {err}\n")
            continue
        print(f"  {answer}")
        print(f"  tokens: {tokens_in} in, {tokens_out} out | {seconds:.2f} s\n")

    if failed:
        print(f"{failed} of {len(MODELS)} models failed. See 'If something goes wrong' in the setup guide.")
        sys.exit(1)
    print("All three answered. Your lab is ready for Week 1.")


if __name__ == "__main__":
    main()