534 lines
15 KiB
TypeScript
534 lines
15 KiB
TypeScript
|
import { createLanguageModelFromResolved } from "@tasks/ai";
|
||
|
|
import { seedDefaultRoleRoutes } from "@tasks/database/model-routing";
|
||
|
|
import {
|
||
|
|
modelProviders,
|
||
|
|
roleModelRoutes,
|
||
|
|
} from "@tasks/database/schema";
|
||
|
|
import {
|
||
|
|
convoyModelPolicySchema,
|
||
|
|
providerTypeSchema,
|
||
|
|
resolveModel,
|
||
|
|
type ModelProvider,
|
||
|
|
type RoleModelRoute,
|
||
|
|
} from "@tasks/shared/model-routing";
|
||
|
|
import { TRPCError } from "@trpc/server";
|
||
|
|
import { generateText } from "ai";
|
||
|
|
import { and, asc, eq } from "drizzle-orm";
|
||
|
|
import { z } from "zod";
|
||
|
|
|
||
|
|
import { rateLimit } from "@/server/lib/rate-limit";
|
||
|
|
import { router, workspaceProcedure } from "@/server/trpc";
|
||
|
|
import type { WorkspaceProcedureContext } from "@/server/trpc";
|
||
|
|
|
||
|
|
type ProviderRow = typeof modelProviders.$inferSelect;
|
||
|
|
type RouteRow = typeof roleModelRoutes.$inferSelect;
|
||
|
|
|
||
|
|
function assertWorkspaceAdmin(ctx: WorkspaceProcedureContext): void {
|
||
|
|
if (ctx.workspace.role !== "owner" && ctx.workspace.role !== "admin") {
|
||
|
|
throw new TRPCError({
|
||
|
|
code: "FORBIDDEN",
|
||
|
|
message: "Only workspace owners and admins can manage model providers.",
|
||
|
|
});
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
function canManageWorkspace(ctx: WorkspaceProcedureContext): boolean {
|
||
|
|
return ctx.workspace.role === "owner" || ctx.workspace.role === "admin";
|
||
|
|
}
|
||
|
|
|
||
|
|
function mapProviderRow(row: ProviderRow): ModelProvider {
|
||
|
|
return {
|
||
|
|
id: row.id,
|
||
|
|
workspaceId: row.workspaceId,
|
||
|
|
alias: row.alias,
|
||
|
|
displayName: row.displayName,
|
||
|
|
providerType: providerTypeSchema.parse(row.providerType),
|
||
|
|
modelId: row.modelId,
|
||
|
|
baseUrl: row.baseUrl,
|
||
|
|
apiKeyEnvVar: row.apiKeyEnvVar,
|
||
|
|
extraHeaders: row.extraHeaders,
|
||
|
|
defaultParams: row.defaultParams,
|
||
|
|
metadata: row.metadata,
|
||
|
|
enabled: row.enabled,
|
||
|
|
};
|
||
|
|
}
|
||
|
|
|
||
|
|
function mapRouteRow(row: RouteRow): RoleModelRoute {
|
||
|
|
return {
|
||
|
|
id: row.id,
|
||
|
|
workspaceId: row.workspaceId,
|
||
|
|
roleId: row.roleId,
|
||
|
|
providerAliasId: row.providerAliasId,
|
||
|
|
enabled: row.enabled,
|
||
|
|
sortOrder: row.sortOrder,
|
||
|
|
};
|
||
|
|
}
|
||
|
|
|
||
|
|
function readApiKey(envVarName: string): string | undefined {
|
||
|
|
return process.env[envVarName]?.trim() || undefined;
|
||
|
|
}
|
||
|
|
|
||
|
|
function validateProviderFields(input: {
|
||
|
|
providerType: z.infer<typeof providerTypeSchema>;
|
||
|
|
baseUrl?: string | null;
|
||
|
|
}): void {
|
||
|
|
if (input.providerType === "openai_compatible" && !input.baseUrl?.trim()) {
|
||
|
|
throw new TRPCError({
|
||
|
|
code: "BAD_REQUEST",
|
||
|
|
message: "baseUrl is required for openai_compatible providers.",
|
||
|
|
});
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
function providerToResolved(row: ProviderRow): ReturnType<typeof resolveModel> {
|
||
|
|
const provider = mapProviderRow(row);
|
||
|
|
const result = resolveModel({
|
||
|
|
roleId: "test",
|
||
|
|
providers: [provider],
|
||
|
|
routes: [],
|
||
|
|
briefRecommendedModel: provider.alias,
|
||
|
|
getApiKey: readApiKey,
|
||
|
|
});
|
||
|
|
return result;
|
||
|
|
}
|
||
|
|
|
||
|
|
const providerFieldsSchema = z.object({
|
||
|
|
alias: z.string().min(1).max(120),
|
||
|
|
displayName: z.string().min(1).max(200),
|
||
|
|
providerType: providerTypeSchema,
|
||
|
|
modelId: z.string().min(1).max(200),
|
||
|
|
baseUrl: z.string().url().nullable().optional(),
|
||
|
|
apiKeyEnvVar: z.string().min(1).max(120).nullable().optional(),
|
||
|
|
extraHeaders: z.record(z.string()).nullable().optional(),
|
||
|
|
defaultParams: z
|
||
|
|
.object({
|
||
|
|
temperature: z.number().min(0).max(2).optional(),
|
||
|
|
maxTokens: z.number().int().positive().optional(),
|
||
|
|
})
|
||
|
|
.nullable()
|
||
|
|
.optional(),
|
||
|
|
metadata: z.record(z.unknown()).nullable().optional(),
|
||
|
|
enabled: z.boolean().optional(),
|
||
|
|
});
|
||
|
|
|
||
|
|
export const aiModelsRouter = router({
|
||
|
|
listProviders: workspaceProcedure.query(async ({ ctx }) => {
|
||
|
|
const rows = await ctx.db
|
||
|
|
.select()
|
||
|
|
.from(modelProviders)
|
||
|
|
.where(eq(modelProviders.workspaceId, ctx.workspace.id))
|
||
|
|
.orderBy(asc(modelProviders.alias));
|
||
|
|
|
||
|
|
return {
|
||
|
|
providers: rows.map(mapProviderRow),
|
||
|
|
canManage: canManageWorkspace(ctx),
|
||
|
|
};
|
||
|
|
}),
|
||
|
|
|
||
|
|
createProvider: workspaceProcedure
|
||
|
|
.input(providerFieldsSchema)
|
||
|
|
.mutation(async ({ ctx, input }) => {
|
||
|
|
assertWorkspaceAdmin(ctx);
|
||
|
|
validateProviderFields(input);
|
||
|
|
|
||
|
|
try {
|
||
|
|
const [created] = await ctx.db
|
||
|
|
.insert(modelProviders)
|
||
|
|
.values({
|
||
|
|
workspaceId: ctx.workspace.id,
|
||
|
|
alias: input.alias.trim(),
|
||
|
|
displayName: input.displayName.trim(),
|
||
|
|
providerType: input.providerType,
|
||
|
|
modelId: input.modelId.trim(),
|
||
|
|
baseUrl: input.baseUrl?.trim() || null,
|
||
|
|
apiKeyEnvVar: input.apiKeyEnvVar?.trim() || null,
|
||
|
|
extraHeaders: input.extraHeaders ?? null,
|
||
|
|
defaultParams: input.defaultParams ?? null,
|
||
|
|
metadata: input.metadata ?? null,
|
||
|
|
enabled: input.enabled ?? true,
|
||
|
|
updatedAt: new Date(),
|
||
|
|
})
|
||
|
|
.returning();
|
||
|
|
|
||
|
|
if (!created) {
|
||
|
|
throw new TRPCError({
|
||
|
|
code: "INTERNAL_SERVER_ERROR",
|
||
|
|
message: "Failed to create provider.",
|
||
|
|
});
|
||
|
|
}
|
||
|
|
|
||
|
|
return mapProviderRow(created);
|
||
|
|
} catch (e) {
|
||
|
|
if (isUniqueViolation(e)) {
|
||
|
|
throw new TRPCError({
|
||
|
|
code: "CONFLICT",
|
||
|
|
message: `Alias "${input.alias}" already exists in this workspace.`,
|
||
|
|
});
|
||
|
|
}
|
||
|
|
throw e;
|
||
|
|
}
|
||
|
|
}),
|
||
|
|
|
||
|
|
updateProvider: workspaceProcedure
|
||
|
|
.input(
|
||
|
|
providerFieldsSchema.partial().extend({
|
||
|
|
id: z.string().uuid(),
|
||
|
|
}),
|
||
|
|
)
|
||
|
|
.mutation(async ({ ctx, input }) => {
|
||
|
|
assertWorkspaceAdmin(ctx);
|
||
|
|
|
||
|
|
const [existing] = await ctx.db
|
||
|
|
.select()
|
||
|
|
.from(modelProviders)
|
||
|
|
.where(
|
||
|
|
and(
|
||
|
|
eq(modelProviders.id, input.id),
|
||
|
|
eq(modelProviders.workspaceId, ctx.workspace.id),
|
||
|
|
),
|
||
|
|
)
|
||
|
|
.limit(1);
|
||
|
|
|
||
|
|
if (!existing) {
|
||
|
|
throw new TRPCError({ code: "NOT_FOUND", message: "Provider not found." });
|
||
|
|
}
|
||
|
|
|
||
|
|
const nextType = input.providerType ?? existing.providerType;
|
||
|
|
const nextBaseUrl =
|
||
|
|
input.baseUrl !== undefined ? input.baseUrl : existing.baseUrl;
|
||
|
|
validateProviderFields({
|
||
|
|
providerType: providerTypeSchema.parse(nextType),
|
||
|
|
baseUrl: nextBaseUrl,
|
||
|
|
});
|
||
|
|
|
||
|
|
try {
|
||
|
|
const [updated] = await ctx.db
|
||
|
|
.update(modelProviders)
|
||
|
|
.set({
|
||
|
|
...(input.alias !== undefined ? { alias: input.alias.trim() } : {}),
|
||
|
|
...(input.displayName !== undefined
|
||
|
|
? { displayName: input.displayName.trim() }
|
||
|
|
: {}),
|
||
|
|
...(input.providerType !== undefined
|
||
|
|
? { providerType: input.providerType }
|
||
|
|
: {}),
|
||
|
|
...(input.modelId !== undefined ? { modelId: input.modelId.trim() } : {}),
|
||
|
|
...(input.baseUrl !== undefined
|
||
|
|
? { baseUrl: input.baseUrl?.trim() || null }
|
||
|
|
: {}),
|
||
|
|
...(input.apiKeyEnvVar !== undefined
|
||
|
|
? { apiKeyEnvVar: input.apiKeyEnvVar?.trim() || null }
|
||
|
|
: {}),
|
||
|
|
...(input.extraHeaders !== undefined
|
||
|
|
? { extraHeaders: input.extraHeaders }
|
||
|
|
: {}),
|
||
|
|
...(input.defaultParams !== undefined
|
||
|
|
? { defaultParams: input.defaultParams }
|
||
|
|
: {}),
|
||
|
|
...(input.metadata !== undefined ? { metadata: input.metadata } : {}),
|
||
|
|
...(input.enabled !== undefined ? { enabled: input.enabled } : {}),
|
||
|
|
updatedAt: new Date(),
|
||
|
|
})
|
||
|
|
.where(
|
||
|
|
and(
|
||
|
|
eq(modelProviders.id, input.id),
|
||
|
|
eq(modelProviders.workspaceId, ctx.workspace.id),
|
||
|
|
),
|
||
|
|
)
|
||
|
|
.returning();
|
||
|
|
|
||
|
|
if (!updated) {
|
||
|
|
throw new TRPCError({ code: "NOT_FOUND", message: "Provider not found." });
|
||
|
|
}
|
||
|
|
|
||
|
|
return mapProviderRow(updated);
|
||
|
|
} catch (e) {
|
||
|
|
if (isUniqueViolation(e)) {
|
||
|
|
throw new TRPCError({
|
||
|
|
code: "CONFLICT",
|
||
|
|
message: "Another provider already uses that alias.",
|
||
|
|
});
|
||
|
|
}
|
||
|
|
throw e;
|
||
|
|
}
|
||
|
|
}),
|
||
|
|
|
||
|
|
deleteProvider: workspaceProcedure
|
||
|
|
.input(z.object({ id: z.string().uuid() }))
|
||
|
|
.mutation(async ({ ctx, input }) => {
|
||
|
|
assertWorkspaceAdmin(ctx);
|
||
|
|
const [deleted] = await ctx.db
|
||
|
|
.delete(modelProviders)
|
||
|
|
.where(
|
||
|
|
and(
|
||
|
|
eq(modelProviders.id, input.id),
|
||
|
|
eq(modelProviders.workspaceId, ctx.workspace.id),
|
||
|
|
),
|
||
|
|
)
|
||
|
|
.returning({ id: modelProviders.id });
|
||
|
|
|
||
|
|
if (!deleted) {
|
||
|
|
throw new TRPCError({ code: "NOT_FOUND", message: "Provider not found." });
|
||
|
|
}
|
||
|
|
|
||
|
|
return { id: deleted.id };
|
||
|
|
}),
|
||
|
|
|
||
|
|
testProvider: workspaceProcedure
|
||
|
|
.input(z.object({ providerId: z.string().uuid() }))
|
||
|
|
.mutation(async ({ ctx, input }) => {
|
||
|
|
assertWorkspaceAdmin(ctx);
|
||
|
|
|
||
|
|
const limit = rateLimit({
|
||
|
|
key: `ai:testProvider:${ctx.workspace.id}`,
|
||
|
|
limit: 10,
|
||
|
|
windowMs: 60_000,
|
||
|
|
});
|
||
|
|
if (!limit.allowed) {
|
||
|
|
throw new TRPCError({
|
||
|
|
code: "TOO_MANY_REQUESTS",
|
||
|
|
message: "Test rate limit exceeded. Try again shortly.",
|
||
|
|
});
|
||
|
|
}
|
||
|
|
|
||
|
|
const [row] = await ctx.db
|
||
|
|
.select()
|
||
|
|
.from(modelProviders)
|
||
|
|
.where(
|
||
|
|
and(
|
||
|
|
eq(modelProviders.id, input.providerId),
|
||
|
|
eq(modelProviders.workspaceId, ctx.workspace.id),
|
||
|
|
),
|
||
|
|
)
|
||
|
|
.limit(1);
|
||
|
|
|
||
|
|
if (!row) {
|
||
|
|
throw new TRPCError({ code: "NOT_FOUND", message: "Provider not found." });
|
||
|
|
}
|
||
|
|
|
||
|
|
const resolved = providerToResolved(row);
|
||
|
|
if (!resolved.ok) {
|
||
|
|
return {
|
||
|
|
ok: false as const,
|
||
|
|
latencyMs: 0,
|
||
|
|
error: resolved.error.message,
|
||
|
|
};
|
||
|
|
}
|
||
|
|
|
||
|
|
const client = createLanguageModelFromResolved(resolved.model);
|
||
|
|
if (!client.ok) {
|
||
|
|
return { ok: false as const, latencyMs: 0, error: client.reason };
|
||
|
|
}
|
||
|
|
|
||
|
|
const started = Date.now();
|
||
|
|
try {
|
||
|
|
await generateText({
|
||
|
|
model: client.model,
|
||
|
|
prompt: "Reply with exactly: OK",
|
||
|
|
maxTokens: 5,
|
||
|
|
temperature: 0,
|
||
|
|
});
|
||
|
|
return {
|
||
|
|
ok: true as const,
|
||
|
|
latencyMs: Date.now() - started,
|
||
|
|
};
|
||
|
|
} catch (e) {
|
||
|
|
const message = e instanceof Error ? e.message : String(e);
|
||
|
|
return {
|
||
|
|
ok: false as const,
|
||
|
|
latencyMs: Date.now() - started,
|
||
|
|
error: message,
|
||
|
|
};
|
||
|
|
}
|
||
|
|
}),
|
||
|
|
|
||
|
|
listRoleRoutes: workspaceProcedure.query(async ({ ctx }) => {
|
||
|
|
const routes = await ctx.db
|
||
|
|
.select({
|
||
|
|
route: roleModelRoutes,
|
||
|
|
providerAlias: modelProviders.alias,
|
||
|
|
providerEnabled: modelProviders.enabled,
|
||
|
|
providerMetadata: modelProviders.metadata,
|
||
|
|
})
|
||
|
|
.from(roleModelRoutes)
|
||
|
|
.innerJoin(
|
||
|
|
modelProviders,
|
||
|
|
eq(roleModelRoutes.providerAliasId, modelProviders.id),
|
||
|
|
)
|
||
|
|
.where(eq(roleModelRoutes.workspaceId, ctx.workspace.id))
|
||
|
|
.orderBy(asc(roleModelRoutes.sortOrder), asc(roleModelRoutes.roleId));
|
||
|
|
|
||
|
|
return {
|
||
|
|
routes: routes.map((row) => ({
|
||
|
|
...mapRouteRow(row.route),
|
||
|
|
providerAlias: row.providerAlias,
|
||
|
|
providerEnabled: row.providerEnabled,
|
||
|
|
providerMetadata: row.providerMetadata,
|
||
|
|
})),
|
||
|
|
canManage: canManageWorkspace(ctx),
|
||
|
|
};
|
||
|
|
}),
|
||
|
|
|
||
|
|
upsertRoleRoute: workspaceProcedure
|
||
|
|
.input(
|
||
|
|
z.object({
|
||
|
|
roleId: z.string().min(1).max(120),
|
||
|
|
providerAliasId: z.string().uuid(),
|
||
|
|
enabled: z.boolean().optional(),
|
||
|
|
sortOrder: z.number().int().optional(),
|
||
|
|
}),
|
||
|
|
)
|
||
|
|
.mutation(async ({ ctx, input }) => {
|
||
|
|
assertWorkspaceAdmin(ctx);
|
||
|
|
|
||
|
|
const [provider] = await ctx.db
|
||
|
|
.select({ id: modelProviders.id })
|
||
|
|
.from(modelProviders)
|
||
|
|
.where(
|
||
|
|
and(
|
||
|
|
eq(modelProviders.id, input.providerAliasId),
|
||
|
|
eq(modelProviders.workspaceId, ctx.workspace.id),
|
||
|
|
),
|
||
|
|
)
|
||
|
|
.limit(1);
|
||
|
|
|
||
|
|
if (!provider) {
|
||
|
|
throw new TRPCError({ code: "NOT_FOUND", message: "Provider not found." });
|
||
|
|
}
|
||
|
|
|
||
|
|
const [existing] = await ctx.db
|
||
|
|
.select({ id: roleModelRoutes.id })
|
||
|
|
.from(roleModelRoutes)
|
||
|
|
.where(
|
||
|
|
and(
|
||
|
|
eq(roleModelRoutes.workspaceId, ctx.workspace.id),
|
||
|
|
eq(roleModelRoutes.roleId, input.roleId),
|
||
|
|
),
|
||
|
|
)
|
||
|
|
.limit(1);
|
||
|
|
|
||
|
|
if (existing) {
|
||
|
|
const [updated] = await ctx.db
|
||
|
|
.update(roleModelRoutes)
|
||
|
|
.set({
|
||
|
|
providerAliasId: input.providerAliasId,
|
||
|
|
enabled: input.enabled ?? true,
|
||
|
|
...(input.sortOrder !== undefined ? { sortOrder: input.sortOrder } : {}),
|
||
|
|
updatedAt: new Date(),
|
||
|
|
})
|
||
|
|
.where(
|
||
|
|
and(
|
||
|
|
eq(roleModelRoutes.id, existing.id),
|
||
|
|
eq(roleModelRoutes.workspaceId, ctx.workspace.id),
|
||
|
|
),
|
||
|
|
)
|
||
|
|
.returning();
|
||
|
|
|
||
|
|
if (!updated) {
|
||
|
|
throw new TRPCError({ code: "NOT_FOUND", message: "Route not found." });
|
||
|
|
}
|
||
|
|
|
||
|
|
return mapRouteRow(updated);
|
||
|
|
}
|
||
|
|
|
||
|
|
const [created] = await ctx.db
|
||
|
|
.insert(roleModelRoutes)
|
||
|
|
.values({
|
||
|
|
workspaceId: ctx.workspace.id,
|
||
|
|
roleId: input.roleId,
|
||
|
|
providerAliasId: input.providerAliasId,
|
||
|
|
enabled: input.enabled ?? true,
|
||
|
|
sortOrder: input.sortOrder ?? 0,
|
||
|
|
updatedAt: new Date(),
|
||
|
|
})
|
||
|
|
.returning();
|
||
|
|
|
||
|
|
if (!created) {
|
||
|
|
throw new TRPCError({
|
||
|
|
code: "INTERNAL_SERVER_ERROR",
|
||
|
|
message: "Failed to create role route.",
|
||
|
|
});
|
||
|
|
}
|
||
|
|
|
||
|
|
return mapRouteRow(created);
|
||
|
|
}),
|
||
|
|
|
||
|
|
deleteRoleRoute: workspaceProcedure
|
||
|
|
.input(z.object({ id: z.string().uuid() }))
|
||
|
|
.mutation(async ({ ctx, input }) => {
|
||
|
|
assertWorkspaceAdmin(ctx);
|
||
|
|
const [deleted] = await ctx.db
|
||
|
|
.delete(roleModelRoutes)
|
||
|
|
.where(
|
||
|
|
and(
|
||
|
|
eq(roleModelRoutes.id, input.id),
|
||
|
|
eq(roleModelRoutes.workspaceId, ctx.workspace.id),
|
||
|
|
),
|
||
|
|
)
|
||
|
|
.returning({ id: roleModelRoutes.id });
|
||
|
|
|
||
|
|
if (!deleted) {
|
||
|
|
throw new TRPCError({ code: "NOT_FOUND", message: "Route not found." });
|
||
|
|
}
|
||
|
|
|
||
|
|
return { id: deleted.id };
|
||
|
|
}),
|
||
|
|
|
||
|
|
resetPipelineDefaults: workspaceProcedure.mutation(async ({ ctx }) => {
|
||
|
|
assertWorkspaceAdmin(ctx);
|
||
|
|
const result = await seedDefaultRoleRoutes(ctx.db, ctx.workspace.id);
|
||
|
|
return result;
|
||
|
|
}),
|
||
|
|
|
||
|
|
resolveForRole: workspaceProcedure
|
||
|
|
.input(
|
||
|
|
z.object({
|
||
|
|
roleId: z.string().min(1).max(120),
|
||
|
|
briefRecommendedModel: z.string().optional(),
|
||
|
|
convoyPolicy: convoyModelPolicySchema.optional(),
|
||
|
|
}),
|
||
|
|
)
|
||
|
|
.query(async ({ ctx, input }) => {
|
||
|
|
const providerRows = await ctx.db
|
||
|
|
.select()
|
||
|
|
.from(modelProviders)
|
||
|
|
.where(eq(modelProviders.workspaceId, ctx.workspace.id));
|
||
|
|
|
||
|
|
const routeRows = await ctx.db
|
||
|
|
.select()
|
||
|
|
.from(roleModelRoutes)
|
||
|
|
.where(eq(roleModelRoutes.workspaceId, ctx.workspace.id));
|
||
|
|
|
||
|
|
const result = resolveModel({
|
||
|
|
roleId: input.roleId,
|
||
|
|
providers: providerRows.map(mapProviderRow),
|
||
|
|
routes: routeRows.map(mapRouteRow),
|
||
|
|
briefRecommendedModel: input.briefRecommendedModel,
|
||
|
|
convoyPolicy: input.convoyPolicy,
|
||
|
|
getApiKey: readApiKey,
|
||
|
|
});
|
||
|
|
|
||
|
|
if (!result.ok) {
|
||
|
|
return result;
|
||
|
|
}
|
||
|
|
|
||
|
|
const { apiKey: _removed, ...publicModel } = result.model;
|
||
|
|
void _removed;
|
||
|
|
return {
|
||
|
|
ok: true as const,
|
||
|
|
model: publicModel,
|
||
|
|
source: result.source,
|
||
|
|
resolvedAlias: result.resolvedAlias,
|
||
|
|
};
|
||
|
|
}),
|
||
|
|
});
|
||
|
|
|
||
|
|
function isUniqueViolation(error: unknown): boolean {
|
||
|
|
return (
|
||
|
|
typeof error === "object" &&
|
||
|
|
error !== null &&
|
||
|
|
"code" in error &&
|
||
|
|
(error as { code?: string }).code === "23505"
|
||
|
|
);
|
||
|
|
}
|