Spaces:
Sleeping
Sleeping
File size: 2,787 Bytes
05c5ed5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 | import { streamObject } from "ai";
import { customModelProvider } from "lib/ai/models";
import { buildAgentGenerationPrompt } from "lib/ai/prompts";
import globalLogger from "logger";
import { ChatModel } from "app-types/chat";
import { getSession } from "auth/server";
import { colorize } from "consola/utils";
import { AgentGenerateSchema } from "app-types/agent";
import { z } from "zod";
import { loadAppDefaultTools } from "../../chat/shared.chat";
import { workflowRepository } from "lib/db/repository";
import { safe } from "ts-safe";
import { objectFlow } from "lib/utils";
import { mcpClientsManager } from "lib/ai/mcp/mcp-manager";
const logger = globalLogger.withDefaults({
message: colorize("blackBright", `Agent Generate API: `),
});
export async function POST(request: Request) {
try {
const json = await request.json();
const { chatModel, message = "hello" } = json as {
chatModel?: ChatModel;
message: string;
};
logger.info(`chatModel: ${chatModel?.provider}/${chatModel?.model}`);
const session = await getSession();
if (!session) {
return new Response("Unauthorized", { status: 401 });
}
const toolNames = new Set<string>();
await safe(loadAppDefaultTools)
.ifOk((appTools) => {
objectFlow(appTools).forEach((_, toolName) => {
toolNames.add(toolName);
});
})
.unwrap();
await safe(mcpClientsManager.tools())
.ifOk((tools) => {
objectFlow(tools).forEach((mcp) => {
toolNames.add(mcp._originToolName);
});
})
.unwrap();
await safe(workflowRepository.selectExecuteAbility(session.user.id))
.ifOk((tools) => {
tools.forEach((tool) => {
toolNames.add(tool.name);
});
})
.unwrap();
const dynamicAgentTable = AgentGenerateSchema.extend({
tools: z
.array(
z.enum(
Array.from(toolNames).length > 0
? ([
Array.from(toolNames)[0],
...Array.from(toolNames).slice(1),
] as [string, ...string[]])
: ([""] as [string]),
),
)
.describe("Agent allowed tools name")
.nullable()
.default([]),
});
const system = buildAgentGenerationPrompt(Array.from(toolNames));
const { getUserPreferences } = await import("@/lib/user/server");
const userPreferences = (await getUserPreferences(session.user.id)) || undefined;
const result = streamObject({
model: customModelProvider.getDynamicModel(chatModel, userPreferences),
system,
prompt: message,
schema: dynamicAgentTable,
});
return result.toTextStreamResponse();
} catch (error) {
logger.error(error);
}
}
|