128 lines
3.2 KiB
TypeScript
128 lines
3.2 KiB
TypeScript
import { PromptTemplate } from "@langchain/core/prompts"
|
|
import { ChatOpenAI } from "@langchain/openai"
|
|
import { CallbackHandler, Langfuse } from "langfuse-langchain"
|
|
import { z } from "zod"
|
|
|
|
import db from "../db"
|
|
|
|
/**
|
|
* 获取Langfuse
|
|
* @returns
|
|
*/
|
|
const getLangfuse = async () => {
|
|
const langfuseParams = {
|
|
publicKey: await db.appConfig.getLangfusePk(),
|
|
secretKey: await db.appConfig.getLangfuseSk(),
|
|
baseUrl: "http://langfuse.c5-cloudml.xiaomi.srv",
|
|
}
|
|
return {
|
|
langfuseHandler: new CallbackHandler(langfuseParams),
|
|
langfuse: new Langfuse(langfuseParams),
|
|
}
|
|
}
|
|
|
|
const modelMap = {
|
|
"deepseek-chat": {
|
|
model: "deepseek-chat",
|
|
apiKey: "xx",
|
|
baseURL: "http://10.38.214.162:8003/v1",
|
|
},
|
|
"qwen2-72b-instruct-int4": {
|
|
model: "qwen2-72b-instruct-int4",
|
|
apiKey: "xx",
|
|
baseURL: "http://10.38.214.206:8000/v1",
|
|
},
|
|
"gpt-4o": {
|
|
model: "gpt-4o",
|
|
apiKey: "sk-EhbBTR0QjhH22iLr9aCb04D2B0F44f88A07c2924Eb54CfA4",
|
|
baseURL: "https://api.gpt.ge/v1",
|
|
},
|
|
"qwen-72b-instruct-int4/v1": {
|
|
model: "qwen-72b-instruct-int4/v1",
|
|
apiKey: "xx",
|
|
baseURL:
|
|
"http://ms-13871-qwen-model-128k-9-1012195754.kscn-tj5-prod2-cloudml.xiaomi.srv/v1",
|
|
},
|
|
}
|
|
|
|
/**
|
|
* 获取模型
|
|
* @param modelName 模型名称
|
|
* @param temperature 温度
|
|
*/
|
|
const getModel = async (modelName: keyof typeof modelMap, temperature = 0) => {
|
|
const { model, apiKey, baseURL } = modelMap[modelName]
|
|
return new ChatOpenAI(
|
|
{ temperature, model, apiKey },
|
|
{
|
|
baseURL,
|
|
}
|
|
)
|
|
}
|
|
|
|
const timeConfig = z.object({
|
|
startTime: z.string().describe("开始时间，格式为 YYYY-MM-DD HH:mm:ss"),
|
|
endTime: z.string().describe("结束时间，格式为 YYYY-MM-DD HH:mm:ss"),
|
|
})
|
|
|
|
/**
|
|
* 解析时间
|
|
* @param userInput 用户输入
|
|
* @returns
|
|
*/
|
|
const parseTime = async (userInput: string) => {
|
|
const model = await getModel("deepseek-chat")
|
|
const structuredLlm = model.withStructuredOutput(timeConfig, { name: "time" })
|
|
return await structuredLlm.invoke(
|
|
`
|
|
当前时间为 ${new Date().toLocaleString("zh-CN", { timeZone: "Asia/Shanghai" })}
|
|
你是一个专业的语义解析工程师，给定以下用户输入，帮我解析出开始时间和结束时间
|
|
如果不包含时间信息，请返回当天的起始时间到当前时间
|
|
用户输入：
|
|
\`\`\`
|
|
${userInput.replaceAll("`", " ")}
|
|
\`\`\`
|
|
`
|
|
)
|
|
}
|
|
|
|
/**
|
|
* 调用LLM模型
|
|
* @param promptName 提示Key
|
|
* @param variables 变量
|
|
* @param temperature 温度
|
|
* @returns
|
|
*/
|
|
const invoke = async (
|
|
promptName: string,
|
|
variables: Record<string, any>,
|
|
temperature = 0
|
|
) => {
|
|
const { langfuse, langfuseHandler } = await getLangfuse()
|
|
const prompt = await langfuse.getPrompt(promptName)
|
|
const config = prompt.config as { modelName: keyof typeof modelMap }
|
|
|
|
const langchainTextPrompt = PromptTemplate.fromTemplate(
|
|
prompt.getLangchainPrompt()
|
|
).withConfig({
|
|
metadata: { langfusePrompt: prompt },
|
|
})
|
|
|
|
const chain = langchainTextPrompt.pipe(
|
|
await getModel(config.modelName, temperature)
|
|
)
|
|
|
|
const { content } = await chain.invoke(variables, {
|
|
callbacks: [langfuseHandler],
|
|
})
|
|
|
|
return content
|
|
}
|
|
|
|
const llm = {
|
|
parseTime,
|
|
invoke,
|
|
}
|
|
|
|
export default llm
|