Skip to content

support sequential tool calls and fix mistral #78

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

support sequential tool calls and fix mistral #78

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 6 additions & 2 deletions packages/workers-ai-provider/src/convert-to-workersai-chat-messages.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,10 @@
import { type LanguageModelV1Prompt, UnsupportedFunctionalityError } from "@ai-sdk/provider";
import type { WorkersAIChatPrompt } from "./workersai-chat-prompt";

export function convertToWorkersAIChatMessages(prompt: LanguageModelV1Prompt): WorkersAIChatPrompt {
export function convertToWorkersAIChatMessages(
prompt: LanguageModelV1Prompt,
excludeContentWithToolCalls = false
): WorkersAIChatPrompt {
const messages: WorkersAIChatPrompt = [];

for (const { role, content } of prompt) {
Expand Down Expand Up @@ -71,7 +74,7 @@ export function convertToWorkersAIChatMessages(prompt: LanguageModelV1Prompt): W

messages.push({
role: "assistant",
content: text,
content: excludeContentWithToolCalls && toolCalls.length > 0 ? undefined : text, // fix for mistral
tool_calls:
toolCalls.length > 0
? toolCalls.map(({ function: { name, arguments: args } }) => ({
Expand All @@ -88,6 +91,7 @@ export function convertToWorkersAIChatMessages(prompt: LanguageModelV1Prompt): W
for (const toolResponse of content) {
messages.push({
role: "tool",
tool_call_id: toolResponse.toolCallId, // required by mistral
name: toolResponse.toolName,
content: JSON.stringify(toolResponse.result),
});
Expand Down
10 changes: 6 additions & 4 deletions packages/workers-ai-provider/src/workersai-chat-language-model.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -83,7 +83,7 @@ export class WorkersAIChatLanguageModel implements LanguageModelV1 {
random_seed: seed,

// messages:
messages: convertToWorkersAIChatMessages(prompt),
messages: convertToWorkersAIChatMessages(prompt, this.modelId.includes('mistral')),
};

switch (type) {
Expand Down Expand Up @@ -139,7 +139,7 @@ export class WorkersAIChatLanguageModel implements LanguageModelV1 {
): Promise<Awaited<ReturnType<LanguageModelV1["doGenerate"]>>> {
const { args, warnings } = this.getArgs(options);

const { gateway, safePrompt, ...passthroughOptions } = this.settings;
const { gateway, safePrompt, sequentialCalls, ...passthroughOptions } = this.settings;

const output = await this.config.binding.run(
args.model,
Expand Down Expand Up @@ -181,11 +181,14 @@ export class WorkersAIChatLanguageModel implements LanguageModelV1 {
options: Parameters<LanguageModelV1["doStream"]>[0],
): Promise<Awaited<ReturnType<LanguageModelV1["doStream"]>>> {
const { args, warnings } = this.getArgs(options);
const { gateway, safePrompt, sequentialCalls, ...passthroughOptions } = this.settings;

// [1] When the latest message is not a tool response, we use the regular generate function
// and simulate it as a streamed response in order to satisfy the AI SDK's interface for
// doStream...
if (args.tools?.length && lastMessageWasUser(args.messages)) {
// To allow a model to chain together tool calls, `sequentialCalls` can be set to `true`.
// This will make all responses be simulated streams.
if (args.tools?.length && (!!sequentialCalls || lastMessageWasUser(args.messages))) {
const response = await this.doGenerate(options);

if (response instanceof ReadableStream) {
Expand Down Expand Up @@ -223,7 +226,6 @@ export class WorkersAIChatLanguageModel implements LanguageModelV1 {
}

// [2] ...otherwise, we just proceed as normal and stream the response directly from the remote model.
const { gateway, ...passthroughOptions } = this.settings;

const response = await this.config.binding.run(
args.model,
Expand Down
6 changes: 6 additions & 0 deletions packages/workers-ai-provider/src/workersai-chat-settings.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,12 @@ export type WorkersAIChatSettings = {
*/
safePrompt?: boolean;

/**
* Whether to allow the model to execute calls when the last message was not a user message. Turns off streaming.
* Defaults to `false`.
*/
sequentialCalls?: boolean;

/**
* Optionally set Cloudflare AI Gateway options.
* @deprecated
Expand Down