calvin.erfmann/pangolin
1
0
Fork 0
mirror of https://github.com/fosrl/pangolin.git synced 2026-08-06 20:51:36 +02:00
Code Issues Packages Projects Releases Wiki Activity

add api capabilities

Browse Source
This commit is contained in:
miloschwartz
2026-08-05 16:55:48 -04:00
parent 2e8bd7a8c7
commit 39e06f2b6d
16 changed files with 715 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files
server/aiGatewayServer.ts
+2 -2
View File
@@ -7,7 +7,7 @@ import {
errorHandlerMiddleware,
notFoundMiddleware
} from "@server/middlewares";
import * as aiGateway from "@server/routers/aiGateway";
import { createAiGatewayRouter } from "@server/routers/aiGateway";
const aiGatewayPort = config.getRawConfig().server.ai_gateway_port;
@@ -23,7 +23,7 @@ export function createAiGatewayServer() {
aiGatewayServer.use(cors());
aiGatewayServer.use(express.json());
aiGatewayServer.post("/chat/completions", aiGateway.chatCompletions);
aiGatewayServer.use(createAiGatewayRouter());
aiGatewayServer.use(notFoundMiddleware);
aiGatewayServer.use(errorHandlerMiddleware);
server/db/pg/schema/schema.ts
+1
View File
@@ -1659,6 +1659,7 @@ export const aiProviders = pgTable("aiProviders", {
.$type<"url" | "target">()
.notNull()
.default("url"),
capabilities: text("capabilities").notNull().default("[]"),
skipTlsVerification: boolean("skipTlsVerification")
.notNull()
.default(false),
server/db/sqlite/schema/schema.ts
+1
View File
@@ -1641,6 +1641,7 @@ export const aiProviders = sqliteTable("aiProviders", {
.$type<"url" | "target">()
.notNull()
.default("url"),
capabilities: text("capabilities").notNull().default("[]"),
skipTlsVerification: integer("skipTlsVerification", { mode: "boolean" })
.notNull()
.default(false),
server/lib/aiCapabilities.ts
+257
View File
@@ -0,0 +1,257 @@
import type { Request } from "express";
import type { AiProviderType } from "@server/lib/aiProviderDefaults";
export const AI_CAPABILITIES = [
"openai_chat",
"openai_responses",
"anthropic_messages",
"gemini_generate_content",
"bedrock_model_invoke",
"google_generate_content",
"google_raw_predict",
"bedrock_converse"
] as const;
export type AiCapability = (typeof AI_CAPABILITIES)[number];
export type AiCapabilityRoute = {
method: "POST";
path: string;
};
export type AiCapabilityDefinition = {
id: AiCapability;
routes: AiCapabilityRoute[];
extractModel: (req: Request) => string | undefined;
resolveUpstreamUrl: (
baseUrl: string,
req: Request,
model: string
) => string;
};
function bodyModel(req: Request): string | undefined {
return typeof req.body?.model === "string" ? req.body.model : undefined;
}
function paramModel(req: Request): string | undefined {
const model = req.params?.model;
return typeof model === "string" && model.length > 0 ? model : undefined;
}
/**
* Join base URL with a path, avoiding double slashes and a duplicated trailing
* /v1 when the inbound path already starts with /v1 and the base ends with /v1.
*/
export function joinUpstreamUrl(baseUrl: string, path: string): string {
const base = baseUrl.replace(/\/+$/, "");
let suffix = path.startsWith("/") ? path : `/${path}`;
if (
base.endsWith("/v1") &&
(suffix === "/v1" || suffix.startsWith("/v1/"))
) {
suffix = suffix.slice("/v1".length) || "/";
}
if (suffix === "/") {
return base;
}
return `${base}${suffix}`;
}
function pathFromRequest(req: Request): string {
// Prefer originalUrl path (includes mounted path) over req.path when available.
const raw =
req.originalUrl?.split("?")[0] || req.url?.split("?")[0] || req.path;
return raw.startsWith("/") ? raw : `/${raw}`;
}
export const AI_CAPABILITY_DEFS: Record<AiCapability, AiCapabilityDefinition> =
{
openai_chat: {
id: "openai_chat",
routes: [
{ method: "POST", path: "/v1/chat/completions" },
{ method: "POST", path: "/chat/completions" }
],
extractModel: bodyModel,
resolveUpstreamUrl: (base, req) =>
joinUpstreamUrl(base, pathFromRequest(req))
},
openai_responses: {
id: "openai_responses",
routes: [{ method: "POST", path: "/v1/responses" }],
extractModel: bodyModel,
resolveUpstreamUrl: (base, req) =>
joinUpstreamUrl(base, pathFromRequest(req))
},
anthropic_messages: {
id: "anthropic_messages",
routes: [{ method: "POST", path: "/v1/messages" }],
extractModel: bodyModel,
resolveUpstreamUrl: (base, req) =>
joinUpstreamUrl(base, pathFromRequest(req))
},
gemini_generate_content: {
id: "gemini_generate_content",
routes: [
{
method: "POST",
path: "/v1beta/models/:model\\:generateContent"
},
{
method: "POST",
path: "/v1beta/models/:model\\:streamGenerateContent"
}
],
extractModel: paramModel,
resolveUpstreamUrl: (base, req) =>
joinUpstreamUrl(base, pathFromRequest(req))
},
google_generate_content: {
id: "google_generate_content",
routes: [
{
method: "POST",
// Vertex publisher model generateContent
path: "/v1/projects/:project/locations/:location/publishers/:publisher/models/:model\\:generateContent"
},
{
method: "POST",
path: "/v1/projects/:project/locations/:location/publishers/:publisher/models/:model\\:streamGenerateContent"
}
],
extractModel: paramModel,
resolveUpstreamUrl: (base, req) =>
joinUpstreamUrl(base, pathFromRequest(req))
},
google_raw_predict: {
id: "google_raw_predict",
routes: [
{
method: "POST",
path: "/v1/projects/:project/locations/:location/publishers/:publisher/models/:model\\:rawPredict"
},
{
method: "POST",
path: "/v1/projects/:project/locations/:location/publishers/:publisher/models/:model\\:streamRawPredict"
}
],
extractModel: paramModel,
resolveUpstreamUrl: (base, req) =>
joinUpstreamUrl(base, pathFromRequest(req))
},
bedrock_model_invoke: {
id: "bedrock_model_invoke",
routes: [
{ method: "POST", path: "/model/:model/invoke" },
{
method: "POST",
path: "/model/:model/invoke-with-response-stream"
}
],
extractModel: paramModel,
resolveUpstreamUrl: (base, req) =>
joinUpstreamUrl(base, pathFromRequest(req))
},
bedrock_converse: {
id: "bedrock_converse",
routes: [
{ method: "POST", path: "/model/:model/converse" },
{ method: "POST", path: "/model/:model/converse-stream" }
],
extractModel: paramModel,
resolveUpstreamUrl: (base, req) =>
joinUpstreamUrl(base, pathFromRequest(req))
}
};
export const AI_PROVIDER_CAPABILITY_DEFAULTS: Record<
Exclude<AiProviderType, "custom">,
readonly AiCapability[]
> = {
openai: ["openai_chat"],
anthropic: ["anthropic_messages"],
googleGemini: ["openai_chat"],
vertexAi: ["google_generate_content"],
bedrock: ["bedrock_converse"],
microsoftFoundry: ["openai_chat"],
openRouter: ["openai_chat"],
vercelAiGateway: ["openai_chat"]
};
export function isAiCapability(value: unknown): value is AiCapability {
return (
typeof value === "string" &&
(AI_CAPABILITIES as readonly string[]).includes(value)
);
}
export function parseCapabilities(raw: unknown): AiCapability[] {
if (raw == null) {
return [];
}
let parsed: unknown = raw;
if (typeof raw === "string") {
const trimmed = raw.trim();
if (!trimmed) {
return [];
}
try {
parsed = JSON.parse(trimmed);
} catch {
return [];
}
}
if (!Array.isArray(parsed)) {
return [];
}
const out: AiCapability[] = [];
const seen = new Set<AiCapability>();
for (const item of parsed) {
if (isAiCapability(item) && !seen.has(item)) {
seen.add(item);
out.push(item);
}
}
return out;
}
export function serializeCapabilities(capabilities: AiCapability[]): string {
return JSON.stringify(capabilities);
}
export function providerHasCapability(
capabilities: AiCapability[] | string | null | undefined,
capability: AiCapability
): boolean {
const list =
typeof capabilities === "string" || capabilities == null
? parseCapabilities(capabilities)
: capabilities;
return list.includes(capability);
}
export function resolveCapabilitiesForCreate(input: {
type: AiProviderType;
capabilities?: AiCapability[] | null;
}): AiCapability[] {
if (input.type === "custom") {
return parseCapabilities(input.capabilities ?? []);
}
return [...AI_PROVIDER_CAPABILITY_DEFAULTS[input.type]];
}
export function defaultsForProviderType(
type: AiProviderType
): readonly AiCapability[] {
if (type === "custom") {
return [];
}
return AI_PROVIDER_CAPABILITY_DEFAULTS[type];
}
server/routers/aiGateway/createAiGatewayRouter.ts
+21
View File
@@ -0,0 +1,21 @@
import { Router } from "express";
import {
AI_CAPABILITY_DEFS,
type AiCapability
} from "@server/lib/aiCapabilities";
import { handleAiGatewayProxy } from "@server/routers/aiGateway/pipeline";
export function createAiGatewayRouter() {
const router = Router();
for (const def of Object.values(AI_CAPABILITY_DEFS)) {
const capability = def.id as AiCapability;
for (const route of def.routes) {
router.post(route.path, (req, res) =>
handleAiGatewayProxy(req, res, capability)
);
}
}
return router;
}
server/routers/aiGateway/index.ts
+2 -1
View File
@@ -1 +1,2 @@
export * from "./chatCompletions";
export { handleAiGatewayProxy } from "./pipeline";
export { createAiGatewayRouter } from "./createAiGatewayRouter";
server/routers/aiGateway/chatCompletions.ts → server/routers/aiGateway/pipeline.ts
+37 -36
View File
@@ -22,6 +22,11 @@ import {
applyAiProviderAuthHeaders,
authTypeRequiresApiKey
} from "@server/lib/aiProviderDefaults";
import {
AI_CAPABILITY_DEFS,
providerHasCapability,
type AiCapability
} from "@server/lib/aiCapabilities";
import {
SESSION_COOKIE_NAME,
validateSessionToken
@@ -45,10 +50,6 @@ const REQUEST_USER_TTL_SEC = 30;
type CachedClient = { clientId: number; userId: string | null } | null;
// The set of CIDRs an exit node manages; client exitNodeSubnets are always
// /32s carved out of one of these ranges. Checking against this small,
// cacheable list lets us skip the (much more frequent) per-IP client lookup
// entirely for traffic that could never match a client anyway.
async function getExitNodeRanges(): Promise<string[]> {
const cached = localCache.get<string[]>(EXIT_NODE_RANGES_CACHE_KEY);
if (cached) {
@@ -96,8 +97,6 @@ type ResolvedTarget = {
siteResourceId: number | null;
orgId: string | null;
attachments: ProviderAttachment[];
// Model IDs on the resource allowlist that belong to allowlist-mode
// providers. Empty when no attached provider uses allowlist mode.
allowlistedModelIds: Set<number>;
};
@@ -150,13 +149,9 @@ async function buildRequestUser(
async function resolveRequestUser(
req: Request,
resourceId: number | null,
_resourceId: number | null,
orgId: string | null
): Promise<RequestUser | null> {
// Public resources behind badger: badger passes the resource session
// cookie through to the backend (same mechanism the browser gateway,
// e.g. the SSH page, relies on), so we can validate it exactly like
// verifySessionUserMiddleware does for the dashboard.
const sessionToken = req.cookies?.[SESSION_COOKIE_NAME];
if (sessionToken) {
const { session, user } = await validateSessionToken(sessionToken);
@@ -167,10 +162,6 @@ async function resolveRequestUser(
// TODO: MAKE SURE THIS CAN NOT BE SPOOFED AND CAN BE TRUSTED AS AN INTERNAL ADDRESS FROM A NODE
// No session cookie - fall back to identifying the caller by source IP.
// A client's exitNodeSubnet is a /32 handed out from one of our exit
// node's address ranges, so an IP that isn't inside any of those ranges
// can never belong to a client and we can skip the DB entirely.
const ip = req.ip;
if (!ip) {
return null;
@@ -193,9 +184,6 @@ async function resolveRequestUser(
}
async function resolveTarget(host: string): Promise<ResolvedTarget | null> {
// TODO: eventually we need to know if it's a private or public resource
// and not just simply check the fullDomain in case there is a private resource with the same fullDomain
const [resourceRow] = await db
.select({
resourceId: resources.resourceId,
@@ -378,7 +366,6 @@ async function selectProvider(
};
}
// One lookup for the requested model key across all attached providers.
const matchingModels = await db
.select({
modelId: aiModels.modelId,
@@ -407,7 +394,6 @@ async function selectProvider(
continue;
}
// allowlist: only models explicitly attached to the resource
if (allowlistedModelIds.has(model.modelId)) {
candidates.push(attachment.provider);
}
@@ -432,11 +418,14 @@ async function selectProvider(
};
}
export async function chatCompletions(
export async function handleAiGatewayProxy(
req: Request,
res: Response
res: Response,
capability: AiCapability
): Promise<any> {
try {
const def = AI_CAPABILITY_DEFS[capability];
const host = (
(req.headers["p-host"] as string | undefined) ||
req.headers.host ||
@@ -448,7 +437,7 @@ export async function chatCompletions(
.json({ error: { message: "Missing Host header" } });
}
logger.info(`AI gateway request for host: ${host}`);
logger.info(`AI gateway ${capability} request for host: ${host}`);
const target = await resolveTarget(host);
if (!target) {
@@ -461,11 +450,6 @@ export async function chatCompletions(
const { attachments, allowlistedModelIds, resourceId, orgId } = target;
logger.debug("+++++ gateway target: ", target);
// Best-effort identity resolution - not yet enforced, but lets us
// start making per-user access decisions (e.g. model/role-based
// restrictions) without another round of plumbing later.
const requestUser = await resolveRequestUser(req, resourceId, orgId);
if (requestUser) {
logger.debug(
@@ -473,11 +457,22 @@ export async function chatCompletions(
);
}
const requestedModel =
typeof req.body?.model === "string" ? req.body.model : undefined;
const capableAttachments = attachments.filter((a) =>
providerHasCapability(a.provider.capabilities, capability)
);
if (capableAttachments.length === 0) {
return res.status(HttpCode.FORBIDDEN).json({
error: {
message: `No AI provider on this resource supports ${capability}`
}
});
}
const requestedModel = def.extractModel(req);
const selection = await selectProvider(
attachments,
capableAttachments,
allowlistedModelIds,
requestedModel
);
@@ -513,11 +508,12 @@ export async function chatCompletions(
apiKey = decrypt(provider.apiKey, secret);
}
const targetUrl = `${upstreamUrl.replace(/\/$/, "")}`;
const targetUrl = def.resolveUpstreamUrl(
upstreamUrl,
req,
requestedModel!
);
// Drop hop-by-hop / proxy-only headers. Forwarding Host especially
// breaks Node fetch (TLS/SNI targets the upstream URL while Host
// still says localhost).
const skipHeaders = new Set([
"p-host",
"host",
@@ -554,6 +550,7 @@ export async function chatCompletions(
const body = JSON.stringify(req.body);
logger.debug("AI gateway upstream request", {
capability,
url: targetUrl,
method: "POST",
headers,
@@ -591,7 +588,11 @@ export async function chatCompletions(
const contentType = upstreamRes.headers.get("content-type") || "";
const isStream =
req.body?.stream === true ||
contentType.includes("text/event-stream");
contentType.includes("text/event-stream") ||
req.path.includes("streamGenerateContent") ||
req.path.includes("streamRawPredict") ||
req.path.includes("converse-stream") ||
req.path.includes("invoke-with-response-stream");
res.status(upstreamRes.status);
res.setHeader("Content-Type", contentType || "application/json");
server/routers/aiProvider/createAiProvider.ts
+12
View File
@@ -14,10 +14,15 @@ import type { CreateOrEditAiProviderResponse } from "@server/routers/aiProvider/
import { toPublicAiProvider } from "@server/routers/aiProvider/types";
import {
aiAuthTypeSchema,
aiCapabilitiesSchema,
aiProviderTypeSchema,
aiRoutingModeSchema,
refineProviderUpstreamFields
} from "@server/routers/aiProvider/validation";
import {
resolveCapabilitiesForCreate,
serializeCapabilities
} from "@server/lib/aiCapabilities";
const paramsSchema = z.strictObject({
orgId: z.string().nonempty()
@@ -31,6 +36,7 @@ const bodySchema = z
apiKey: z.string().optional(),
authType: aiAuthTypeSchema.optional(),
routingMode: aiRoutingModeSchema.optional(),
capabilities: aiCapabilitiesSchema.optional(),
skipTlsVerification: z.boolean().optional(),
enabled: z.boolean().optional()
})
@@ -94,6 +100,7 @@ export async function createAiProvider(
apiKey,
authType,
routingMode,
capabilities,
skipTlsVerification,
enabled
} = parsedBody.data;
@@ -108,6 +115,10 @@ export async function createAiProvider(
authType,
routingMode
});
const resolvedCapabilities = resolveCapabilitiesForCreate({
type,
capabilities
});
const [provider] = await db
.insert(aiProviders)
@@ -120,6 +131,7 @@ export async function createAiProvider(
apiKeyLastChars,
authType: resolved.authType,
routingMode: resolved.routingMode,
capabilities: serializeCapabilities(resolvedCapabilities),
skipTlsVerification: skipTlsVerification ?? false,
enabled: enabled ?? true,
createdAt: now,
server/routers/aiProvider/types.ts
+12 -2
View File
@@ -1,12 +1,17 @@
import type { AiModel, AiProvider } from "@server/db";
import type { PaginatedResponse } from "@server/types/Pagination";
import type { AiProviderAuthType } from "@server/lib/aiProviderDefaults";
import {
parseCapabilities,
type AiCapability
} from "@server/lib/aiCapabilities";
import { decrypt } from "@server/lib/crypto";
import config from "@server/lib/config";
export type AiProviderPublic = Omit<AiProvider, "apiKey"> & {
export type AiProviderPublic = Omit<AiProvider, "apiKey" | "capabilities"> & {
/** Decrypted API key. Only included on get/create/update of a single provider. */
apiKey?: string | null;
capabilities: AiCapability[];
effectiveUpstreamUrl: string | null;
effectiveAuthType: AiProviderAuthType;
};
@@ -39,7 +44,11 @@ export function toPublicAiProvider(
provider: AiProvider,
options?: { includeApiKey?: boolean }
): AiProviderPublic {
const { apiKey: encryptedApiKey, ...rest } = provider;
const {
apiKey: encryptedApiKey,
capabilities: rawCapabilities,
...rest
} = provider;
let apiKey: string | null | undefined;
if (options?.includeApiKey) {
@@ -56,6 +65,7 @@ export function toPublicAiProvider(
return {
...rest,
...(options?.includeApiKey ? { apiKey } : {}),
capabilities: parseCapabilities(rawCapabilities),
effectiveUpstreamUrl: provider.upstreamUrl,
effectiveAuthType: provider.authType as AiProviderAuthType
};
server/routers/aiProvider/updateAiProvider.ts
+29 -2
View File
@@ -14,6 +14,7 @@ import type { CreateOrEditAiProviderResponse } from "@server/routers/aiProvider/
import { toPublicAiProvider } from "@server/routers/aiProvider/types";
import {
aiAuthTypeSchema,
aiCapabilitiesSchema,
aiProviderTypeSchema,
aiRoutingModeSchema,
refineProviderUpstreamFields
@@ -23,6 +24,10 @@ import type {
AiProviderRoutingMode,
AiProviderType
} from "@server/lib/aiProviderDefaults";
import {
parseCapabilities,
serializeCapabilities
} from "@server/lib/aiCapabilities";
const paramsSchema = z.strictObject({
providerId: z.coerce.number().int().positive()
@@ -34,6 +39,7 @@ const bodySchema = z.strictObject({
apiKey: z.string().optional(),
authType: aiAuthTypeSchema.optional(),
routingMode: aiRoutingModeSchema.optional(),
capabilities: aiCapabilitiesSchema.optional(),
skipTlsVerification: z.boolean().optional(),
enabled: z.boolean().optional()
});
@@ -122,19 +128,37 @@ export async function updateAiProvider(
? body.authType
: (existing.authType as AiProviderAuthType);
if (body.capabilities !== undefined && providerType !== "custom") {
return next(
createHttpError(
HttpCode.BAD_REQUEST,
"Capabilities can only be updated for custom providers"
)
);
}
const nextCapabilities =
providerType === "custom"
? body.capabilities !== undefined
? body.capabilities
: parseCapabilities(existing.capabilities)
: parseCapabilities(existing.capabilities);
const validation = z
.object({
type: aiProviderTypeSchema,
upstreamUrl: z.string().nullable().optional(),
authType: aiAuthTypeSchema,
routingMode: aiRoutingModeSchema.optional()
routingMode: aiRoutingModeSchema.optional(),
capabilities: aiCapabilitiesSchema.optional()
})
.superRefine((data, ctx) => refineProviderUpstreamFields(data, ctx))
.safeParse({
type: providerType,
upstreamUrl: nextUpstreamUrl,
authType: nextAuthType,
routingMode: nextRoutingMode
routingMode: nextRoutingMode,
capabilities: nextCapabilities
});
if (!validation.success) {
@@ -168,6 +192,9 @@ export async function updateAiProvider(
if (body.authType !== undefined) {
updateData.authType = body.authType;
}
if (providerType === "custom" && body.capabilities !== undefined) {
updateData.capabilities = serializeCapabilities(body.capabilities);
}
if (body.apiKey !== undefined) {
const key = config.getRawConfig().server.secret!;
server/routers/aiProvider/validation.ts
+18
View File
@@ -6,6 +6,7 @@ import {
type AiProviderRoutingMode,
type AiProviderType
} from "@server/lib/aiProviderDefaults";
import { AI_CAPABILITIES } from "@server/lib/aiCapabilities";
export const aiProviderTypeSchema = z.enum([
"openai",
@@ -23,12 +24,17 @@ export const aiAuthTypeSchema = z.enum(AI_PROVIDER_AUTH_TYPES);
export const aiRoutingModeSchema = z.enum(["url", "target"]);
export const aiCapabilitySchema = z.enum(AI_CAPABILITIES);
export const aiCapabilitiesSchema = z.array(aiCapabilitySchema);
export function refineProviderUpstreamFields(
data: {
type: AiProviderType;
upstreamUrl?: string | null;
authType?: AiProviderAuthType | null;
routingMode?: AiProviderRoutingMode | null;
capabilities?: z.infer<typeof aiCapabilitiesSchema> | null;
},
ctx: z.RefinementCtx
) {
@@ -52,4 +58,16 @@ export function refineProviderUpstreamFields(
path: ["upstreamUrl"]
});
}
if (data.type === "custom") {
const caps = data.capabilities;
if (!caps || caps.length === 0) {
ctx.addIssue({
code: "custom",
message:
"At least one capability is required for custom providers",
path: ["capabilities"]
});
}
}
}