Tool Calling using Vercel's AI SDK not working as intended

[kldzj]

System Info

/info Output:

{
  "model_id": "casperhansen/llama-3.3-70b-instruct-awq",
  "model_sha": "64d255621f40b42adaf6d1f32a47e1d4534c0f14",
  "model_pipeline_tag": "text-generation",
  "max_concurrent_requests": 128,
  "max_best_of": 2,
  "max_stop_sequences": 4,
  "max_input_tokens": 8191,
  "max_total_tokens": 8192,
  "validation_workers": 2,
  "max_client_batch_size": 4,
  "router": "text-generation-router",
  "version": "3.0.2-dev0",
  "sha": "23bc38b10d06f8cc271d086c26270976faf67cc2",
  "docker_label": "sha-23bc38b"
}

Information

• Docker
• The CLI directly

Tasks

• An officially supported command
• My own modifications

Reproduction

1. Install dependencies: npm i ai @ai-sdk/openai zod

Tool Calls

2. Run the following script (e.g. node repro_generate.js):

const { createOpenAI } = require("@ai-sdk/openai");
const { generateObject } = require("ai");
const { z } = require("zod");

const openai = createOpenAI({
  baseURL: "http://localhost:8080/v1",
  apiKey: "-",
});

async function main() {
  const result = await generateObject({
    model: openai("casperhansen/llama-3.3-70b-instruct-awq"),
    system:
      "You are a summarization expert. Summarize the following text into a short title and a short description.",
    prompt:
      "Lorem ipsum dolor sit amet, consectetur adipiscing elit. Sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum.",
    schema: z.object({
      title: z.string(),
      content: z.string(),
    }),
    temperature: 0.7,
    maxTokens: 200,
  });

  console.log(result.object);
}

main().catch(console.error);

This will result in an error from Vercels AI SDK:

Invalid JSON response
Error message: [
  {
    "code": "invalid_type",
    "expected": "string",
    "received": "object",
    "path": [
      "choices",
      0,
      "message",
      "tool_calls",
      0,
      "function",
      "arguments"
    ],
    "message": "Expected string, received object"
  }
]

For this issue, there's a simple workaround of rewriting the response and stringifying the function arguments:

// ...
const openai = createOpenAI({
  baseURL: "http://localhost:8080/v1",
  apiKey: "-",
  fetch: async (url, options) => {
    const response = await fetch(url, {
      ...options,
      body: JSON.stringify({
        ...JSON.parse(options.body),
      }),
    });

    const body = await response.json();
    const choices = body.choices.map((choice) => {
      if (choice.message.tool_calls) {
        choice.message.tool_calls = choice.message.tool_calls.map(
          (toolCall) => {
            return {
              ...toolCall,
              arguments: JSON.stringify(toolCall.function.arguments),
              function: {
                ...toolCall.function,
                arguments: JSON.stringify(toolCall.function.arguments),
              },
            };
          }
        );

        return {
          ...choice,
          message: {
            ...choice.message,
            tool_calls: choice.message.tool_calls,
          },
        };
      }
    });

    return new Response(JSON.stringify({ ...body, choices }), {
      status: response.status,
      headers: response.headers,
    });
  },
});
// ...

Tool Streaming

3. Run the following script (e.g. node repro_stream.js):

const { createOpenAI } = require("@ai-sdk/openai");
const { generateObject, streamObject } = require("ai");
const { z } = require("zod");

const openai = createOpenAI({
  baseURL: "http://localhost:8080/v1",
  apiKey: "-",
});

async function main() {
  const result = await streamObject({
    model: openai("casperhansen/llama-3.3-70b-instruct-awq"),
    system:
      "You are a summarization expert. Summarize the following text into a short title and a short description.",
    prompt:
      "Lorem ipsum dolor sit amet, consectetur adipiscing elit. Sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum.",
    schema: z.object({
      title: z.string(),
      content: z.string(),
    }),
    temperature: 0.7,
    maxTokens: 200,
  });

  for await (const partialObject of result.partialObjectStream) {
    console.log(partialObject);
  }
}

main().catch(console.error);

This will result in an error from Vercels AI SDK:

Type validation failed: Value: {"object":"chat.completion.chunk","id":"","created":1734971997,"model":"casperhansen/llama-3.3-70b-instruct-awq","system_fingerprint":"3.0.2-dev0-sha-23bc38b","choices":[{"index":0,"delta":{"role":"assistant","tool_calls":{"index":0,"id":"","type":"function","function":{"name":null,"arguments":"{\""}}},"logprobs":null,"finish_reason":null}],"usage":null}.
Error message: [
  {
    "code": "invalid_union",
    "unionErrors": [
      {
        "issues": [
          {
            "code": "invalid_type",
            "expected": "array",
            "received": "object",
            "path": [
              "choices",
              0,
              "delta",
              "tool_calls"
            ],
            "message": "Expected array, received object"
          }
        ],
        "name": "ZodError"
      },
      {
        "issues": [
          {
            "code": "invalid_type",
            "expected": "object",
            "received": "undefined",
            "path": [
              "error"
            ],
            "message": "Required"
          }
        ],
        "name": "ZodError"
      }
    ],
    "path": [],
    "message": "Invalid input"
  }
]

Edit: There's another one, the client sometimes sends response_format without value property:

requestBodyValues: {
  model: 'casperhansen/llama-3.3-70b-instruct-awq',
  logit_bias: undefined,
  logprobs: undefined,
  top_logprobs: undefined,
  user: undefined,
  parallel_tool_calls: undefined,
  max_tokens: 300,
  temperature: 0.7,
  top_p: undefined,
  frequency_penalty: undefined,
  presence_penalty: undefined,
  stop: undefined,
  seed: undefined,
  max_completion_tokens: undefined,
  store: undefined,
  metadata: undefined,
  response_format: { type: 'json_object' },
  messages: [ [Object], [Object] ]
}

which results in the following error: Failed to deserialize the JSON body into the target type: response_format: missing field value at line 1 column 126.

Expected behavior

Being able to call tools/generate objects using the OpenAI compatible clients like Vercels AI SDK.

[kldzj]

I'm not well versed in Rust, but I'll try to aggregate the code issues and perhaps try to fix them when I find the time.

(Some) Problematic lines:

• main/router/src/server.rs:1441 should be stringified
• main/router/src/lib.rs:205 if I understand correctly this forces the object to have a value property

[kldzj]

I think I got it mostly working with a few request and response overrides:

import { getStopTokens } from "@/util";
import { createOpenAI } from "@ai-sdk/openai";
import { generateObject, generateText, streamObject, streamText, tool } from "ai";
import { z } from "zod";

function getRuntime() {
  const modelName = process.env.MODEL_NAME!;
  const authToken = process.env.OPENAI_API_KEY!;

  const openai = createOpenAI({
    baseURL: process.env.OPENAI_BASE_URL!,
    apiKey: authToken,
    compatibility: "strict",
    fetch: async (url, options) => {
      const requestBody = JSON.parse(options!.body! as string);
      if (requestBody.response_format && requestBody.response_format.type === "json_schema") {
        requestBody.response_format.type = "json";
        requestBody.response_format.value = requestBody.response_format.json_schema.schema;
        delete requestBody.response_format.json_schema;
      }

      const response = await fetch(url, {
        ...options,
        body: JSON.stringify({
          stop: getStopTokens(modelName),
          ...requestBody,
        }),
      });

      const contentType = response.headers.get("content-type");
      if (response.ok && contentType?.includes("application/json")) {
        const responseBody = await response.json();
        const choices = responseBody.choices.map((choice: any) => {
          if (choice.message?.tool_calls) {
            choice.message.tool_calls = choice.message.tool_calls.map((toolCall: any) => {
              return {
                ...toolCall,
                arguments: JSON.stringify(toolCall.function.arguments),
                function: {
                  ...toolCall.function,
                  arguments: JSON.stringify(toolCall.function.arguments),
                },
              };
            });

            return {
              ...choice,
              message: {
                ...choice.message,
                tool_calls: choice.message.tool_calls,
              },
            };
          } else {
            return choice;
          }
        });

        return new Response(JSON.stringify({ ...responseBody, choices }), {
          status: response.status,
          headers: response.headers,
        });
      }

      return response;
    },
  });

  return openai.chat(modelName, { structuredOutputs: true });
}

async function main() {
  await test("generateText", testGenerateText);
  await test("streamText", testStreamText);
  await test("generateObject", testGenerateObject);
  await test("streamObject", testStreamObject);
  await test("toolCall", testToolCall);
}

main().catch(console.error);

async function test(what: string, call: () => Promise<void>) {
  console.log(`${what}...`);
  try {
    await call();
    console.log("✅", what);
  } catch (error) {
    console.error("❌", what, error);
    console.log();
  }
}

async function testGenerateText() {
  const model = getRuntime();
  const response = await generateText({
    model,
    system: "You are a helpful assistant.",
    prompt: "What is the capital of Germany?",
  });

  console.log(response.text);
}

async function testStreamText() {
  const model = getRuntime();
  const stream = await streamText({
    model,
    system: "You are a helpful assistant.",
    prompt: "What is the capital of Germany?",
  });

  for await (const chunk of stream.textStream) {
    process.stdout.write(chunk);
  }

  process.stdout.write("\n");
}

async function testGenerateObject() {
  const model = getRuntime();
  const response = await generateObject({
    model,
    schema: z.object({
      text: z.string(),
    }),
    system: "You are a helpful assistant.",
    prompt: "What is the capital of Germany?",
  });

  console.log(response.object);
}

async function testStreamObject() {
  const model = getRuntime();
  const stream = await streamObject({
    model,
    schema: z.object({ text: z.string() }),
    system: "You are a helpful assistant.",
    prompt: "What is the capital of Germany?",
  });

  for await (const chunk of stream.partialObjectStream) {
    // ...
  }

  console.log(await stream.object);
}

async function testToolCall() {
  const model = getRuntime();
  const response = await generateText({
    model,
    system:
      "You are a helpful assistant. You have a tool to get the capital of a country. Return the result of the tool call.",
    prompt: "What is the capital of Germany?",
    tools: {
      getCapital: tool({
        parameters: z.object({
          country: z.string(),
        }),
        execute: async ({ country }) => {
          return country === "Germany" ? "Berlin" : "Unknown";
        },
      }),
    },
  });

  console.log(response.text || response.toolResults.find((r) => r.toolName === "getCapital")?.result);
}

Note the rewriting of the requests function calls and the responses response_format and the settings for the model, so it doesn't use tool calls for structured output (openai.chat(modelName, { structuredOutputs: true })).

Also note, that unlike other inference framework, at the end of the tool calls there will be no text generated. But perhaps that's an issue on my end.