You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 
 

228 lines
5.8 KiB

import { dbGlobal } from "drizzle-pkg/lib/db";
import { llmProviders, llmModels } from "drizzle-pkg/lib/schema/llm";
import { eq, desc, like, and, sql } from "drizzle-orm";
import type { LlmParseMode, LlmProviderStatus, LlmModelType } from "drizzle-pkg/lib/schema/llm";
export interface LlmProviderRow {
id: number;
name: string;
slug: string;
baseUrl: string | null;
parseMode: LlmParseMode;
apiKey: string | null;
status: LlmProviderStatus;
description: string | null;
createdAt: Date;
updatedAt: Date;
}
export interface LlmModelRow {
id: number;
providerId: number;
name: string;
modelId: string;
type: LlmModelType;
enabled: number;
description: string | null;
maxTokens: number | null;
createdAt: Date;
updatedAt: Date;
}
export interface LlmProviderWithModels extends Omit<LlmProviderRow, "apiKey"> {
apiKeySet: boolean;
models: LlmModelRow[];
}
export async function listProviders(
page = 1,
pageSize = 20,
search?: string,
status?: LlmProviderStatus,
) {
const conditions = [];
if (search) {
conditions.push(like(llmProviders.name, `%${search}%`));
}
if (status) {
conditions.push(eq(llmProviders.status, status));
}
const where = conditions.length > 0 ? and(...conditions) : undefined;
const [countResult] = await dbGlobal
.select({ count: sql<number>`count(*)` })
.from(llmProviders)
.where(where);
const total = countResult?.count ?? 0;
const totalPages = Math.ceil(total / pageSize);
const rows = await dbGlobal
.select()
.from(llmProviders)
.where(where)
.orderBy(desc(llmProviders.createdAt))
.limit(pageSize)
.offset((page - 1) * pageSize);
const sanitized = rows.map(({ apiKey, ...rest }) => ({
...rest,
apiKeySet: !!apiKey,
}));
return { list: sanitized, total, page, pageSize, totalPages };
}
export async function getProviderById(id: number) {
const [row] = await dbGlobal
.select()
.from(llmProviders)
.where(eq(llmProviders.id, id))
.limit(1);
return row ?? null;
}
export async function getProviderWithModels(id: number): Promise<LlmProviderWithModels | null> {
const provider = await getProviderById(id);
if (!provider) return null;
const models = await dbGlobal
.select()
.from(llmModels)
.where(eq(llmModels.providerId, id))
.orderBy(desc(llmModels.createdAt));
const { apiKey, ...rest } = provider;
return {
...rest,
apiKeySet: !!apiKey,
models,
};
}
export async function createProvider(data: {
name: string;
slug: string;
baseUrl?: string;
parseMode: LlmParseMode;
apiKey?: string;
status?: LlmProviderStatus;
description?: string;
}) {
const [inserted] = await dbGlobal
.insert(llmProviders)
.values({
name: data.name,
slug: data.slug,
baseUrl: data.baseUrl || null,
parseMode: data.parseMode,
apiKey: data.apiKey || null,
status: data.status || "active",
description: data.description || null,
})
.returning({ id: llmProviders.id });
return inserted!;
}
export async function updateProvider(
id: number,
data: {
name?: string;
slug?: string;
baseUrl?: string;
parseMode?: LlmParseMode;
apiKey?: string;
status?: LlmProviderStatus;
description?: string;
},
) {
const updates: Record<string, any> = {};
if (data.name !== undefined) updates.name = data.name;
if (data.slug !== undefined) updates.slug = data.slug;
if (data.baseUrl !== undefined) updates.baseUrl = data.baseUrl;
if (data.parseMode !== undefined) updates.parseMode = data.parseMode;
if (data.status !== undefined) updates.status = data.status;
if (data.description !== undefined) updates.description = data.description;
if (data.apiKey !== undefined) {
if (data.apiKey === "") {
updates.apiKey = null;
} else if (data.apiKey !== "__unchanged__") {
updates.apiKey = data.apiKey;
}
}
if (Object.keys(updates).length === 0) return;
await dbGlobal
.update(llmProviders)
.set(updates)
.where(eq(llmProviders.id, id));
}
export async function deleteProvider(id: number) {
await dbGlobal.delete(llmProviders).where(eq(llmProviders.id, id));
}
export async function listModelsByProvider(providerId: number) {
return dbGlobal
.select()
.from(llmModels)
.where(eq(llmModels.providerId, providerId))
.orderBy(desc(llmModels.createdAt));
}
export async function createModel(data: {
providerId: number;
name: string;
modelId: string;
type: LlmModelType;
enabled?: number;
description?: string;
maxTokens?: number;
}) {
const [inserted] = await dbGlobal
.insert(llmModels)
.values({
providerId: data.providerId,
name: data.name,
modelId: data.modelId,
type: data.type,
enabled: data.enabled ?? 1,
description: data.description || null,
maxTokens: data.maxTokens || null,
})
.returning({ id: llmModels.id });
return inserted!;
}
export async function updateModel(
id: number,
data: {
name?: string;
modelId?: string;
type?: LlmModelType;
enabled?: number;
description?: string;
maxTokens?: number;
},
) {
const updates: Record<string, any> = {};
if (data.name !== undefined) updates.name = data.name;
if (data.modelId !== undefined) updates.modelId = data.modelId;
if (data.type !== undefined) updates.type = data.type;
if (data.enabled !== undefined) updates.enabled = data.enabled;
if (data.description !== undefined) updates.description = data.description;
if (data.maxTokens !== undefined) updates.maxTokens = data.maxTokens;
if (Object.keys(updates).length === 0) return;
await dbGlobal
.update(llmModels)
.set(updates)
.where(eq(llmModels.id, id));
}
export async function deleteModel(id: number) {
await dbGlobal.delete(llmModels).where(eq(llmModels.id, id));
}