Add possibility to add a reference for telemetry

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`.

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,16 +44,20 @@ 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
 
       assert_raise RuntimeError, fn ->
         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