feat: added allowed tools (#5377)

Co-authored-by: Khan Life <[email protected]>

Author: 2 people

.changeset/shiny-points-call.md

@@ -0,0 +1,5 @@
+---
+"@effect/ai": minor
+---
+
+Added "oneOf" property for generateText and streamText "toolChoice" to be able to pick a subset of tools, which are going to be passed to LLM

packages/ai/ai/src/AiLanguageModel.ts

@@ -61,12 +61,15 @@ export interface IdentifiedSchema<A, I, R> extends Schema.Schema<A, I, R> {
  * - `required`: The model **must** call a tool but can decide which tool will be called.
  * - `none`: The model **must not** call a tool.
  * - `{ tool: <tool_name> }`: The model must call the specified tool.
- *
+ * - `{ mode?: "auto" (default) | "required", oneOf: [<tool_names>] }`: The model is restricted to a subset of tools. When `mode` is `"auto"` or omitted, the model can decide whether to call a tool from the allowed subset. When `mode` is `"required"`, the model must call one tool from the allowed subset.
  * @since 1.0.0
  * @category Models
  */
 export type ToolChoice<Tool extends AiTool.Any> = "auto" | "none" | "required" | {
   readonly tool: Tool["name"]
+} | {
+  readonly mode?: "auto" | "required"
+  readonly oneOf: Array<Tool["name"]>
 }
 
 /**
@@ -390,8 +393,12 @@ export const make: (
             return response
           }
           modelOptions.toolChoice = toolChoice
+          const hasUnallowedTools = typeof toolChoice === "object" && "oneOf" in toolChoice
           const actualToolkit = Effect.isEffect(toolkit) ? yield* toolkit : toolkit
           for (const tool of actualToolkit.tools) {
+            if (hasUnallowedTools && !toolChoice.oneOf.includes(tool.name)) {
+              continue
+            }
             modelOptions.tools.push(convertTool(tool))
           }
           const response = yield* opts.generateText(modelOptions)
@@ -426,7 +433,11 @@ export const make: (
       const actualToolkit = Effect.isEffect(toolkit)
         ? yield* (toolkit as Effect.Effect<AiToolkit.ToHandler<any>>)
         : toolkit
+      const hasUnallowedTools = typeof toolChoice === "object" && "oneOf" in toolChoice
       for (const tool of actualToolkit.tools) {
+        if (hasUnallowedTools && !toolChoice.oneOf.includes(tool.name)) {
+          continue
+        }
         modelOptions.tools.push(convertTool(tool))
       }
       const stream = opts.streamText(modelOptions)

packages/ai/amazon-bedrock/src/AmazonBedrockLanguageModel.ts

@@ -151,6 +151,15 @@ export const make = Effect.fnUntraced(function*(options: {
       toolConfig.toolChoice = { tool: { name: options.tools[0].name } }
       return toolConfig
     }
+    const hasUnallowedTools = typeof options.toolChoice === "object" && "oneOf" in options.toolChoice
+    if (hasUnallowedTools) {
+      if (options.toolChoice.mode === "required") {
+        toolConfig.toolChoice = { any: {} }
+      } else {
+        toolConfig.toolChoice = { auto: {} }
+      }
+      return toolConfig
+    }
     if (options.toolChoice === "auto") {
       toolConfig.toolChoice = { auto: {} }
     } else if (options.toolChoice === "none") {

packages/ai/anthropic/src/AnthropicLanguageModel.ts

@@ -142,9 +142,16 @@ export const make = Effect.fnUntraced(function*(options: {
     function*(method: string, { prompt, system, toolChoice, tools }: AiLanguageModel.AiLanguageModelOptions) {
       const context = yield* Effect.context<never>()
       const useStructured = tools.length === 1 && tools[0].structured
+      const hasUnallowedTools = typeof toolChoice === "object" && "oneOf" in toolChoice
       let tool_choice: typeof Generated.ToolChoice.Encoded | undefined = undefined
       if (useStructured) {
         tool_choice = { type: "tool", name: tools[0].name }
+      } else if (hasUnallowedTools) {
+        if (toolChoice.mode === "required") {
+          tool_choice = { type: "any" }
+        } else {
+          tool_choice = { type: "auto" }
+        }
       } else if (tools.length > 0) {
         if (toolChoice === "required") {
           tool_choice = { type: "any" }

packages/ai/google/src/GoogleAiLanguageModel.ts

@@ -158,7 +158,8 @@ export const make = Effect.fnUntraced(function*(options: {
         ? "application/json"
         : (options.config?.generationConfig?.responseMimeType ?? perRequestConfig?.generationConfig?.responseMimeType)
       let toolConfig: typeof Generated.ToolConfig.Encoded | undefined = options.config?.toolConfig
-      if (Predicate.isNotUndefined(toolChoice) && !useStructured && tools.length > 0) {
+      const hasUnallowedTools = typeof toolChoice === "object" && "oneOf" in toolChoice
+      if (Predicate.isNotUndefined(toolChoice) && !useStructured && !hasUnallowedTools && tools.length > 0) {
         if (toolChoice === "none") {
           toolConfig = { functionCallingConfig: { ...toolConfig?.functionCallingConfig, mode: "NONE" } }
         } else if (toolChoice === "auto") {
@@ -168,6 +169,12 @@ export const make = Effect.fnUntraced(function*(options: {
         } else {
           toolConfig = { functionCallingConfig: { allowedFunctionNames: [toolChoice.tool], mode: "ANY" } }
         }
+      } else if (hasUnallowedTools && !useStructured) {
+        if (toolChoice.mode === "required") {
+          toolConfig = { functionCallingConfig: { ...toolConfig?.functionCallingConfig, mode: "ANY" } }
+        } else {
+          toolConfig = { functionCallingConfig: { ...toolConfig?.functionCallingConfig, mode: "AUTO" } }
+        }
       }
       const contents = yield* makeContents(method, prompt)
       return {

packages/ai/openai/src/OpenAiLanguageModel.ts

@@ -169,12 +169,15 @@ export const make = Effect.fnUntraced(function*(options: {
       const context = yield* Effect.context<never>()
       const useStructured = tools.length === 1 && tools[0].structured
       let tool_choice: typeof Generated.ChatCompletionToolChoiceOption.Encoded | undefined = undefined
-      if (Predicate.isNotUndefined(toolChoice) && !useStructured && tools.length > 0) {
+      const hasUnallowedTools = typeof toolChoice === "object" && "oneOf" in toolChoice
+      if (Predicate.isNotUndefined(toolChoice) && !useStructured && !hasUnallowedTools && tools.length > 0) {
         if (toolChoice === "auto" || toolChoice === "required") {
           tool_choice = toolChoice
         } else if (typeof toolChoice === "object") {
           tool_choice = { type: "function", function: { name: toolChoice.tool } }
         }
+      } else if (hasUnallowedTools && !useStructured) {
+        tool_choice = toolChoice.mode ?? "auto"
       }
       const messages = yield* makeMessages(method, system, prompt)
       return {