diff --git a/lib/chat_models/chat_vertex_ai.ex b/lib/chat_models/chat_vertex_ai.ex index b3e3147e..d3ae0b1a 100644 --- a/lib/chat_models/chat_vertex_ai.ex +++ b/lib/chat_models/chat_vertex_ai.ex @@ -116,6 +116,14 @@ defmodule LangChain.ChatModels.ChatVertexAI do # Refer to `https://hexdocs.pm/req/Req.html#new/1-options` for # `Req.new` supported set of options. field :req_config, :map, default: %{} + + # The safety settings for the model, specified as a list of maps. Each map + # should contain a `category` and a `threshold` for that category. + # e.g. [%{"category" => "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold" => "BLOCK_ONLY_HIGH"}] + # see https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/configure-safety-filters + # for the list of categories and thresholds. Defaults to `[]` (the + # provider's own defaults apply). + field :safety_settings, {:array, :map}, default: [] end @type t :: %ChatVertexAI{} @@ -132,7 +140,8 @@ defmodule LangChain.ChatModels.ChatVertexAI do :json_response, :json_schema, :stream, - :req_config + :req_config, + :safety_settings ] @required_fields [ :endpoint, @@ -209,6 +218,7 @@ defmodule LangChain.ChatModels.ChatVertexAI do "generationConfig" => generation_config_params } |> Utils.conditionally_add_to_map("system_instruction", for_api(sys_instructions)) + |> Utils.conditionally_add_to_map("safetySettings", vertex_ai.safety_settings) if functions && not Enum.empty?(functions) do req diff --git a/test/chat_models/chat_vertex_ai_test.exs b/test/chat_models/chat_vertex_ai_test.exs index 0ebc3359..49ad79ea 100644 --- a/test/chat_models/chat_vertex_ai_test.exs +++ b/test/chat_models/chat_vertex_ai_test.exs @@ -102,6 +102,29 @@ defmodule ChatModels.ChatVertexAITest do assert %{"temperature" => 1.0, "topK" => 1.0, "topP" => 1.0} = config end + test "adds safety settings to the request if present" do + settings = [ + %{"category" => "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold" => "BLOCK_ONLY_HIGH"} + ] + + vertex_ai = + ChatVertexAI.new!(%{ + model: @test_model, + endpoint: "http://localhost:1234/", + safety_settings: settings + }) + + data = ChatVertexAI.for_api(vertex_ai, [], []) + assert %{"safetySettings" => ^settings} = data + end + + test "does not add safety settings to the request when the list is empty", %{ + vertex_ai: vertex_ai + } do + data = ChatVertexAI.for_api(vertex_ai, [], []) + refute Map.has_key?(data, "safetySettings") + end + test "generate a map containing a text, inline image, and image url parts", %{ vertex_ai: google_ai } do