add api capabilities

server/routers/aiProvider/createAiProvider.ts

@@ -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

@@ -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

@@ -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

@@ -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"]
+            });
+        }
+    }
 }