From d521e6faeba38ece81cce3a586894c9e69c8e60f Mon Sep 17 00:00:00 2001 From: Georgy Shabunin Date: Tue, 25 Aug 2026 14:55:12 +0100 Subject: [PATCH] feat: stream OpenAI and Gemini tokens into :messages mode Graph streaming already forwarded :on_token as {:message_delta, ...}, but only Anthropic honored it. SSE-stream GPT and Gemini the same way so every built-in provider yields content chunks, while batch JSON stays the default when neither :on_token nor :stream is set. Co-authored-by: Cursor --- AGENTS.md | 8 +- CHANGELOG.md | 13 ++ lib/lang_ex/llm/gemini.ex | 101 +++++++++++++- lib/lang_ex/llm/gemini/sse.ex | 195 +++++++++++++++++++++++++++ lib/lang_ex/llm/openai.ex | 105 ++++++++++++++- lib/lang_ex/llm/openai/sse.ex | 162 ++++++++++++++++++++++ test/lang_ex/llm/gemini_sse_test.exs | 129 ++++++++++++++++++ test/lang_ex/llm/gemini_test.exs | 74 ++++++++++ test/lang_ex/llm/openai_sse_test.exs | 101 ++++++++++++++ test/lang_ex/llm/openai_test.exs | 121 +++++++++++++++++ 10 files changed, 997 insertions(+), 12 deletions(-) create mode 100644 lib/lang_ex/llm/gemini/sse.ex create mode 100644 lib/lang_ex/llm/openai/sse.ex create mode 100644 test/lang_ex/llm/gemini_sse_test.exs create mode 100644 test/lang_ex/llm/openai_sse_test.exs diff --git a/AGENTS.md b/AGENTS.md index d4687d6..15813a6 100644 --- a/AGENTS.md +++ b/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 diff --git a/CHANGELOG.md b/CHANGELOG.md index fa70a1d..16324dc 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/lib/lang_ex/llm/gemini.ex b/lib/lang_ex/llm/gemini.ex index 3e475da..4881712 100644 --- a/lib/lang_ex/llm/gemini.ex +++ b/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,6 +19,12 @@ 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) """ @@ -24,6 +32,7 @@ defmodule LangEx.LLM.Gemini do @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 diff --git a/lib/lang_ex/llm/gemini/sse.ex b/lib/lang_ex/llm/gemini/sse.ex new file mode 100644 index 0000000..1070b9d --- /dev/null +++ b/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 diff --git a/lib/lang_ex/llm/openai.ex b/lib/lang_ex/llm/openai.ex index 97fe3c7..2cc3a22 100644 --- a/lib/lang_ex/llm/openai.ex +++ b/lib/lang_ex/llm/openai.ex @@ -15,11 +15,24 @@ defmodule LangEx.LLM.OpenAI do model: "gpt-4o-mini", tools: [%LangEx.Tool{name: "get_weather", ...}] ) + + ## Options + + - `:on_token` — `fn(text_delta) -> any()` callback invoked per streamed + content token (used by graph streaming's `:messages` mode). Implies SSE + streaming. Tool-call argument fragments are assembled, not emitted. + - `: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 tool use: `:auto` (default), `:required`/`:any` + (must call some tool), or `{:tool, name}` (must call that tool) + - `:base_url` — override the API root (OpenRouter and other compatible hosts) """ @behaviour LangEx.LLM alias LangEx.Config + alias LangEx.LLM.OpenAI.SSE alias LangEx.Message alias LangEx.Tool @@ -41,16 +54,32 @@ defmodule LangEx.LLM.OpenAI do model = Config.model(:openai, opts) tools = Keyword.get(opts, :tools, []) base_url = Keyword.get(opts, :base_url, @base_url) + stream? = stream_requested?(opts) %{model: model, messages: Enum.map(messages, &format_message/1)} |> put_present(:temperature, opts[:temperature]) |> put_present(:max_tokens, opts[:max_tokens]) |> put_tools(tools) |> put_tool_choice(Keyword.get(opts, :tool_choice)) - |> send_request(api_key, base_url) - |> handle_response() + |> put_stream(stream?) + |> send_request(api_key, base_url, 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_stream(body, true), + do: body |> Map.put(:stream, true) |> Map.put(:stream_options, %{include_usage: true}) + + defp put_stream(body, false), do: body + defp put_tool_choice(body, nil), do: body defp put_tool_choice(body, choice), do: Map.put(body, :tool_choice, format_tool_choice(choice)) @@ -114,14 +143,80 @@ defmodule LangEx.LLM.OpenAI do defp parse_decoded({:ok, parsed}), do: parsed defp parse_decoded(_), do: %{} - defp send_request(body, api_key, base_url) do - Req.post("#{base_url}/chat/completions", + defp send_request(body, api_key, base_url, callbacks, stream?) do + [ json: body, headers: [ {"authorization", "Bearer #{api_key}"}, {"content-type", "application/json"} ] - ) + ] + |> add_stream_timeouts(stream?) + |> dispatch_request(callbacks, stream?, "#{base_url}/chat/completions") + end + + 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 format_message(%Message.Human{content: c}), do: %{role: "user", content: c} diff --git a/lib/lang_ex/llm/openai/sse.ex b/lib/lang_ex/llm/openai/sse.ex new file mode 100644 index 0000000..7b99da2 --- /dev/null +++ b/lib/lang_ex/llm/openai/sse.ex @@ -0,0 +1,162 @@ +defmodule LangEx.LLM.OpenAI.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()}, + usage: map(), + line_buffer: String.t() + } + + @spec initial_state() :: state() + def initial_state do + %{text: "", tools: %{}, 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("[DONE]", 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_token(event, callbacks.on_token) + updated + end + + defp apply_event(_, acc, _callbacks), do: acc + + defp emit_token(%{"choices" => [%{"delta" => %{"content" => text}} | _]}, on_token) + when is_binary(text) and text != "" and is_function(on_token, 1), + do: on_token.(text) + + defp emit_token(_, _), do: :ok + + defp handle_event(%{"choices" => [%{"delta" => delta} | _]} = event, state) do + state + |> append_text(delta) + |> merge_tool_calls(delta) + |> put_usage(event) + end + + defp handle_event(event, state), do: put_usage(state, event) + + defp append_text(state, %{"content" => text}) when is_binary(text), + do: %{state | text: state.text <> text} + + defp append_text(state, _), do: state + + defp merge_tool_calls(state, %{"tool_calls" => calls}) when is_list(calls), + do: Enum.reduce(calls, state, &merge_tool_call/2) + + defp merge_tool_calls(state, _), do: state + + defp merge_tool_call(%{"index" => idx} = call, state), + do: update_in(state, [:tools, idx], &merge_tool(&1, call)) + + defp merge_tool_call(_, state), do: state + + defp merge_tool(nil, call) do + %{id: call["id"], name: function_field(call, "name"), args: args_fragment(call)} + end + + defp merge_tool(existing, call) do + %{ + id: call["id"] || existing.id, + name: function_field(call, "name") || existing.name, + args: existing.args <> args_fragment(call) + } + end + + defp function_field(%{"function" => function}, key) when is_map(function), do: function[key] + defp function_field(_, _), do: nil + + defp args_fragment(call) do + call + |> function_field("arguments") + |> binary_or_empty() + end + + defp binary_or_empty(args) when is_binary(args), do: args + defp binary_or_empty(_), do: "" + + defp put_usage(state, %{"usage" => usage}) when is_map(usage), do: %{state | usage: usage} + 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: decode_args(tc.args)} + end) + end + + defp decode_args(args) when is_binary(args) do + args + |> Jason.decode() + |> parsed_args() + end + + defp decode_args(args) when is_map(args), do: args + defp decode_args(_), do: %{} + + defp parsed_args({:ok, parsed}), do: parsed + defp parsed_args(_), do: %{} + + defp extract_usage(%{"prompt_tokens" => inp, "completion_tokens" => out}), + do: %{input_tokens: inp, output_tokens: out} + + defp extract_usage(_), do: %{input_tokens: 0, output_tokens: 0} + + defp presence(""), do: nil + defp presence(text), do: text +end diff --git a/test/lang_ex/llm/gemini_sse_test.exs b/test/lang_ex/llm/gemini_sse_test.exs new file mode 100644 index 0000000..982b508 --- /dev/null +++ b/test/lang_ex/llm/gemini_sse_test.exs @@ -0,0 +1,129 @@ +defmodule LangEx.LLM.Gemini.SSETest do + use ExUnit.Case, async: true + + alias LangEx.LLM.Gemini.SSE + alias LangEx.Message + + # streamGenerateContent?alt=sse emits one GenerateContentResponse per data: + # line. Text lives in candidates[0].content.parts[].text; last chunk carries + # usageMetadata.promptTokenCount / candidatesTokenCount. No [DONE] sentinel. + @sse_body """ + data: {"candidates":[{"content":{"role":"model","parts":[{"text":"Hel"}]},"index":0}]} + + data: {"candidates":[{"content":{"role":"model","parts":[{"text":"lo"}]},"finishReason":"STOP","index":0}],"usageMetadata":{"promptTokenCount":5,"candidatesTokenCount":2,"totalTokenCount":7}} + """ + + test "on_token receives each content delta and the message assembles" do + test_pid = self() + callbacks = SSE.callbacks(&send(test_pid, {:token, &1})) + + assert {:ok, %Message.AI{content: "Hello", tool_calls: []}, + %{input_tokens: 5, output_tokens: 2}} = + SSE.parse_response(@sse_body, callbacks) + + assert_received {:token, "Hel"} + assert_received {:token, "lo"} + end + + test "without callbacks the same body parses silently" do + assert {:ok, %Message.AI{content: "Hello"}, %{input_tokens: 5, output_tokens: 2}} = + SSE.parse_response(@sse_body, SSE.callbacks(nil)) + + refute_received {:token, _} + end + + test "thought parts are not content and do not fire on_token" do + body = """ + data: {"candidates":[{"content":{"role":"model","parts":[{"thought":true,"text":"pondering"}]}}]} + + data: {"candidates":[{"content":{"role":"model","parts":[{"text":"Hello"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":5,"candidatesTokenCount":1,"totalTokenCount":6}} + """ + + test_pid = self() + + assert {:ok, %Message.AI{content: "Hello", tool_calls: []}, + %{input_tokens: 5, output_tokens: 1}} = + SSE.parse_response(body, SSE.callbacks(&send(test_pid, {:token, &1}))) + + assert_received {:token, "Hello"} + refute_received {:token, "pondering"} + end + + @tool_body """ + data: {"candidates":[{"content":{"role":"model","parts":[{"functionCall":{"name":"get_weather","args":{"location":"Paris"}}}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":8,"candidatesTokenCount":12,"totalTokenCount":20}} + """ + + test "a whole functionCall assembles one ToolCall and does not fire on_token" do + test_pid = self() + callbacks = SSE.callbacks(&send(test_pid, {:token, &1})) + + assert {:ok, %Message.AI{content: nil, tool_calls: [call]}, + %{input_tokens: 8, output_tokens: 12}} = + SSE.parse_response(@tool_body, callbacks) + + assert %Message.ToolCall{name: "get_weather", id: nil, args: %{"location" => "Paris"}} = call + refute_received {:token, _} + end + + test "functionCall id is kept when the API sends one" do + body = """ + data: {"candidates":[{"content":{"parts":[{"functionCall":{"id":"fc_1","name":"get_weather","args":{"location":"Paris"}}}]}}]} + """ + + assert {:ok, %Message.AI{tool_calls: [call]}, _} = + SSE.parse_response(body, SSE.callbacks(nil)) + + assert %Message.ToolCall{name: "get_weather", id: "fc_1", args: %{"location" => "Paris"}} = + call + end + + test "chunks with the same functionCall id merge args" do + body = """ + data: {"candidates":[{"content":{"parts":[{"functionCall":{"id":"fc_1","name":"get_weather","args":{"location":"Paris"}}}]}}]} + + data: {"candidates":[{"content":{"parts":[{"functionCall":{"id":"fc_1","name":"get_weather","args":{"units":"celsius"}}}]}}]} + """ + + assert {:ok, %Message.AI{tool_calls: [call]}, _} = + SSE.parse_response(body, SSE.callbacks(nil)) + + assert %Message.ToolCall{ + name: "get_weather", + id: "fc_1", + args: %{"location" => "Paris", "units" => "celsius"} + } = call + end + + @split_tool_body """ + data: {"candidates":[{"content":{"role":"model","parts":[{"functionCall":{"name":"get_weather","args":{"location":"Par"}}}]}}]} + + data: {"candidates":[{"content":{"role":"model","parts":[{"functionCall":{"name":"get_weather","args":{"location":"Paris","units":"celsius"}}}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":8,"candidatesTokenCount":12,"totalTokenCount":20}} + """ + + test "the same functionCall name across chunks merges args objects" do + assert {:ok, %Message.AI{tool_calls: [call]}, %{input_tokens: 8, output_tokens: 12}} = + SSE.parse_response(@split_tool_body, SSE.callbacks(nil)) + + assert %Message.ToolCall{ + name: "get_weather", + args: %{"location" => "Paris", "units" => "celsius"} + } = call + end + + test "two functionCall parts in one chunk are two ToolCalls" do + body = """ + data: {"candidates":[{"content":{"parts":[{"functionCall":{"name":"get_weather","args":{"location":"Paris"}}},{"functionCall":{"name":"get_time","args":{"tz":"UTC"}}}]}}]} + """ + + assert {:ok, %Message.AI{tool_calls: [weather, time]}, _} = + SSE.parse_response(body, SSE.callbacks(nil)) + + assert %Message.ToolCall{name: "get_weather", args: %{"location" => "Paris"}} = weather + assert %Message.ToolCall{name: "get_time", args: %{"tz" => "UTC"}} = time + end + + test "usageMetadata on the last chunk becomes token counts" do + assert {:ok, %Message.AI{content: "Hello"}, %{input_tokens: 5, output_tokens: 2}} = + SSE.parse_response(@sse_body, SSE.callbacks(nil)) + end +end diff --git a/test/lang_ex/llm/gemini_test.exs b/test/lang_ex/llm/gemini_test.exs index c04d613..1a12ab8 100644 --- a/test/lang_ex/llm/gemini_test.exs +++ b/test/lang_ex/llm/gemini_test.exs @@ -130,4 +130,78 @@ defmodule LangEx.LLM.GeminiTest do ) end end + + describe "streaming" do + @sse_hello """ + data: {"candidates":[{"content":{"role":"model","parts":[{"text":"Hel"}]},"index":0}]} + + data: {"candidates":[{"content":{"role":"model","parts":[{"text":"lo"}]},"finishReason":"STOP","index":0}],"usageMetadata":{"promptTokenCount":3,"candidatesTokenCount":2,"totalTokenCount":5}} + """ + + test "on_token hits streamGenerateContent with :into set" do + test_pid = self() + + expect(Req, :post, fn url, opts -> + assert url =~ "streamGenerateContent" + assert url =~ "alt=sse" + assert is_function(opts[:into], 2) + refute url =~ ":generateContent?" + refute Map.has_key?(opts[:json], :stream) + assert {"x-goog-api-key", "test"} in opts[:headers] + + {:ok, %{status: 200, body: @sse_hello}} + end) + + assert {:ok, %Message.AI{content: "Hello"}, %{input_tokens: 3, output_tokens: 2}} = + LangEx.LLM.Gemini.chat_with_usage( + [Message.human("hi")], + model: "gemini-2.0-flash", + api_key: "test", + on_token: &send(test_pid, {:token, &1}) + ) + + assert_received {:token, "Hel"} + assert_received {:token, "lo"} + end + + test "stream: true uses the :into callback accumulator" do + expect(Req, :post, fn url, opts -> + assert url =~ "streamGenerateContent" + assert is_function(opts[:into], 2) + {:cont, _} = opts[:into].({:data, @sse_hello}, {nil, nil}) + {:ok, %{status: 200, body: ""}} + end) + + assert {:ok, %Message.AI{content: "Hello"}} = + LangEx.LLM.Gemini.chat( + [Message.human("hi")], + model: "gemini-2.0-flash", + api_key: "test", + stream: true + ) + end + + test "batch still hits generateContent" do + expect(Req, :post, fn url, opts -> + assert url =~ ":generateContent" + refute url =~ "streamGenerateContent" + assert is_nil(opts[:into]) + + {:ok, + %{ + status: 200, + body: %{ + "candidates" => [%{"content" => %{"parts" => [%{"text" => "ok"}]}}] + } + }} + end) + + assert {:ok, %Message.AI{content: "ok"}} = + LangEx.LLM.Gemini.chat( + [Message.human("hi")], + model: "gemini-2.0-flash", + api_key: "test" + ) + end + end end diff --git a/test/lang_ex/llm/openai_sse_test.exs b/test/lang_ex/llm/openai_sse_test.exs new file mode 100644 index 0000000..aafb132 --- /dev/null +++ b/test/lang_ex/llm/openai_sse_test.exs @@ -0,0 +1,101 @@ +defmodule LangEx.LLM.OpenAI.SSETest do + use ExUnit.Case, async: true + + alias LangEx.LLM.OpenAI.SSE + alias LangEx.Message + + # Official CreateChatCompletionStreamResponse shape: first chunk often has + # role + empty content, usage is null until the final choices:[] chunk, then + # data: [DONE]. See OpenAPI ChatCompletionStreamOptions.include_usage. + @sse_body """ + data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1694268190,"model":"gpt-4o-mini","choices":[{"index":0,"delta":{"role":"assistant","content":""},"logprobs":null,"finish_reason":null}],"usage":null} + + data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1694268190,"model":"gpt-4o-mini","choices":[{"index":0,"delta":{"content":"Hel"},"logprobs":null,"finish_reason":null}],"usage":null} + + data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1694268190,"model":"gpt-4o-mini","choices":[{"index":0,"delta":{"content":"lo"},"logprobs":null,"finish_reason":null}],"usage":null} + + data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1694268190,"model":"gpt-4o-mini","choices":[{"index":0,"delta":{},"logprobs":null,"finish_reason":"stop"}],"usage":null} + + data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1694268190,"model":"gpt-4o-mini","choices":[],"usage":{"prompt_tokens":4,"completion_tokens":2,"total_tokens":6}} + + data: [DONE] + """ + + test "on_token receives each content delta and the message assembles" do + test_pid = self() + callbacks = SSE.callbacks(&send(test_pid, {:token, &1})) + + assert {:ok, %Message.AI{content: "Hello", tool_calls: []}, + %{input_tokens: 4, output_tokens: 2}} = + SSE.parse_response(@sse_body, callbacks) + + assert_received {:token, "Hel"} + assert_received {:token, "lo"} + refute_received {:token, ""} + end + + test "without callbacks the same body parses silently" do + assert {:ok, %Message.AI{content: "Hello"}, %{input_tokens: 4, output_tokens: 2}} = + SSE.parse_response(@sse_body, SSE.callbacks(nil)) + + refute_received {:token, _} + end + + test "null content and empty first-chunk content are not tokens" do + body = """ + data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1,"model":"gpt-4o-mini","choices":[{"index":0,"delta":{"role":"assistant","content":null},"finish_reason":null}],"usage":null} + + data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1,"model":"gpt-4o-mini","choices":[{"index":0,"delta":{"content":"Hi"},"finish_reason":null}],"usage":null} + + data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1,"model":"gpt-4o-mini","choices":[],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}} + + data: [DONE] + """ + + test_pid = self() + + assert {:ok, %Message.AI{content: "Hi"}, %{input_tokens: 1, output_tokens: 1}} = + SSE.parse_response(body, SSE.callbacks(&send(test_pid, {:token, &1}))) + + assert_received {:token, "Hi"} + refute_received {:token, _} + end + + @tool_body """ + data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1,"model":"gpt-4o-mini","choices":[{"index":0,"delta":{"role":"assistant","tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"get_weather","arguments":""}}]},"finish_reason":null}],"usage":null} + + data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1,"model":"gpt-4o-mini","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\\"location\\":"}}]},"finish_reason":null}],"usage":null} + + data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1,"model":"gpt-4o-mini","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\\"Paris\\"}"}}]},"finish_reason":null}],"usage":null} + + data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1,"model":"gpt-4o-mini","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}],"usage":null} + + data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1,"model":"gpt-4o-mini","choices":[],"usage":{"prompt_tokens":10,"completion_tokens":8,"total_tokens":18}} + + data: [DONE] + """ + + test "tool-call argument fragments assemble one ToolCall and never fire on_token" do + test_pid = self() + callbacks = SSE.callbacks(&send(test_pid, {:token, &1})) + + assert {:ok, %Message.AI{content: nil, tool_calls: [call]}, + %{input_tokens: 10, output_tokens: 8}} = + SSE.parse_response(@tool_body, callbacks) + + assert %Message.ToolCall{name: "get_weather", id: "call_1", args: %{"location" => "Paris"}} = + call + + refute_received {:token, _} + end + + test "[DONE] is ignored even when it is the only payload" do + assert {:ok, %Message.AI{content: nil, tool_calls: []}, %{input_tokens: 0, output_tokens: 0}} = + SSE.parse_response("data: [DONE]\n", SSE.callbacks(nil)) + end + + test "usage is taken from the final choices:[] chunk" do + assert {:ok, %Message.AI{content: "Hello"}, %{input_tokens: 4, output_tokens: 2}} = + SSE.parse_response(@sse_body, SSE.callbacks(nil)) + end +end diff --git a/test/lang_ex/llm/openai_test.exs b/test/lang_ex/llm/openai_test.exs index 81e3192..0fe8efe 100644 --- a/test/lang_ex/llm/openai_test.exs +++ b/test/lang_ex/llm/openai_test.exs @@ -2,6 +2,8 @@ defmodule LangEx.LLM.OpenAITest do use ExUnit.Case, async: false use Mimic + alias LangEx.Graph + alias LangEx.LLM.ChatModel alias LangEx.Message alias LangEx.Tool @@ -72,4 +74,123 @@ defmodule LangEx.LLM.OpenAITest do ) end end + + describe "streaming" do + @sse_hello """ + data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1,"model":"gpt-4o-mini","choices":[{"index":0,"delta":{"role":"assistant","content":""},"finish_reason":null}],"usage":null} + + data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1,"model":"gpt-4o-mini","choices":[{"index":0,"delta":{"content":"Hel"},"finish_reason":null}],"usage":null} + + data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1,"model":"gpt-4o-mini","choices":[{"index":0,"delta":{"content":"lo"},"finish_reason":null}],"usage":null} + + data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1,"model":"gpt-4o-mini","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":null} + + data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1,"model":"gpt-4o-mini","choices":[],"usage":{"prompt_tokens":3,"completion_tokens":2,"total_tokens":5}} + + data: [DONE] + """ + + test "on_token sends stream: true, include_usage, and :into" do + test_pid = self() + + expect(Req, :post, fn url, opts -> + assert url =~ "/chat/completions" + assert opts[:json].stream == true + assert opts[:json].stream_options == %{include_usage: true} + refute Map.has_key?(opts[:json].stream_options, :include_obfuscation) + assert is_function(opts[:into], 2) + + {:ok, %{status: 200, body: @sse_hello}} + end) + + assert {:ok, %Message.AI{content: "Hello"}, %{input_tokens: 3, output_tokens: 2}} = + LangEx.LLM.OpenAI.chat_with_usage( + [Message.human("hi")], + model: "gpt-4o-mini", + api_key: "test", + on_token: &send(test_pid, {:token, &1}) + ) + + assert_received {:token, "Hel"} + assert_received {:token, "lo"} + end + + test "stream: true uses the :into callback accumulator" do + expect(Req, :post, fn _url, opts -> + assert opts[:json].stream == true + assert is_function(opts[:into], 2) + {:cont, _} = opts[:into].({:data, @sse_hello}, {nil, nil}) + {:ok, %{status: 200, body: ""}} + end) + + assert {:ok, %Message.AI{content: "Hello"}} = + LangEx.LLM.OpenAI.chat( + [Message.human("hi")], + model: "gpt-4o-mini", + api_key: "test", + stream: true + ) + end + + test "without stream opts the request stays batch" do + expect(Req, :post, fn _url, opts -> + refute Map.has_key?(opts[:json], :stream) + refute Map.has_key?(opts[:json], :stream_options) + assert is_nil(opts[:into]) + + {:ok, + %{ + status: 200, + body: %{"choices" => [%{"message" => %{"content" => "ok"}}]} + }} + end) + + assert {:ok, %Message.AI{content: "ok"}} = + LangEx.LLM.OpenAI.chat( + [Message.human("hi")], + model: "gpt-4o-mini", + api_key: "test" + ) + end + + test "streaming honors base_url" do + expect(Req, :post, fn url, opts -> + assert url == "https://openrouter.example/v1/chat/completions" + assert opts[:json].stream == true + {:ok, %{status: 200, body: @sse_hello}} + end) + + assert {:ok, %Message.AI{content: "Hello"}} = + LangEx.LLM.OpenAI.chat( + [Message.human("hi")], + model: "gpt-4o-mini", + api_key: "test", + stream: true, + base_url: "https://openrouter.example/v1" + ) + end + + test "ChatModel.node yields message deltas under :messages stream mode" do + expect(Req, :post, fn _url, opts -> + assert opts[:json].stream == true + assert is_function(opts[:into], 2) + {:ok, %{status: 200, body: @sse_hello}} + end) + + events = + Graph.new(messages: {[], &Message.add_messages/2}) + |> Graph.add_node(:llm, ChatModel.node(model: "gpt-4o", api_key: "test")) + |> Graph.add_edge(:__start__, :llm) + |> Graph.add_edge(:llm, :__end__) + |> Graph.compile() + |> LangEx.stream(%{messages: [Message.human("hi")]}, modes: [:messages]) + |> Enum.to_list() + + assert [ + {:message_delta, %{node: :llm, kind: :content, text: "Hel"}}, + {:message_delta, %{node: :llm, kind: :content, text: "lo"}}, + {:done, {:ok, _}} + ] = events + end + end end