@hebo-ai/gateway 0.9.4 → 0.10.1

package/dist/endpoints/messages/handler.js

@@ -0,0 +1,142 @@
+import { generateText, Output, streamText, wrapLanguageModel, } from "ai";
+import * as z from "zod";
+import { GatewayError } from "../../errors/gateway";
+import { winterCgHandler } from "../../lifecycle";
+import { logger } from "../../logger";
+import { modelMiddlewareMatcher } from "../../middleware/matcher";
+import { resolveProvider } from "../../providers/registry";
+import { getGenAiGeneralAttributes, recordTimePerOutputToken, recordTimeToFirstToken, recordTokenUsage, } from "../../telemetry/gen-ai";
+import { addSpanEvent, setSpanAttributes } from "../../telemetry/span";
+import { parseRequestBody } from "../../utils/body";
+import { prepareForwardHeaders } from "../../utils/request";
+import { convertToTextCallOptions, toMessages, toMessagesStream } from "./converters";
+import { getMessagesRequestAttributes, getMessagesResponseAttributes } from "./otel";
+import { MessagesBodySchema } from "./schema";
+export const messages = (config) => {
+    const hooks = config.hooks;
+    const handler = async (ctx, cfg) => {
+        const start = performance.now();
+        ctx.operation = "messages";
+        setSpanAttributes({ "gen_ai.operation.name": ctx.operation });
+        addSpanEvent("hebo.handler.started");
+        if (!ctx.request || ctx.request.method !== "POST") {
+            throw new GatewayError("Method Not Allowed", 405);
+        }
+        // Parse + validate input (handles Content-Encoding decompression + body size limits).
+        ctx.body = (await parseRequestBody(ctx.request, cfg.maxBodySize));
+        logger.trace({ requestId: ctx.requestId, body: ctx.body }, "[messages] MessagesBody");
+        addSpanEvent("hebo.request.deserialized");
+        const parsed = MessagesBodySchema.safeParse(ctx.body);
+        if (!parsed.success) {
+            // FUTURE: consider adding body shape to metadata
+            throw new GatewayError(z.prettifyError(parsed.error), 400, undefined, parsed.error);
+        }
+        ctx.body = parsed.data;
+        addSpanEvent("hebo.request.parsed");
+        if (hooks?.before) {
+            ctx.body = (await hooks.before(ctx)) ?? ctx.body;
+            addSpanEvent("hebo.hooks.before.completed");
+        }
+        ctx.modelId = ctx.body.model;
+        ctx.resolvedModelId =
+            (await hooks?.resolveModelId?.(ctx)) ?? ctx.modelId;
+        logger.debug(`[messages] resolved ${ctx.modelId} to ${ctx.resolvedModelId}`);
+        addSpanEvent("hebo.model.resolved");
+        const override = await hooks?.resolveProvider?.(ctx);
+        ctx.provider =
+            override ??
+                resolveProvider({
+                    providers: ctx.providers,
+                    models: ctx.models,
+                    modelId: ctx.resolvedModelId,
+                    operation: ctx.operation,
+                });
+        const languageModel = ctx.provider.languageModel(ctx.resolvedModelId);
+        ctx.resolvedProviderId = languageModel.provider;
+        logger.debug(`[messages] using ${languageModel.provider} for ${ctx.resolvedModelId}`);
+        addSpanEvent("hebo.provider.resolved");
+        const genAiSignalLevel = cfg.telemetry?.signals?.gen_ai;
+        const genAiGeneralAttrs = getGenAiGeneralAttributes(ctx, genAiSignalLevel);
+        setSpanAttributes(genAiGeneralAttrs);
+        const { model: _model, stream, ...inputs } = ctx.body;
+        const textOptions = convertToTextCallOptions(inputs);
+        logger.trace({ requestId: ctx.requestId, options: textOptions }, "[messages] AI SDK options");
+        addSpanEvent("hebo.options.prepared");
+        setSpanAttributes(getMessagesRequestAttributes(ctx.body, genAiSignalLevel));
+        const languageModelWithMiddleware = wrapLanguageModel({
+            model: languageModel,
+            middleware: modelMiddlewareMatcher.for(ctx.resolvedModelId, languageModel.provider),
+        });
+        if (stream) {
+            addSpanEvent("hebo.ai-sdk.started");
+            let ttft = 0;
+            const result = streamText({
+                model: languageModelWithMiddleware,
+                headers: prepareForwardHeaders(ctx.request),
+                abortSignal: ctx.request.signal,
+                timeout: {
+                    totalMs: cfg.timeouts.normal,
+                },
+                onAbort: () => {
+                    throw new DOMException("The operation was aborted.", "AbortError");
+                },
+                onError: () => { },
+                onChunk: () => {
+                    if (!ttft) {
+                        ttft = performance.now() - start;
+                        recordTimeToFirstToken(ttft, genAiGeneralAttrs, genAiSignalLevel);
+                    }
+                },
+                onFinish: (res) => {
+                    addSpanEvent("hebo.ai-sdk.completed");
+                    const streamResult = toMessages(res, ctx.resolvedModelId);
+                    logger.trace({ requestId: ctx.requestId, result: streamResult }, "[messages] Messages");
+                    addSpanEvent("hebo.result.transformed");
+                    const genAiResponseAttrs = getMessagesResponseAttributes(streamResult, genAiSignalLevel, res.finishReason);
+                    setSpanAttributes(genAiResponseAttrs);
+                    recordTokenUsage(genAiResponseAttrs, genAiGeneralAttrs, genAiSignalLevel);
+                    recordTimePerOutputToken(start, ttft, genAiResponseAttrs, genAiGeneralAttrs, genAiSignalLevel);
+                },
+                experimental_include: {
+                    requestBody: false,
+                },
+                includeRawChunks: false,
+                ...textOptions,
+            });
+            ctx.result = toMessagesStream(result, ctx.resolvedModelId);
+            if (hooks?.after) {
+                ctx.result = (await hooks.after(ctx)) ?? ctx.result;
+                addSpanEvent("hebo.hooks.after.completed");
+            }
+            return ctx.result;
+        }
+        addSpanEvent("hebo.ai-sdk.started");
+        const result = await generateText({
+            model: languageModelWithMiddleware,
+            headers: prepareForwardHeaders(ctx.request),
+            abortSignal: ctx.request.signal,
+            timeout: cfg.timeouts.normal,
+            experimental_include: {
+                requestBody: false,
+                responseBody: false,
+            },
+            ...textOptions,
+        });
+        logger.trace({ requestId: ctx.requestId, result }, "[messages] AI SDK result");
+        addSpanEvent("hebo.ai-sdk.completed");
+        recordTimeToFirstToken(performance.now() - start, genAiGeneralAttrs, genAiSignalLevel);
+        ctx.result = toMessages(result, ctx.resolvedModelId);
+        logger.trace({ requestId: ctx.requestId, result: ctx.result }, "[messages] Messages");
+        addSpanEvent("hebo.result.transformed");
+        const genAiResponseAttrs = getMessagesResponseAttributes(ctx.result, genAiSignalLevel, result.finishReason);
+        setSpanAttributes(genAiResponseAttrs);
+        recordTokenUsage(genAiResponseAttrs, genAiGeneralAttrs, genAiSignalLevel);
+        if (hooks?.after) {
+            ctx.result = (await hooks.after(ctx)) ?? ctx.result;
+            addSpanEvent("hebo.hooks.after.completed");
+        }
+        recordTimePerOutputToken(start, 0, genAiResponseAttrs, genAiGeneralAttrs, genAiSignalLevel);
+        return ctx.result;
+    };
+    return { handler: winterCgHandler(handler, config) };
+};

package/dist/endpoints/messages/index.d.ts

@@ -0,0 +1,4 @@
+export * from "./converters";
+export * from "./handler";
+export * from "./schema";
+export * from "./otel";

package/dist/endpoints/messages/index.js

@@ -0,0 +1,4 @@
+export * from "./converters";
+export * from "./handler";
+export * from "./schema";
+export * from "./otel";

package/dist/endpoints/messages/otel.d.ts

@@ -0,0 +1,6 @@
+import type { Attributes } from "@opentelemetry/api";
+import type { FinishReason } from "ai";
+import { type TelemetrySignalLevel } from "../../types";
+import type { Messages, MessagesBody } from "./schema";
+export declare const getMessagesRequestAttributes: (body: MessagesBody, signalLevel?: TelemetrySignalLevel) => Attributes;
+export declare const getMessagesResponseAttributes: (response: Messages, signalLevel?: TelemetrySignalLevel, finishReason?: FinishReason) => Attributes;

package/dist/endpoints/messages/otel.js

@@ -0,0 +1,171 @@
+import {} from "../../types";
+const toBlobPart = (modality, mimeType) => {
+    const part = {
+        type: "blob",
+        modality,
+        content: "[REDACTED_BINARY_DATA]",
+    };
+    if (mimeType)
+        part["mime_type"] = mimeType;
+    return part;
+};
+const toUserBlockParts = (block) => {
+    switch (block.type) {
+        case "text":
+            return { type: "text", content: block.text };
+        case "image":
+            if (block.source.type === "base64") {
+                return toBlobPart("image", block.source.media_type);
+            }
+            return { type: "uri", modality: "image", uri: block.source.url };
+        case "tool_result":
+            return {
+                type: "tool_call_response",
+                id: block.tool_use_id,
+                response: typeof block.content === "string"
+                    ? block.content
+                    : block.content
+                        ? block.content.map((p) => (p.type === "text" ? p.text : "")).join("")
+                        : "",
+            };
+        case "document":
+            if (block.source.type === "base64") {
+                return toBlobPart("file", block.source.media_type);
+            }
+            if (block.source.type === "url") {
+                return { type: "uri", modality: "file", uri: block.source.url };
+            }
+            return { type: "text", content: block.source.data };
+        default:
+            return { type: block.type, content: "[UNHANDLED_CONTENT_BLOCK]" };
+    }
+};
+const toMessageParts = (message) => {
+    if (typeof message.content === "string") {
+        return [{ type: "text", content: message.content }];
+    }
+    const parts = [];
+    for (const block of message.content) {
+        if (message.role === "user") {
+            parts.push(toUserBlockParts(block));
+        }
+        else {
+            const assistantBlock = block;
+            switch (assistantBlock.type) {
+                case "text":
+                    parts.push({ type: "text", content: assistantBlock.text });
+                    break;
+                case "tool_use":
+                    parts.push({
+                        type: "tool_call",
+                        id: assistantBlock.id,
+                        name: assistantBlock.name,
+                        arguments: typeof assistantBlock.input === "string"
+                            ? assistantBlock.input
+                            : JSON.stringify(assistantBlock.input),
+                    });
+                    break;
+                case "thinking":
+                    parts.push({ type: "reasoning", content: assistantBlock.thinking });
+                    break;
+                case "redacted_thinking":
+                    parts.push({ type: "reasoning", content: "[ENCRYPTED_REASONING]" });
+                    break;
+            }
+        }
+    }
+    return parts;
+};
+const toResponseBlockPart = (block) => {
+    switch (block.type) {
+        case "text":
+            return { type: "text", content: block.text };
+        case "tool_use":
+            return {
+                type: "tool_call",
+                id: block.id,
+                name: block.name,
+                arguments: typeof block.input === "string" ? block.input : JSON.stringify(block.input),
+            };
+        case "thinking":
+            return { type: "reasoning", content: block.thinking };
+        case "redacted_thinking":
+            return { type: "reasoning", content: "[ENCRYPTED_REASONING]" };
+        default:
+            return { type: "unknown", content: "[UNHANDLED_RESPONSE_BLOCK]" };
+    }
+};
+export const getMessagesRequestAttributes = (body, signalLevel) => {
+    if (!signalLevel || signalLevel === "off")
+        return {};
+    const attrs = {};
+    if (signalLevel !== "required") {
+        Object.assign(attrs, {
+            "gen_ai.request.stream": body.stream,
+            "gen_ai.request.service_tier": body.service_tier,
+            "gen_ai.request.max_tokens": body.max_tokens,
+            "gen_ai.request.temperature": body.temperature,
+            "gen_ai.request.top_p": body.top_p,
+        });
+        if (body.metadata) {
+            for (const key in body.metadata) {
+                attrs[`gen_ai.request.metadata.${key}`] = body.metadata[key];
+            }
+        }
+    }
+    if (signalLevel === "full") {
+        const inputMessages = [];
+        // System prompt
+        if (body.system) {
+            const systemText = typeof body.system === "string" ? body.system : body.system.map((b) => b.text).join("");
+            inputMessages.push(JSON.stringify({
+                role: "system",
+                parts: [{ type: "text", content: systemText }],
+            }));
+        }
+        // Messages
+        for (const message of body.messages) {
+            inputMessages.push(JSON.stringify({
+                role: message.role,
+                parts: toMessageParts(message),
+            }));
+        }
+        Object.assign(attrs, {
+            "gen_ai.input.messages": inputMessages,
+            "gen_ai.tool.definitions": body.tools?.map((toolDef) => JSON.stringify(toolDef)),
+        });
+    }
+    return attrs;
+};
+export const getMessagesResponseAttributes = (response, signalLevel, finishReason) => {
+    if (!signalLevel || signalLevel === "off")
+        return {};
+    const attrs = {
+        "gen_ai.response.id": response.id,
+    };
+    if (signalLevel !== "required") {
+        Object.assign(attrs, {
+            "gen_ai.response.finish_reasons": finishReason
+                ? [finishReason]
+                : response.stop_reason
+                    ? [response.stop_reason]
+                    : [],
+            "gen_ai.response.service_tier": response.service_tier,
+            "gen_ai.usage.input_tokens": response.usage?.input_tokens,
+            "gen_ai.usage.output_tokens": response.usage?.output_tokens,
+            "gen_ai.usage.cache_read.input_tokens": response.usage?.cache_read_input_tokens,
+            "gen_ai.usage.cache_creation.input_tokens": response.usage?.cache_creation_input_tokens,
+        });
+    }
+    if (signalLevel === "full") {
+        Object.assign(attrs, {
+            "gen_ai.output.messages": [
+                JSON.stringify({
+                    role: "assistant",
+                    parts: response.content.map(toResponseBlockPart),
+                }),
+            ],
+        });
+    }
+    return attrs;
+};