diff --git a/lib/chat_models/chat_ollama_ai.ex b/lib/chat_models/chat_ollama_ai.ex index 3a178504..b3b308bc 100644 --- a/lib/chat_models/chat_ollama_ai.ex +++ b/lib/chat_models/chat_ollama_ai.ex @@ -37,7 +37,11 @@ defmodule LangChain.ChatModels.ChatOllamaAI do alias LangChain.ChatModels.ChatModel alias LangChain.ChatModels.ChatOpenAI alias LangChain.Message + alias LangChain.Message.ToolCall + alias LangChain.Message.ToolResult alias LangChain.MessageDelta + alias LangChain.Function + alias LangChain.FunctionParam alias LangChain.LangChainError alias LangChain.Utils @@ -185,14 +189,24 @@ defmodule LangChain.ChatModels.ChatOllamaAI do |> validate_number(:mirostat_eta, greater_than_or_equal_to: 0.0, less_than_or_equal_to: 1.0) end + defp messages_for_api(messages) do + Enum.reduce(messages, [], fn m, acc -> + case for_api(m) do + data when is_map(data) -> [data | acc] + data when is_list(data) -> Enum.reverse(data) ++ acc + end + end) + |> Enum.reverse() + end + @doc """ Return the params formatted for an API request. """ - def for_api(%ChatOllamaAI{} = model, messages, _functions) do + def for_api(%ChatOllamaAI{} = model, messages, tools) do %{ model: model.model, temperature: model.temperature, - messages: messages |> Enum.map(&ChatOpenAI.for_api/1), + messages: messages_for_api(messages), stream: model.stream, seed: model.seed, num_ctx: model.num_ctx, @@ -211,15 +225,94 @@ defmodule LangChain.ChatModels.ChatOllamaAI do top_k: model.top_k, top_p: model.top_p } + |> Utils.conditionally_add_to_map(:tools, get_tools_for_api(tools)) + end + + def for_api(%Message{role: :assistant, tool_calls: tool_calls} = msg) + when is_list(tool_calls) do + %{ + "role" => :assistant, + "content" => msg.content + } + |> Utils.conditionally_add_to_map("tool_calls", Enum.map(tool_calls, &for_api(&1))) + end + + # ToolCall support + def for_api(%ToolCall{type: :function} = fun) do + %{ + "id" => fun.call_id, + "type" => "function", + "function" => %{ + "name" => fun.name, + "arguments" => fun.arguments + } + } + end + + # Function support + def for_api(%Function{} = fun) do + %{ + "name" => fun.name, + "parameters" => get_parameters(fun) + } + |> Utils.conditionally_add_to_map("description", fun.description) + end + + def for_api(%Message{role: :tool, tool_results: tool_results}) when is_list(tool_results) do + Enum.map(tool_results, &for_api/1) + end + + def for_api(%ToolResult{content: content}) do + %{ + "role" => :tool, + "content" => content + } + end + + def for_api(%Message{content: content} = msg) when is_binary(content) do + %{ + "role" => msg.role, + "content" => msg.content + } + |> Utils.conditionally_add_to_map("name", msg.name) + end + + def for_api(%Message{role: :user, content: content} = msg) when is_list(content) do + %{ + "role" => msg.role, + "content" => Enum.map(content, &for_api(&1)) + } + |> Utils.conditionally_add_to_map("name", msg.name) + end + + defp get_tools_for_api(nil), do: [] + + defp get_tools_for_api(tools) do + Enum.map(tools, fn %Function{} = function -> + %{"type" => "function", "function" => for_api(function)} + end) + end + + defp get_parameters(%Function{parameters: [], parameters_schema: nil} = _fun) do + %{ + "type" => "object", + "properties" => %{} + } + end + + defp get_parameters(%Function{parameters: [], parameters_schema: schema} = _fun) + when is_map(schema) do + schema + end + + defp get_parameters(%Function{parameters: params} = _fun) do + FunctionParam.to_parameters_schema(params) end @doc """ Calls the Ollama Chat Completion API struct with configuration, plus either a simple message or the list of messages to act as the prompt. - **NOTE:** This API as of right now does not support functions. More - information here: https://github.com/jmorganca/ollama/issues/1729 - **NOTE:** This function *can* be used directly, but the primary interface should be through `LangChain.Chains.LLMChain`. The `ChatOllamaAI` module is more focused on translating the `LangChain` data structures to and from the Ollama API. @@ -232,21 +325,20 @@ defmodule LangChain.ChatModels.ChatOllamaAI do """ @impl ChatModel - def call(ollama_ai, prompt, functions \\ []) + def call(ollama_ai, prompt, tools \\ []) - def call(%ChatOllamaAI{} = ollama_ai, prompt, functions) when is_binary(prompt) do + def call(%ChatOllamaAI{} = ollama_ai, prompt, tools) when is_binary(prompt) do messages = [ Message.new_system!(), Message.new_user!(prompt) ] - call(ollama_ai, messages, functions) + call(ollama_ai, messages, tools) end - def call(%ChatOllamaAI{} = ollama_ai, messages, functions) - when is_list(messages) do + def call(%ChatOllamaAI{} = ollama_ai, messages, tools) when is_list(messages) do try do - case do_api_request(ollama_ai, messages, functions) do + case __MODULE__.do_api_request(ollama_ai, messages, tools) do {:error, reason} -> {:error, reason} @@ -288,13 +380,13 @@ defmodule LangChain.ChatModels.ChatOllamaAI do def do_api_request( %ChatOllamaAI{stream: false} = ollama_ai, messages, - functions, + tools, retry_count ) do req = Req.new( url: ollama_ai.endpoint, - json: for_api(ollama_ai, messages, functions), + json: for_api(ollama_ai, messages, tools), receive_timeout: ollama_ai.receive_timeout, retry: :transient, max_retries: 3, @@ -320,7 +412,7 @@ defmodule LangChain.ChatModels.ChatOllamaAI do {:error, %Req.TransportError{reason: :closed}} -> # Force a retry by making a recursive call decrementing the counter Logger.debug(fn -> "Mint connection closed: retry count = #{inspect(retry_count)}" end) - do_api_request(ollama_ai, messages, functions, retry_count - 1) + do_api_request(ollama_ai, messages, tools, retry_count - 1) other -> Logger.error("Unexpected and unhandled API response! #{inspect(other)}") @@ -331,12 +423,12 @@ defmodule LangChain.ChatModels.ChatOllamaAI do def do_api_request( %ChatOllamaAI{stream: true} = ollama_ai, messages, - functions, + tools, retry_count ) do Req.new( url: ollama_ai.endpoint, - json: for_api(ollama_ai, messages, functions), + json: for_api(ollama_ai, messages, tools), inet6: true, receive_timeout: ollama_ai.receive_timeout ) @@ -362,7 +454,7 @@ defmodule LangChain.ChatModels.ChatOllamaAI do {:error, %Req.TransportError{reason: :closed}} -> # Force a retry by making a recursive call decrementing the counter Logger.debug(fn -> "Mint connection closed: retry count = #{inspect(retry_count)}" end) - do_api_request(ollama_ai, messages, functions, retry_count - 1) + do_api_request(ollama_ai, messages, tools, retry_count - 1) other -> Logger.error( @@ -378,6 +470,18 @@ defmodule LangChain.ChatModels.ChatOllamaAI do create_message(message, :complete, MessageDelta) end + def do_process_response(model, %{ + "message" => %{"tool_calls" => calls} = message, + "done" => true + }) + when calls != [] do + message + |> Map.merge(%{ + "tool_calls" => Enum.map(calls, &do_process_response(model, &1)) + }) + |> create_message(:complete, Message) + end + def do_process_response(_model, %{"message" => message, "done" => true}) do create_message(message, :complete, Message) end @@ -391,6 +495,50 @@ defmodule LangChain.ChatModels.ChatOllamaAI do {:error, LangChainError.exception(message: reason)} end + def do_process_response(_model, %{ + "function" => %{ + "arguments" => args, + "name" => name + } + }) do + case ToolCall.new(%{ + call_id: Ecto.UUID.generate(), + type: :function, + name: name, + arguments: args + }) do + {:ok, %ToolCall{} = call} -> + call + + {:error, changeset} -> + reason = Utils.changeset_error_to_string(changeset) + Logger.error("Failed to process ToolCall for a function. Reason: #{reason}") + {:error, reason} + end + end + + def do_process_response(_model, %{ + "function" => %{ + "arguments" => args, + "name" => name + } + }) do + case ToolCall.new(%{ + call_id: Ecto.UUID.generate(), + type: :function, + name: name, + arguments: args + }) do + {:ok, %ToolCall{} = call} -> + call + + {:error, changeset} -> + reason = Utils.changeset_error_to_string(changeset) + Logger.error("Failed to process ToolCall for a function. Reason: #{reason}") + {:error, reason} + end + end + defp create_message(message, status, message_type) do case message_type.new(Map.merge(message, %{"status" => status})) do {:ok, new_message} -> diff --git a/test/chat_models/chat_ollama_ai_test.exs b/test/chat_models/chat_ollama_ai_test.exs index b1fbd3f1..fd90691c 100644 --- a/test/chat_models/chat_ollama_ai_test.exs +++ b/test/chat_models/chat_ollama_ai_test.exs @@ -2,7 +2,12 @@ defmodule ChatModels.ChatOllamaAITest do use LangChain.BaseCase doctest LangChain.ChatModels.ChatOllamaAI + alias LangChain.ChatModels.ChatOllamaAI + alias LangChain.Function + alias LangChain.FunctionParam + + use Mimic setup do model = ChatOllamaAI.new!(%{"model" => "llama2:latest"}) @@ -124,6 +129,285 @@ defmodule ChatModels.ChatOllamaAITest do assert system_msg["content"] == "You are a weather man" assert user_msg["content"] == "What color is the sky?" end + + test "generates a map for an API call with a tool", %{ollama_ai: ollama_ai} do + fun = + Function.new!(%{ + name: "give_greeting", + description: "Gives a friendly greeting for the given subject", + parameters_schema: %{ + type: "object", + properties: %{ + name: %{ + type: "string", + description: "The subject to greet" + } + }, + required: ["name"] + }, + function: fn %{"name" => name} = _arguments, _context -> {:ok, "Hello, #{name}!"} end + }) + + data = ChatOllamaAI.for_api(ollama_ai, [], [fun]) + + assert [%{} = result_data] = data.tools + + assert %{ + "type" => "function", + "function" => %{ + "description" => "Gives a friendly greeting for the given subject", + "name" => "give_greeting", + "parameters" => %{ + type: "object", + required: ["name"], + properties: %{ + name: %{ + type: "string", + description: "The subject to greet" + } + } + } + } + } = result_data + end + + test "generates a map for an API call with a tool using FunctionParams", %{ + ollama_ai: ollama_ai + } do + fun = + Function.new!(%{ + name: "give_greeting", + description: "Gives a friendly greeting for the given subject", + parameters: [ + FunctionParam.new!(%{name: "name", type: :string, required: true}) + ], + function: fn %{"name" => name} = _arguments, _context -> {:ok, "Hello, #{name}!"} end + }) + + data = ChatOllamaAI.for_api(ollama_ai, [], [fun]) + + assert [%{} = result_data] = data.tools + + assert %{ + "function" => %{ + "description" => "Gives a friendly greeting for the given subject", + "name" => "give_greeting", + "parameters" => %{ + "properties" => %{"name" => %{"type" => "string"}}, + "required" => ["name"], + "type" => "object" + } + }, + "type" => "function" + } = result_data + end + + test "generates a map for an API call with a tool without parameters", %{ollama_ai: ollama_ai} do + fun = + Function.new!(%{ + name: "greet_the_world", + description: "Be friendly to the world", + function: fn _arguments, _context -> {:ok, "Hello, world!"} end + }) + + data = ChatOllamaAI.for_api(ollama_ai, [], [fun]) + + assert [%{} = result_data] = data.tools + + assert %{ + "function" => %{ + "description" => "Be friendly to the world", + "name" => "greet_the_world", + "parameters" => %{"properties" => %{}, "type" => "object"} + }, + "type" => "function" + } = result_data + end + + test "generates a map for an API call without tools", %{ollama_ai: ollama_ai} do + data = ChatOllamaAI.for_api(ollama_ai, [], nil) + + assert data[:tools] == nil + end + + test "for assistant message with non-empty tool_calls, generates an assistant message with a list of ToolCall messages", + %{ollama_ai: ollama_ai} do + tool_call = + Message.ToolCall.new!(%{ + call_id: "call_123", + name: "give_greeting", + arguments: %{"name" => "world"} + }) + + data = + ChatOllamaAI.for_api(ollama_ai, [Message.new_assistant!(%{tool_calls: [tool_call]})], []) + + assert [%{"role" => :assistant} = assistant_msg] = data.messages + assert [%{"id" => "call_123"}] = assistant_msg["tool_calls"] + end + + test "for assistant message empty tool-calls, generates an assistant message", %{ + ollama_ai: ollama_ai + } do + data = ChatOllamaAI.for_api(ollama_ai, [Message.new_assistant!("Hello, world!")], []) + + assert [%{"role" => :assistant} = assistant_msg] = data.messages + assert assistant_msg["content"] == "Hello, world!" + end + + test "for tool call, generate expected structure" do + tool_call = + Message.ToolCall.new!(%{ + call_id: "call_123", + name: "give_greeting", + arguments: %{"name" => "world"} + }) + + expected = %{ + "function" => %{"arguments" => %{"name" => "world"}, "name" => "give_greeting"}, + "id" => "call_123", + "type" => "function" + } + + assert ChatOllamaAI.for_api(tool_call) == expected + end + + test "for function, return expected structure" do + function = + Function.new!(%{ + name: "give_greeting", + description: "Gives a friendly greeting to the given recipient", + parameters: [ + FunctionParam.new!(%{name: "name", type: :string, required: true}) + ], + function: fn %{"name" => name} = _args, _context -> + {:ok, "Hello, #{name}!"} + end + }) + + expected = %{ + "description" => "Gives a friendly greeting to the given recipient", + "name" => "give_greeting", + "parameters" => %{ + "properties" => %{ + "name" => %{"type" => "string"} + }, + "required" => ["name"], + "type" => "object" + } + } + + assert ChatOllamaAI.for_api(function) == expected + end + + test "for message with a list of tool results, generate expected structure" do + tool_result = + Message.ToolResult.new!(%{ + type: :function, + tool_call_id: "call_123", + name: "give_greeting", + content: "Hello, world!", + display_text: nil, + is_error: false + }) + + message = + Message.new_tool_result!(%{ + tool_results: [tool_result] + }) + + expected = [%{"role" => :tool, "content" => "Hello, world!"}] + assert ChatOllamaAI.for_api(message) == expected + end + + test "for user message, generate expected structure" do + message = Message.new_user!("Hello!") + expected = %{"role" => :user, "content" => "Hello!"} + + assert ChatOllamaAI.for_api(message) == expected + end + + test "for nested messages, handle them all", %{ollama_ai: ollama_ai} do + messages = [ + %LangChain.Message{ + content: "Where is the hairbrush located?", + processed_content: nil, + index: nil, + status: :complete, + role: :user, + name: nil, + tool_calls: [], + tool_results: nil + }, + %LangChain.Message{ + content: nil, + processed_content: nil, + index: nil, + status: :complete, + role: :assistant, + name: nil, + tool_calls: [ + %LangChain.Message.ToolCall{ + status: :complete, + type: :function, + call_id: "54836033-8394-4a97-abc5-34c2d4b9fdbf", + name: "custom", + arguments: %{"thing" => "hairbrush"}, + index: nil + } + ], + tool_results: nil + }, + %LangChain.Message{ + content: nil, + processed_content: nil, + index: nil, + status: :complete, + role: :tool, + name: nil, + tool_calls: [], + tool_results: [ + %LangChain.Message.ToolResult{ + type: :function, + tool_call_id: "54836033-8394-4a97-abc5-34c2d4b9fdbf", + name: "custom", + content: "drawer", + display_text: nil, + is_error: false + } + ] + }, + %LangChain.Message{ + content: "The hairbrush is located in the drawer.", + processed_content: nil, + index: nil, + status: :complete, + role: :assistant, + name: nil, + tool_calls: [], + tool_results: nil + } + ] + + expected = [ + %{"content" => "Where is the hairbrush located?", "role" => :user}, + %{ + "content" => nil, + "role" => :assistant, + "tool_calls" => [ + %{ + "function" => %{"arguments" => %{"thing" => "hairbrush"}, "name" => "custom"}, + "id" => "54836033-8394-4a97-abc5-34c2d4b9fdbf", + "type" => "function" + } + ] + }, + %{"content" => "drawer", "role" => :tool}, + %{"content" => "The hairbrush is located in the drawer.", "role" => :assistant} + ] + + assert %{messages: ^expected} = ChatOllamaAI.for_api(ollama_ai, messages, nil) + end end describe "call/2" do @@ -190,6 +474,174 @@ defmodule ChatModels.ChatOllamaAITest do assert reason == "model '#{invalid_model}' not found, try pulling it first" end + + @tag live_call: true, live_ollama_ai: true + test "provided tool is not necessarily used", %{ + models: %{llama31: model}, + tools: %{locator: locator} + } do + {:ok, chat} = ChatOllamaAI.new(model) + {:ok, msg} = Message.new_user("Good morning") + {:ok, %{tool_calls: calls} = _message} = ChatOllamaAI.call(chat, [msg], [locator]) + + assert [] == calls + end + + @tag live_call: true, live_ollama_ai: true + test "provided tool is called (online)", %{ + models: %{llama31: model}, + tools: %{locator: locator} + } do + {:ok, chat} = ChatOllamaAI.new(model) + {:ok, msg} = Message.new_user("Where is the hairbrush located?") + {:ok, %{tool_calls: calls} = _message} = ChatOllamaAI.call(chat, [msg], [locator]) + + assert [%Message.ToolCall{name: "locator", arguments: %{"thing" => "hairbrush"}}] = calls + end + + @tag live_ollama_ai: true + test "provided tool is called", %{models: %{llama31: model}, tools: %{locator: locator}} do + {:ok, chat} = ChatOllamaAI.new(model) + {:ok, msg} = Message.new_user("Where is the hairbrush located?") + + expect(ChatOllamaAI, :do_api_request, fn _model, _msgs, _tools -> + %LangChain.Message{ + content: nil, + processed_content: nil, + index: nil, + status: :complete, + role: :assistant, + name: nil, + tool_calls: [ + %LangChain.Message.ToolCall{ + status: :complete, + type: :function, + call_id: "4806e4e4-b1fd-48a4-b969-b6ae0045bb90", + name: "locator", + arguments: %{"thing" => "hairbrush"}, + index: nil + } + ], + tool_results: nil + } + end) + + {:ok, %{tool_calls: calls} = _message} = ChatOllamaAI.call(chat, [msg], [locator]) + + assert [%Message.ToolCall{name: "locator", arguments: %{"thing" => "hairbrush"}}] = calls + end + + setup do + locator = + Function.new!(%{ + name: "locator", + description: "Returns the location of the requested element or item.", + parameters: [ + FunctionParam.new!(%{ + name: "thing", + type: :string, + description: "the thing whose location is being request" + }) + ], + function: fn %{"thing" => thing} = _arguments, context -> + # our context is a pretend item/location location map + {:ok, context[thing]} + end + }) + + llama31 = %{ + model: "llama3.1:latest", + temperature: 1, + seed: 0, + stream: false + } + + {:ok, %{models: %{llama31: llama31}, tools: %{locator: locator}}} + end + + @tag live_call: true, live_ollama_ai: true + test "provided tool is not necessarily used", %{ + models: %{llama31: model}, + tools: %{locator: locator} + } do + {:ok, chat} = ChatOllamaAI.new(model) + {:ok, msg} = Message.new_user("Good morning") + {:ok, %{tool_calls: calls} = _message} = ChatOllamaAI.call(chat, [msg], [locator]) + + assert [] == calls + end + + @tag live_call: true, live_ollama_ai: true + test "provided tool is called (online)", %{ + models: %{llama31: model}, + tools: %{locator: locator} + } do + {:ok, chat} = ChatOllamaAI.new(model) + {:ok, msg} = Message.new_user("Where is the hairbrush located?") + {:ok, %{tool_calls: calls} = _message} = ChatOllamaAI.call(chat, [msg], [locator]) + + assert [%Message.ToolCall{name: "locator", arguments: %{"thing" => "hairbrush"}}] = calls + end + + @tag live_ollama_ai: true + test "provided tool is called", %{models: %{llama31: model}, tools: %{locator: locator}} do + {:ok, chat} = ChatOllamaAI.new(model) + {:ok, msg} = Message.new_user("Where is the hairbrush located?") + + expect(ChatOllamaAI, :do_api_request, fn _model, _msgs, _tools -> + %LangChain.Message{ + content: nil, + processed_content: nil, + index: nil, + status: :complete, + role: :assistant, + name: nil, + tool_calls: [ + %LangChain.Message.ToolCall{ + status: :complete, + type: :function, + call_id: "4806e4e4-b1fd-48a4-b969-b6ae0045bb90", + name: "locator", + arguments: %{"thing" => "hairbrush"}, + index: nil + } + ], + tool_results: nil + } + end) + + {:ok, %{tool_calls: calls} = _message} = ChatOllamaAI.call(chat, [msg], [locator]) + + assert [%Message.ToolCall{name: "locator", arguments: %{"thing" => "hairbrush"}}] = calls + end + + setup do + locator = + Function.new!(%{ + name: "locator", + description: "Returns the location of the requested element or item.", + parameters: [ + FunctionParam.new!(%{ + name: "thing", + type: :string, + description: "the thing whose location is being request" + }) + ], + function: fn %{"thing" => thing} = _arguments, context -> + # our context is a pretend item/location location map + {:ok, context[thing]} + end + }) + + llama31 = %{ + model: "llama3.1:latest", + temperature: 1, + seed: 0, + stream: false + } + + {:ok, %{models: %{llama31: llama31}, tools: %{locator: locator}}} + end end describe "do_process_response/1" do @@ -232,6 +684,90 @@ defmodule ChatModels.ChatOllamaAITest do assert struct.content == "Gre" assert struct.status == :incomplete end + + test "handles receiving a tool call request response", %{model: model} do + response = %{ + "created_at" => "2024-08-05T09:13:24.222066Z", + "done" => true, + "done_reason" => "stop", + "eval_count" => 17, + "eval_duration" => 303_049_000, + "load_duration" => 12_754_875, + "message" => %{ + "content" => "", + "role" => "assistant", + "tool_calls" => [ + %{ + "function" => %{ + "arguments" => %{"thing" => "hairbrush"}, + "name" => "custom" + } + } + ] + }, + "model" => "llama3.1", + "prompt_eval_count" => 160, + "prompt_eval_duration" => 441_402_000, + "total_duration" => 757_930_875 + } + + assert %Message{} = msg = ChatOllamaAI.do_process_response(model, response) + assert msg.role == :assistant + assert msg.content == nil + assert msg.index == nil + + assert [ + %LangChain.Message.ToolCall{ + status: :complete, + type: :function, + name: "custom", + arguments: %{"thing" => "hairbrush"}, + index: nil + } + ] = msg.tool_calls + end + + test "handles receiving a tool call request response", %{model: model} do + response = %{ + "created_at" => "2024-08-05T09:13:24.222066Z", + "done" => true, + "done_reason" => "stop", + "eval_count" => 17, + "eval_duration" => 303_049_000, + "load_duration" => 12_754_875, + "message" => %{ + "content" => "", + "role" => "assistant", + "tool_calls" => [ + %{ + "function" => %{ + "arguments" => %{"thing" => "hairbrush"}, + "name" => "custom" + } + } + ] + }, + "model" => "llama3.1", + "prompt_eval_count" => 160, + "prompt_eval_duration" => 441_402_000, + "total_duration" => 757_930_875 + } + + assert %Message{} = msg = ChatOllamaAI.do_process_response(model, response) + assert msg.role == :assistant + assert msg.content == nil + assert msg.index == nil + + assert [ + %LangChain.Message.ToolCall{ + status: :complete, + type: :function, + name: "custom", + arguments: %{"thing" => "hairbrush"}, + index: nil + } + ] = msg.tool_calls + end end describe "serialize_config/2" do diff --git a/test/test_helper.exs b/test/test_helper.exs index 0b50cb7e..152bc8aa 100644 --- a/test/test_helper.exs +++ b/test/test_helper.exs @@ -24,6 +24,7 @@ Mimic.copy(LangChain.ChatModels.ChatOpenAI) Mimic.copy(LangChain.ChatModels.ChatAnthropic) Mimic.copy(LangChain.ChatModels.ChatMistralAI) Mimic.copy(LangChain.ChatModels.ChatBumblebee) +Mimic.copy(LangChain.ChatModels.ChatOllamaAI) Mimic.copy(LangChain.Images.OpenAIImage) ExUnit.configure(capture_log: true, exclude: [live_call: true])