app.py

"""
Agent 101 — a side-by-side demo of a plain LLM vs. an LLM-with-tools agent.

Based on the Week 12 CS 203 lab at IIT Gandhinagar:
"Building AI Agents from Scratch with Gemma 4".

Left panel: a plain LLM answers from memory alone.
Right panel: the same model runs an agent loop — it can call a calculator,
weather lookup, unit converter, course-notes search, or CS/ML dictionary,
and we show every tool call + result as it happens.

Inference is delegated to the HF Inference Providers API, so no GPU is
needed on the Space. Set HF_TOKEN as a Space secret.
"""

from __future__ import annotations

import datetime as _dt
import inspect
import json
import os
import re
import textwrap

import gradio as gr
from huggingface_hub import InferenceClient

# --- Model registry ------------------------------------------------------
# Models known to handle OpenAI-style tool calling on HF Inference Providers.
# 70B/72B-class models are MUCH more reliable at chaining multiple tool
# calls than 8B — try the smaller one to see the failure modes.
MODELS = [
    "google/gemma-3-27b-it",
    "google/gemma-3-12b-it",
    # Gemma 4 E2B and Gemma 3n E2B are tiny on-device models. They usually
    # aren't served on Inference Providers, so picking them may return
    # `model_not_supported` — listed here for experimentation.
    "google/gemma-3n-E2B-it",
    "google/gemma-4-e2b-it",
    "meta-llama/Llama-3.3-70B-Instruct",
    "Qwen/Qwen2.5-7B-Instruct",
    "meta-llama/Llama-3.1-8B-Instruct",
]
DEFAULT_MODEL = os.environ.get("MODEL_ID", MODELS[0])

client = InferenceClient(token=os.environ.get("HF_TOKEN"))


# --- System prompt -------------------------------------------------------

SYSTEM_PROMPT = """You are a helpful assistant with access to tools.

Rules you must follow:
1. For ANY arithmetic — even simple multiplication like 5 * 7 — call the
   `calculate` tool. Never compute numbers in your head.
2. For real-time facts like weather, current time, or stock prices,
   call the matching tool (`get_weather`, `get_time`, `get_stock_price`).
   Do not guess or say you don't know.
3. For unit conversions, call `convert_units` — don't approximate.
4. For currency conversion, call `get_exchange_rate`.
5. For distances between cities, call `get_distance`. For city
   populations, call `get_population`.
6. For questions about the CS 203 course ("what week did we cover X?"),
   call `search_notes`.
7. For definitions of CS / ML terms, call `define_word`.
8. Multi-step questions need multiple tool calls in sequence. Examples:
   - "5 km/day for a week in miles" → `convert_units` then `calculate`.
   - "Delhi-Mumbai distance in miles" → `get_distance` then
     `convert_units`.
   - "Shares of AAPL for 5000 INR" → `get_exchange_rate` →
     `get_stock_price` → `calculate`.
9. After every tool result, decide: do I need another tool, or can I
   write the final answer? Only answer once you have ALL the numbers.
10. If the question genuinely doesn't need any tool (e.g. "capital of
    France"), answer directly.

CRITICAL: when you ARE calling a tool, use the provider's structured
tool_calls interface — do NOT write tool calls as Python-style text like
`[get_weather(city="Delhi")]` or as JSON in your reply. If you want to
call a tool, emit a real tool_call; otherwise write a final natural-
English answer with the numbers spelled out.
"""


# --- Tools ---------------------------------------------------------------

def calculate(expression: str) -> str:
    allowed = set("0123456789+-*/.() ")
    if not all(c in allowed for c in expression):
        return "Error: only numbers and +-*/() are allowed"
    try:
        return str(round(eval(expression), 10))
    except Exception as e:
        return f"Error: {e}"


def get_weather(city: str) -> str:
    data = {
        "gandhinagar": {"temp_c": 38, "condition": "Sunny",  "humidity": 25},
        "mumbai":      {"temp_c": 32, "condition": "Humid",  "humidity": 80},
        "bangalore":   {"temp_c": 28, "condition": "Rainy",  "humidity": 65},
        "delhi":       {"temp_c": 40, "condition": "Haze",   "humidity": 30},
        "chennai":     {"temp_c": 35, "condition": "Cloudy", "humidity": 70},
        "kolkata":     {"temp_c": 33, "condition": "Humid",  "humidity": 75},
        "paris":       {"temp_c": 18, "condition": "Cloudy", "humidity": 60},
        "new york":    {"temp_c": 22, "condition": "Clear",  "humidity": 45},
    }
    return json.dumps(data.get(city.lower(), {"error": f"No data for {city}"}))


def convert_units(value: float, from_unit: str, to_unit: str) -> str:
    conversions = {
        ("celsius", "fahrenheit"):  lambda v: v * 9 / 5 + 32,
        ("fahrenheit", "celsius"):  lambda v: (v - 32) * 5 / 9,
        ("kg", "pounds"):           lambda v: v * 2.20462,
        ("pounds", "kg"):           lambda v: v / 2.20462,
        ("km", "miles"):            lambda v: v * 0.621371,
        ("miles", "km"):            lambda v: v / 0.621371,
        ("meters", "feet"):         lambda v: v * 3.28084,
        ("feet", "meters"):         lambda v: v / 3.28084,
        ("liters", "gallons"):      lambda v: v * 0.264172,
        ("gallons", "liters"):      lambda v: v / 0.264172,
    }
    try:
        value = float(value)
    except (TypeError, ValueError):
        return f"Error: value '{value}' is not a number"
    fn = conversions.get((str(from_unit).lower(), str(to_unit).lower()))
    if fn is None:
        return f"Cannot convert {from_unit} to {to_unit}"
    return f"{fn(value):.4f} {to_unit}"


def search_notes(query: str) -> str:
    topics = {
        "data drift":          "Week 10: detecting distribution shift with KS test, PSI, chi-squared test",
        "profiling":           "Week 11: cProfile, timeit, finding bottlenecks in ML code",
        "quantization":        "Week 11: INT8 / FP16, dynamic quantization, ONNX, model compression",
        "pruning":             "Week 11: unstructured and structured pruning, removing weights",
        "distillation":        "Week 11: teacher-student training, soft labels, knowledge transfer",
        "docker":              "Week 10: containerization, Dockerfiles, reproducible environments",
        "fastapi":             "Week 12: building REST APIs, Pydantic validation, /predict endpoints",
        "agents":              "Week 12: tool calling, Gemma 4, the agent loop, function calling",
        "git":                 "Week 9: version control, commits, branches, merge conflicts",
        "experiment tracking": "Week 8: MLflow, Weights & Biases, hyperparameter tuning",
        "cross validation":    "Week 7: k-fold CV, train/val/test splits, bias-variance tradeoff",
        "gradio":              "Week 12: building demo UIs for ML models, share=True for public link",
        "streamlit":           "Week 12: building dashboard-style ML apps with Python",
        "batching":            "Week 11: processing multiple inputs at once for throughput",
        "onnx":                "Week 11: portable model format, export once run anywhere",
    }
    q = query.lower()
    matches = {k: v for k, v in topics.items() if q in k or q in v.lower()}
    if matches:
        return json.dumps(matches, indent=2)
    return json.dumps({"message": f"No results for '{query}'. Try: {', '.join(list(topics.keys())[:5])}"})


def get_time(timezone: str = "UTC") -> str:
    """Return the current wall-clock time in a named timezone or UTC offset."""
    offsets = {
        "utc": 0, "ist": 5.5, "gmt": 0, "bst": 1,
        "edt": -4, "est": -5, "pdt": -7, "pst": -8,
        "jst": 9, "kst": 9, "cst": 8, "cet": 1, "eet": 2,
        "sgt": 8, "hkt": 8, "aest": 10,
        "gandhinagar": 5.5, "mumbai": 5.5, "delhi": 5.5, "bangalore": 5.5,
        "tokyo": 9, "seoul": 9, "singapore": 8, "london": 0, "paris": 1,
        "new york": -4, "san francisco": -7, "sydney": 10,
    }
    key = timezone.lower().strip()
    offset_h = offsets.get(key)
    if offset_h is None:
        m = re.match(r"^utc([+-])(\d+(?:\.\d+)?)$", key)
        if m:
            offset_h = float(m.group(2)) * (1 if m.group(1) == "+" else -1)
    if offset_h is None:
        return json.dumps({"error": f"Unknown timezone '{timezone}'. Try UTC, IST, JST, 'Tokyo', 'New York', etc."})
    now_utc = _dt.datetime.utcnow()
    local = now_utc + _dt.timedelta(hours=offset_h)
    return json.dumps({
        "timezone": timezone,
        "utc_offset_hours": offset_h,
        "iso": local.strftime("%Y-%m-%d %H:%M:%S"),
        "day_of_week": local.strftime("%A"),
    })


def get_exchange_rate(from_currency: str, to_currency: str, amount: float = 1.0) -> str:
    """Mock currency converter — rates are approximate and fixed for the demo."""
    # Rates expressed as "1 unit of X in USD".
    usd_rates = {
        "usd": 1.0, "eur": 1.08, "gbp": 1.27, "jpy": 0.0067,
        "inr": 0.012, "cny": 0.14, "aud": 0.66, "cad": 0.73,
        "sgd": 0.74, "chf": 1.13, "krw": 0.00075,
    }
    try:
        amount = float(amount)
    except (TypeError, ValueError):
        return json.dumps({"error": f"amount '{amount}' is not a number"})
    f = str(from_currency).lower().strip()
    t = str(to_currency).lower().strip()
    if f not in usd_rates or t not in usd_rates:
        return json.dumps({
            "error": f"Unsupported currency pair {from_currency}->{to_currency}",
            "supported": sorted(usd_rates.keys()),
        })
    usd = amount * usd_rates[f]
    converted = usd / usd_rates[t]
    return json.dumps({
        "from": f.upper(), "to": t.upper(),
        "amount": amount, "converted": round(converted, 2),
        "note": "rates are approximate demo values, not live market data",
    })


def get_distance(city_a: str, city_b: str) -> str:
    """Great-circle distance between two cities in kilometres (demo data)."""
    # Stored as city -> (lat, lon) so we can list any pair.
    coords = {
        "gandhinagar": (23.22, 72.65), "mumbai": (19.08, 72.88),
        "delhi":       (28.61, 77.21), "bangalore": (12.97, 77.59),
        "chennai":     (13.08, 80.27), "kolkata": (22.57, 88.36),
        "hyderabad":   (17.38, 78.49), "pune": (18.52, 73.86),
        "tokyo":       (35.68, 139.69), "seoul": (37.57, 126.98),
        "singapore":   (1.35, 103.82),  "hong kong": (22.32, 114.17),
        "london":      (51.51, -0.13),  "paris": (48.86, 2.35),
        "new york":    (40.71, -74.00), "san francisco": (37.77, -122.42),
        "sydney":      (-33.87, 151.21),
    }
    a = coords.get(city_a.lower().strip())
    b = coords.get(city_b.lower().strip())
    if a is None or b is None:
        return json.dumps({"error": f"Unknown city. Supported: {sorted(coords.keys())}"})
    import math
    lat1, lon1 = map(math.radians, a)
    lat2, lon2 = map(math.radians, b)
    dlat, dlon = lat2 - lat1, lon2 - lon1
    h = math.sin(dlat / 2) ** 2 + math.cos(lat1) * math.cos(lat2) * math.sin(dlon / 2) ** 2
    km = 2 * 6371.0 * math.asin(math.sqrt(h))
    return json.dumps({"from": city_a, "to": city_b, "distance_km": round(km, 1)})


def get_population(city: str) -> str:
    """Approximate metro-area population (demo data, in millions)."""
    data = {
        "gandhinagar": 0.36, "mumbai": 20.96, "delhi": 32.07, "bangalore": 13.61,
        "chennai": 11.56,    "kolkata": 15.33, "hyderabad": 10.53, "pune": 6.81,
        "tokyo": 37.44,      "seoul": 9.97,    "singapore": 5.94,  "hong kong": 7.49,
        "london": 9.54,      "paris": 11.21,   "new york": 18.87,  "san francisco": 3.32,
        "sydney": 5.37,
    }
    v = data.get(city.lower().strip())
    if v is None:
        return json.dumps({"error": f"No population data for {city}"})
    return json.dumps({"city": city, "population_millions": v})


def get_stock_price(ticker: str) -> str:
    """Mock share price in USD for a handful of tickers."""
    prices = {
        "aapl": 228.45, "msft": 417.12, "googl": 178.32, "amzn": 212.91,
        "meta": 574.80, "tsla": 249.07, "nvda": 135.60, "amd": 162.33,
        "ibm": 224.18,  "intc":  21.40, "nflx": 745.20, "orcl": 176.55,
    }
    p = prices.get(ticker.lower().strip())
    if p is None:
        return json.dumps({"error": f"Unknown ticker '{ticker}'. Try: {sorted(prices.keys())}"})
    return json.dumps({"ticker": ticker.upper(), "price_usd": p, "note": "demo price, not live market data"})


def define_word(word: str) -> str:
    definitions = {
        "overfitting":    "When a model learns training data too well (including noise) and performs poorly on new data.",
        "gradient":       "The vector of partial derivatives of the loss w.r.t. each parameter.",
        "epoch":          "One full pass through the entire training dataset.",
        "batch size":     "The number of training examples in one forward/backward pass.",
        "learning rate":  "A hyperparameter controlling how much to adjust weights per update.",
        "regularization": "Techniques to prevent overfitting, e.g. L1/L2 penalties, dropout.",
        "transformer":    "A neural-network architecture based on self-attention (GPT, BERT, etc.).",
        "tokenizer":      "Converts text into a sequence of integer IDs a model can process.",
    }
    r = definitions.get(word.lower())
    if r:
        return json.dumps({"word": word, "definition": r})
    return json.dumps({"error": f"'{word}' not found. Available: {', '.join(definitions.keys())}"})


TOOL_SCHEMAS = [
    {
        "type": "function",
        "function": {
            "name": "calculate",
            "description": "Evaluate a math expression. Use Python syntax: + - * / ** (). Example: '4729 * 8314'. ALWAYS use this for arithmetic.",
            "parameters": {
                "type": "object",
                "properties": {"expression": {"type": "string", "description": "e.g. '2 + 3 * 4'"}},
                "required": ["expression"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "get_weather",
            "description": "Current weather (temperature in Celsius, condition, humidity) for a city. Cities: Gandhinagar, Mumbai, Bangalore, Delhi, Chennai, Kolkata, Paris, New York.",
            "parameters": {
                "type": "object",
                "properties": {"city": {"type": "string", "description": "City name, e.g. 'Mumbai'"}},
                "required": ["city"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "convert_units",
            "description": "Convert a value between units. Supports celsius/fahrenheit, kg/pounds, km/miles, meters/feet, liters/gallons.",
            "parameters": {
                "type": "object",
                "properties": {
                    "value": {"type": "number", "description": "The numeric value"},
                    "from_unit": {"type": "string", "description": "Source unit, e.g. 'celsius'"},
                    "to_unit": {"type": "string", "description": "Target unit, e.g. 'fahrenheit'"},
                },
                "required": ["value", "from_unit", "to_unit"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "search_notes",
            "description": "Search the CS 203 (IIT Gandhinagar) course notes by keyword. Returns which week a topic was covered.",
            "parameters": {
                "type": "object",
                "properties": {"query": {"type": "string", "description": "Topic keyword, e.g. 'docker'"}},
                "required": ["query"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "define_word",
            "description": "Look up the definition of a common CS / ML term.",
            "parameters": {
                "type": "object",
                "properties": {"word": {"type": "string", "description": "Term, e.g. 'overfitting'"}},
                "required": ["word"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "get_time",
            "description": "Get the current wall-clock time in a given timezone or city. Accepts UTC, IST, JST, EST, PDT, 'Tokyo', 'New York', etc.",
            "parameters": {
                "type": "object",
                "properties": {"timezone": {"type": "string", "description": "Timezone code or city, e.g. 'IST' or 'Tokyo'"}},
                "required": ["timezone"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "get_exchange_rate",
            "description": "Convert a currency amount using fixed demo rates. Supports USD, EUR, GBP, JPY, INR, CNY, AUD, CAD, SGD, CHF, KRW.",
            "parameters": {
                "type": "object",
                "properties": {
                    "from_currency": {"type": "string", "description": "Source ISO code, e.g. 'USD'"},
                    "to_currency":   {"type": "string", "description": "Target ISO code, e.g. 'INR'"},
                    "amount":        {"type": "number", "description": "Amount to convert"},
                },
                "required": ["from_currency", "to_currency", "amount"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "get_distance",
            "description": "Great-circle distance in kilometres between two cities. Supports major Indian and global cities (Gandhinagar, Mumbai, Delhi, Bangalore, Chennai, Kolkata, Hyderabad, Pune, Tokyo, Seoul, Singapore, Hong Kong, London, Paris, New York, San Francisco, Sydney).",
            "parameters": {
                "type": "object",
                "properties": {
                    "city_a": {"type": "string", "description": "First city, e.g. 'Delhi'"},
                    "city_b": {"type": "string", "description": "Second city, e.g. 'Mumbai'"},
                },
                "required": ["city_a", "city_b"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "get_population",
            "description": "Approximate metro-area population (in millions). Same city list as get_distance.",
            "parameters": {
                "type": "object",
                "properties": {"city": {"type": "string", "description": "City name"}},
                "required": ["city"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "get_stock_price",
            "description": "Current share price in USD for a handful of tickers (AAPL, MSFT, GOOGL, AMZN, META, TSLA, NVDA, AMD, IBM, INTC, NFLX, ORCL).",
            "parameters": {
                "type": "object",
                "properties": {"ticker": {"type": "string", "description": "Ticker symbol, e.g. 'AAPL'"}},
                "required": ["ticker"],
            },
        },
    },
]

TOOL_FUNCTIONS = {
    "calculate": calculate,
    "get_weather": get_weather,
    "convert_units": convert_units,
    "search_notes": search_notes,
    "define_word": define_word,
    "get_time": get_time,
    "get_exchange_rate": get_exchange_rate,
    "get_distance": get_distance,
    "get_population": get_population,
    "get_stock_price": get_stock_price,
}


# --- Inference plumbing --------------------------------------------------

def _chat(messages, model, tools=None, max_tokens=512):
    kwargs = dict(
        model=model,
        messages=messages,
        max_tokens=max_tokens,
        temperature=0.0,
    )
    if tools:
        kwargs["tools"] = tools
        kwargs["tool_choice"] = "auto"
    return client.chat.completions.create(**kwargs)


def _friendly_error(e: Exception, model: str) -> str:
    s = str(e)
    if "model_not_supported" in s or "is not supported by any provider" in s:
        return (
            f"**`{model}` isn't available on any Inference Provider enabled for your account.**\n\n"
            f"Pick a different model from the dropdown, or enable a provider that hosts it at "
            f"[hf.co/settings/inference-providers](https://huggingface.co/settings/inference-providers)."
        )
    if "tool_use_failed" in s:
        return (
            f"**`{model}` generated a malformed tool call that its provider rejected.**\n\n"
            f"This is a known model/provider parser issue (common with Llama 3.x on Groq/Cerebras "
            f"when the model tries to nest calls). Try `google/gemma-3-27b-it` instead — "
            f"it handles tool syntax more cleanly."
        )
    if "401" in s or "Unauthorized" in s or "authentication" in s.lower():
        return (
            "**No `HF_TOKEN` set or the token is invalid.**\n\n"
            "Set the `HF_TOKEN` secret under the Space's Settings → Variables and secrets."
        )
    return f"**Error calling model:** `{s}`"


def _parse_args(raw):
    if isinstance(raw, dict):
        return raw
    if isinstance(raw, str):
        try:
            return json.loads(raw)
        except Exception:
            return {}
    return {}


# Fallback: some models emit tool calls as plain text instead of as
# structured `tool_calls`. We scrape a few common shapes so the loop
# doesn't silently give up:
#   - JSON     : {"name": "foo", "arguments": {...}}
#   - Python   : [foo(k=v, k="v")]   (Gemma 3)
#   - Llama    : <function=foo {"k": "v"}>
_JSON_LEAK_RE = re.compile(
    r'\{[^{}]*"name"\s*:\s*"(?P<name>[a-zA-Z_]+)"[^{}]*"arguments"\s*:\s*(?P<args>\{[^{}]*\})[^{}]*\}'
)
_PY_LEAK_RE = re.compile(r"\[?\s*(?P<name>[a-zA-Z_]\w*)\s*\((?P<args>[^)]*)\)\s*\]?")
_LLAMA_LEAK_RE = re.compile(r"<function=(?P<name>[a-zA-Z_]\w*)\s*(?P<args>\{.*?\})\s*>", re.DOTALL)
# Gemma 3 sometimes emits: `get_weather, city="Delhi"` (no parens, just `name, k=v`)
_KWARG_LEAK_RE = re.compile(
    r"(?P<name>[a-zA-Z_]\w*)\s*,\s*"
    r"(?P<args>\w+\s*=\s*(?:\"[^\"]*\"|'[^']*'|[^,\s]+)"
    r"(?:\s*,\s*\w+\s*=\s*(?:\"[^\"]*\"|'[^']*'|[^,\s]+))*)"
)


def _coerce_scalar(s: str):
    s = s.strip()
    if not s:
        return s
    if (s[0] == s[-1]) and s[0] in {'"', "'"}:
        return s[1:-1]
    try:
        if "." in s:
            return float(s)
        return int(s)
    except ValueError:
        return s


def _parse_py_args(arg_str: str) -> dict:
    """Parse a loose k=v, k=v string used in Python-like tool-call leaks."""
    args = {}
    # Split on commas that are not inside quotes.
    parts = re.findall(r'(?:[^,"\']|"[^"]*"|\'[^\']*\')+', arg_str)
    for part in parts:
        if "=" not in part:
            continue
        k, v = part.split("=", 1)
        args[k.strip()] = _coerce_scalar(v)
    return args


def _extract_leaked_tool_calls(content: str):
    if not content:
        return []
    calls = []
    for m in _JSON_LEAK_RE.finditer(content):
        name = m.group("name")
        try:
            args = json.loads(m.group("args"))
        except Exception:
            continue
        if name in TOOL_FUNCTIONS:
            calls.append({"name": name, "arguments": args})
    for m in _LLAMA_LEAK_RE.finditer(content):
        name = m.group("name")
        try:
            args = json.loads(m.group("args"))
        except Exception:
            continue
        if name in TOOL_FUNCTIONS:
            calls.append({"name": name, "arguments": args})
    if not calls:
        for m in _PY_LEAK_RE.finditer(content):
            name = m.group("name")
            if name not in TOOL_FUNCTIONS:
                continue
            args = _parse_py_args(m.group("args"))
            calls.append({"name": name, "arguments": args})
    if not calls:
        for m in _KWARG_LEAK_RE.finditer(content):
            name = m.group("name")
            if name not in TOOL_FUNCTIONS:
                continue
            args = _parse_py_args(m.group("args"))
            if args:
                calls.append({"name": name, "arguments": args})
    return calls


# --- Inference entry points ----------------------------------------------

def answer_without_tools(question: str, model: str) -> str:
    messages = [{"role": "user", "content": question}]
    try:
        resp = _chat(messages, model=model)
        return (resp.choices[0].message.content or "(empty response)").strip()
    except Exception as e:
        return _friendly_error(e, model)


def answer_with_tools(question: str, model: str, max_steps: int = 5):
    """Run the agent loop. Yields (trace_markdown, final_answer_markdown)."""
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": question},
    ]
    trace = [f"**Question:** {question}\n"]
    final = "_Still thinking..._"
    last_tool_result = None
    last_tool_name = None

    for step in range(1, max_steps + 1):
        trace.append(f"### Step {step}")
        try:
            resp = _chat(messages, model=model, tools=TOOL_SCHEMAS)
        except Exception as e:
            msg = _friendly_error(e, model)
            trace.append(msg)
            yield "\n\n".join(trace), msg
            return

        msg = resp.choices[0].message
        tool_calls = getattr(msg, "tool_calls", None) or []

        # Fallback: did the model leak tool calls into content?
        leaked = []
        if not tool_calls and msg.content:
            leaked = _extract_leaked_tool_calls(msg.content)
            if leaked:
                trace.append(
                    "_Model emitted tool calls as text instead of structured "
                    "`tool_calls` — parsing them out of `content`. "
                    "(This is a small-model symptom; try a 70B+ model.)_"
                )

        if not tool_calls and not leaked:
            final_text = (msg.content or "").strip()
            if not final_text and last_tool_result is not None:
                # Model went quiet after a tool call — surface the tool's result.
                final_text = (
                    f"The model returned nothing after calling `{last_tool_name}`. "
                    f"Showing its result: **{last_tool_result}**"
                )
                trace.append("_Model returned empty content — using the last tool result._")
            elif not final_text:
                final_text = "(empty response)"
            else:
                trace.append("No tool call — model answered directly.")
            final = final_text
            yield "\n\n".join(trace), f"### ✅ Final answer\n\n{final}"
            return

        # Record the assistant turn so the model sees its own tool calls.
        if tool_calls:
            messages.append({
                "role": "assistant",
                "content": msg.content or "",
                "tool_calls": [
                    {
                        "id": tc.id,
                        "type": "function",
                        "function": {
                            "name": tc.function.name,
                            "arguments": tc.function.arguments
                            if isinstance(tc.function.arguments, str)
                            else json.dumps(tc.function.arguments),
                        },
                    }
                    for tc in tool_calls
                ],
            })
            iter_calls = [
                (tc.id, tc.function.name, _parse_args(tc.function.arguments))
                for tc in tool_calls
            ]
        else:
            # Synthesised from leaked content; there are no tool_call ids,
            # so we just append user-role tool results as plain text.
            messages.append({"role": "assistant", "content": msg.content or ""})
            iter_calls = [(None, c["name"], c["arguments"]) for c in leaked]

        for call_id, name, args in iter_calls:
            trace.append(f"- **Tool call:** `{name}({json.dumps(args)})`")

            fn = TOOL_FUNCTIONS.get(name)
            if fn is None:
                tool_result = f"Error: unknown tool '{name}'"
            else:
                try:
                    tool_result = fn(**args)
                except TypeError as e:
                    tool_result = f"Error: bad arguments to {name}: {e}"
                except Exception as e:
                    tool_result = f"Error running {name}: {e}"

            trace.append(f"- **Result:** `{tool_result}`")
            last_tool_result = tool_result
            last_tool_name = name
            if call_id is not None:
                messages.append({
                    "role": "tool",
                    "tool_call_id": call_id,
                    "content": str(tool_result),
                })
            else:
                messages.append({
                    "role": "user",
                    "content": f"(Tool {name} returned: {tool_result})",
                })

        yield "\n\n".join(trace), final

    trace.append(f"\n**Stopped:** hit max_steps = {max_steps}.")
    yield "\n\n".join(trace), "_Stopped before producing a final answer (max steps)._"


def run_both(question: str, model: str):
    if not question or not question.strip():
        yield "_Type a question first._", "", ""
        return
    plain = answer_without_tools(question, model)
    yield plain, "_Running agent loop..._", ""
    for trace, final in answer_with_tools(question, model):
        yield plain, trace, final


# --- "How it works" tab content -----------------------------------------

def _source_of(fn) -> str:
    return textwrap.dedent(inspect.getsource(fn)).strip()


TOOLS_MARKDOWN = "\n\n".join(
    f"### `{schema['function']['name']}`\n\n"
    f"**Schema** (what the model sees):\n\n"
    f"```json\n{json.dumps(schema, indent=2)}\n```\n\n"
    f"**Python implementation** (what actually runs):\n\n"
    f"```python\n{_source_of(TOOL_FUNCTIONS[schema['function']['name']])}\n```"
    for schema in TOOL_SCHEMAS
)


HOW_IT_WORKS = f"""
## The three pieces of an agent

| Piece | What it does |
|:--|:--|
| **LLM** | Decides what to do next |
| **Tools** | Python functions the LLM can call |
| **Loop** | Keep asking the LLM until it's done |

## Pseudocode

```python
while True:
    response = llm.chat(messages, tools=schemas)
    if not response.tool_calls:
        return response.text          # done
    for call in response.tool_calls:
        result = TOOLS[call.name](**call.args)
        messages.append(result)       # feed back, loop again
```

The LLM never runs code. It just *names* a tool and *your* code runs it.
"""


TOOLS_DETAIL = f"""
## System prompt

```text
{SYSTEM_PROMPT}
```

## Tools the agent has

{TOOLS_MARKDOWN}
"""


# Hand-rolled inline SVG flowchart — actual arrows, no runtime deps.
FLOWCHART_HTML = """
<div style="padding: .5rem 0 1rem 0;">
  <h2>Without tools vs. with tools</h2>
  <p style="color:var(--muted);max-width:680px;">
    Same model on both sides. The only difference is the loop on the right —
    it can ask for a tool, run it in your code, feed the result back, and try again.
  </p>

  <svg viewBox="0 0 900 620" xmlns="http://www.w3.org/2000/svg"
       style="width:100%;max-width:900px;height:auto;display:block;margin-top:.75rem;font-family:'Work Sans',system-ui,sans-serif;">
    <defs>
      <marker id="arr"  viewBox="0 0 10 10" refX="9" refY="5" markerWidth="7" markerHeight="7" orient="auto-start-reverse">
        <path d="M0,0 L10,5 L0,10 Z" fill="#5a6070"/>
      </marker>
      <marker id="arrr" viewBox="0 0 10 10" refX="9" refY="5" markerWidth="7" markerHeight="7" orient="auto-start-reverse">
        <path d="M0,0 L10,5 L0,10 Z" fill="#c44536"/>
      </marker>
    </defs>

    <!-- Column backgrounds -->
    <rect x="20"  y="20" width="320" height="580" rx="6" fill="#fff" stroke="#e5e5e5"/>
    <rect x="20"  y="20" width="320" height="3"   fill="#9a6b6b"/>
    <rect x="380" y="20" width="500" height="580" rx="6" fill="#fff" stroke="#e5e5e5"/>
    <rect x="380" y="20" width="500" height="3"   fill="#c44536"/>

    <!-- Titles -->
    <text x="40"  y="55" font-family="'Playfair Display',serif" font-size="18" font-weight="600" fill="#1a1a2e">Without tools</text>
    <text x="400" y="55" font-family="'Playfair Display',serif" font-size="18" font-weight="600" fill="#1a1a2e">With tools (agent loop)</text>

    <!-- ============ LEFT COLUMN ============ -->
    <g font-size="13" fill="#1a1a2e">
      <rect x="100" y="90"  width="160" height="40" rx="4" fill="#faf7f2" stroke="#ddd"/>
      <text x="180" y="115" text-anchor="middle">User question</text>
      <path d="M180 130 L180 188" stroke="#5a6070" stroke-width="1.5" fill="none" marker-end="url(#arr)"/>

      <rect x="100" y="190" width="160" height="40" rx="4" fill="#faf7f2" stroke="#ddd"/>
      <text x="180" y="215" text-anchor="middle">LLM</text>
      <path d="M180 230 L180 288" stroke="#5a6070" stroke-width="1.5" fill="none" marker-end="url(#arr)"/>

      <rect x="80"  y="290" width="200" height="56" rx="4" fill="#fff" stroke="#9a6b6b" stroke-width="1.5"/>
      <text x="180" y="313" text-anchor="middle">Answer</text>
      <text x="180" y="330" text-anchor="middle" font-size="10" fill="#5a6070">(may hallucinate)</text>
    </g>

    <!-- ============ RIGHT COLUMN (x=380..880) ============ -->
    <g font-size="13" fill="#1a1a2e">
      <rect x="550" y="90"  width="160" height="40" rx="4" fill="#faf7f2" stroke="#ddd"/>
      <text x="630" y="115" text-anchor="middle">User question</text>
      <path d="M630 130 L630 178" stroke="#5a6070" stroke-width="1.5" fill="none" marker-end="url(#arr)"/>

      <rect x="535" y="180" width="190" height="42" rx="4" fill="#fff" stroke="#c44536" stroke-width="1.5"/>
      <text x="630" y="206" text-anchor="middle" fill="#c44536" font-weight="500">LLM + tool menu</text>
      <path d="M630 222 L630 268" stroke="#5a6070" stroke-width="1.5" fill="none" marker-end="url(#arr)"/>

      <!-- Diamond decision -->
      <polygon points="630,270 720,315 630,360 540,315" fill="#fff" stroke="#c44536" stroke-width="1.5"/>
      <text x="630" y="311" text-anchor="middle" font-size="12" fill="#c44536" font-weight="500">Tool call</text>
      <text x="630" y="327" text-anchor="middle" font-size="12" fill="#c44536" font-weight="500">needed?</text>

      <!-- NO branch: left + down. Final-answer box must stay inside the column (x >= 380). -->
      <path d="M540 315 L480 315 L480 405" stroke="#5a6070" stroke-width="1.5" fill="none" marker-end="url(#arr)"/>
      <text x="495" y="308" font-size="11" fill="#5a6070" font-weight="600">NO</text>
      <rect x="410" y="407" width="140" height="56" rx="4" fill="#fff" stroke="#c44536"/>
      <text x="480" y="430" text-anchor="middle">Final answer</text>
      <text x="480" y="447" text-anchor="middle" font-size="10" fill="#5a6070">(grounded in data)</text>

      <!-- YES branch: right + down -->
      <path d="M720 315 L780 315 L780 397" stroke="#5a6070" stroke-width="1.5" fill="none" marker-end="url(#arr)"/>
      <text x="737" y="308" font-size="11" fill="#5a6070" font-weight="600">YES</text>

      <rect x="710" y="400" width="140" height="38" rx="4" fill="#faf7f2" stroke="#ddd"/>
      <text x="780" y="423" text-anchor="middle">Run tool</text>
      <path d="M780 438 L780 468" stroke="#5a6070" stroke-width="1.5" fill="none" marker-end="url(#arr)"/>

      <rect x="710" y="470" width="140" height="38" rx="4" fill="#faf7f2" stroke="#ddd"/>
      <text x="780" y="493" text-anchor="middle">Append result</text>

      <!-- Loop back: dashed red, routed along the right edge (x=860, inside right column x=880). -->
      <path d="M780 508 L780 560 L860 560 L860 155 L630 155 L630 176"
            stroke="#c44536" stroke-width="1.5" fill="none"
            stroke-dasharray="5 3" marker-end="url(#arrr)"/>
      <text x="790" y="553" font-size="11" fill="#c44536" font-style="italic">loop back to LLM</text>
    </g>
  </svg>
</div>
"""


CUSTOM_CSS = """
@import url('https://fonts.googleapis.com/css2?family=Playfair+Display:wght@500;600;700&family=Work+Sans:wght@300;400;500;600&display=swap');

:root {
  --ink: #1a1a2e;
  --muted: #5a6070;
  --paper: #ffffff;
  --cream: #faf7f2;
  --accent: #c44536;
  --accent-hover: #a63c30;
  --line: rgba(26, 26, 46, 0.1);
  --shadow-sm: 0 2px 8px rgba(26, 26, 46, 0.06);
  --shadow-md: 0 8px 30px rgba(26, 26, 46, 0.12);
}

.gradio-container, .gradio-container * {
  font-family: "Work Sans", system-ui, sans-serif !important;
  color: var(--ink);
}

.gradio-container { background: var(--paper) !important; }

.gradio-container h1, .gradio-container h2 {
  font-family: "Playfair Display", serif !important;
  font-weight: 600 !important;
  letter-spacing: -0.01em;
}

.gradio-container h1, .gradio-container h2, .gradio-container h3,
.gradio-container p, .gradio-container label {
  text-align: left !important;
}

/* Header card — replaces the purple/pink gradient */
#title-card {
  background: var(--cream);
  border: 1px solid var(--line);
  border-left: 4px solid var(--accent);
  padding: 1.5rem 1.75rem;
  border-radius: 4px;
  margin-bottom: 1.25rem;
  box-shadow: var(--shadow-sm);
}
#title-card h1 {
  font-family: "Playfair Display", serif;
  color: var(--ink);
  margin: 0 0 .35rem 0;
  font-size: 2rem;
  font-weight: 600;
}
#title-card p {
  color: var(--muted);
  margin: 0;
  font-size: 1rem;
  line-height: 1.6;
}
#title-card a { color: var(--accent); text-decoration: none; }
#title-card a:hover { color: var(--accent-hover); text-decoration: underline; }

/* Tabs */
.gradio-container .tab-nav button {
  font-family: "Work Sans", sans-serif !important;
  font-weight: 500 !important;
  color: var(--muted) !important;
  background: transparent !important;
  border: none !important;
  border-bottom: 2px solid transparent !important;
  border-radius: 0 !important;
  padding: .6rem 1.1rem !important;
}
.gradio-container .tab-nav button.selected {
  color: var(--accent) !important;
  border-bottom-color: var(--accent) !important;
}

/* Question textbox */
#question-box textarea {
  font-family: "Work Sans", sans-serif !important;
  font-size: 1.05rem !important;
  border-radius: 4px !important;
  padding: .75rem 1rem !important;
  border: 1px solid var(--line) !important;
  background: var(--paper) !important;
  transition: border-color .15s;
}
#question-box textarea:focus {
  border-color: var(--accent) !important;
  outline: none !important;
  box-shadow: 0 0 0 3px rgba(196, 69, 54, .1) !important;
}

/* Primary button (Run both) */
.gradio-container button.primary {
  background: var(--accent) !important;
  border: none !important;
  border-radius: 4px !important;
  font-weight: 500 !important;
  letter-spacing: .02em !important;
  transition: background .15s;
}
.gradio-container button.primary:hover {
  background: var(--accent-hover) !important;
}

/* Pill-style example buttons */
#pill-row { gap: .5rem !important; flex-wrap: wrap !important; }
#pill-row .gr-button, #pill-row button {
  border-radius: 999px !important;
  padding: .4rem 1rem !important;
  font-family: "Work Sans", sans-serif !important;
  font-size: .88rem !important;
  font-weight: 400 !important;
  background: var(--paper) !important;
  color: var(--ink) !important;
  border: 1px solid var(--line) !important;
  min-width: 0 !important;
  width: auto !important;
  text-align: left !important;
  white-space: normal !important;
  line-height: 1.3 !important;
  flex: 0 0 auto !important;
  transition: all .15s ease;
}
#pill-row .gr-button:hover, #pill-row button:hover {
  background: var(--cream) !important;
  border-color: var(--accent) !important;
  color: var(--accent) !important;
}

/* Result panels */
.panel {
  border-radius: 4px;
  padding: 1.1rem 1.25rem;
  min-height: 220px;
  background: var(--paper);
  box-shadow: var(--shadow-sm);
}
.panel-plain {
  background: var(--cream);
  border: 1px solid var(--line);
  border-top: 3px solid #9a6b6b;
}
.panel-agent {
  background: var(--paper);
  border: 1px solid var(--line);
  border-top: 3px solid var(--accent);
}
.panel h3 {
  font-family: "Playfair Display", serif !important;
  margin-top: 0 !important;
  margin-bottom: .25rem !important;
  color: var(--ink) !important;
}

.final-box {
  background: var(--cream);
  border-left: 3px solid var(--accent);
  padding: .75rem 1rem;
  border-radius: 2px;
  margin-bottom: .75rem;
}

/* Dropdown */
.gradio-container .gr-dropdown,
.gradio-container .gr-dropdown input {
  border-radius: 4px !important;
  border-color: var(--line) !important;
  font-family: "Work Sans", sans-serif !important;
}

/* Code blocks keep mono font */
.gradio-container code, .gradio-container pre,
.gradio-container pre code {
  font-family: "JetBrains Mono", "SF Mono", Menlo, monospace !important;
  font-size: .88rem !important;
}

/* Links */
.gradio-container a {
  color: var(--accent);
  text-decoration: none;
}
.gradio-container a:hover {
  color: var(--accent-hover);
  text-decoration: underline;
}

"""


# --- UI ------------------------------------------------------------------


# --- UI ------------------------------------------------------------------

EXAMPLES = [
    # --- Hard arithmetic (plain LLMs routinely get these wrong) ---
    "What is 4729 times 8314?",
    "Compute (127 ** 3) + (49 ** 4).",
    "What is 31.5% of 128,400?",
    "Evaluate 2847 * 9183 - 17^5.",
    # --- Single-tool: real-time / local data ---
    "What's the temperature in Gandhinagar in Fahrenheit?",
    "What time is it in Tokyo right now?",
    "What's the current price of NVDA?",
    # --- Two-tool chains ---
    "How much hotter is Delhi than Bangalore right now, in degrees Celsius?",
    "If I run 5 km every day for a week, how many miles is that total?",
    "I have 1000 USD. How much is that in INR, and what time is it in Delhi?",
    # --- Three-tool chains (real compositional reasoning) ---
    "What's the distance from Delhi to Mumbai in miles, rounded to the nearest 10?",
    "If Apple stock is $X, how many whole AAPL shares can I buy with 5000 INR? Convert first, then tell me.",
    "Mumbai's population density if its metro area is 603 sq km — and how does that compare to Delhi's metro at 1484 sq km?",
    "I earn 2500 EUR/month. After converting to INR, what is my annual salary?",
    "Convert the distance from London to New York to kilometres, then tell me how long it would take at 900 km/h.",
    # --- Four-tool / open-ended chain ---
    "I want to video-call a colleague in Tokyo at 9 AM their time. What time is that in Delhi, and roughly how far apart are the two cities?",
    "If Paris is 18 C and New York is 22 C, what's the average in Fahrenheit?",
    # --- Course knowledge base ---
    "Which week of CS 203 covered Docker?",
    "What does 'overfitting' mean?",
    # --- No-tool question (sanity check) ---
    "What is the capital of France?",
]

TITLE_CARD = """
<div id="title-card">
  <h1>Agent 101</h1>
  <p>LLM vs. LLM-with-tools — same model on both sides, the only difference is the agent loop on the right. Watch what happens on questions the LLM can't answer from memory.</p>
  <p style="margin-top:.5rem;font-size:.88rem;color:var(--muted);">Built for CS 203 at IIT Gandhinagar · <a href="https://github.com/nipunbatra/stt-ai-teaching/blob/master/lecture-demos/week12/colab-notebooks/01-agents-from-scratch.ipynb">companion Colab</a></p>
</div>
"""


with gr.Blocks(
    title="Agent 101",
    theme=gr.themes.Default(
        primary_hue=gr.themes.colors.red,
        font=[gr.themes.GoogleFont("Work Sans"), "system-ui", "sans-serif"],
        font_mono=[gr.themes.GoogleFont("JetBrains Mono"), "monospace"],
    ),
    css=CUSTOM_CSS,
) as demo:
    gr.HTML(TITLE_CARD)

    with gr.Tabs():
        with gr.Tab("Demo"):
            with gr.Row():
                question = gr.Textbox(
                    label="Ask something",
                    placeholder="e.g. I have 1000 USD. How much is that in INR, and what time is it in Delhi?",
                    lines=2,
                    scale=4,
                    elem_id="question-box",
                )
                model_choice = gr.Dropdown(
                    label="Model",
                    choices=MODELS,
                    value=DEFAULT_MODEL,
                    scale=2,
                )
            go = gr.Button("Run both", variant="primary", size="lg")

            gr.Markdown("**Try an example** — click a pill to auto-run it:")
            with gr.Row(elem_id="pill-row"):
                pill_buttons = [gr.Button(ex, size="sm") for ex in EXAMPLES]

            latex_delims = [
                {"left": "$$", "right": "$$", "display": True},
                {"left": "$",  "right": "$",  "display": False},
                {"left": "\\(", "right": "\\)", "display": False},
                {"left": "\\[", "right": "\\]", "display": True},
            ]

            with gr.Row():
                with gr.Column(elem_classes="panel panel-plain"):
                    gr.Markdown("### 🚫 Without tools\n_Plain LLM — no calculator, no weather, no notes._")
                    out_plain = gr.Markdown(latex_delimiters=latex_delims)
                with gr.Column(elem_classes="panel panel-agent"):
                    gr.Markdown("### 🛠️ With tools (agent loop)\n_Same model, but it can call functions._")
                    out_final = gr.Markdown(elem_classes="final-box", latex_delimiters=latex_delims)
                    with gr.Accordion("Step-by-step trace", open=True):
                        out_trace = gr.Markdown(latex_delimiters=latex_delims)

            go.click(run_both, inputs=[question, model_choice], outputs=[out_plain, out_trace, out_final])
            question.submit(run_both, inputs=[question, model_choice], outputs=[out_plain, out_trace, out_final])

            for btn, ex in zip(pill_buttons, EXAMPLES):
                btn.click(
                    lambda e=ex: e,
                    outputs=question,
                ).then(
                    run_both,
                    inputs=[question, model_choice],
                    outputs=[out_plain, out_trace, out_final],
                )

        with gr.Tab("Flowchart"):
            gr.HTML(FLOWCHART_HTML)

        with gr.Tab("How it works"):
            gr.Markdown(HOW_IT_WORKS)
            with gr.Accordion("System prompt + tool definitions", open=False):
                gr.Markdown(TOOLS_DETAIL)

if __name__ == "__main__":
    demo.queue().launch()