|
|
@@ -1,4 +1,5 @@
|
|
|
import { t } from "i18next"
|
|
|
+import { FunctionCallingConfigMode } from "@google/genai"
|
|
|
|
|
|
import { GeminiHandler } from "../gemini"
|
|
|
import type { ApiHandlerOptions } from "../../../shared/api"
|
|
|
@@ -141,4 +142,152 @@ describe("GeminiHandler backend support", () => {
|
|
|
}).rejects.toThrow(t("common:errors.gemini.generate_stream", { error: "API rate limit exceeded" }))
|
|
|
})
|
|
|
})
|
|
|
+
|
|
|
+ describe("allowedFunctionNames support", () => {
|
|
|
+ const testTools = [
|
|
|
+ {
|
|
|
+ type: "function" as const,
|
|
|
+ function: {
|
|
|
+ name: "read_file",
|
|
|
+ description: "Read a file",
|
|
|
+ parameters: { type: "object", properties: {} },
|
|
|
+ },
|
|
|
+ },
|
|
|
+ {
|
|
|
+ type: "function" as const,
|
|
|
+ function: {
|
|
|
+ name: "write_to_file",
|
|
|
+ description: "Write to a file",
|
|
|
+ parameters: { type: "object", properties: {} },
|
|
|
+ },
|
|
|
+ },
|
|
|
+ {
|
|
|
+ type: "function" as const,
|
|
|
+ function: {
|
|
|
+ name: "execute_command",
|
|
|
+ description: "Execute a command",
|
|
|
+ parameters: { type: "object", properties: {} },
|
|
|
+ },
|
|
|
+ },
|
|
|
+ ]
|
|
|
+
|
|
|
+ it("should pass allowedFunctionNames to toolConfig when provided", async () => {
|
|
|
+ const options = {
|
|
|
+ apiProvider: "gemini",
|
|
|
+ } as ApiHandlerOptions
|
|
|
+ const handler = new GeminiHandler(options)
|
|
|
+ const stub = vi.fn().mockReturnValue((async function* () {})())
|
|
|
+ // @ts-ignore access private client
|
|
|
+ handler["client"].models.generateContentStream = stub
|
|
|
+
|
|
|
+ await handler
|
|
|
+ .createMessage("test", [] as any, {
|
|
|
+ taskId: "test-task",
|
|
|
+ tools: testTools,
|
|
|
+ allowedFunctionNames: ["read_file", "write_to_file"],
|
|
|
+ })
|
|
|
+ .next()
|
|
|
+
|
|
|
+ const config = stub.mock.calls[0][0].config
|
|
|
+ expect(config.toolConfig).toEqual({
|
|
|
+ functionCallingConfig: {
|
|
|
+ mode: FunctionCallingConfigMode.ANY,
|
|
|
+ allowedFunctionNames: ["read_file", "write_to_file"],
|
|
|
+ },
|
|
|
+ })
|
|
|
+ })
|
|
|
+
|
|
|
+ it("should include all tools but restrict callable functions via allowedFunctionNames", async () => {
|
|
|
+ const options = {
|
|
|
+ apiProvider: "gemini",
|
|
|
+ } as ApiHandlerOptions
|
|
|
+ const handler = new GeminiHandler(options)
|
|
|
+ const stub = vi.fn().mockReturnValue((async function* () {})())
|
|
|
+ // @ts-ignore access private client
|
|
|
+ handler["client"].models.generateContentStream = stub
|
|
|
+
|
|
|
+ await handler
|
|
|
+ .createMessage("test", [] as any, {
|
|
|
+ taskId: "test-task",
|
|
|
+ tools: testTools,
|
|
|
+ allowedFunctionNames: ["read_file"],
|
|
|
+ })
|
|
|
+ .next()
|
|
|
+
|
|
|
+ const config = stub.mock.calls[0][0].config
|
|
|
+ // All tools should be passed to the model
|
|
|
+ expect(config.tools[0].functionDeclarations).toHaveLength(3)
|
|
|
+ // But only read_file should be allowed to be called
|
|
|
+ expect(config.toolConfig.functionCallingConfig.allowedFunctionNames).toEqual(["read_file"])
|
|
|
+ })
|
|
|
+
|
|
|
+ it("should take precedence over tool_choice when allowedFunctionNames is provided", async () => {
|
|
|
+ const options = {
|
|
|
+ apiProvider: "gemini",
|
|
|
+ } as ApiHandlerOptions
|
|
|
+ const handler = new GeminiHandler(options)
|
|
|
+ const stub = vi.fn().mockReturnValue((async function* () {})())
|
|
|
+ // @ts-ignore access private client
|
|
|
+ handler["client"].models.generateContentStream = stub
|
|
|
+
|
|
|
+ await handler
|
|
|
+ .createMessage("test", [] as any, {
|
|
|
+ taskId: "test-task",
|
|
|
+ tools: testTools,
|
|
|
+ tool_choice: "auto",
|
|
|
+ allowedFunctionNames: ["read_file"],
|
|
|
+ })
|
|
|
+ .next()
|
|
|
+
|
|
|
+ const config = stub.mock.calls[0][0].config
|
|
|
+ // allowedFunctionNames should take precedence - mode should be ANY, not AUTO
|
|
|
+ expect(config.toolConfig.functionCallingConfig.mode).toBe(FunctionCallingConfigMode.ANY)
|
|
|
+ expect(config.toolConfig.functionCallingConfig.allowedFunctionNames).toEqual(["read_file"])
|
|
|
+ })
|
|
|
+
|
|
|
+ it("should fall back to tool_choice when allowedFunctionNames is empty", async () => {
|
|
|
+ const options = {
|
|
|
+ apiProvider: "gemini",
|
|
|
+ } as ApiHandlerOptions
|
|
|
+ const handler = new GeminiHandler(options)
|
|
|
+ const stub = vi.fn().mockReturnValue((async function* () {})())
|
|
|
+ // @ts-ignore access private client
|
|
|
+ handler["client"].models.generateContentStream = stub
|
|
|
+
|
|
|
+ await handler
|
|
|
+ .createMessage("test", [] as any, {
|
|
|
+ taskId: "test-task",
|
|
|
+ tools: testTools,
|
|
|
+ tool_choice: "auto",
|
|
|
+ allowedFunctionNames: [],
|
|
|
+ })
|
|
|
+ .next()
|
|
|
+
|
|
|
+ const config = stub.mock.calls[0][0].config
|
|
|
+ // Empty allowedFunctionNames should fall back to tool_choice behavior
|
|
|
+ expect(config.toolConfig.functionCallingConfig.mode).toBe(FunctionCallingConfigMode.AUTO)
|
|
|
+ expect(config.toolConfig.functionCallingConfig.allowedFunctionNames).toBeUndefined()
|
|
|
+ })
|
|
|
+
|
|
|
+ it("should not set toolConfig when allowedFunctionNames is undefined and no tool_choice", async () => {
|
|
|
+ const options = {
|
|
|
+ apiProvider: "gemini",
|
|
|
+ } as ApiHandlerOptions
|
|
|
+ const handler = new GeminiHandler(options)
|
|
|
+ const stub = vi.fn().mockReturnValue((async function* () {})())
|
|
|
+ // @ts-ignore access private client
|
|
|
+ handler["client"].models.generateContentStream = stub
|
|
|
+
|
|
|
+ await handler
|
|
|
+ .createMessage("test", [] as any, {
|
|
|
+ taskId: "test-task",
|
|
|
+ tools: testTools,
|
|
|
+ })
|
|
|
+ .next()
|
|
|
+
|
|
|
+ const config = stub.mock.calls[0][0].config
|
|
|
+ // No toolConfig should be set when neither allowedFunctionNames nor tool_choice is provided
|
|
|
+ expect(config.toolConfig).toBeUndefined()
|
|
|
+ })
|
|
|
+ })
|
|
|
})
|