@hebo-ai/gateway 0.1.2 → 0.2.0

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

@@ -1,36 +1,35 @@
-import { generateText, streamText } from "ai";
+import { generateText, streamText, wrapLanguageModel } from "ai";
 import * as z from "zod/mini";
 
-import type { GatewayConfig, Endpoint } from "../../types";
+import type { GatewayConfig, Endpoint, GatewayContext } from "../../types";
 
-import { parseConfig } from "../../config";
+import { withLifecycle } from "../../lifecycle";
+import { modelMiddlewareMatcher } from "../../middleware/matcher";
 import { resolveProvider } from "../../providers/registry";
 import { createErrorResponse } from "../../utils/errors";
-import { withHooks } from "../../utils/hooks";
 import {
-  transformCompletionsInputs,
-  createCompletionsResponse,
-  createCompletionsStreamResponse,
+  convertToTextCallOptions,
+  toChatCompletionsResponse,
+  toChatCompletionsStreamResponse,
 } from "./converters";
-import { CompletionsBodySchema } from "./schema";
+import { ChatCompletionsBodySchema } from "./schema";
 
 export const chatCompletions = (config: GatewayConfig): Endpoint => {
-  const { providers, models, hooks } = parseConfig(config);
+  const hooks = config.hooks;
 
-  const handler = async (req: Request): Promise<Response> => {
-    if (req.method !== "POST") {
+  const handler = async (ctx: GatewayContext): Promise<Response> => {
+    if (!ctx.request || ctx.request.method !== "POST") {
       return createErrorResponse("METHOD_NOT_ALLOWED", "Method Not Allowed", 405);
     }
 
     let body;
     try {
-      body = await req.json();
+      body = await ctx.request.json();
     } catch {
       return createErrorResponse("BAD_REQUEST", "Invalid JSON", 400);
     }
 
-    const parsed = CompletionsBodySchema.safeParse(body);
-
+    const parsed = ChatCompletionsBodySchema.safeParse(body);
     if (!parsed.success) {
       return createErrorResponse(
         "UNPROCESSABLE_ENTITY",
@@ -39,65 +38,72 @@ export const chatCompletions = (config: GatewayConfig): Endpoint => {
         z.prettifyError(parsed.error),
       );
     }
+    ctx.body = parsed.data;
 
-    const { model: modelId, stream, ...inputs } = parsed.data;
+    let stream, inputs;
+    ({ model: ctx.modelId, stream, ...inputs } = parsed.data);
 
-    let resolvedModelId;
     try {
-      resolvedModelId = (await hooks?.resolveModelId?.({ modelId })) ?? modelId;
+      ctx.resolvedModelId = (await hooks?.resolveModelId?.(ctx)) ?? ctx.modelId;
     } catch (error) {
       return createErrorResponse("BAD_REQUEST", error, 400);
     }
 
-    let textOptions;
+    ctx.operation = "text";
     try {
-      textOptions = transformCompletionsInputs(inputs);
+      const override = await hooks?.resolveProvider?.(ctx);
+      ctx.provider =
+        override ??
+        resolveProvider({
+          providers: ctx.providers,
+          models: ctx.models,
+          modelId: ctx.resolvedModelId,
+          operation: ctx.operation,
+        });
     } catch (error) {
       return createErrorResponse("BAD_REQUEST", error, 400);
     }
 
-    let provider;
+    const languageModel = ctx.provider.languageModel(ctx.resolvedModelId);
+
+    let textOptions;
     try {
-      const args = {
-        providers,
-        models,
-        modelId: resolvedModelId,
-        operation: "text" as const,
-      };
-      const override = await hooks?.resolveProvider?.(args);
-      provider = override ?? resolveProvider(args);
+      textOptions = convertToTextCallOptions(inputs);
     } catch (error) {
       return createErrorResponse("BAD_REQUEST", error, 400);
     }
 
-    const languageModel = provider.languageModel(resolvedModelId);
+    const languageModelWithMiddleware = wrapLanguageModel({
+      model: languageModel,
+      middleware: modelMiddlewareMatcher.for(ctx.resolvedModelId, languageModel.provider),
+    });
 
     if (stream) {
       let result;
       try {
         result = streamText({
-          model: languageModel,
+          model: languageModelWithMiddleware,
           ...textOptions,
         });
       } catch (error) {
         return createErrorResponse("INTERNAL_SERVER_ERROR", error, 500);
       }
 
-      return createCompletionsStreamResponse(result, modelId);
+      return toChatCompletionsStreamResponse(result, ctx.modelId);
     }
 
     let result;
     try {
       result = await generateText({
-        model: languageModel,
+        model: languageModelWithMiddleware,
         ...textOptions,
       });
     } catch (error) {
       return createErrorResponse("INTERNAL_SERVER_ERROR", error, 500);
     }
 
-    return createCompletionsResponse(result, modelId);
+    return toChatCompletionsResponse(result, ctx.modelId);
   };
 
-  return { handler: withHooks(hooks, handler) };
+  return { handler: withLifecycle(handler, config) };
 };

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

@@ -1,33 +1,34 @@
-import * as z from "zod/mini";
+import * as z from "zod";
 
-export const CompletionsContentPartTextSchema = z.object({
+export const ChatCompletionsContentPartTextSchema = z.object({
   type: z.literal("text"),
   text: z.string(),
 });
 
-export const CompletionsContentPartImageSchema = z.object({
+export const ChatCompletionsContentPartImageSchema = z.object({
   type: z.literal("image_url"),
   image_url: z.object({
     url: z.string(),
-    detail: z.optional(z.union([z.literal("low"), z.literal("high"), z.literal("auto")])),
+    detail: z.union([z.literal("low"), z.literal("high"), z.literal("auto")]).optional(),
   }),
 });
 
-export const CompletionsContentPartFileSchema = z.object({
+export const ChatCompletionsContentPartFileSchema = z.object({
   type: z.literal("file"),
   file: z.object({
     data: z.string(),
     media_type: z.string(),
-    filename: z.string(),
+    filename: z.string().optional(),
   }),
 });
 
-export type CompletionsContentPart =
-  | z.infer<typeof CompletionsContentPartTextSchema>
-  | z.infer<typeof CompletionsContentPartImageSchema>
-  | z.infer<typeof CompletionsContentPartFileSchema>;
+// FUTURE: missing ContentPartAudio
+export type ChatCompletionsContentPart =
+  | z.infer<typeof ChatCompletionsContentPartTextSchema>
+  | z.infer<typeof ChatCompletionsContentPartImageSchema>
+  | z.infer<typeof ChatCompletionsContentPartFileSchema>;
 
-export const CompletionsToolCallSchema = z.object({
+export const ChatCompletionsToolCallSchema = z.object({
   type: z.literal("function"),
   id: z.string(),
   function: z.object({
@@ -35,67 +36,75 @@ export const CompletionsToolCallSchema = z.object({
     name: z.string(),
   }),
 });
-export type CompletionsToolCall = z.infer<typeof CompletionsToolCallSchema>;
+export type ChatCompletionsToolCall = z.infer<typeof ChatCompletionsToolCallSchema>;
 
-export const CompletionsSystemMessageSchema = z.object({
+export const ChatCompletionsSystemMessageSchema = z.object({
   role: z.literal("system"),
   content: z.string(),
+  name: z.string().optional(),
 });
-export type CompletionsSystemMessage = z.infer<typeof CompletionsSystemMessageSchema>;
+export type ChatCompletionsSystemMessage = z.infer<typeof ChatCompletionsSystemMessageSchema>;
 
-export const CompletionsUserMessageSchema = z.object({
+export const ChatCompletionsUserMessageSchema = z.object({
   role: z.literal("user"),
   content: z.union([
     z.string(),
     z.array(
       z.union([
-        CompletionsContentPartTextSchema,
-        CompletionsContentPartImageSchema,
-        CompletionsContentPartFileSchema,
+        ChatCompletionsContentPartTextSchema,
+        ChatCompletionsContentPartImageSchema,
+        ChatCompletionsContentPartFileSchema,
       ]),
     ),
   ]),
+  name: z.string().optional(),
 });
-export type CompletionsUserMessage = z.infer<typeof CompletionsUserMessageSchema>;
+export type ChatCompletionsUserMessage = z.infer<typeof ChatCompletionsUserMessageSchema>;
 
-export const CompletionsAssistantMessageSchema = z.object({
+export const ChatCompletionsAssistantMessageSchema = z.object({
   role: z.literal("assistant"),
-  content: z.union([z.string(), z.null()]),
-  tool_calls: z.optional(z.array(CompletionsToolCallSchema)),
-  reasoning: z.optional(z.string()),
-  reasoning_content: z.optional(z.string()),
+  // FUTURE: this should support arrays of TextContentPart and RefusalContentPart
+  content: z.union([z.string(), z.null()]).optional(),
+  name: z.string().optional(),
+  // FUTURE: This should also support Custom Tool Calls
+  tool_calls: z.array(ChatCompletionsToolCallSchema).optional(),
+  // Extensions
+  reasoning_content: z.string().optional().meta({ extension: true }),
 });
-export type CompletionsAssistantMessage = z.infer<typeof CompletionsAssistantMessageSchema>;
+export type ChatCompletionsAssistantMessage = z.infer<typeof ChatCompletionsAssistantMessageSchema>;
 
-export const CompletionsToolMessageSchema = z.object({
+export const ChatCompletionsToolMessageSchema = z.object({
   role: z.literal("tool"),
+  // FUTURE: this should also support arrays of TextContentParts
   content: z.string(),
   tool_call_id: z.string(),
 });
-export type CompletionsToolMessage = z.infer<typeof CompletionsToolMessageSchema>;
+export type ChatCompletionsToolMessage = z.infer<typeof ChatCompletionsToolMessageSchema>;
 
-export const CompletionsMessageSchema = z.union([
-  CompletionsSystemMessageSchema,
-  CompletionsUserMessageSchema,
-  CompletionsAssistantMessageSchema,
-  CompletionsToolMessageSchema,
+export const ChatCompletionsMessageSchema = z.union([
+  ChatCompletionsSystemMessageSchema,
+  ChatCompletionsUserMessageSchema,
+  ChatCompletionsAssistantMessageSchema,
+  ChatCompletionsToolMessageSchema,
 ]);
-export type CompletionsMessage = z.infer<typeof CompletionsMessageSchema>;
+export type ChatCompletionsMessage = z.infer<typeof ChatCompletionsMessageSchema>;
 
-export const CompletionsToolSchema = z.object({
+export const ChatCompletionsToolSchema = z.object({
   type: z.literal("function"),
   function: z.object({
     name: z.string(),
-    description: z.optional(z.string()),
+    description: z.string().optional(),
     parameters: z.record(z.string(), z.any()),
+    // Missing strict parameter
   }),
 });
-export type CompletionsTool = z.infer<typeof CompletionsToolSchema>;
+export type ChatCompletionsTool = z.infer<typeof ChatCompletionsToolSchema>;
 
-export const CompletionsToolChoiceSchema = z.union([
+export const ChatCompletionsToolChoiceSchema = z.union([
   z.literal("none"),
   z.literal("auto"),
   z.literal("required"),
+  // FUTURE: missing AllowedTools and CustomToolChoice
   z.object({
     type: z.literal("function"),
     function: z.object({
@@ -103,70 +112,134 @@ export const CompletionsToolChoiceSchema = z.union([
     }),
   }),
 ]);
-export type CompletionsToolChoice = z.infer<typeof CompletionsToolChoiceSchema>;
+export type ChatCompletionsToolChoice = z.infer<typeof ChatCompletionsToolChoiceSchema>;
 
-export const CompletionsInputsSchema = z.object({
-  messages: z.array(CompletionsMessageSchema),
-  tools: z.optional(z.array(CompletionsToolSchema)),
-  tool_choice: z.optional(CompletionsToolChoiceSchema),
-  temperature: z.optional(z.number()),
+export const ChatCompletionsReasoningEffortSchema = z.union([
+  z.literal("none"),
+  z.literal("minimal"),
+  z.literal("low"),
+  z.literal("medium"),
+  z.literal("high"),
+  z.literal("xhigh"),
+]);
+export type ChatCompletionsReasoningEffort = z.infer<typeof ChatCompletionsReasoningEffortSchema>;
+
+export const ChatCompletionsReasoningConfigSchema = z.object({
+  enabled: z.optional(z.boolean()),
+  effort: z.optional(ChatCompletionsReasoningEffortSchema),
+  max_tokens: z.optional(z.number()),
+  exclude: z.optional(z.boolean()),
+});
+export type ChatCompletionsReasoningConfig = z.infer<typeof ChatCompletionsReasoningConfigSchema>;
+
+const ChatCompletionsInputsSchema = z.object({
+  messages: z.array(ChatCompletionsMessageSchema),
+  tools: z
+    .array(
+      // FUTURE: Missing CustomTool
+      ChatCompletionsToolSchema,
+    )
+    .optional(),
+  tool_choice: ChatCompletionsToolChoiceSchema.optional(),
+  temperature: z.number().min(0).max(2).optional(),
+  max_tokens: z.int().nonnegative().optional(),
+  max_completion_tokens: z.int().nonnegative().optional(),
+  frequency_penalty: z.number().min(-2.0).max(2.0).optional(),
+  presence_penalty: z.number().min(-2.0).max(2.0).optional(),
+  seed: z.int().optional(),
+  stop: z.union([z.string(), z.array(z.string())]).optional(),
+  top_p: z.number().min(0).max(1.0).optional(),
+  reasoning_effort: ChatCompletionsReasoningEffortSchema.optional(),
+  // Extensions
+  reasoning: ChatCompletionsReasoningConfigSchema.optional().meta({ extension: true }),
 });
-export type CompletionsInputs = z.infer<typeof CompletionsInputsSchema>;
+export type ChatCompletionsInputs = z.infer<typeof ChatCompletionsInputsSchema>;
 
-export const CompletionsBodySchema = z.extend(CompletionsInputsSchema, {
+export const ChatCompletionsBodySchema = z.looseObject({
   model: z.string(),
-  stream: z.optional(z.boolean()),
+  stream: z.boolean().optional(),
+  ...ChatCompletionsInputsSchema.shape,
 });
-export type CompletionsBody = z.infer<typeof CompletionsBodySchema>;
+export type ChatCompletionsBody = z.infer<typeof ChatCompletionsBodySchema>;
 
-export const CompletionsFinishReasonSchema = z.union([
+export const ChatCompletionsFinishReasonSchema = z.union([
   z.literal("stop"),
   z.literal("length"),
   z.literal("content_filter"),
   z.literal("tool_calls"),
 ]);
-export type CompletionsFinishReason = z.infer<typeof CompletionsFinishReasonSchema>;
-
-export const CompletionsChoiceSchema = z.object({
-  index: z.number(),
-  message: CompletionsAssistantMessageSchema,
-  finish_reason: CompletionsFinishReasonSchema,
-  logprobs: z.optional(z.any()),
-});
-export type CompletionsChoice = z.infer<typeof CompletionsChoiceSchema>;
-
-export const CompletionsUsageSchema = z.object({
-  prompt_tokens: z.number(),
-  completion_tokens: z.number(),
-  total_tokens: z.number(),
-  completion_tokens_details: z.optional(
-    z.object({
-      reasoning_tokens: z.optional(z.number()),
-    }),
-  ),
-  prompt_tokens_details: z.optional(
-    z.object({
-      cached_tokens: z.optional(z.number()),
-    }),
-  ),
+export type ChatCompletionsFinishReason = z.infer<typeof ChatCompletionsFinishReasonSchema>;
+
+export const ChatCompletionsChoiceSchema = z.object({
+  index: z.int().nonnegative(),
+  message: ChatCompletionsAssistantMessageSchema,
+  finish_reason: ChatCompletionsFinishReasonSchema,
+  // FUTURE: model this out
+  logprobs: z.any().optional(),
 });
-export type CompletionsUsage = z.infer<typeof CompletionsUsageSchema>;
+export type ChatCompletionsChoice = z.infer<typeof ChatCompletionsChoiceSchema>;
+
+export const ChatCompletionsUsageSchema = z.object({
+  prompt_tokens: z.int().nonnegative().optional(),
+  completion_tokens: z.int().nonnegative().optional(),
+  total_tokens: z.int().nonnegative().optional(),
+  completion_tokens_details: z
+    .object({
+      // FUTURE: add missing properties
+      reasoning_tokens: z.int().nonnegative().optional(),
+    })
+    .optional(),
+  prompt_tokens_details: z
+    .object({
+      // FUTURE: add missing properties
+      cached_tokens: z.int().nonnegative().optional(),
+    })
+    .optional(),
+});
+export type ChatCompletionsUsage = z.infer<typeof ChatCompletionsUsageSchema>;
 
-export const CompletionsSchema = z.object({
+export const ChatCompletionsSchema = z.object({
   id: z.string(),
   object: z.literal("chat.completion"),
-  created: z.number(),
+  created: z.int().nonnegative(),
   model: z.string(),
-  choices: z.array(CompletionsChoiceSchema),
-  usage: z.optional(CompletionsUsageSchema),
-  system_fingerprint: z.optional(z.string()),
-  providerMetadata: z.optional(z.any()),
-});
-export type Completions = z.infer<typeof CompletionsSchema>;
-
-export type CompletionsToolCallDelta = {
-  id: string;
-  index: number;
-  type: "function";
-  function: { name: string; arguments: string };
-};
+  choices: z.array(ChatCompletionsChoiceSchema),
+  usage: ChatCompletionsUsageSchema.nullable(),
+  // Extensions
+  provider_metadata: z.any().optional().meta({ extension: true }),
+});
+export type ChatCompletions = z.infer<typeof ChatCompletionsSchema>;
+
+export const ChatCompletionsToolCallDeltaSchema = ChatCompletionsToolCallSchema.partial().extend({
+  index: z.int().nonnegative(),
+});
+export type ChatCompletionsToolCallDelta = z.infer<typeof ChatCompletionsToolCallDeltaSchema>;
+
+export const ChatCompletionsAssistantMessageDeltaSchema =
+  ChatCompletionsAssistantMessageSchema.partial().extend({
+    tool_calls: z.array(ChatCompletionsToolCallDeltaSchema).optional(),
+  });
+export type ChatCompletionsAssistantMessageDelta = z.infer<
+  typeof ChatCompletionsAssistantMessageDeltaSchema
+>;
+
+export const ChatCompletionsChoiceDeltaSchema = z.object({
+  index: z.int().nonnegative(),
+  delta: ChatCompletionsAssistantMessageDeltaSchema,
+  finish_reason: ChatCompletionsFinishReasonSchema.nullable(),
+  // FUTURE: model this out
+  logprobs: z.any().optional(),
+});
+export type ChatCompletionsChoiceDelta = z.infer<typeof ChatCompletionsChoiceDeltaSchema>;
+
+export const ChatCompletionsChunkSchema = z.object({
+  id: z.string(),
+  object: z.literal("chat.completion.chunk"),
+  created: z.int().nonnegative(),
+  model: z.string(),
+  choices: z.array(ChatCompletionsChoiceDeltaSchema),
+  usage: ChatCompletionsUsageSchema.nullable(),
+  // Extensions
+  provider_metadata: z.any().optional().meta({ extension: true }),
+});
+export type ChatCompletionsChunk = z.infer<typeof ChatCompletionsChunkSchema>;

package/src/endpoints/embeddings/converters.ts

@@ -1,20 +1,22 @@
-import type { ProviderOptions } from "@ai-sdk/provider-utils";
+import type { JSONObject, SharedV3ProviderOptions } from "@ai-sdk/provider";
 import type { EmbedManyResult } from "ai";
 
-import type { EmbeddingsInputs, Embeddings, EmbeddingsData, EmbeddingsUsage } from "./schema";
+import type { EmbeddingsInputs, EmbeddingsData, EmbeddingsUsage, Embeddings } from "./schema";
+
+import { mergeResponseInit } from "../../utils/response";
 
 export type EmbedCallOptions = {
   values: string[];
-  providerOptions: ProviderOptions;
+  providerOptions: SharedV3ProviderOptions;
 };
 
-export function transformEmbeddingsInputs(params: EmbeddingsInputs): EmbedCallOptions {
+export function convertToEmbedCallOptions(params: EmbeddingsInputs): EmbedCallOptions {
   const { input, ...rest } = params;
 
   return {
     values: Array.isArray(input) ? input : [input],
     providerOptions: {
-      openAICompat: rest,
+      unknown: rest as JSONObject,
     },
   };
 }
@@ -27,8 +29,8 @@ export function toEmbeddings(embedManyResult: EmbedManyResult, modelId: string):
   }));
 
   const usage: EmbeddingsUsage = {
-    prompt_tokens: embedManyResult.usage?.tokens || 0,
-    total_tokens: embedManyResult.usage?.tokens || 0,
+    prompt_tokens: embedManyResult.usage.tokens,
+    total_tokens: embedManyResult.usage.tokens,
   };
 
   return {
@@ -36,15 +38,17 @@ export function toEmbeddings(embedManyResult: EmbedManyResult, modelId: string):
     data,
     model: modelId,
     usage,
-    providerMetadata: embedManyResult.providerMetadata,
+    provider_metadata: embedManyResult.providerMetadata,
   };
 }
 
 export function createEmbeddingsResponse(
   embedManyResult: EmbedManyResult,
   modelId: string,
+  responseInit?: ResponseInit,
 ): Response {
-  return new Response(JSON.stringify(toEmbeddings(embedManyResult, modelId)), {
-    headers: { "Content-Type": "application/json" },
-  });
+  return new Response(
+    JSON.stringify(toEmbeddings(embedManyResult, modelId)),
+    mergeResponseInit({ "Content-Type": "application/json" }, responseInit),
+  );
 }

package/src/endpoints/embeddings/handler.test.ts

@@ -1,9 +1,7 @@
-import { createProviderRegistry } from "ai";
 import { MockEmbeddingModelV3, MockProviderV3 } from "ai/test";
 import { describe, expect, test } from "bun:test";
 
 import { parseResponse, postJson } from "../../../test/helpers/http";
-import { createModelCatalog } from "../../models/catalog";
 import { embeddings } from "./handler";
 
 const baseUrl = "http://localhost/embeddings";
@@ -20,45 +18,44 @@ const expectedEmbeddingResponse = (count: number) => ({
     prompt_tokens: count * 10,
     total_tokens: count * 10,
   },
-  providerMetadata: {
-    openai: {
+  provider_metadata: {
+    provider: {
       key: "value",
     },
   },
 });
 
 describe("Embeddings Handler", () => {
-  const registry = createProviderRegistry({
-    openai: new MockProviderV3({
-      embeddingModels: {
-        "text-embedding-3-small": new MockEmbeddingModelV3({
-          // eslint-disable-next-line require-await
-          doEmbed: async (options) => ({
-            embeddings: options.values.map(() => [0.1, 0.2, 0.3]),
-            usage: { tokens: 10 },
-            providerMetadata: { openai: { key: "value" } },
-            warnings: [],
+  const endpoint = embeddings({
+    providers: {
+      openai: new MockProviderV3({
+        embeddingModels: {
+          "text-embedding-3-small": new MockEmbeddingModelV3({
+            // eslint-disable-next-line require-await
+            doEmbed: async (options) => ({
+              embeddings: options.values.map(() => [0.1, 0.2, 0.3]),
+              usage: { tokens: 10 },
+              providerMetadata: { provider: { key: "value" } },
+              warnings: [],
+            }),
           }),
-        }),
-      },
-    }),
-  });
-
-  const catalog = createModelCatalog({
-    "text-embedding-3-small": {
-      name: "OpenAI Embedding Model",
-      modalities: { input: ["text"], output: ["embeddings"] },
-      providers: ["openai"],
+        },
+      }),
     },
-    "gpt-oss-20b": {
-      name: "GPT-OSS 20B",
-      modalities: { input: ["text"], output: ["text"] },
-      providers: ["openai"],
+    models: {
+      "text-embedding-3-small": {
+        name: "OpenAI Embedding Model",
+        modalities: { input: ["text"], output: ["embeddings"] },
+        providers: ["openai"],
+      },
+      "gpt-oss-20b": {
+        name: "GPT-OSS 20B",
+        modalities: { input: ["text"], output: ["text"] },
+        providers: ["openai"],
+      },
     },
   });
 
-  const endpoint = embeddings({ providers: registry, models: catalog });
-
   test("should return 400 if model does not support embeddings", async () => {
     const request = postJson(baseUrl, {
       model: "gpt-oss-20b",