Streaming and Prompt Caching: Time to First Token, SSE and Cheaper Prompts

"""Streaming and prompt caching for Pennant Outdoor's support assistant."""
import json
import time

import requests

POLICY = open("policy.md", encoding="utf-8").read()


# ---------- Step 1: stream and time it ----------
def stream_chat(client, messages, max_tokens=120, model="sim"):
    """Stream one completion. Returns {"text", "ttft" (seconds to the first text), "total" (seconds to the end),
    "chunks" (pieces of text received), "usage" (the final usage as a dict, or None)}."""
    # TODO (Step 1): create the completion with stream=True and stream_options={"include_usage": True}.
    # For each chunk: if chunk.choices and chunk.choices[0].delta.content, record the time of the first one (ttft)
    # and collect the text; if chunk.usage, keep chunk.usage.model_dump(). Time everything from before the call.
    raise NotImplementedError("Step 1: write stream_chat()")


def is_buffered(result):
    """True when something between you and the model held the stream back: the first text arrived at or after 80 %
    of the total time."""
    # TODO (Step 1): True when result["ttft"] is at least 80 % of result["total"]
    raise NotImplementedError("Step 1: write is_buffered()")


# ---------- Step 2: read Server-Sent Events yourself ----------
def parse_sse(lines):
    """Yield the JSON payload of every "data:" line, in order. Skip blank lines and comment lines (starting with
    ":"), accept "data:" with or without a space after it, and stop at "data: [DONE]"."""
    raise NotImplementedError("parse_sse() arrives in Step 2")


def raw_stream_text(base_url, messages, max_tokens=60):
    """POST the request with "stream": true using requests (no SDK), read the response line by line with
    iter_lines(decode_unicode=True), and return the concatenated delta text."""
    raise NotImplementedError("raw_stream_text() arrives in Step 2")


# ---------- Step 3: prompts that hit the cache ----------
INSTRUCTIONS = ("You are Pennant Outdoor's support assistant. Answer from the policy below, quote the return window "
                "or warranty that applies, and keep answers under 120 words.")


def naive_messages(customer, question, now):
    """How the prototype built its prompt: greeting and time first, then everything else."""
    return [{"role": "system", "content": f"Current time: {now}. You are talking to {customer['name']} "
                                          f"(member: {customer['member']}).\n{INSTRUCTIONS}\n\n{POLICY}"},
            {"role": "user", "content": question}]


def build_messages(customer, question, now):
    """The same information, ordered so that every request starts with the same tokens: the system message holds
    only INSTRUCTIONS and POLICY; the time and the customer's name and membership go in the user message, before
    the question."""
    return naive_messages(customer, question, now)  # Step 3 reorders this for the cache


# ---------- Step 4: stop when the user does ----------
def stream_until(client, messages, should_stop, max_tokens=400, model="sim"):
    """Stream the answer, and as soon as should_stop(text_so_far) is True, close the stream so the provider stops
    generating (and billing). Returns the text received."""
    raise NotImplementedError("stream_until() arrives in Step 4")


# ---------- Step 5: relay the stream to your own users ----------
def relay_events(client, question):
    """Generator for the /ask endpoint in relay.py: stream the answer to `question` (build_messages with a guest
    customer) and yield each piece of text as an SSE event, 'data: {"text": "..."}\\n\\n', the moment it arrives.
    End with 'data: [DONE]\\n\\n'."""
    raise NotImplementedError("relay_events() arrives in Step 5")