@lobehub/chat
Version:
Lobe Chat - an open-source, high-performance chatbot framework that supports speech synthesis, multimodal, and extensible Function Call plugin system. Supports one-click free deployment of your private ChatGPT/LLM web application.
139 lines (113 loc) • 4.81 kB
text/typescript
import { TRPCError } from '@trpc/server';
import OpenAI from 'openai';
import { z } from 'zod';
import { chainAnswerWithContext } from '@/chains/answerWithContext';
import { DEFAULT_EMBEDDING_MODEL, DEFAULT_MODEL } from '@/const/settings';
import { ChunkModel } from '@/database/models/chunk';
import { EmbeddingModel } from '@/database/models/embedding';
import { FileModel } from '@/database/models/file';
import {
EvalDatasetRecordModel,
EvalEvaluationModel,
EvaluationRecordModel,
} from '@/database/server/models/ragEval';
import { ModelProvider } from '@/libs/model-runtime';
import { asyncAuthedProcedure, asyncRouter as router } from '@/libs/trpc/async';
import { initAgentRuntimeWithUserPayload } from '@/server/modules/AgentRuntime';
import { ChunkService } from '@/server/services/chunk';
import { AsyncTaskError } from '@/types/asyncTask';
import { EvalEvaluationStatus } from '@/types/eval';
const ragEvalProcedure = asyncAuthedProcedure.use(async (opts) => {
const { ctx } = opts;
return opts.next({
ctx: {
chunkModel: new ChunkModel(ctx.serverDB, ctx.userId),
chunkService: new ChunkService(ctx.userId),
datasetRecordModel: new EvalDatasetRecordModel(ctx.userId),
embeddingModel: new EmbeddingModel(ctx.serverDB, ctx.userId),
evalRecordModel: new EvaluationRecordModel(ctx.userId),
evaluationModel: new EvalEvaluationModel(ctx.userId),
fileModel: new FileModel(ctx.serverDB, ctx.userId),
},
});
});
export const ragEvalRouter = router({
runRecordEvaluation: ragEvalProcedure
.input(
z.object({
evalRecordId: z.number(),
}),
)
.mutation(async ({ ctx, input }) => {
const evalRecord = await ctx.evalRecordModel.findById(input.evalRecordId);
if (!evalRecord) {
throw new TRPCError({ code: 'BAD_REQUEST', message: 'Evaluation not found' });
}
const now = Date.now();
try {
const agentRuntime = await initAgentRuntimeWithUserPayload(
ModelProvider.OpenAI,
ctx.jwtPayload,
);
const { question, languageModel, embeddingModel } = evalRecord;
let questionEmbeddingId = evalRecord.questionEmbeddingId;
let context = evalRecord.context;
// 如果不存在 questionEmbeddingId,那么就需要做一次 embedding
if (!questionEmbeddingId) {
const embeddings = await agentRuntime.embeddings({
dimensions: 1024,
input: question,
model: !!embeddingModel ? embeddingModel : DEFAULT_EMBEDDING_MODEL,
});
const embeddingId = await ctx.embeddingModel.create({
embeddings: embeddings?.[0],
model: embeddingModel,
});
await ctx.evalRecordModel.update(evalRecord.id, {
questionEmbeddingId: embeddingId,
});
questionEmbeddingId = embeddingId;
}
// 如果不存在 context,那么就需要做一次检索
if (!context || context.length === 0) {
const datasetRecord = await ctx.datasetRecordModel.findById(evalRecord.datasetRecordId);
const embeddingItem = await ctx.embeddingModel.findById(questionEmbeddingId);
const chunks = await ctx.chunkModel.semanticSearchForChat({
embedding: embeddingItem!.embeddings!,
fileIds: datasetRecord!.referenceFiles!,
query: evalRecord.question,
});
context = chunks.map((item) => item.text).filter(Boolean) as string[];
await ctx.evalRecordModel.update(evalRecord.id, { context });
}
// 做一次生成 LLM 答案生成
const { messages } = chainAnswerWithContext({ context, knowledge: [], question });
const response = await agentRuntime.chat({
messages: messages!,
model: !!languageModel ? languageModel : DEFAULT_MODEL,
responseMode: 'json',
stream: false,
temperature: 1,
});
const data = (await response.json()) as OpenAI.ChatCompletion;
const answer = data.choices[0].message.content;
await ctx.evalRecordModel.update(input.evalRecordId, {
answer,
duration: Date.now() - now,
languageModel,
status: EvalEvaluationStatus.Success,
});
return { success: true };
} catch (e) {
await ctx.evalRecordModel.update(input.evalRecordId, {
error: new AsyncTaskError((e as Error).name, (e as Error).message),
status: EvalEvaluationStatus.Error,
});
await ctx.evaluationModel.update(evalRecord.evaluationId, {
status: EvalEvaluationStatus.Error,
});
console.error('[RAGEvaluation] error', e);
return { success: false };
}
}),
});