ton/pi_harness
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add deferred tools support for kimi in openai-completions api

Browse Source
This commit is contained in:
David Brailovsky
2026-07-16 12:26:56 +00:00
parent 36db3fa385
commit f16b4e0cda
10 changed files with 223 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/ai/scripts/generate-models.ts
+2 -1
View File
@@ -373,8 +373,9 @@ const OPENAI_COMPLETIONS_DEFAULT_COMPAT = {
supportsStrictMode: true,
sendSessionAffinityHeaders: false,
supportsLongCacheRetention: true,
} satisfies Required<Omit<OpenAICompletionsCompat, "cacheControlFormat">> & {
} satisfies Required<Omit<OpenAICompletionsCompat, "cacheControlFormat" | "deferredToolsMode">> & {
cacheControlFormat?: OpenAICompletionsCompat["cacheControlFormat"];
deferredToolsMode?: OpenAICompletionsCompat["deferredToolsMode"];
};
type OpenAICompletionsResolvedCompat = typeof OPENAI_COMPLETIONS_DEFAULT_COMPAT & {
packages/ai/src/api/openai-completions.ts
+56 -3
View File
@@ -77,6 +77,26 @@ function hasToolHistory(messages: Message[]): boolean {
return false;
}
function getDeferredToolNames(messages: Message[]): Set<string> {
const names = new Set<string>();
for (const message of messages) {
if (message.role === "toolResult") {
for (const name of message.addedToolNames ?? []) {
names.add(name);
}
}
}
return names;
}
function getToolsByName(tools: Tool[] | undefined, names: Iterable<string>): Tool[] {
if (!tools) return [];
const toolsByName = new Map(tools.map((tool) => [tool.name, tool]));
return Array.from(names)
.map((name) => toolsByName.get(name))
.filter((tool): tool is Tool => tool !== undefined);
}
function isTextContentBlock(block: { type: string }): block is TextContent {
return block.type === "text";
}
@@ -117,14 +137,23 @@ interface OpenAICompatCacheControl {
ttl?: string;
}
type ResolvedOpenAICompletionsCompat = Omit<Required<OpenAICompletionsCompat>, "cacheControlFormat"> & {
type ResolvedOpenAICompletionsCompat = Omit<
Required<OpenAICompletionsCompat>,
"cacheControlFormat" | "deferredToolsMode"
> & {
cacheControlFormat?: OpenAICompletionsCompat["cacheControlFormat"];
deferredToolsMode?: OpenAICompletionsCompat["deferredToolsMode"];
};
type ResolvedChatTemplateKwargValue = string | number | boolean | null;
type ChatCompletionInstructionMessageParam = ChatCompletionDeveloperMessageParam | ChatCompletionSystemMessageParam;
type KimiToolSystemMessageParam = {
role: "system";
tools: OpenAI.Chat.Completions.ChatCompletionTool[];
};
type OpenAIEncryptedReasoningDetail = {
type: "reasoning.encrypted";
id: string;
@@ -585,8 +614,11 @@ function buildParams(
params.temperature = options.temperature;
}
if (context.tools && context.tools.length > 0) {
params.tools = convertTools(context.tools, compat);
const deferredToolNames =
compat.deferredToolsMode === "kimi" ? getDeferredToolNames(context.messages) : new Set<string>();
const activeTools = context.tools?.filter((tool) => !deferredToolNames.has(tool.name));
if (activeTools && activeTools.length > 0) {
params.tools = convertTools(activeTools, compat);
if (compat.zaiToolStream) {
(params as any).tool_stream = true;
}
@@ -1025,6 +1057,7 @@ export function convertMessages(
params.push(assistantMsg);
} else if (msg.role === "toolResult") {
const imageBlocks: Array<{ type: "image_url"; image_url: { url: string } }> = [];
const deferredToolNames = new Set<string>();
let j = i;
for (; j < transformedMessages.length && transformedMessages[j].role === "toolResult"; j++) {
@@ -1051,6 +1084,12 @@ export function convertMessages(
}
params.push(toolResultMsg);
if (compat.deferredToolsMode === "kimi") {
for (const name of toolMsg.addedToolNames ?? []) {
deferredToolNames.add(name);
}
}
if (hasImages && model.input.includes("image")) {
for (const block of toolMsg.content) {
if (isImageContentBlock(block)) {
@@ -1089,6 +1128,18 @@ export function convertMessages(
} else {
lastRole = "toolResult";
}
if (deferredToolNames.size > 0) {
const deferredTools = getToolsByName(context.tools, deferredToolNames);
if (deferredTools.length > 0) {
const kimiToolMessage: KimiToolSystemMessageParam = {
role: "system",
tools: convertTools(deferredTools, compat),
};
// Kimi accepts a system message with tools but omits the standard content field.
params.push(kimiToolMessage as unknown as ChatCompletionMessageParam);
}
}
continue;
}
@@ -1256,6 +1307,7 @@ function detectCompat(model: Model<"openai-completions">): ResolvedOpenAIComplet
supportsStrictMode: !isMoonshot && !isTogether && !isCloudflareAiGateway && !isNvidia,
cacheControlFormat,
sendSessionAffinityHeaders: false,
deferredToolsMode: undefined,
sessionAffinityFormat: isOpenRouter ? "openrouter" : "openai",
supportsLongCacheRetention: !(
isTogether ||
@@ -1296,6 +1348,7 @@ function getCompat(model: Model<"openai-completions">): ResolvedOpenAICompletion
supportsStrictMode: model.compat.supportsStrictMode ?? detected.supportsStrictMode,
cacheControlFormat: model.compat.cacheControlFormat ?? detected.cacheControlFormat,
sendSessionAffinityHeaders: model.compat.sendSessionAffinityHeaders ?? detected.sendSessionAffinityHeaders,
deferredToolsMode: model.compat.deferredToolsMode ?? detected.deferredToolsMode,
sessionAffinityFormat: model.compat.sessionAffinityFormat ?? detected.sessionAffinityFormat,
supportsLongCacheRetention: model.compat.supportsLongCacheRetention ?? detected.supportsLongCacheRetention,
};
packages/ai/src/types.ts
+2
View File
@@ -524,6 +524,8 @@ export interface OpenAICompletionsCompat {
cacheControlFormat?: "anthropic";
/** Whether to send session-affinity data from `options.sessionId`. Default: false. */
sendSessionAffinityHeaders?: boolean;
/** Provider-specific deferred tool serialization mode. */
deferredToolsMode?: "kimi";
/** Session-affinity header format: `openai` sends `session_id`, `x-client-request-id`, and `x-session-affinity`; `openai-nosession` sends `x-client-request-id` and `x-session-affinity`; `openrouter` sends `x-session-id`. Does not affect the `prompt_cache_key` body param, which is governed by cache retention. Default: auto-detected. */
sessionAffinityFormat?: SessionAffinityFormat;
/** Whether the provider supports long prompt cache retention (`prompt_cache_retention: "24h"` or Anthropic-style `cache_control.ttl: "1h"`, depending on format). Default: true. */
packages/ai/test/deferred-tools.test.ts
+94
View File
@@ -1,5 +1,6 @@
import { Type } from "typebox";
import { describe, expect, it } from "vitest";
import { convertMessages } from "../src/api/openai-completions.ts";
import { getModel, streamSimple } from "../src/compat.ts";
import type { Api, AssistantMessage, Context, Model, Tool, ToolResultMessage, UserMessage } from "../src/types.ts";
import { estimateContextTokens } from "../src/utils/estimate.ts";
@@ -49,6 +50,26 @@ interface OpenAIPayload {
input?: Array<OpenAIToolSearchCall | OpenAIToolSearchOutput | { type?: string }>;
}
interface KimiTool {
type: "function";
function: {
name: string;
description?: string;
parameters?: Record<string, unknown>;
};
}
interface KimiMessage {
role: string;
content?: unknown;
tools?: KimiTool[];
}
interface KimiPayload {
tools?: KimiTool[];
messages: KimiMessage[];
}
class PayloadCaptured extends Error {}
function makeTool(name: string): Tool {
@@ -102,6 +123,22 @@ function makeContext(tools: Tool[], addedToolNames = ["late_tool"]): Context {
};
}
function makeKimiModel(deferredToolsMode?: "kimi"): Model<"openai-completions"> {
return {
id: "deferred-tools-model",
name: "Deferred Tools Model",
api: "openai-completions",
provider: "moonshotai",
baseUrl: "http://127.0.0.1:9/v1",
reasoning: false,
input: ["text"],
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
contextWindow: 128000,
maxTokens: 4096,
compat: deferredToolsMode ? { deferredToolsMode } : undefined,
};
}
async function capturePayload<T>(model: Model<Api>, context: Context, apiKey = "fake-key"): Promise<T> {
let captured: T | undefined;
const stream = streamSimple({ ...model, baseUrl: "http://127.0.0.1:9" }, context, {
@@ -297,6 +334,63 @@ describe("deferred tools", () => {
expect(payload.tools?.find((tool) => tool.name === "late_tool")?.defer_loading).toBe(true);
});
it("serializes Kimi deferred tools as system tool definitions", async () => {
const context = makeContext([makeTool("base_tool"), makeTool("late_tool")]);
const payload = await capturePayload<KimiPayload>(makeKimiModel("kimi"), context);
expect(payload.tools?.map((tool) => tool.function.name)).toEqual(["base_tool"]);
const toolResultIndex = payload.messages.findIndex((message) => message.role === "tool");
const systemToolIndex = payload.messages.findIndex((message) => message.tools !== undefined);
expect(toolResultIndex).toBeGreaterThanOrEqual(0);
expect(systemToolIndex).toBeGreaterThan(toolResultIndex);
expect(payload.messages[systemToolIndex]?.tools?.map((tool) => tool.function.name)).toEqual(["late_tool"]);
});
it("emits Kimi deferred schemas after all tool results in a batch", () => {
const context = makeContext([makeTool("base_tool"), makeTool("late_tool"), makeTool("later_tool")]);
context.messages.splice(3, 0, {
...makeToolResult(["later_tool"]),
toolCallId: "call_2",
});
const messages = convertMessages(makeKimiModel("kimi"), context, {
supportsStore: false,
supportsDeveloperRole: false,
supportsReasoningEffort: false,
supportsUsageInStreaming: true,
maxTokensField: "max_tokens",
requiresToolResultName: false,
requiresAssistantAfterToolResult: false,
requiresThinkingAsText: false,
requiresReasoningContentOnAssistantMessages: false,
thinkingFormat: "openai",
openRouterRouting: {},
vercelGatewayRouting: {},
chatTemplateKwargs: {},
zaiToolStream: false,
supportsStrictMode: false,
cacheControlFormat: undefined,
sendSessionAffinityHeaders: false,
deferredToolsMode: "kimi",
sessionAffinityFormat: "openai",
supportsLongCacheRetention: false,
});
expect(messages.map((message) => message.role)).toEqual(["user", "assistant", "tool", "tool", "system", "user"]);
expect((messages[4] as { tools?: KimiTool[] }).tools?.map((tool) => tool.function.name)).toEqual([
"late_tool",
"later_tool",
]);
});
it("leaves OpenAI Completions tools unchanged without Kimi mode", async () => {
const context = makeContext([makeTool("base_tool"), makeTool("late_tool")]);
const payload = await capturePayload<KimiPayload>(makeKimiModel(), context);
expect(payload.tools?.map((tool) => tool.function.name)).toEqual(["base_tool", "late_tool"]);
expect(payload.messages.some((message) => message.tools !== undefined)).toBe(false);
});
it("loads an OpenAI Responses tool through client tool search", async () => {
const context = makeContext([makeTool("base_tool"), makeTool("late_tool")]);
const payload = await capturePayload<OpenAIPayload>(getModel("openai", "gpt-5.4"), context);
packages/ai/test/openai-completions-thinking-as-text.test.ts
+2 -1
View File
@@ -41,8 +41,9 @@ const compat = {
sendSessionAffinityHeaders: false,
sessionAffinityFormat: "openai",
supportsLongCacheRetention: true,
} satisfies Required<Omit<OpenAICompletionsCompat, "cacheControlFormat">> & {
} satisfies Omit<Required<OpenAICompletionsCompat>, "cacheControlFormat" | "deferredToolsMode"> & {
cacheControlFormat?: OpenAICompletionsCompat["cacheControlFormat"];
deferredToolsMode?: OpenAICompletionsCompat["deferredToolsMode"];
};
function buildModel(baseUrl = "http://127.0.0.1:1"): Model<"openai-completions"> {
packages/ai/test/openai-completions-tool-result-images.test.ts
+3 -1
View File
@@ -19,7 +19,9 @@ const emptyUsage: Usage = {
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
};
const compat: Required<OpenAICompletionsCompat> = {
const compat: Omit<Required<OpenAICompletionsCompat>, "deferredToolsMode"> & {
deferredToolsMode?: OpenAICompletionsCompat["deferredToolsMode"];
} = {
supportsStore: true,
supportsDeveloperRole: true,
supportsReasoningEffort: true,