Lab 1 — The machine

"""GatewayMeter, stage 1: how many events per model, every ten seconds.

This is the whole job for module 1 — a source, a parse, a key, a window, a print. It bills nothing
and remembers nothing across a restart, on purpose: checkpointing is **off by default** in Flink,
and module 7 is where we turn it on and earn the right to make a durability claim.

Run it from the JobManager container:

    docker compose exec jobmanager flink run -py /labs/01-the-machine/gatewaymeter/job.py

Then watch http://localhost:8081 and the TaskManager logs (`docker compose logs -f taskmanager`),
which is where `print()` lands — a Flink job's stdout is the cluster's, not your shell's.
"""
from __future__ import annotations

import sys

sys.path.insert(0, "/labs")

from pyflink.common import Time, Types, WatermarkStrategy  # noqa: E402
from pyflink.common.serialization import SimpleStringSchema  # noqa: E402
from pyflink.datastream import StreamExecutionEnvironment  # noqa: E402
from pyflink.datastream.connectors.kafka import (  # noqa: E402
    KafkaOffsetsInitializer, KafkaSource,
)
from pyflink.datastream.window import TumblingProcessingTimeWindows  # noqa: E402

from common.events import parse_event  # noqa: E402

BOOTSTRAP = "kafka:19092"
TOPIC = "gateway-events"
WINDOW_SECONDS = 10


def model_and_one(raw: str):
    """One JSON line -> (model, 1). `parse_event` is the shared envelope; the count is this module's."""
    return parse_event(raw).model, 1


def build_job(env: StreamExecutionEnvironment):
    """Wire the job. Every module in this course adds to this function and nothing else."""
    source = (
        KafkaSource.builder()
        .set_bootstrap_servers(BOOTSTRAP)
        .set_topics(TOPIC)
        .set_group_id("gatewaymeter-m01")
        # Read only what arrives from now on. `earliest()` would replay the whole topic every
        # time you submit — useful in module 7, confusing here.
        .set_starting_offsets(KafkaOffsetsInitializer.latest())
        .set_value_only_deserializer(SimpleStringSchema())
        .build()
    )

    # no_watermarks(): this stage has no notion of event time yet. The records carry a `ts` and the
    # job ignores it — module 4 is where that stops being acceptable.
    events = env.from_source(source, WatermarkStrategy.no_watermarks(), "gateway-events")

    counts = (
        events
        .map(model_and_one, output_type=Types.TUPLE([Types.STRING(), Types.INT()]))
        .key_by(lambda t: t[0], key_type=Types.STRING())
        .window(TumblingProcessingTimeWindows.of(Time.seconds(WINDOW_SECONDS)))
        .reduce(lambda a, b: (a[0], a[1] + b[1]))
    )
    counts.print()
    return counts


def main() -> None:
    env = StreamExecutionEnvironment.get_execution_environment()
    env.set_parallelism(4)
    # A Python UDF does not run in this process. It runs in a Python worker beside each TaskManager,
    # and that worker only imports what the job shipped to it — so `common` has to be declared, not
    # merely importable here. Forget this and the job submits cleanly, then dies on the first record
    # with ModuleNotFoundError deep inside the SDK harness.
    env.add_python_file("/labs/common")
    build_job(env)
    env.execute("GatewayMeter - module 1 - events per model")


if __name__ == "__main__":
    main()

"""The envelope is the one thing every lab depends on, so it gets a test that needs no cluster."""
import json
import pathlib
import sys

sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[2]))

from common.events import EVENT_TYPES, REQUEST_COMPLETED, Event, parse_event  # noqa: E402


def test_round_trip():
    e = Event(REQUEST_COMPLETED, "r-1", "acme", "gpt-5.5", 1789000000000,
              {"status": "ok", "completion_tokens": 207, "latency_ms": 486})
    assert parse_event(e.to_json()) == e


def test_ts_is_an_int_even_when_the_gateway_sends_a_string():
    raw = json.dumps({"event_type": "ModelError", "request_id": "r-2", "tenant_id": "globex",
                      "model": "gpt-5.5", "ts": "1789000000000", "payload": {"code": 429}})
    assert parse_event(raw).ts == 1789000000000


def test_missing_payload_is_empty_not_none():
    raw = json.dumps({"event_type": "RequestStarted", "request_id": "r-3", "tenant_id": "hooli",
                      "model": "gpt-5.5", "ts": 1})
    assert parse_event(raw).payload == {}


def test_producer_emits_only_known_event_types():
    import subprocess
    root = pathlib.Path(__file__).resolve().parents[2]
    out = subprocess.run([sys.executable, str(root / "common/produce.py"), "--rate", "50",
                          "--seconds", "1", "--seed", "1"], capture_output=True, text=True, timeout=60)
    lines = [l for l in out.stdout.splitlines() if l.strip()]
    assert lines, out.stderr
    for line in lines:
        key, _, raw = line.partition("\t")
        ev = parse_event(raw)
        assert ev.event_type in EVENT_TYPES
        assert ev.request_id == key      # the Kafka key is the request id: one request, one partition

"""The event envelope every lab shares.

The gateway emits one JSON object per thing that happens to a request. `ts` is **event time** —
the millisecond the gateway says the thing happened — and it is the only clock the job may bill by.
Kafka's append time is not it, and neither is the moment an operator sees the record.
"""
from __future__ import annotations

import json
from dataclasses import dataclass, field
from typing import Any, Dict

# The six things the gateway reports. Frozen: a lab that needs a new fact adds a payload field,
# never a seventh type.
REQUEST_STARTED = "RequestStarted"          # payload: prompt_tokens
TOKEN_GENERATED = "TokenGenerated"          # payload: n
TOOL_CALLED = "ToolCalled"                  # payload: tool, latency_ms
AGENT_STEP_COMPLETED = "AgentStepCompleted"  # payload: step, tokens_in, tokens_out
MODEL_ERROR = "ModelError"                  # payload: code
REQUEST_COMPLETED = "RequestCompleted"      # payload: status, completion_tokens, latency_ms

EVENT_TYPES = (
    REQUEST_STARTED, TOKEN_GENERATED, TOOL_CALLED,
    AGENT_STEP_COMPLETED, MODEL_ERROR, REQUEST_COMPLETED,
)


@dataclass
class Event:
    event_type: str
    request_id: str
    tenant_id: str
    model: str
    ts: int                                   # event time, ms since epoch
    payload: Dict[str, Any] = field(default_factory=dict)

    def to_json(self) -> str:
        return json.dumps({
            "event_type": self.event_type, "request_id": self.request_id,
            "tenant_id": self.tenant_id, "model": self.model, "ts": self.ts,
            "payload": self.payload,
        })


def parse_event(raw: str) -> Event:
    """JSON line -> Event. A plain map; no state, no side effects, no clock.

    Malformed records raise. Module 2 is where the job decides what to do about that; until then
    a bad record should stop you, not vanish.
    """
    d = json.loads(raw)
    return Event(
        event_type=d["event_type"], request_id=d["request_id"], tenant_id=d["tenant_id"],
        model=d["model"], ts=int(d["ts"]), payload=d.get("payload") or {},
    )