diff --git a/lib/rag/generation.ex b/lib/rag/generation.ex index 6ddd35b..a5ac7a5 100644 --- a/lib/rag/generation.ex +++ b/lib/rag/generation.ex @@ -21,7 +21,8 @@ defmodule Rag.Generation do response: String.t() | nil, evaluations: %{optional(atom()) => any()}, halted?: boolean(), - errors: list(any()) + errors: list(any()), + ref: any() } @enforce_keys [:query] @@ -34,13 +35,17 @@ defmodule Rag.Generation do response: nil, evaluations: %{}, halted?: false, - errors: [] + errors: [], + ref: nil @doc """ Creates a new generation struct from a query. """ - @spec new(String.t()) :: t() - def new(query) when is_binary(query), do: %Generation{query: query} + @spec new(String.t(), opts :: keyword()) :: t() + def new(query, opts \\ []) when is_binary(query) do + ref = Keyword.get(opts, :ref) + %Generation{query: query, ref: ref} + end @doc """ Puts `query_embedding` in `generation.query_embedding`. diff --git a/test/rag/generation_test.exs b/test/rag/generation_test.exs index 4bbc031..3dcfae4 100644 --- a/test/rag/generation_test.exs +++ b/test/rag/generation_test.exs @@ -3,6 +3,18 @@ defmodule Rag.GenerationTest do alias Rag.Generation + describe "new/1" do + test "builds a new generation" do + assert %Generation{query: "query", ref: nil} = Generation.new("query") + end + end + + describe "new/2" do + test "allows to pass a reference value" do + assert %Generation{query: "query", ref: "foo"} = Generation.new("query", ref: "foo") + end + end + describe "generate_response/2" do test "calls response_fn with a prompt to generate a response" do generation = %Generation{query: "query", prompt: "a prompt"} @@ -20,7 +32,7 @@ defmodule Rag.GenerationTest do end test "emits start, stop, and exception telemetry events" do - generation = %Generation{query: "query", prompt: "a prompt"} + generation = %Generation{query: "query", prompt: "a prompt", ref: "test-reference"} response_fn = fn "a prompt", _opts -> {:ok, "a response"} end ref = @@ -32,8 +44,11 @@ defmodule Rag.GenerationTest do Generation.generate_response(generation, response_fn) - assert_received {[:rag, :generate_response, :start], ^ref, _measurement, _meta} - assert_received {[:rag, :generate_response, :stop], ^ref, _measurement, _meta} + assert_received {[:rag, :generate_response, :start], ^ref, _measurement, + %{generation: %Generation{ref: "test-reference"}}} + + assert_received {[:rag, :generate_response, :stop], ^ref, _measurement, + %{generation: %Generation{ref: "test-reference"}}} crashing_response_fn = fn _prompt, _opts -> raise "boom" end @@ -41,7 +56,8 @@ defmodule Rag.GenerationTest do Generation.generate_response(generation, crashing_response_fn) end - assert_received {[:rag, :generate_response, :exception], ^ref, _measurement, _meta} + assert_received {[:rag, :generate_response, :exception], ^ref, _measurement, + %{generation: %Generation{ref: "test-reference"}}} end test "halts and sets error when response_fn returns error tuple" do