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.
 
 
 
 

224 lines
7.2 KiB

import { dbGlobal } from "drizzle-pkg/lib/db";
import { agents, agentToolAssociations } from "drizzle-pkg/lib/schema/agent";
import { agentTools } from "drizzle-pkg/lib/schema/agent-tool";
import { eq, and, asc, sql } from "drizzle-orm";
import type { AgentToolRow } from "#server/service/agent-tool";
export type AgentRow = typeof agents.$inferSelect;
export type AgentTitleStrategy = "llm" | "first-line" | "none";
export interface AgentWithTools extends AgentRow {
tools: AgentToolRow[];
}
export interface CreateAgentInput {
slug: string;
name: string;
description?: string;
systemPrompt: string;
defaultModelId?: number | null;
titleModelId?: number | null;
titleStrategy?: AgentTitleStrategy;
enableThinking?: number;
enableTools?: number;
maxStepCount?: number;
isDefault?: number;
isCallable?: number;
sortOrder?: number;
enabled?: number;
toolIds?: string[];
}
export interface UpdateAgentInput {
slug?: string;
name?: string;
description?: string;
systemPrompt?: string;
defaultModelId?: number | null;
titleModelId?: number | null;
titleStrategy?: AgentTitleStrategy;
enableThinking?: number;
enableTools?: number;
maxStepCount?: number;
isDefault?: number;
isCallable?: number;
sortOrder?: number;
enabled?: number;
toolIds?: string[];
}
export async function getAgentBySlug(slug: string): Promise<AgentRow | null> {
const [row] = await dbGlobal
.select()
.from(agents)
.where(and(eq(agents.slug, slug), eq(agents.enabled, 1)))
.limit(1);
return row ?? null;
}
export async function getAgentBySlugAny(slug: string): Promise<AgentRow | null> {
const [row] = await dbGlobal
.select()
.from(agents)
.where(eq(agents.slug, slug))
.limit(1);
return row ?? null;
}
export async function getAgentById(id: number): Promise<AgentRow | null> {
const [row] = await dbGlobal.select().from(agents).where(eq(agents.id, id)).limit(1);
return row ?? null;
}
export async function getDefaultAgent(): Promise<AgentRow | null> {
const [row] = await dbGlobal
.select()
.from(agents)
.where(and(eq(agents.isDefault, 1), eq(agents.enabled, 1)))
.limit(1);
if (row) return row;
const [fallback] = await dbGlobal
.select()
.from(agents)
.where(eq(agents.enabled, 1))
.orderBy(asc(agents.sortOrder))
.limit(1);
return fallback ?? null;
}
export async function listAgents(onlyEnabled = true): Promise<AgentRow[]> {
const where = onlyEnabled ? eq(agents.enabled, 1) : sql`1=1`;
return dbGlobal.select().from(agents).where(where).orderBy(asc(agents.sortOrder), asc(agents.id));
}
export async function getAgentWithTools(slug: string): Promise<AgentWithTools | null> {
const agent = await getAgentBySlug(slug);
if (!agent) return null;
const tools = await getAgentTools(agent.id);
return { ...agent, tools };
}
export async function getAgentTools(agentId: number): Promise<AgentToolRow[]> {
const rows = await dbGlobal
.select({
tool: agentTools,
sortOrder: agentToolAssociations.sortOrder,
})
.from(agentToolAssociations)
.innerJoin(agentTools, eq(agentToolAssociations.toolId, agentTools.id))
.where(and(eq(agentToolAssociations.agentId, agentId), eq(agentTools.enabled, 1)))
.orderBy(asc(agentToolAssociations.sortOrder), asc(agentTools.sortOrder));
return rows.map((r) => r.tool);
}
export async function getAgentToolIds(agentId: number): Promise<string[]> {
const tools = await getAgentTools(agentId);
return tools.map((t) => t.id);
}
export async function createAgent(input: CreateAgentInput): Promise<AgentRow> {
const existing = await getAgentBySlugAny(input.slug);
if (existing) {
throw new Error(`slug 已存在: ${input.slug}`);
}
const [row] = await dbGlobal
.insert(agents)
.values({
slug: input.slug,
name: input.name,
description: input.description ?? null,
systemPrompt: input.systemPrompt,
defaultModelId: input.defaultModelId ?? null,
titleModelId: input.titleModelId ?? null,
titleStrategy: input.titleStrategy ?? "llm",
enableThinking: input.enableThinking ?? 0,
enableTools: input.enableTools ?? 1,
maxStepCount: input.maxStepCount ?? 6,
isDefault: input.isDefault ?? 0,
isCallable: input.isCallable ?? 0,
sortOrder: input.sortOrder ?? 0,
enabled: input.enabled ?? 1,
})
.returning();
if (!row) throw new Error("创建 agent 失败");
if (input.isDefault === 1) {
await dbGlobal
.update(agents)
.set({ isDefault: 0 })
.where(and(sql`${agents.id} != ${row.id}`));
}
if (input.toolIds && input.toolIds.length > 0) {
for (let i = 0; i < input.toolIds.length; i++) {
const toolId = input.toolIds[i]!;
await dbGlobal.insert(agentToolAssociations).values({
agentId: row.id,
toolId,
sortOrder: i,
});
}
}
return row;
}
export async function updateAgent(id: number, input: UpdateAgentInput): Promise<AgentRow | null> {
const existing = await getAgentById(id);
if (!existing) return null;
if (input.slug && input.slug !== existing.slug) {
const slugConflict = await getAgentBySlugAny(input.slug);
if (slugConflict) {
throw new Error(`slug 已存在: ${input.slug}`);
}
}
const updates: Partial<AgentRow> = {};
if (input.slug !== undefined) updates.slug = input.slug;
if (input.name !== undefined) updates.name = input.name;
if (input.description !== undefined) updates.description = input.description;
if (input.systemPrompt !== undefined) updates.systemPrompt = input.systemPrompt;
if (input.defaultModelId !== undefined) updates.defaultModelId = input.defaultModelId;
if (input.titleModelId !== undefined) updates.titleModelId = input.titleModelId;
if (input.titleStrategy !== undefined) updates.titleStrategy = input.titleStrategy;
if (input.enableThinking !== undefined) updates.enableThinking = input.enableThinking;
if (input.enableTools !== undefined) updates.enableTools = input.enableTools;
if (input.maxStepCount !== undefined) updates.maxStepCount = input.maxStepCount;
if (input.isCallable !== undefined) updates.isCallable = input.isCallable;
if (input.sortOrder !== undefined) updates.sortOrder = input.sortOrder;
if (input.enabled !== undefined) updates.enabled = input.enabled;
if (input.isDefault !== undefined) {
updates.isDefault = input.isDefault;
if (input.isDefault === 1) {
await dbGlobal
.update(agents)
.set({ isDefault: 0 })
.where(and(sql`${agents.id} != ${id}`));
}
}
if (Object.keys(updates).length > 0) {
await dbGlobal.update(agents).set(updates).where(eq(agents.id, id));
}
if (input.toolIds !== undefined) {
await dbGlobal.delete(agentToolAssociations).where(eq(agentToolAssociations.agentId, id));
for (let i = 0; i < input.toolIds.length; i++) {
const toolId = input.toolIds[i]!;
await dbGlobal.insert(agentToolAssociations).values({
agentId: id,
toolId,
sortOrder: i,
});
}
}
return getAgentById(id);
}
export async function deleteAgent(id: number): Promise<void> {
await dbGlobal.delete(agentToolAssociations).where(eq(agentToolAssociations.agentId, id));
await dbGlobal.delete(agents).where(eq(agents.id, id));
}