build prompt with builder function

Author: joelpaulkoch

lib/rag/generation.ex

@@ -10,6 +10,7 @@ defmodule Rag.Generation do
   @type response_function :: (String.t(), keyword() -> String.t())
   @type context_builder_function :: (Generation.t(), keyword() -> String.t())
   @type context_sources_builder_function :: (Generation.t(), keyword() -> list(String.t()))
+  @type prompt_builder_function :: (Generation.t(), keyword() -> String.t())
 
   @typedoc """
   Represents a generation, the main datastructure in `rag`.
@@ -184,4 +185,16 @@ defmodule Rag.Generation do
 
     Generation.put_context_sources(generation, context_sources)
   end
+
+  @doc """
+  Passes `generation` and `opts` to `prompt_builder_function` to determine the prompt.
+  Then, puts the prompt in `generation.prompt`.
+  """
+  @spec build_prompt(t(), prompt_builder_function(), keyword()) :: t()
+  def build_prompt(%Generation{} = generation, prompt_builder_function, opts \\ [])
+      when is_function(prompt_builder_function, 2) do
+    prompt = prompt_builder_function.(generation, opts)
+
+    Generation.put_prompt(generation, prompt)
+  end
 end

test/rag/generation_test.exs

@@ -90,4 +90,15 @@ defmodule Rag.GenerationTest do
                Generation.build_context_sources(generation, builder_fn, [])
     end
   end
+
+  describe "build_prompt/3" do
+    test "builds prompt according to builder function" do
+      generation = %Generation{query: "query"}
+
+      builder_fn = fn _generation, _opts -> "this is the prompt" end
+
+      assert %{prompt: "this is the prompt"} =
+               Generation.build_prompt(generation, builder_fn, [])
+    end
+  end
 end