Stream OpenAI and Gemini tokens into :messages mode

AGENTS.md

@@ -39,8 +39,10 @@ LangEx (facade: invoke/3, stream/3, get_state/2, get_state_history/2, update_sta
 │   ├── LLM.Anthropic                 # Claude adapter (streaming SSE)
 │   │   ├── Anthropic.SSE             # SSE state machine (internal)
 │   │   └── Anthropic.Formatter       # Message wire format (internal)
-│   ├── LLM.OpenAI                    # GPT adapter
-│   ├── LLM.Gemini                    # Gemini adapter
+│   ├── LLM.OpenAI                    # GPT adapter (streaming SSE)
+│   │   └── OpenAI.SSE                # SSE state machine (internal)
+│   ├── LLM.Gemini                    # Gemini adapter (streaming SSE)
+│   │   └── Gemini.SSE                # SSE state machine (internal)
 │   ├── LLM.Resilient                 # Retry wrapper with backoff
 │   ├── LLM.ChatModel                 # Graph node helper for LLM calls
 │   └── LLM.Registry                  # Provider resolution by model string
@@ -143,7 +145,9 @@ lib/lang_ex/
 │   ├── anthropic/sse.ex             → LangEx.LLM.Anthropic.SSE (@moduledoc false)
 │   ├── anthropic/formatter.ex       → LangEx.LLM.Anthropic.Formatter (@moduledoc false)
 │   ├── openai.ex                    → LangEx.LLM.OpenAI
+│   ├── openai/sse.ex                → LangEx.LLM.OpenAI.SSE (@moduledoc false)
 │   ├── gemini.ex                    → LangEx.LLM.Gemini
+│   ├── gemini/sse.ex                → LangEx.LLM.Gemini.SSE (@moduledoc false)
 │   ├── resilient.ex                 → LangEx.LLM.Resilient
 │   ├── chat_model.ex                → LangEx.LLM.ChatModel
 │   └── chat_models.ex               → LangEx.LLM.Registry

CHANGELOG.md

@@ -2,6 +2,19 @@
 
 ## Unreleased
 
+### Token streaming for OpenAI and Gemini
+
+- `LangEx.LLM.OpenAI` and `LangEx.LLM.Gemini` honor `:on_token` and
+  `:stream`, so `LangEx.stream(..., modes: [:messages])` yields
+  `{:message_delta, ...}` content chunks for every built-in provider.
+  Without those opts the adapters still send a single JSON completion —
+  existing batch tests and `invoke/3` are unchanged.
+- OpenAI requests `stream_options.include_usage` so the final SSE chunk
+  carries token counts. Gemini uses `streamGenerateContent?alt=sse`.
+  Tool-call / function-call payloads are assembled into the final
+  `Message.AI` and are not emitted as content deltas. Gemini thought
+  parts (`thought: true`) are excluded from content.
+
 ### Documentation
 
 - README rewritten as a product tour. Options, edge cases, and behaviours

lib/lang_ex/llm/gemini.ex

@@ -3,6 +3,8 @@ defmodule LangEx.LLM.Gemini do
   Google Gemini chat adapter with function calling support.
 
   Supports Gemini models via the `/v1beta/models/{model}:generateContent` endpoint.
+  Streaming uses `:streamGenerateContent?alt=sse` and the same `x-goog-api-key`
+  header.
 
   ## Tool Calling
 
@@ -17,13 +19,20 @@ defmodule LangEx.LLM.Gemini do
 
   ## Options
 
+  - `:on_token` — `fn(text_delta) -> any()` callback invoked per streamed
+    content token (used by graph streaming's `:messages` mode). Implies SSE
+    streaming. Function-call payloads are assembled, not emitted as tokens.
+  - `:stream` — use SSE streaming (`true` / `false`, default `false`). Also
+    streams when `:on_token` is set. The final return stays
+    `{:ok, %Message.AI{}, usage}` — streaming is how the body arrives.
   - `:tool_choice` — force function calling: `:auto` (default), `:required`/
     `:any` (must call some function), or `{:tool, name}` (must call that one)
   """
 
   @behaviour LangEx.LLM
 
   alias LangEx.Config
+  alias LangEx.LLM.Gemini.SSE
   alias LangEx.Message
   alias LangEx.Tool
 
@@ -44,6 +53,7 @@ defmodule LangEx.LLM.Gemini do
     api_key = Config.api_key!(:gemini, opts)
     model = Config.model(:gemini, opts)
     tools = Keyword.get(opts, :tools, [])
+    stream? = stream_requested?(opts)
 
     {system_instruction, contents} = extract_system(messages)
 
@@ -52,10 +62,19 @@ defmodule LangEx.LLM.Gemini do
     |> put_generation_config(opts)
     |> put_tools(tools)
     |> put_tool_choice(Keyword.get(opts, :tool_choice))
-    |> send_request(api_key, model)
-    |> handle_response()
+    |> send_request(api_key, model, SSE.callbacks(Keyword.get(opts, :on_token)), stream?)
+  end
+
+  defp stream_requested?(opts) do
+    opts
+    |> Keyword.get(:stream, false)
+    |> stream_enabled?(Keyword.get(opts, :on_token))
   end
 
+  defp stream_enabled?(true, _), do: true
+  defp stream_enabled?(_, on_token) when is_function(on_token, 1), do: true
+  defp stream_enabled?(_, _), do: false
+
   defp put_tool_choice(body, nil), do: body
   defp put_tool_choice(body, choice), do: Map.put(body, :tool_config, format_tool_choice(choice))
 
@@ -127,14 +146,86 @@ defmodule LangEx.LLM.Gemini do
   defp to_text(%{"text" => text}), do: {:text, text}
   defp to_text(_), do: nil
 
-  defp send_request(body, api_key, model) do
-    Req.post("#{@base_url}/models/#{model}:generateContent",
+  defp send_request(body, api_key, model, callbacks, stream?) do
+    [
       json: body,
       headers: [
         {"x-goog-api-key", api_key},
         {"content-type", "application/json"}
       ]
-    )
+    ]
+    |> add_stream_timeouts(stream?)
+    |> dispatch_request(callbacks, stream?, endpoint(model, stream?))
+  end
+
+  defp endpoint(model, true),
+    do: "#{@base_url}/models/#{model}:streamGenerateContent?alt=sse"
+
+  defp endpoint(model, false),
+    do: "#{@base_url}/models/#{model}:generateContent"
+
+  defp add_stream_timeouts(opts, true),
+    do: opts |> Keyword.put(:receive_timeout, 300_000) |> Keyword.put(:pool_timeout, 60_000)
+
+  defp add_stream_timeouts(opts, false), do: opts
+
+  defp dispatch_request(req_opts, callbacks, true, url),
+    do: stream_request(req_opts, callbacks, url)
+
+  defp dispatch_request(req_opts, _callbacks, false, url),
+    do: batch_request(req_opts, url)
+
+  defp stream_request(req_opts, callbacks, url) do
+    pkey = {__MODULE__, make_ref()}
+    Process.put(pkey, SSE.initial_state())
+
+    callback = fn {:data, chunk}, {req, resp} ->
+      pkey
+      |> Process.get()
+      |> SSE.process_chunk(callbacks, chunk)
+      |> then(&Process.put(pkey, &1))
+
+      {:cont, {req, resp}}
+    end
+
+    result =
+      req_opts
+      |> Keyword.put(:into, callback)
+      |> then(&Req.post(url, &1))
+      |> handle_streaming_response(pkey, callbacks)
+
+    Process.delete(pkey)
+    result
+  end
+
+  defp handle_streaming_response({:ok, %{status: 200, body: ""}}, pkey, _callbacks),
+    do: SSE.build_message(Process.get(pkey))
+
+  defp handle_streaming_response(
+         {:ok, %{status: 200, body: %Req.Response.Async{}}},
+         pkey,
+         _callbacks
+       ),
+       do: SSE.build_message(Process.get(pkey))
+
+  defp handle_streaming_response({:ok, %{status: 200, body: raw}}, _pkey, callbacks)
+       when is_binary(raw) and byte_size(raw) > 0,
+       do: SSE.parse_response(raw, callbacks)
+
+  defp handle_streaming_response({:ok, %{status: 200, body: response}}, _pkey, _callbacks)
+       when is_map(response),
+       do: handle_response({:ok, %{status: 200, body: response}})
+
+  defp handle_streaming_response({:ok, %{status: status, body: resp_body}}, _pkey, _callbacks),
+    do: {:error, {status, resp_body}}
+
+  defp handle_streaming_response({:error, reason}, _pkey, _callbacks),
+    do: {:error, reason}
+
+  defp batch_request(req_opts, url) do
+    url
+    |> Req.post(req_opts)
+    |> handle_response()
   end
 
   defp extract_system(messages) do

lib/lang_ex/llm/gemini/sse.ex

@@ -0,0 +1,195 @@
+defmodule LangEx.LLM.Gemini.SSE do
+  @moduledoc false
+
+  alias LangEx.Message
+
+  @type callbacks :: %{on_token: (String.t() -> any()) | nil}
+  @type state :: %{
+          text: String.t(),
+          tools: %{optional(non_neg_integer()) => map()},
+          tool_index: non_neg_integer(),
+          usage: map(),
+          line_buffer: String.t()
+        }
+
+  @spec initial_state() :: state()
+  def initial_state do
+    %{text: "", tools: %{}, tool_index: 0, usage: %{}, line_buffer: ""}
+  end
+
+  @spec callbacks((String.t() -> any()) | nil) :: callbacks()
+  def callbacks(on_token), do: %{on_token: on_token}
+
+  @spec process_chunk(state(), callbacks(), String.t()) :: state()
+  def process_chunk(state, callbacks, chunk) do
+    {lines, remainder} =
+      (state.line_buffer <> chunk)
+      |> split_buffer()
+
+    Enum.reduce(lines, %{state | line_buffer: remainder}, &reduce_line(&1, &2, callbacks))
+  end
+
+  @spec parse_response(String.t(), callbacks()) :: {:ok, Message.AI.t(), map()}
+  def parse_response(raw, callbacks) do
+    raw
+    |> String.split("\n")
+    |> Enum.reduce(initial_state(), &reduce_line(&1, &2, callbacks))
+    |> build_message()
+  end
+
+  @spec build_message(state()) :: {:ok, Message.AI.t(), map()}
+  def build_message(state) do
+    {:ok, Message.ai(presence(state.text), tool_calls: tool_calls(state.tools)),
+     extract_usage(state.usage)}
+  end
+
+  defp split_buffer(buffer) do
+    buffer
+    |> String.split("\n")
+    |> split_lines()
+  end
+
+  defp split_lines([single]), do: {[], single}
+  defp split_lines(parts), do: {Enum.slice(parts, 0..-2//1), List.last(parts)}
+
+  defp reduce_line("data: " <> payload, acc, callbacks) do
+    payload
+    |> String.trim()
+    |> apply_payload(acc, callbacks)
+  end
+
+  defp reduce_line(_line, acc, _callbacks), do: acc
+
+  defp apply_payload(json_str, acc, callbacks) do
+    json_str
+    |> Jason.decode()
+    |> apply_event(acc, callbacks)
+  end
+
+  defp apply_event({:ok, event}, acc, callbacks) do
+    updated = handle_event(event, acc)
+    emit_tokens(event, callbacks.on_token)
+    updated
+  end
+
+  defp apply_event(_, acc, _callbacks), do: acc
+
+  defp emit_tokens(event, on_token) do
+    event
+    |> parts()
+    |> Enum.each(&emit_part_token(&1, on_token))
+  end
+
+  defp emit_part_token(%{"thought" => true}, _), do: :ok
+
+  defp emit_part_token(%{"text" => text}, on_token)
+       when is_binary(text) and text != "" and is_function(on_token, 1),
+       do: on_token.(text)
+
+  defp emit_part_token(_, _), do: :ok
+
+  defp handle_event(event, state) do
+    event
+    |> parts()
+    |> Enum.reduce(state, &apply_part/2)
+    |> put_usage(event)
+  end
+
+  defp parts(%{"candidates" => [%{"content" => %{"parts" => parts}} | _]}) when is_list(parts),
+    do: parts
+
+  defp parts(_), do: []
+
+  defp apply_part(%{"thought" => true}, state), do: state
+
+  defp apply_part(%{"text" => text}, state) when is_binary(text),
+    do: %{state | text: state.text <> text}
+
+  defp apply_part(%{"functionCall" => call}, state) when is_map(call),
+    do: apply_function_call(call, state)
+
+  defp apply_part(_, state), do: state
+
+  defp apply_function_call(%{"id" => id} = call, state) when is_binary(id) do
+    state.tools
+    |> Enum.find(&match?({_idx, %{id: ^id}}, &1))
+    |> merge_found(call, state)
+  end
+
+  defp apply_function_call(call, state), do: apply_named_call(call, state)
+
+  defp merge_found({idx, _tool}, call, state),
+    do: update_in(state, [:tools, idx], &merge_tool(&1, call))
+
+  defp merge_found(nil, call, state), do: apply_named_call(call, state)
+
+  defp apply_named_call(%{"name" => name} = call, %{tool_index: idx} = state)
+       when is_binary(name) do
+    state.tools
+    |> Map.get(idx - 1)
+    |> merge_named(name, call, state)
+  end
+
+  defp apply_named_call(_call, state), do: state
+
+  defp merge_named(
+         %{name: name, id: last_id},
+         name,
+         %{"id" => id} = call,
+         %{tool_index: idx} = state
+       )
+       when is_binary(last_id) and is_binary(id) and last_id != id,
+       do: append_tool(state, idx, name, call)
+
+  defp merge_named(%{name: name}, name, call, %{tool_index: idx} = state),
+    do: update_in(state, [:tools, idx - 1], &merge_tool(&1, call))
+
+  defp merge_named(_last, name, call, %{tool_index: idx} = state),
+    do: append_tool(state, idx, name, call)
+
+  defp append_tool(state, idx, name, call) do
+    state
+    |> put_in([:tools, idx], new_tool(name, call))
+    |> Map.put(:tool_index, idx + 1)
+  end
+
+  defp new_tool(name, call) do
+    %{name: name, id: call["id"], args: args_map(call)}
+  end
+
+  defp merge_tool(tool, call) do
+    %{
+      tool
+      | id: call["id"] || tool.id,
+        name: call["name"] || tool.name,
+        args: Map.merge(tool.args, args_map(call))
+    }
+  end
+
+  defp args_map(%{"args" => args}) when is_map(args), do: args
+  defp args_map(_), do: %{}
+
+  defp put_usage(state, %{"usageMetadata" => meta}) when is_map(meta),
+    do: %{state | usage: meta}
+
+  defp put_usage(state, _), do: state
+
+  defp tool_calls(tools) do
+    tools
+    |> Enum.sort_by(&elem(&1, 0))
+    |> Enum.map(fn {_idx, tc} ->
+      %Message.ToolCall{name: tc.name, id: tc.id, args: tc.args}
+    end)
+  end
+
+  defp extract_usage(%{"promptTokenCount" => inp, "candidatesTokenCount" => out}),
+    do: %{input_tokens: inp, output_tokens: out}
+
+  defp extract_usage(%{"promptTokenCount" => inp}),
+    do: %{input_tokens: inp, output_tokens: 0}
+
+  defp extract_usage(_), do: %{input_tokens: 0, output_tokens: 0}
+
+  defp presence(""), do: nil
+  defp presence(text), do: text
+end