add deferred tools support for kimi in openai-completions api

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
+2 -1
View File
@@ -373,8 +373,9 @@ const OPENAI_COMPLETIONS_DEFAULT_COMPAT = {
supportsStrictMode: true, supportsStrictMode: true,
sendSessionAffinityHeaders: false, sendSessionAffinityHeaders: false,
supportsLongCacheRetention: true, supportsLongCacheRetention: true,
} satisfies Required<Omit<OpenAICompletionsCompat, "cacheControlFormat">> & { } satisfies Required<Omit<OpenAICompletionsCompat, "cacheControlFormat" | "deferredToolsMode">> & {
cacheControlFormat?: OpenAICompletionsCompat["cacheControlFormat"]; cacheControlFormat?: OpenAICompletionsCompat["cacheControlFormat"];
deferredToolsMode?: OpenAICompletionsCompat["deferredToolsMode"];
}; };
type OpenAICompletionsResolvedCompat = typeof OPENAI_COMPLETIONS_DEFAULT_COMPAT & { type OpenAICompletionsResolvedCompat = typeof OPENAI_COMPLETIONS_DEFAULT_COMPAT & {
+56 -3
View File
@@ -77,6 +77,26 @@ function hasToolHistory(messages: Message[]): boolean {
return false; 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 { function isTextContentBlock(block: { type: string }): block is TextContent {
return block.type === "text"; return block.type === "text";
} }
@@ -117,14 +137,23 @@ interface OpenAICompatCacheControl {
ttl?: string; ttl?: string;
} }
type ResolvedOpenAICompletionsCompat = Omit<Required<OpenAICompletionsCompat>, "cacheControlFormat"> & { type ResolvedOpenAICompletionsCompat = Omit<
Required<OpenAICompletionsCompat>,
"cacheControlFormat" | "deferredToolsMode"
> & {
cacheControlFormat?: OpenAICompletionsCompat["cacheControlFormat"]; cacheControlFormat?: OpenAICompletionsCompat["cacheControlFormat"];
deferredToolsMode?: OpenAICompletionsCompat["deferredToolsMode"];
}; };
type ResolvedChatTemplateKwargValue = string | number | boolean | null; type ResolvedChatTemplateKwargValue = string | number | boolean | null;
type ChatCompletionInstructionMessageParam = ChatCompletionDeveloperMessageParam | ChatCompletionSystemMessageParam; type ChatCompletionInstructionMessageParam = ChatCompletionDeveloperMessageParam | ChatCompletionSystemMessageParam;
type KimiToolSystemMessageParam = {
role: "system";
tools: OpenAI.Chat.Completions.ChatCompletionTool[];
};
type OpenAIEncryptedReasoningDetail = { type OpenAIEncryptedReasoningDetail = {
type: "reasoning.encrypted"; type: "reasoning.encrypted";
id: string; id: string;
@@ -585,8 +614,11 @@ function buildParams(
params.temperature = options.temperature; params.temperature = options.temperature;
} }
if (context.tools && context.tools.length > 0) { const deferredToolNames =
params.tools = convertTools(context.tools, compat); 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) { if (compat.zaiToolStream) {
(params as any).tool_stream = true; (params as any).tool_stream = true;
} }
@@ -1025,6 +1057,7 @@ export function convertMessages(
params.push(assistantMsg); params.push(assistantMsg);
} else if (msg.role === "toolResult") { } else if (msg.role === "toolResult") {
const imageBlocks: Array<{ type: "image_url"; image_url: { url: string } }> = []; const imageBlocks: Array<{ type: "image_url"; image_url: { url: string } }> = [];
const deferredToolNames = new Set<string>();
let j = i; let j = i;
for (; j < transformedMessages.length && transformedMessages[j].role === "toolResult"; j++) { for (; j < transformedMessages.length && transformedMessages[j].role === "toolResult"; j++) {
@@ -1051,6 +1084,12 @@ export function convertMessages(
} }
params.push(toolResultMsg); params.push(toolResultMsg);
if (compat.deferredToolsMode === "kimi") {
for (const name of toolMsg.addedToolNames ?? []) {
deferredToolNames.add(name);
}
}
if (hasImages && model.input.includes("image")) { if (hasImages && model.input.includes("image")) {
for (const block of toolMsg.content) { for (const block of toolMsg.content) {
if (isImageContentBlock(block)) { if (isImageContentBlock(block)) {
@@ -1089,6 +1128,18 @@ export function convertMessages(
} else { } else {
lastRole = "toolResult"; 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; continue;
} }
@@ -1256,6 +1307,7 @@ function detectCompat(model: Model<"openai-completions">): ResolvedOpenAIComplet
supportsStrictMode: !isMoonshot && !isTogether && !isCloudflareAiGateway && !isNvidia, supportsStrictMode: !isMoonshot && !isTogether && !isCloudflareAiGateway && !isNvidia,
cacheControlFormat, cacheControlFormat,
sendSessionAffinityHeaders: false, sendSessionAffinityHeaders: false,
deferredToolsMode: undefined,
sessionAffinityFormat: isOpenRouter ? "openrouter" : "openai", sessionAffinityFormat: isOpenRouter ? "openrouter" : "openai",
supportsLongCacheRetention: !( supportsLongCacheRetention: !(
isTogether || isTogether ||
@@ -1296,6 +1348,7 @@ function getCompat(model: Model<"openai-completions">): ResolvedOpenAICompletion
supportsStrictMode: model.compat.supportsStrictMode ?? detected.supportsStrictMode, supportsStrictMode: model.compat.supportsStrictMode ?? detected.supportsStrictMode,
cacheControlFormat: model.compat.cacheControlFormat ?? detected.cacheControlFormat, cacheControlFormat: model.compat.cacheControlFormat ?? detected.cacheControlFormat,
sendSessionAffinityHeaders: model.compat.sendSessionAffinityHeaders ?? detected.sendSessionAffinityHeaders, sendSessionAffinityHeaders: model.compat.sendSessionAffinityHeaders ?? detected.sendSessionAffinityHeaders,
deferredToolsMode: model.compat.deferredToolsMode ?? detected.deferredToolsMode,
sessionAffinityFormat: model.compat.sessionAffinityFormat ?? detected.sessionAffinityFormat, sessionAffinityFormat: model.compat.sessionAffinityFormat ?? detected.sessionAffinityFormat,
supportsLongCacheRetention: model.compat.supportsLongCacheRetention ?? detected.supportsLongCacheRetention, supportsLongCacheRetention: model.compat.supportsLongCacheRetention ?? detected.supportsLongCacheRetention,
}; };
+2
View File
@@ -524,6 +524,8 @@ export interface OpenAICompletionsCompat {
cacheControlFormat?: "anthropic"; cacheControlFormat?: "anthropic";
/** Whether to send session-affinity data from `options.sessionId`. Default: false. */ /** Whether to send session-affinity data from `options.sessionId`. Default: false. */
sendSessionAffinityHeaders?: boolean; 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. */ /** 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; 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. */ /** Whether the provider supports long prompt cache retention (`prompt_cache_retention: "24h"` or Anthropic-style `cache_control.ttl: "1h"`, depending on format). Default: true. */
+94
View File
@@ -1,5 +1,6 @@
import { Type } from "typebox"; import { Type } from "typebox";
import { describe, expect, it } from "vitest"; import { describe, expect, it } from "vitest";
import { convertMessages } from "../src/api/openai-completions.ts";
import { getModel, streamSimple } from "../src/compat.ts"; import { getModel, streamSimple } from "../src/compat.ts";
import type { Api, AssistantMessage, Context, Model, Tool, ToolResultMessage, UserMessage } from "../src/types.ts"; import type { Api, AssistantMessage, Context, Model, Tool, ToolResultMessage, UserMessage } from "../src/types.ts";
import { estimateContextTokens } from "../src/utils/estimate.ts"; import { estimateContextTokens } from "../src/utils/estimate.ts";
@@ -49,6 +50,26 @@ interface OpenAIPayload {
input?: Array<OpenAIToolSearchCall | OpenAIToolSearchOutput | { type?: string }>; 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 {} class PayloadCaptured extends Error {}
function makeTool(name: string): Tool { 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> { async function capturePayload<T>(model: Model<Api>, context: Context, apiKey = "fake-key"): Promise<T> {
let captured: T | undefined; let captured: T | undefined;
const stream = streamSimple({ ...model, baseUrl: "http://127.0.0.1:9" }, context, { 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); 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 () => { it("loads an OpenAI Responses tool through client tool search", async () => {
const context = makeContext([makeTool("base_tool"), makeTool("late_tool")]); const context = makeContext([makeTool("base_tool"), makeTool("late_tool")]);
const payload = await capturePayload<OpenAIPayload>(getModel("openai", "gpt-5.4"), context); const payload = await capturePayload<OpenAIPayload>(getModel("openai", "gpt-5.4"), context);
@@ -41,8 +41,9 @@ const compat = {
sendSessionAffinityHeaders: false, sendSessionAffinityHeaders: false,
sessionAffinityFormat: "openai", sessionAffinityFormat: "openai",
supportsLongCacheRetention: true, supportsLongCacheRetention: true,
} satisfies Required<Omit<OpenAICompletionsCompat, "cacheControlFormat">> & { } satisfies Omit<Required<OpenAICompletionsCompat>, "cacheControlFormat" | "deferredToolsMode"> & {
cacheControlFormat?: OpenAICompletionsCompat["cacheControlFormat"]; cacheControlFormat?: OpenAICompletionsCompat["cacheControlFormat"];
deferredToolsMode?: OpenAICompletionsCompat["deferredToolsMode"];
}; };
function buildModel(baseUrl = "http://127.0.0.1:1"): Model<"openai-completions"> { function buildModel(baseUrl = "http://127.0.0.1:1"): Model<"openai-completions"> {
@@ -19,7 +19,9 @@ const emptyUsage: Usage = {
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, 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, supportsStore: true,
supportsDeveloperRole: true, supportsDeveloperRole: true,
supportsReasoningEffort: true, supportsReasoningEffort: true,
+1
View File
@@ -449,6 +449,7 @@ For providers with partial OpenAI compatibility, use the `compat` field.
| `sendSessionAffinityHeaders` | For `openai-completions`, send session-affinity headers from the session id when caching is enabled. Default: `false`. | | `sendSessionAffinityHeaders` | For `openai-completions`, send session-affinity headers from the session id when caching is enabled. Default: `false`. |
| `sessionAffinityFormat` | For `openai-completions` and `openai-responses`, the session-affinity header format: `openai` sends `session_id`/`x-client-request-id` (completions also `x-session-affinity`), `openai-nosession` omits the underscore-containing `session_id` header, `openrouter` sends `x-session-id`. Does not affect the `prompt_cache_key` body param. Default: auto-detected. | | `sessionAffinityFormat` | For `openai-completions` and `openai-responses`, the session-affinity header format: `openai` sends `session_id`/`x-client-request-id` (completions also `x-session-affinity`), `openai-nosession` omits the underscore-containing `session_id` header, `openrouter` sends `x-session-id`. Does not affect the `prompt_cache_key` body param. Default: auto-detected. |
| `supportsStrictMode` | Include the `strict` field in tool definitions | | `supportsStrictMode` | Include the `strict` field in tool definitions |
| `deferredToolsMode` | Use provider-specific deferred tool serialization. Currently only `"kimi"` is supported for Kimi's OpenAI-compatible Chat Completions format. |
| `supportsLongCacheRetention` | Whether the provider accepts long cache retention when cache retention is `long`: `prompt_cache_retention: "24h"` for OpenAI prompt caching, or `cache_control.ttl: "1h"` when `cacheControlFormat` is `anthropic`. Default: `true`. | | `supportsLongCacheRetention` | Whether the provider accepts long cache retention when cache retention is `long`: `prompt_cache_retention: "24h"` for OpenAI prompt caching, or `cache_control.ttl: "1h"` when `cacheControlFormat` is `anthropic`. Default: `true`. |
| `openRouterRouting` | OpenRouter provider routing preferences. This object is sent as-is in the `provider` field of the [OpenRouter API request](https://openrouter.ai/docs/guides/routing/provider-selection). | | `openRouterRouting` | OpenRouter provider routing preferences. This object is sent as-is in the `provider` field of the [OpenRouter API request](https://openrouter.ai/docs/guides/routing/provider-selection). |
| `vercelGatewayRouting` | Vercel AI Gateway routing config for provider selection (`only`, `order`) | | `vercelGatewayRouting` | Vercel AI Gateway routing config for provider selection (`only`, `order`) |
@@ -36,6 +36,7 @@ cp permission-gate.ts ~/.pi/agent/extensions/
| `questionnaire.ts` | Multi-question input with tab bar navigation between questions | | `questionnaire.ts` | Multi-question input with tab bar navigation between questions |
| `tool-override.ts` | Override built-in tools (e.g., add logging/access control to `read`) | | `tool-override.ts` | Override built-in tools (e.g., add logging/access control to `read`) |
| `dynamic-tools.ts` | Register tools after startup (`session_start`) and at runtime via command, with prompt snippets and tool-specific prompt guidelines | | `dynamic-tools.ts` | Register tools after startup (`session_start`) and at runtime via command, with prompt snippets and tool-specific prompt guidelines |
| `kimi-deferred-tools.ts` | Search for and progressively activate tools for Kimi's deferred-tool loading protocol |
| `structured-output.ts` | Final structured-output tool that returns `terminate: true` so the agent can end on the tool call | | `structured-output.ts` | Final structured-output tool that returns `terminate: true` so the agent can end on the tool call |
| `built-in-tool-renderer.ts` | Custom compact rendering for built-in tools (read, bash, edit, write) while keeping original behavior | | `built-in-tool-renderer.ts` | Custom compact rendering for built-in tools (read, bash, edit, write) while keeping original behavior |
| `minimal-mode.ts` | Override built-in tool rendering for minimal display (only tool calls, no output in collapsed mode) | | `minimal-mode.ts` | Override built-in tool rendering for minimal display (only tool calls, no output in collapsed mode) |
@@ -0,0 +1,61 @@
/**
* Minimal Kimi deferred-tool loading demo.
*
* pi -e ./kimi-deferred-tools.ts
* example prompt: Use the available tools to calculate 100 + 500. Do not calculate it yourself.
*/
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { Type } from "typebox";
function calculate(_expr: string): string {
return "42";
}
export default function (pi: ExtensionAPI): void {
pi.registerTool({
name: "Calculator",
label: "Calculator",
description: "Evaluate a simple arithmetic expression.",
parameters: Type.Object({
expr: Type.String({ description: "An expression such as 100 + 500" }),
}),
async execute(_toolCallId, params) {
return {
content: [{ type: "text", text: calculate(params.expr) }],
details: {},
};
},
});
pi.registerTool({
name: "tool_search",
label: "Tool Search",
description: "Find and activate tools for a capability.",
promptSnippet: "Search for additional tools when the active tools cannot perform the task",
parameters: Type.Object({
query: Type.String({ description: "Capability to search for" }),
}),
async execute(_toolCallId, params) {
if (!params.query.toLowerCase().includes("calc")) {
return {
content: [{ type: "text", text: "The relevant tools do not exist." }],
details: { matches: [], added: [] },
};
}
const active = pi.getActiveTools();
const added = active.includes("Calculator") ? [] : ["Calculator"];
if (added.length > 0) pi.setActiveTools([...active, ...added]);
return {
content: [{ type: "text", text: "Success. Found 1 matching tool(s)" }],
details: { matches: ["Calculator"], added },
};
},
});
pi.on("session_start", () => {
pi.setActiveTools(["tool_search"]);
});
}
@@ -99,6 +99,7 @@ const OpenAICompletionsCompatSchema = Type.Object({
vercelGatewayRouting: Type.Optional(VercelGatewayRoutingSchema), vercelGatewayRouting: Type.Optional(VercelGatewayRoutingSchema),
supportsStrictMode: Type.Optional(Type.Boolean()), supportsStrictMode: Type.Optional(Type.Boolean()),
sendSessionAffinityHeaders: Type.Optional(Type.Boolean()), sendSessionAffinityHeaders: Type.Optional(Type.Boolean()),
deferredToolsMode: Type.Optional(Type.Literal("kimi")),
sessionAffinityFormat: Type.Optional( sessionAffinityFormat: Type.Optional(
Type.Union([Type.Literal("openai"), Type.Literal("openai-nosession"), Type.Literal("openrouter")]), Type.Union([Type.Literal("openai"), Type.Literal("openai-nosession"), Type.Literal("openrouter")]),
), ),