@hebo-ai/gateway 0.4.1 → 0.5.0-beta.0

package/src/endpoints/chat-completions/converters.ts

@@ -11,16 +11,16 @@ import type {
   UserContent,
   AssistantContent,
   LanguageModelUsage,
-  Output,
   TextStreamPart,
   ReasoningOutput,
+  JSONValue,
   AssistantModelMessage,
   ToolModelMessage,
   UserModelMessage,
 } from "ai";
 
-import { convertBase64ToUint8Array } from "@ai-sdk/provider-utils";
-import { jsonSchema, tool } from "ai";
+import { Output, jsonSchema, tool } from "ai";
+import { z } from "zod";
 
 import type {
   ChatCompletionsToolCall,
@@ -44,6 +44,8 @@ import type {
   ChatCompletionsReasoningEffort,
   ChatCompletionsReasoningConfig,
   ChatCompletionsReasoningDetail,
+  ChatCompletionsResponseFormat,
+  ChatCompletionsContentPartText,
 } from "./schema";
 
 import { GatewayError } from "../../errors/gateway";
@@ -54,6 +56,7 @@ export type TextCallOptions = {
   messages: ModelMessage[];
   tools?: ToolSet;
   toolChoice?: ToolChoice<ToolSet>;
+  output?: Output.Output;
   temperature?: number;
   maxOutputTokens?: number;
   frequencyPenalty?: number;
@@ -74,6 +77,7 @@ export function convertToTextCallOptions(params: ChatCompletionsInputs): TextCal
     temperature,
     max_tokens,
     max_completion_tokens,
+    response_format,
     reasoning_effort,
     reasoning,
     frequency_penalty,
@@ -90,6 +94,7 @@ export function convertToTextCallOptions(params: ChatCompletionsInputs): TextCal
     messages: convertToModelMessages(messages),
     tools: convertToToolSet(tools),
     toolChoice: convertToToolChoice(tool_choice),
+    output: convertToOutput(response_format),
     temperature,
     maxOutputTokens: max_completion_tokens ?? max_tokens,
     frequencyPenalty: frequency_penalty,
@@ -103,6 +108,19 @@ export function convertToTextCallOptions(params: ChatCompletionsInputs): TextCal
   };
 }
 
+function convertToOutput(responseFormat: ChatCompletionsResponseFormat | undefined) {
+  if (!responseFormat || responseFormat.type === "text") {
+    return;
+  }
+
+  const { name, description, schema } = responseFormat.json_schema;
+  return Output.object({
+    name,
+    description,
+    schema: jsonSchema(schema),
+  });
+}
+
 export function convertToModelMessages(messages: ChatCompletionsMessage[]): ModelMessage[] {
   const modelMessages: ModelMessage[] = [];
   const toolById = indexToolMessages(messages);
@@ -154,65 +172,73 @@ export function fromChatCompletionsAssistantMessage(
 
   const parts: AssistantContent = [];
 
-  if (Array.isArray(parts)) {
-    if (reasoning_details?.length) {
-      for (const detail of reasoning_details) {
-        if (detail.text && detail.type === "reasoning.text") {
-          parts.push({
-            type: "reasoning",
-            text: detail.text,
-            providerOptions: detail.signature
-              ? {
-                  unknown: {
-                    signature: detail.signature,
-                  },
-                }
-              : undefined,
-          });
-        } else if (detail.type === "reasoning.encrypted" && detail.data) {
-          parts.push({
-            type: "reasoning",
-            text: "",
-            providerOptions: {
-              unknown: {
-                redactedData: detail.data,
-              },
+  if (reasoning_details?.length) {
+    for (const detail of reasoning_details) {
+      if (detail.text && detail.type === "reasoning.text") {
+        parts.push({
+          type: "reasoning",
+          text: detail.text,
+          providerOptions: detail.signature
+            ? {
+                unknown: {
+                  signature: detail.signature,
+                },
+              }
+            : undefined,
+        });
+      } else if (detail.type === "reasoning.encrypted" && detail.data) {
+        parts.push({
+          type: "reasoning",
+          text: "",
+          providerOptions: {
+            unknown: {
+              redactedData: detail.data,
             },
-          });
-        }
+          },
+        });
       }
     }
+  }
 
-    if (tool_calls?.length) {
-      for (const tc of tool_calls) {
-        // eslint-disable-next-line no-shadow
-        const { id, function: fn, extra_content } = tc;
-        const out: ToolCallPart = {
-          type: "tool-call",
-          toolCallId: id,
-          toolName: fn.name,
-          input: parseToolOutput(fn.arguments).value,
-        };
-        if (extra_content) {
-          out.providerOptions = extra_content;
-        }
-        parts.push(out);
+  if (content !== undefined && content !== null) {
+    const inputContent =
+      typeof content === "string"
+        ? ([{ type: "text", text: content }] as ChatCompletionsContentPartText[])
+        : content;
+    for (const part of inputContent) {
+      if (part.type === "text") {
+        parts.push({
+          type: "text",
+          text: part.text,
+        });
       }
-    } else if (content !== undefined && content !== null) {
-      parts.push({
-        type: "text",
-        text: content,
-      });
+    }
+  }
+
+  if (tool_calls?.length) {
+    for (const tc of tool_calls) {
+      // eslint-disable-next-line no-shadow
+      const { id, function: fn, extra_content } = tc;
+      const out: ToolCallPart = {
+        type: "tool-call",
+        toolCallId: id,
+        toolName: fn.name,
+        input: parseJsonOrText(fn.arguments).value,
+      };
+      if (extra_content) {
+        out.providerOptions = extra_content as SharedV3ProviderOptions;
+      }
+      parts.push(out);
     }
   }
 
   const out: AssistantModelMessage = {
-    role: role,
-    content: Array.isArray(parts) && parts.length > 0 ? parts : (content ?? ""),
+    role,
+    content: parts.length > 0 ? parts : (content ?? ""),
   };
 
   if (extra_content) {
-    out.providerOptions = extra_content;
+    out.providerOptions = extra_content as SharedV3ProviderOptions;
   }
 
   return out;
@@ -234,7 +260,7 @@ export function fromChatCompletionsToolResultMessage(
       type: "tool-result",
       toolCallId: tc.id,
       toolName: tc.function.name,
-      output: parseToolOutput(toolMsg.content),
+      output: parseToolResult(toolMsg.content),
     });
   }
 
@@ -243,48 +269,48 @@ export function fromChatCompletionsToolResultMessage(
 
 export function fromChatCompletionsContent(content: ChatCompletionsContentPart[]): UserContent {
   return content.map((part) => {
-    if (part.type === "image_url") {
-      const url = part.image_url.url;
-      if (url.startsWith("data:")) {
-        const { mimeType, base64Data } = parseDataUrl(url);
-
-        return mimeType.startsWith("image/")
-          ? {
-              type: "image" as const,
-              image: convertBase64ToUint8Array(base64Data),
-              mediaType: mimeType,
-            }
-          : {
-              type: "file" as const,
-              data: convertBase64ToUint8Array(base64Data),
-              mediaType: mimeType,
-            };
-      }
-
-      return {
-        type: "image" as const,
-        image: new URL(url),
-      };
+    switch (part.type) {
+      case "image_url":
+        return fromImageUrlPart(part.image_url.url);
+      case "file":
+        return fromFilePart(part.file.data, part.file.media_type, part.file.filename);
+      case "input_audio":
+        return fromFilePart(part.input_audio.data, `audio/${part.input_audio.format}`);
+      default:
+        return part;
     }
-    if (part.type === "file") {
-      let { data, media_type, filename } = part.file;
-      return media_type.startsWith("image/")
-        ? {
-            type: "image" as const,
-            image: convertBase64ToUint8Array(data),
-            mediaType: media_type,
-          }
-        : {
-            type: "file" as const,
-            data: convertBase64ToUint8Array(data),
-            filename,
-            mediaType: media_type,
-          };
-    }
-    return part;
   });
 }
 
+function fromImageUrlPart(url: string) {
+  if (url.startsWith("data:")) {
+    const { mimeType, base64Data } = parseDataUrl(url);
+    return fromFilePart(base64Data, mimeType);
+  }
+
+  return {
+    type: "image" as const,
+    image: new URL(url),
+  };
+}
+
+function fromFilePart(base64Data: string, mediaType: string, filename?: string) {
+  if (mediaType.startsWith("image/")) {
+    return {
+      type: "image" as const,
+      image: z.util.base64ToUint8Array(base64Data),
+      mediaType,
+    };
+  }
+
+  return {
+    type: "file" as const,
+    data: z.util.base64ToUint8Array(base64Data),
+    filename,
+    mediaType,
+  };
+}
+
 export const convertToToolSet = (tools: ChatCompletionsTool[] | undefined): ToolSet | undefined => {
   if (!tools) {
     return;
@@ -311,17 +337,39 @@ export const convertToToolChoice = (
     return toolChoice;
   }
 
+  // FUTURE: this is right now google specific, which is not supported by AI SDK, until then, we temporarily map it to auto for now https://docs.cloud.google.com/vertex-ai/generative-ai/docs/migrate/openai/overview
+  if (toolChoice === "validated") {
+    return "auto";
+  }
+
   return {
     type: "tool",
     toolName: toolChoice.function.name,
   };
 };
 
-function parseToolOutput(content: string) {
+function parseToolResult(
+  content: string | ChatCompletionsContentPartText[],
+): ToolResultPart["output"] {
+  if (Array.isArray(content)) {
+    return {
+      type: "content",
+      value: content.map((part) => ({
+        type: "text",
+        text: part.text,
+      })),
+    };
+  }
+  return parseJsonOrText(content);
+}
+
+function parseJsonOrText(
+  content: string,
+): { type: "json"; value: JSONValue } | { type: "text"; value: string } {
   try {
-    return { type: "json" as const, value: JSON.parse(content) };
+    return { type: "json", value: JSON.parse(content) };
   } catch {
-    return { type: "text" as const, value: content };
+    return { type: "text", value: content };
   }
 }
 
@@ -379,8 +427,6 @@ export function toChatCompletions(
   result: GenerateTextResult<ToolSet, Output.Output>,
   model: string,
 ): ChatCompletions {
-  const finish_reason = toChatCompletionsFinishReason(result.finishReason);
-
   return {
     id: "chatcmpl-" + crypto.randomUUID(),
     object: "chat.completion",
@@ -390,7 +436,7 @@ export function toChatCompletions(
       {
         index: 0,
         message: toChatCompletionsAssistantMessage(result),
-        finish_reason,
+        finish_reason: toChatCompletionsFinishReason(result.finishReason),
       } satisfies ChatCompletionsChoice,
     ],
     usage: result.totalUsage ? toChatCompletionsUsage(result.totalUsage) : null,
@@ -430,6 +476,7 @@ export class ChatCompletionsStream<E extends boolean = false> extends TransformS
     const creationTime = Math.floor(Date.now() / 1000);
     let toolCallIndexCounter = 0;
     const reasoningIdToIndex = new Map<string, number>();
+    let finishProviderMetadata: SharedV3ProviderMetadata | undefined;
 
     const createChunk = (
       delta: ChatCompletionsAssistantMessageDelta,
@@ -512,14 +559,7 @@ export class ChatCompletionsStream<E extends boolean = false> extends TransformS
           }
 
           case "finish-step": {
-            controller.enqueue(
-              createChunk(
-                {},
-                part.providerMetadata,
-                toChatCompletionsFinishReason(part.finishReason),
-                toChatCompletionsUsage(part.usage),
-              ),
-            );
+            finishProviderMetadata = part.providerMetadata;
             break;
           }
 
@@ -527,7 +567,7 @@ export class ChatCompletionsStream<E extends boolean = false> extends TransformS
             controller.enqueue(
               createChunk(
                 {},
-                undefined,
+                finishProviderMetadata,
                 toChatCompletionsFinishReason(part.finishReason),
                 toChatCompletionsUsage(part.totalUsage),
               ),

package/src/endpoints/chat-completions/handler.test.ts

@@ -12,6 +12,7 @@ describe("Chat Completions Handler", () => {
   const mockLanguageModel = new MockLanguageModelV3({
     // eslint-disable-next-line require-await
     doGenerate: async (options) => {
+      const isStructuredOutput = options.responseFormat?.type === "json";
       const isToolCall = options.tools && options.tools.length > 0;
 
       if (isToolCall) {
@@ -34,6 +35,24 @@ describe("Chat Completions Handler", () => {
         };
       }
 
+      if (isStructuredOutput) {
+        return {
+          finishReason: { unified: "stop", raw: "stop" },
+          usage: {
+            inputTokens: { total: 10, noCache: 10, cacheRead: 20, cacheWrite: 0 },
+            outputTokens: { total: 20, text: 20, reasoning: 10 },
+          },
+          content: [
+            {
+              type: "text",
+              text: '{"city":"San Francisco","temp_c":18}',
+            },
+          ],
+          providerMetadata: { provider: { key: "value" } },
+          warnings: [],
+        };
+      }
+
       return {
         finishReason: { unified: "stop", raw: "stop" },
         usage: {
@@ -184,6 +203,31 @@ describe("Chat Completions Handler", () => {
     });
   });
 
+  test("should accept input_audio content parts", async () => {
+    const request = postJson(baseUrl, {
+      model: "openai/gpt-oss-20b",
+      messages: [
+        {
+          role: "user",
+          content: [
+            {
+              type: "input_audio",
+              input_audio: {
+                data: "aGVsbG8=",
+                format: "wav",
+              },
+            },
+          ],
+        },
+      ],
+    });
+
+    const res = await endpoint.handler(request);
+    expect(res.status).toBe(200);
+    const data = await parseResponse(res);
+    expect(data.model).toBe("openai/gpt-oss-20b");
+  });
+
   test("should generate completion with tool calls successfully", async () => {
     const request = postJson(baseUrl, {
       model: "openai/gpt-oss-20b",
@@ -299,4 +343,47 @@ describe("Chat Completions Handler", () => {
     const data = await parseResponse(res);
     expect(data.model).toBe("openai/gpt-oss-20b");
   });
+
+  test("should generate non-streaming structured output", async () => {
+    const request = postJson(baseUrl, {
+      model: "openai/gpt-oss-20b",
+      messages: [{ role: "user", content: "Return weather as JSON" }],
+      response_format: {
+        type: "json_schema",
+        json_schema: {
+          name: "weather",
+          schema: {
+            type: "object",
+            properties: {
+              city: { type: "string" },
+              temp_c: { type: "number" },
+            },
+            required: ["city", "temp_c"],
+            additionalProperties: false,
+          },
+          strict: true,
+        },
+      },
+    });
+
+    const res = await endpoint.handler(request);
+    expect(res.status).toBe(200);
+    const data = await parseResponse(res);
+    expect(data.choices[0].message.content).toBe('{"city":"San Francisco","temp_c":18}');
+  });
+
+  test('should accept response_format type "text"', async () => {
+    const request = postJson(baseUrl, {
+      model: "openai/gpt-oss-20b",
+      messages: [{ role: "user", content: "Say hi" }],
+      response_format: {
+        type: "text",
+      },
+    });
+
+    const res = await endpoint.handler(request);
+    expect(res.status).toBe(200);
+    const data = await parseResponse(res);
+    expect(data.choices[0].message.content).toBe("Hello from AI");
+  });
 });

package/src/endpoints/chat-completions/handler.ts

@@ -29,7 +29,6 @@ import {
   recordTokenUsage,
 } from "../../telemetry/gen-ai";
 import { addSpanEvent, setSpanAttributes } from "../../telemetry/span";
-import { resolveRequestId } from "../../utils/headers";
 import { prepareForwardHeaders } from "../../utils/request";
 import { convertToTextCallOptions, toChatCompletions, toChatCompletionsStream } from "./converters";
 import {
@@ -52,8 +51,6 @@ export const chatCompletions = (config: GatewayConfig): Endpoint => {
       throw new GatewayError("Method Not Allowed", 405);
     }
 
-    const requestId = resolveRequestId(ctx.request);
-
     // Parse + validate input.
     try {
       ctx.body = await ctx.request.json();
@@ -107,7 +104,7 @@ export const chatCompletions = (config: GatewayConfig): Endpoint => {
     const textOptions = convertToTextCallOptions(inputs);
     logger.trace(
       {
-        requestId,
+        requestId: ctx.requestId,
         options: textOptions,
       },
       "[chat] AI SDK options",
@@ -178,7 +175,7 @@ export const chatCompletions = (config: GatewayConfig): Endpoint => {
       },
       ...textOptions,
     });
-    logger.trace({ requestId, result }, "[chat] AI SDK result");
+    logger.trace({ requestId: ctx.requestId, result }, "[chat] AI SDK result");
     addSpanEvent("hebo.ai-sdk.completed");
 
     // Transform result.

package/src/endpoints/chat-completions/otel.ts

@@ -41,6 +41,12 @@ const toMessageParts = (message: ChatCompletionsMessage): Record<string, unknown
           parts.push(toTextPart(part.text));
         } else if (part.type === "image_url") {
           parts.push({ type: "image", content: part.image_url.url });
+        } else if (part.type === "input_audio") {
+          parts.push({
+            type: "audio",
+            content: "[REDACTED_BINARY_DATA]",
+            format: part.input_audio.format,
+          });
         } else {
           parts.push({
             type: "file",