import { Effect, Schema } from "effect";import type { AppContext } from "../../app-context";import { findObservedInferenceProcess } from "../../core/function-observability";import { decodeJsonBody } from "../../core/validation";import { defineRoutes, effectRoute, mergeRoutes } from "../../http/route-registrar";import { fetchInference } from "../../http/local-fetch"; const CountTokensRequestSchema = Schema.Struct({ text: Schema.optional(Schema.String), model: Schema.optional(Schema.String),}); const TokenizeChatRequestSchema = Schema.Struct({ messages: Schema.optional(Schema.Array(Schema.Unknown)), tools: Schema.optional(Schema.Array(Schema.Unknown)), model: Schema.optional(Schema.String),}); const TokenizeResponseSchema = Schema.Struct({ tokens: Schema.optional(Schema.Array(Schema.Unknown)),}); const TextPartSchema = Schema.Struct({ type: Schema.String, text: Schema.optional(Schema.String) });const MessageSchema = Schema.Struct({ content: Schema.optional(Schema.Union([Schema.String, Schema.Array(TextPartSchema)])),}); const responseTokens = (response: Response): Effect.Effect<number, unknown> => Effect.tryPromise({ try: () => response.json(), catch: (source) => source }).pipe( Effect.flatMap(Schema.decodeUnknownEffect(TokenizeResponseSchema)), Effect.map((payload) => payload.tokens?.length ?? 0), ); const tokenize = ( context: AppContext, model: string, prompt: string,): Effect.Effect<number, unknown> => fetchInference(context, "/tokenize", { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify({ model, prompt }), }).pipe( Effect.flatMap((response) => response.ok ? responseTokens(response) : Effect.fail(new Error(`Tokenize failed: ${response.status}`)), ), ); const messageText = (messages: readonly unknown[]): string => messages .flatMap((message) => { const decoded = Schema.decodeUnknownOption(MessageSchema)(message); if (decoded._tag === "None" || decoded.value.content === undefined) return []; return typeof decoded.value.content === "string" ? [decoded.value.content] : decoded.value.content.flatMap((part) => part.type === "text" && part.text ? [part.text] : [], ); }) .join("\n"); export const registerTokenizationRoutes = defineRoutes((app, context) => { return mergeRoutes( effectRoute(app.post, "/v1/count-tokens", (ctx) => Effect.gen(function* () { const current = yield* findObservedInferenceProcess(context, "countTokens"); if (!current) return ctx.json({ error: "No model running", num_tokens: 0 }); const body = yield* decodeJsonBody(ctx, CountTokensRequestSchema); const model = body.model ?? current.served_model_name ?? "default"; return yield* tokenize(context, model, body.text ?? "").pipe( Effect.map((numberTokens) => ctx.json({ num_tokens: numberTokens, model })), Effect.catch((error) => Effect.succeed(ctx.json({ error: String(error), num_tokens: 0 })), ), ); }), ), effectRoute(app.post, "/v1/tokenize-chat-completions", (ctx) => Effect.gen(function* () { const current = yield* findObservedInferenceProcess(context, "tokenizeChatCompletions"); if (!current) return ctx.json({ error: "No model running", input_tokens: 0 }); const body = yield* decodeJsonBody(ctx, TokenizeChatRequestSchema); const messages = body.messages ?? []; const tools = body.tools ?? []; const model = body.model ?? current.served_model_name ?? "default"; const messagesTokens = yield* tokenize(context, model, messageText(messages)).pipe( Effect.orElseSucceed(() => 0), ); const toolsTokens = tools.length > 0 ? yield* tokenize(context, model, JSON.stringify(tools)).pipe( Effect.orElseSucceed(() => 0), ) : 0; const overhead = messages.length * 4; return ctx.json({ input_tokens: messagesTokens + toolsTokens + overhead, breakdown: { messages: messagesTokens + overhead, tools: toolsTokens }, model, }); }), ), );});
import { Effect, Schema } from "effect";
import type { AppContext } from "../../app-context";
import { findObservedInferenceProcess } from "../../core/function-observability";
import { decodeJsonBody } from "../../core/validation";
import { defineRoutes, effectRoute, mergeRoutes } from "../../http/route-registrar";
import { fetchInference } from "../../http/local-fetch";
const CountTokensRequestSchema = Schema.Struct({
text: Schema.optional(Schema.String),
model: Schema.optional(Schema.String),
});
const TokenizeChatRequestSchema = Schema.Struct({
messages: Schema.optional(Schema.Array(Schema.Unknown)),
tools: Schema.optional(Schema.Array(Schema.Unknown)),
model: Schema.optional(Schema.String),
});
const TokenizeResponseSchema = Schema.Struct({
tokens: Schema.optional(Schema.Array(Schema.Unknown)),
});
const TextPartSchema = Schema.Struct({ type: Schema.String, text: Schema.optional(Schema.String) });
const MessageSchema = Schema.Struct({
content: Schema.optional(Schema.Union([Schema.String, Schema.Array(TextPartSchema)])),
});
const responseTokens = (response: Response): Effect.Effect<number, unknown> =>
Effect.tryPromise({ try: () => response.json(), catch: (source) => source }).pipe(
Effect.flatMap(Schema.decodeUnknownEffect(TokenizeResponseSchema)),
Effect.map((payload) => payload.tokens?.length ?? 0),
);
const tokenize = (
context: AppContext,
model: string,
prompt: string,
): Effect.Effect<number, unknown> =>
fetchInference(context, "/tokenize", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({ model, prompt }),
}).pipe(
Effect.flatMap((response) =>
response.ok
? responseTokens(response)
: Effect.fail(new Error(`Tokenize failed: ${response.status}`)),
),
);
const messageText = (messages: readonly unknown[]): string =>
messages
.flatMap((message) => {
const decoded = Schema.decodeUnknownOption(MessageSchema)(message);
if (decoded._tag === "None" || decoded.value.content === undefined) return [];
return typeof decoded.value.content === "string"
? [decoded.value.content]
: decoded.value.content.flatMap((part) =>
part.type === "text" && part.text ? [part.text] : [],
);
})
.join("\n");
export const registerTokenizationRoutes = defineRoutes((app, context) => {
return mergeRoutes(
effectRoute(app.post, "/v1/count-tokens", (ctx) =>
Effect.gen(function* () {
const current = yield* findObservedInferenceProcess(context, "countTokens");
if (!current) return ctx.json({ error: "No model running", num_tokens: 0 });
const body = yield* decodeJsonBody(ctx, CountTokensRequestSchema);
const model = body.model ?? current.served_model_name ?? "default";
return yield* tokenize(context, model, body.text ?? "").pipe(
Effect.map((numberTokens) => ctx.json({ num_tokens: numberTokens, model })),
Effect.catch((error) =>
Effect.succeed(ctx.json({ error: String(error), num_tokens: 0 })),
),
);
}),
),
effectRoute(app.post, "/v1/tokenize-chat-completions", (ctx) =>
Effect.gen(function* () {
const current = yield* findObservedInferenceProcess(context, "tokenizeChatCompletions");
if (!current) return ctx.json({ error: "No model running", input_tokens: 0 });
const body = yield* decodeJsonBody(ctx, TokenizeChatRequestSchema);
const messages = body.messages ?? [];
const tools = body.tools ?? [];
const model = body.model ?? current.served_model_name ?? "default";
const messagesTokens = yield* tokenize(context, model, messageText(messages)).pipe(
Effect.orElseSucceed(() => 0),
);
const toolsTokens =
tools.length > 0
? yield* tokenize(context, model, JSON.stringify(tools)).pipe(
Effect.orElseSucceed(() => 0),
)
: 0;
const overhead = messages.length * 4;
return ctx.json({
input_tokens: messagesTokens + toolsTokens + overhead,
breakdown: { messages: messagesTokens + overhead, tools: toolsTokens },
model,
});
}),
),
);
});