feat: apply chat template from callback (#231)

* Enable applying chat template from callback

* Add typespec and docs

lib/utils/chat_templates.ex

@@ -79,11 +79,16 @@ defmodule LangChain.Utils.ChatTemplates do
 
   Note: The `:zephyr` format supports specific system messages.
 
+  ## Template callback
+
+  It's possible to pass a callback as a template.
+  The function receives the list of messages as first argument and `opts` as second and must return a string.
   """
   alias LangChain.LangChainError
   alias LangChain.Message
 
-  @type chat_format :: :inst | :im_start | :llama_2 | :llama_3 | :zephyr
+  @type template_callback :: ([Message.t()], Keyword.t() -> String.t())
+  @type chat_format :: :inst | :im_start | :llama_2 | :llama_3 | :zephyr | template_callback()
 
   # Option:
   # - `add_generation_prompt`: boolean. Defaults to False.
@@ -338,6 +343,10 @@ defmodule LangChain.Utils.ChatTemplates do
     )
   end
 
+  def apply_chat_template!(messages, template_callback, opts)
+      when is_function(template_callback, 2),
+      do: template_callback.(messages, opts)
+
   # return the desired true/false value. Only set to true when the last message
   # is a user prompt.
   defp default_add_generation_prompt_value(messages) do

test/utils/chat_templates_test.exs

@@ -479,4 +479,30 @@ defmodule LangChain.Utils.ChatTemplatesTest do
       assert result == expected
     end
   end
+
+  describe "apply_chat_template!/3 - with template callback" do
+    test "formats according to template callback" do
+      messages = [
+        Message.new_system!("system_message"),
+        Message.new_user!("user_prompt"),
+        Message.new_assistant!("assistant_response"),
+        Message.new_user!("user_2nd")
+      ]
+
+      format =
+        "<|start_of_template|><%= for message <- @messages do %><%= message.role %>\n<%= message.content %>\n\n<% end %><|end_of_template|>"
+
+      template_callback = fn messages, _opts ->
+        EEx.eval_string(format,
+          assigns: [messages: messages]
+        )
+      end
+
+      expected =
+        "<|start_of_template|>system\nsystem_message\n\nuser\nuser_prompt\n\nassistant\nassistant_response\n\nuser\nuser_2nd\n\n<|end_of_template|>"
+
+      result = ChatTemplates.apply_chat_template!(messages, template_callback)
+      assert result == expected
+    end
+  end
 end