bikoret/vav_pipeline/lib/vav_pipeline/gate/stream.ex

defmodule VavPipeline.Gate.Stream do
  @moduledoc """
  The gate for a streaming client (Construction §3.10, "streaming").

  The model's events are relayed to the caller as they arrive — bytes
  unchanged — while the text is gathered. When the model's stream ends,
  the gathered text is examined like any reply, and before the terminal
  events are released the fringe is sent as one more text block:
  Anthropic `content_block_start` / `content_block_delta` /
  `content_block_stop`; OpenAI one more chunk with `delta.content`. The
  A forged fringe in a streamed reply (Tablet 53) has already reached
  the client by the time it is read; it cannot be stripped, so the real
  fringe names it (form forged_fringe) and the last mark is the gate's.
  The `bikoret` object rides on the terminal event the client already parses
  (`message_delta`; the chunk with `finish_reason`), where an unknown
  field is ignored by every SDK.

  Because the status line and headers are sent before the model has
  spoken, a streaming reply carries no `X-Bikoret-*` headers and cannot
  be rebuked (nothing streamed can be taken back); `X-Bikoret-Rebuke`
  with `stream: true` is refused at the door. If the examination itself
  fails after the stream has begun, the trailer says so — the reply is
  delivered, marked unexamined, never silently.
  """

  alias VavPipeline.{Anchors, Bedikah, Gate}
  alias VavPipeline.Gate.{Shape, Upstream}

  @idle_ms 120_000
  @max_text 65_536

  @type callbacks :: %{
          start: (list() -> term()),
          chunk: (term(), iodata() -> term())
        }

  @doc """
  Stream `body` to `upstream` under `key`. `callbacks.start` is called
  once with the upstream's response headers when the model answers 200
  and must return an emitter; `callbacks.chunk` is called with the
  emitter and each piece of SSE text to deliver. Returns
  `{:ok, %{status: 200, bedikah: result | nil, ...}}` after the stream,
  `{:ok, %{status: s, reply: map}}` for a non-200 (nothing started), or
  `{:error, reason}` / `{:waiting, why}`.
  """
  def pass(upstream, body, key, callbacks, opts \\ []) do
    extra = Keyword.get(opts, :extra_headers, [])
    examine_opts = Keyword.drop(opts, [:extra_headers])
    body = Shape.unfringed(upstream, body)

    with :ok <- Gate.gather_call() do
      case Upstream.open_stream(upstream, body, key, extra) do
        {:ok, ref} ->
          state = %{
            upstream: upstream,
            body: body,
            ref: ref,
            buffer: "",
            text: [],
            text_len: 0,
            held: [],
            max_index: -1,
            model: nil,
            emitter: nil,
            callbacks: callbacks,
            examine_opts: examine_opts
          }

          loop(state)

        {:error, reason} ->
          {:error, :upstream_unavailable, reason}
      end
    end
  end

  # ------------------------------------------------------------
  # The relay
  # ------------------------------------------------------------

  defp loop(%{ref: ref} = state) do
    receive do
      {:http, {^ref, :stream_start, headers}} ->
        loop(%{state | emitter: state.callbacks.start.(headers)})

      {:http, {^ref, :stream, part}} ->
        loop(absorb(state, part))

      {:http, {^ref, :stream_end, _headers}} ->
        finish(state)

      {:http, {^ref, {{_, status, _}, _headers, reply}}} ->
        # A non-200 arrives whole; nothing has been sent to the caller yet.
        case Jason.decode(reply) do
          {:ok, decoded} when is_map(decoded) -> {:ok, %{status: status, reply: decoded, bedikah: nil, bikoret: nil}}
          _ -> {:error, :upstream_unavailable, "upstream #{state.upstream} answered #{status} with a body that is not a JSON object"}
        end

      {:http, {^ref, {:error, reason}}} ->
        if state.emitter, do: finish(state, "upstream stream failed: #{inspect(reason)}"), else: {:error, :upstream_unavailable, "upstream #{state.upstream} unreachable: #{inspect(reason)}"}
    after
      @idle_ms ->
        :httpc.cancel_request(ref)
        if state.emitter, do: finish(state, "upstream stream idle for #{div(@idle_ms, 1000)} s"), else: {:error, :upstream_unavailable, "upstream #{state.upstream} idle for #{div(@idle_ms, 1000)} s"}
    end
  end

  # Split the buffer into complete SSE events; relay the ordinary ones, hold the terminal ones.
  defp absorb(state, part) do
    {events, rest} = split_events(state.buffer <> part)
    state = %{state | buffer: rest}
    Enum.reduce(events, state, &handle_event(&2, &1))
  end

  defp split_events(buffer) do
    parts = String.split(buffer, ~r/\r?\n\r?\n/)
    {complete, [rest]} = Enum.split(parts, length(parts) - 1)
    {Enum.reject(complete, &(&1 == "")), rest}
  end

  defp handle_event(state, raw) do
    {event, data} = parse_event(raw)

    cond do
      terminal?(state.upstream, event, data) ->
        %{state | held: state.held ++ [raw]}

      true ->
        state = gather(state, event, data)
        emit(state, raw <> "\n\n")
    end
  end

  defp parse_event(raw) do
    lines = String.split(raw, ~r/\r?\n/)
    event = lines |> Enum.find_value(fn "event:" <> e -> String.trim(e); _ -> nil end)

    data =
      lines
      |> Enum.flat_map(fn "data:" <> d -> [String.trim_leading(d)]; _ -> [] end)
      |> Enum.join("\n")

    {event, data}
  end

  # Anthropic ends with message_delta then message_stop; OpenAI with the finish_reason chunk then [DONE].
  defp terminal?(:anthropic, event, _data), do: event in ["message_delta", "message_stop"]
  defp terminal?(:openai, _event, "[DONE]"), do: true

  defp terminal?(:openai, _event, data) do
    case Jason.decode(data) do
      {:ok, %{"choices" => [%{"finish_reason" => r} | _]}} when not is_nil(r) -> true
      _ -> false
    end
  end

  defp gather(state, "message_start", data) do
    case Jason.decode(data) do
      {:ok, %{"message" => %{"model" => m}}} -> %{state | model: m}
      _ -> state
    end
  end

  defp gather(%{upstream: :anthropic} = state, "content_block_start", data) do
    case Jason.decode(data) do
      {:ok, %{"index" => i}} -> %{state | max_index: max(state.max_index, i)}
      _ -> state
    end
  end

  defp gather(%{upstream: :anthropic} = state, "content_block_delta", data) do
    case Jason.decode(data) do
      {:ok, %{"delta" => %{"type" => "text_delta", "text" => t}}} -> add_text(state, t)
      _ -> state
    end
  end

  defp gather(%{upstream: :openai} = state, _event, data) do
    case Jason.decode(data) do
      {:ok, %{"choices" => [%{"delta" => %{"content" => t}} | _]} = chunk} when is_binary(t) ->
        state = if state.model, do: state, else: %{state | model: chunk["model"]}
        add_text(state, t)

      {:ok, %{"model" => m}} when is_binary(m) and is_nil(state.model) ->
        %{state | model: m}

      _ ->
        state
    end
  end

  defp gather(state, _event, _data), do: state

  defp add_text(state, t) when state.text_len + byte_size(t) > @max_text, do: state
  defp add_text(state, t), do: %{state | text: [state.text | [t]], text_len: state.text_len + byte_size(t)}

  defp emit(%{emitter: nil} = state, _data), do: state
  defp emit(state, data), do: %{state | emitter: state.callbacks.chunk.(state.emitter, data)}

  # ------------------------------------------------------------
  # The end of the stream: examine, fringe, release
  # ------------------------------------------------------------

  defp finish(state, failure \\ nil) do
    text = state.text |> IO.iodata_to_binary()
    examined = if failure || String.trim(text) == "", do: nil, else: examine(state, text)

    {trailer, bikoret} =
      case examined do
        {:ok, result} ->
          anchors = Anchors.for_violations(result.violations)

          {Shape.trailer(result, anchors, nil),
           %{
             bedikah_ref: result.bedikah_ref,
             verdict: Atom.to_string(result.verdict),
             violations: result.violations,
             anchors: anchors,
             fired: Shape.fired(result, anchors),
             sugya: VavPipeline.Sugya.for_verdict(result),
             sources: result.sources,
             model: state.model,
             streamed: true
           }}

        nil when not is_nil(failure) ->
          {"\n\n— bikoret: not examined (#{failure})", %{verdict: "none", reason: failure, streamed: true}}

        nil ->
          # A turn with no text (tool use only) is not a reply to a reader: no block, the object says so.
          {nil, %{verdict: "none", reason: "no text", streamed: true}}

        {:waiting, why} ->
          {"\n\n— bikoret: not examined (manna: #{why})", %{verdict: "none", reason: "manna #{why}", streamed: true}}

        {:error, kind, reason} ->
          {"\n\n— bikoret: not examined (#{kind}: #{reason})", %{verdict: "none", reason: "#{kind}: #{reason}", streamed: true}}
      end

    state = if trailer, do: inject(state, trailer, bikoret), else: state
    state = Enum.reduce(state.held, state, fn raw, st -> emit(st, with_object(st.upstream, raw, bikoret) <> "\n\n") end)

    result = case examined, do: ({:ok, r} -> r; _ -> nil)
    {:ok, %{status: 200, bedikah: result, bikoret: bikoret, emitter: state.emitter}}
  end

  defp examine(state, text) do
    provenance = Shape.provenance(state.upstream)
    query = Shape.query(state.upstream, state.body)
    sources = Shape.sources(state.upstream, state.body)
    Bedikah.examine(text, provenance, query, Keyword.put(state.examine_opts, :sources, sources))
  end

  # The fringe as one more block of the kind the client is already reading.
  defp inject(%{upstream: :anthropic} = state, trailer, _bikoret) do
    i = state.max_index + 1

    state
    |> emit(sse("content_block_start", %{type: "content_block_start", index: i, content_block: %{type: "text", text: ""}}))
    |> emit(sse("content_block_delta", %{type: "content_block_delta", index: i, delta: %{type: "text_delta", text: trailer}}))
    |> emit(sse("content_block_stop", %{type: "content_block_stop", index: i}))
  end

  defp inject(%{upstream: :openai} = state, trailer, _bikoret) do
    chunk = %{id: "bikoret", object: "chat.completion.chunk", model: state.model, choices: [%{index: 0, delta: %{content: trailer}, finish_reason: nil}]}
    emit(state, "data: " <> Jason.encode!(chunk) <> "\n\n")
  end

  defp sse(event, map), do: "event: #{event}\ndata: " <> Jason.encode!(map) <> "\n\n"

  # The object rides on the terminal event: message_delta (Anthropic), the finish chunk (OpenAI).
  defp with_object(:anthropic, raw, bikoret) do
    {event, data} = parse_event(raw)

    if event == "message_delta" do
      case Jason.decode(data) do
        {:ok, map} -> sse("message_delta", Map.put(map, "bikoret", bikoret)) |> String.trim_trailing()
        _ -> raw
      end
    else
      raw
    end
  end

  defp with_object(:openai, raw, bikoret) do
    {_event, data} = parse_event(raw)

    case Jason.decode(data) do
      {:ok, %{"choices" => [%{"finish_reason" => r} | _]} = map} when not is_nil(r) ->
        "data: " <> Jason.encode!(Map.put(map, "bikoret", bikoret))

      _ ->
        raw
    end
  end
end