mirror of
https://github.com/toeverything/AFFiNE.git
synced 2026-07-17 10:06:17 +08:00
feat: add session impl (#6254)
This commit is contained in:
darkskygit
@@ -2,10 +2,11 @@ import { ServerFeature } from '../../core/config';
|
||||
import { Plugin } from '../registry';
|
||||
import { PromptService } from './prompt';
|
||||
import { assertProvidersConfigs, CopilotProviderService } from './providers';
|
||||
import { ChatSessionService } from './session';
|
||||
|
||||
@Plugin({
|
||||
name: 'copilot',
|
||||
providers: [PromptService, CopilotProviderService],
|
||||
providers: [ChatSessionService, PromptService, CopilotProviderService],
|
||||
contributesTo: ServerFeature.Copilot,
|
||||
if: config => {
|
||||
if (config.flavor.graphql) {
|
||||
|
||||
@@ -1,7 +1,124 @@
|
||||
import { Injectable } from '@nestjs/common';
|
||||
import { PrismaClient } from '@prisma/client';
|
||||
import { AiPrompt, PrismaClient } from '@prisma/client';
|
||||
import Mustache from 'mustache';
|
||||
import { Tiktoken } from 'tiktoken';
|
||||
|
||||
import { ChatMessage } from './types';
|
||||
import {
|
||||
getTokenEncoder,
|
||||
PromptMessage,
|
||||
PromptMessageSchema,
|
||||
PromptParams,
|
||||
} from './types';
|
||||
|
||||
// disable escaping
|
||||
Mustache.escape = (text: string) => text;
|
||||
|
||||
function extractMustacheParams(template: string) {
|
||||
const regex = /\{\{\s*([^{}]+)\s*\}\}/g;
|
||||
const params = [];
|
||||
let match;
|
||||
|
||||
while ((match = regex.exec(template)) !== null) {
|
||||
params.push(match[1]);
|
||||
}
|
||||
|
||||
return Array.from(new Set(params));
|
||||
}
|
||||
|
||||
export class ChatPrompt {
|
||||
public readonly encoder?: Tiktoken;
|
||||
private readonly promptTokenSize: number;
|
||||
private readonly templateParamKeys: string[] = [];
|
||||
private readonly templateParams: PromptParams = {};
|
||||
|
||||
static createFromPrompt(
|
||||
options: Omit<AiPrompt, 'id' | 'createdAt'> & {
|
||||
messages: PromptMessage[];
|
||||
}
|
||||
) {
|
||||
return new ChatPrompt(
|
||||
options.name,
|
||||
options.action,
|
||||
options.model,
|
||||
options.messages
|
||||
);
|
||||
}
|
||||
|
||||
constructor(
|
||||
public readonly name: string,
|
||||
public readonly action: string | null,
|
||||
public readonly model: string | null,
|
||||
private readonly messages: PromptMessage[]
|
||||
) {
|
||||
this.encoder = getTokenEncoder(model);
|
||||
this.promptTokenSize =
|
||||
this.encoder?.encode_ordinary(messages.map(m => m.content).join('') || '')
|
||||
.length || 0;
|
||||
this.templateParamKeys = extractMustacheParams(
|
||||
messages.map(m => m.content).join('')
|
||||
);
|
||||
this.templateParams = messages.reduce(
|
||||
(acc, m) => Object.assign(acc, m.params),
|
||||
{} as PromptParams
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* get prompt token size
|
||||
*/
|
||||
get tokens() {
|
||||
return this.promptTokenSize;
|
||||
}
|
||||
|
||||
/**
|
||||
* get prompt param keys in template
|
||||
*/
|
||||
get paramKeys() {
|
||||
return this.templateParamKeys.slice();
|
||||
}
|
||||
|
||||
/**
|
||||
* get prompt params
|
||||
*/
|
||||
get params() {
|
||||
return { ...this.templateParams };
|
||||
}
|
||||
|
||||
encode(message: string) {
|
||||
return this.encoder?.encode_ordinary(message).length || 0;
|
||||
}
|
||||
|
||||
private checkParams(params: PromptParams) {
|
||||
const selfParams = this.templateParams;
|
||||
for (const key of Object.keys(selfParams)) {
|
||||
const options = selfParams[key];
|
||||
const income = params[key];
|
||||
if (
|
||||
typeof income !== 'string' ||
|
||||
(Array.isArray(options) && !options.includes(income))
|
||||
) {
|
||||
throw new Error(`Invalid param: ${key}`);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* render prompt messages with params
|
||||
* @param params record of params, e.g. { name: 'Alice' }
|
||||
* @returns e.g. [{ role: 'system', content: 'Hello, {{name}}' }] => [{ role: 'system', content: 'Hello, Alice' }]
|
||||
*/
|
||||
finish(params: PromptParams) {
|
||||
this.checkParams(params);
|
||||
return this.messages.map(m => ({
|
||||
...m,
|
||||
content: Mustache.render(m.content, params),
|
||||
}));
|
||||
}
|
||||
|
||||
free() {
|
||||
this.encoder?.free();
|
||||
}
|
||||
}
|
||||
|
||||
@Injectable()
|
||||
export class PromptService {
|
||||
@@ -22,51 +139,74 @@ export class PromptService {
|
||||
* @param name prompt name
|
||||
* @returns prompt messages
|
||||
*/
|
||||
async get(name: string): Promise<ChatMessage[]> {
|
||||
return this.db.aiPrompt.findMany({
|
||||
where: {
|
||||
name,
|
||||
},
|
||||
select: {
|
||||
role: true,
|
||||
content: true,
|
||||
},
|
||||
orderBy: {
|
||||
idx: 'asc',
|
||||
},
|
||||
});
|
||||
async get(name: string): Promise<ChatPrompt | null> {
|
||||
return this.db.aiPrompt
|
||||
.findUnique({
|
||||
where: {
|
||||
name,
|
||||
},
|
||||
select: {
|
||||
name: true,
|
||||
action: true,
|
||||
model: true,
|
||||
messages: {
|
||||
select: {
|
||||
role: true,
|
||||
content: true,
|
||||
params: true,
|
||||
},
|
||||
orderBy: {
|
||||
idx: 'asc',
|
||||
},
|
||||
},
|
||||
},
|
||||
})
|
||||
.then(p => {
|
||||
const messages = PromptMessageSchema.array().safeParse(p?.messages);
|
||||
if (p && messages.success) {
|
||||
return ChatPrompt.createFromPrompt({ ...p, messages: messages.data });
|
||||
}
|
||||
return null;
|
||||
});
|
||||
}
|
||||
|
||||
async set(name: string, messages: ChatMessage[]) {
|
||||
return this.db.$transaction(async tx => {
|
||||
const prompts = await tx.aiPrompt.count({ where: { name } });
|
||||
if (prompts > 0) {
|
||||
return 0;
|
||||
}
|
||||
return tx.aiPrompt
|
||||
.createMany({
|
||||
data: messages.map((m, idx) => ({ name, idx, ...m })),
|
||||
})
|
||||
.then(ret => ret.count);
|
||||
});
|
||||
async set(name: string, messages: PromptMessage[]) {
|
||||
return await this.db.aiPrompt
|
||||
.create({
|
||||
data: {
|
||||
name,
|
||||
messages: {
|
||||
create: messages.map((m, idx) => ({
|
||||
idx,
|
||||
...m,
|
||||
params: m.params || undefined,
|
||||
})),
|
||||
},
|
||||
},
|
||||
})
|
||||
.then(ret => ret.id);
|
||||
}
|
||||
|
||||
async update(name: string, messages: ChatMessage[]) {
|
||||
return this.db.$transaction(async tx => {
|
||||
await tx.aiPrompt.deleteMany({ where: { name } });
|
||||
return tx.aiPrompt
|
||||
.createMany({
|
||||
data: messages.map((m, idx) => ({ name, idx, ...m })),
|
||||
})
|
||||
.then(ret => ret.count);
|
||||
});
|
||||
async update(name: string, messages: PromptMessage[]) {
|
||||
return this.db.aiPrompt
|
||||
.update({
|
||||
where: { name },
|
||||
data: {
|
||||
messages: {
|
||||
// cleanup old messages
|
||||
deleteMany: {},
|
||||
create: messages.map((m, idx) => ({
|
||||
idx,
|
||||
...m,
|
||||
params: m.params || undefined,
|
||||
})),
|
||||
},
|
||||
},
|
||||
})
|
||||
.then(ret => ret.id);
|
||||
}
|
||||
|
||||
async delete(name: string) {
|
||||
return this.db.aiPrompt
|
||||
.deleteMany({
|
||||
where: { name },
|
||||
})
|
||||
.then(ret => ret.count);
|
||||
return this.db.aiPrompt.delete({ where: { name } }).then(ret => ret.id);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,203 @@
|
||||
import { randomUUID } from 'node:crypto';
|
||||
|
||||
import { Injectable, Logger } from '@nestjs/common';
|
||||
import { PrismaClient } from '@prisma/client';
|
||||
|
||||
import { ChatPrompt, PromptService } from './prompt';
|
||||
import {
|
||||
ChatMessage,
|
||||
ChatMessageSchema,
|
||||
PromptMessage,
|
||||
PromptParams,
|
||||
} from './types';
|
||||
|
||||
export interface ChatSessionOptions {
|
||||
userId: string;
|
||||
workspaceId: string;
|
||||
docId: string;
|
||||
promptName: string;
|
||||
}
|
||||
|
||||
export interface ChatSessionState
|
||||
extends Omit<ChatSessionOptions, 'promptName'> {
|
||||
// connect ids
|
||||
sessionId: string;
|
||||
// states
|
||||
prompt: ChatPrompt;
|
||||
messages: ChatMessage[];
|
||||
}
|
||||
|
||||
export class ChatSession implements AsyncDisposable {
|
||||
constructor(
|
||||
private readonly state: ChatSessionState,
|
||||
private readonly dispose?: (state: ChatSessionState) => Promise<void>,
|
||||
private readonly maxTokenSize = 3840
|
||||
) {}
|
||||
|
||||
get model() {
|
||||
return this.state.prompt.model;
|
||||
}
|
||||
|
||||
push(message: ChatMessage) {
|
||||
this.state.messages.push(message);
|
||||
}
|
||||
|
||||
pop() {
|
||||
this.state.messages.pop();
|
||||
}
|
||||
|
||||
private takeMessages(): ChatMessage[] {
|
||||
if (this.state.prompt.action) {
|
||||
const messages = this.state.messages;
|
||||
return messages.slice(messages.length - 1);
|
||||
}
|
||||
const ret = [];
|
||||
const messages = this.state.messages.slice();
|
||||
|
||||
let size = this.state.prompt.tokens;
|
||||
while (messages.length) {
|
||||
const message = messages.pop();
|
||||
if (!message) break;
|
||||
|
||||
size += this.state.prompt.encode(message.content);
|
||||
if (size > this.maxTokenSize) {
|
||||
break;
|
||||
}
|
||||
ret.push(message);
|
||||
}
|
||||
ret.reverse();
|
||||
|
||||
return ret;
|
||||
}
|
||||
|
||||
finish(params: PromptParams): PromptMessage[] {
|
||||
const messages = this.takeMessages();
|
||||
return [...this.state.prompt.finish(params), ...messages];
|
||||
}
|
||||
|
||||
async save() {
|
||||
await this.dispose?.(this.state);
|
||||
}
|
||||
|
||||
async [Symbol.asyncDispose]() {
|
||||
this.state.prompt.free();
|
||||
await this.save?.();
|
||||
}
|
||||
}
|
||||
|
||||
@Injectable()
|
||||
export class ChatSessionService {
|
||||
private readonly logger = new Logger(ChatSessionService.name);
|
||||
constructor(
|
||||
private readonly db: PrismaClient,
|
||||
private readonly prompt: PromptService
|
||||
) {}
|
||||
|
||||
private async setSession(state: ChatSessionState): Promise<void> {
|
||||
await this.db.aiSession.upsert({
|
||||
where: {
|
||||
id: state.sessionId,
|
||||
},
|
||||
update: {
|
||||
messages: {
|
||||
create: state.messages.map((m, idx) => ({ idx, ...m })),
|
||||
},
|
||||
},
|
||||
create: {
|
||||
id: state.sessionId,
|
||||
messages: { create: state.messages },
|
||||
// connect
|
||||
user: { connect: { id: state.userId } },
|
||||
workspace: { connect: { id: state.workspaceId } },
|
||||
doc: {
|
||||
connect: {
|
||||
id_workspaceId: {
|
||||
id: state.docId,
|
||||
workspaceId: state.workspaceId,
|
||||
},
|
||||
},
|
||||
},
|
||||
prompt: { connect: { name: state.prompt.name } },
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
private async getSession(
|
||||
sessionId: string
|
||||
): Promise<ChatSessionState | undefined> {
|
||||
return await this.db.aiSession
|
||||
.findUnique({
|
||||
where: { id: sessionId },
|
||||
select: {
|
||||
id: true,
|
||||
userId: true,
|
||||
workspaceId: true,
|
||||
docId: true,
|
||||
messages: true,
|
||||
prompt: {
|
||||
select: {
|
||||
name: true,
|
||||
action: true,
|
||||
model: true,
|
||||
messages: {
|
||||
select: {
|
||||
role: true,
|
||||
content: true,
|
||||
},
|
||||
orderBy: {
|
||||
idx: 'asc',
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
})
|
||||
.then(async session => {
|
||||
if (!session) return;
|
||||
const messages = ChatMessageSchema.array().safeParse(session.messages);
|
||||
|
||||
return {
|
||||
sessionId: session.id,
|
||||
userId: session.userId,
|
||||
workspaceId: session.workspaceId,
|
||||
docId: session.docId,
|
||||
prompt: ChatPrompt.createFromPrompt(session.prompt),
|
||||
messages: messages.success ? messages.data : [],
|
||||
};
|
||||
});
|
||||
}
|
||||
|
||||
async create(options: ChatSessionOptions): Promise<string> {
|
||||
const sessionId = randomUUID();
|
||||
const prompt = await this.prompt.get(options.promptName);
|
||||
if (!prompt) {
|
||||
this.logger.error(`Prompt not found: ${options.promptName}`);
|
||||
throw new Error('Prompt not found');
|
||||
}
|
||||
await this.setSession({ ...options, sessionId, prompt, messages: [] });
|
||||
return sessionId;
|
||||
}
|
||||
|
||||
/**
|
||||
* usage:
|
||||
* ``` typescript
|
||||
* {
|
||||
* // allocate a session, can be reused chat in about 12 hours with same session
|
||||
* await using session = await session.get(sessionId);
|
||||
* session.push(message);
|
||||
* copilot.generateText(session.finish(), model);
|
||||
* }
|
||||
* // session will be disposed after the block
|
||||
* @param sessionId session id
|
||||
* @returns
|
||||
*/
|
||||
async get(sessionId: string): Promise<ChatSession | null> {
|
||||
const state = await this.getSession(sessionId);
|
||||
if (state) {
|
||||
return new ChatSession(state, async state => {
|
||||
await this.setSession(state);
|
||||
});
|
||||
}
|
||||
return null;
|
||||
}
|
||||
}
|
||||
@@ -1,5 +1,11 @@
|
||||
import { AiPromptRole } from '@prisma/client';
|
||||
import type { ClientOptions as OpenAIClientOptions } from 'openai';
|
||||
import {
|
||||
encoding_for_model,
|
||||
get_encoding,
|
||||
Tiktoken,
|
||||
TiktokenModel,
|
||||
} from 'tiktoken';
|
||||
import { z } from 'zod';
|
||||
|
||||
export interface CopilotConfig {
|
||||
@@ -9,6 +15,76 @@ export interface CopilotConfig {
|
||||
};
|
||||
}
|
||||
|
||||
export enum AvailableModels {
|
||||
// text to text
|
||||
Gpt4VisionPreview = 'gpt-4-vision-preview',
|
||||
Gpt4TurboPreview = 'gpt-4-turbo-preview',
|
||||
Gpt35Turbo = 'gpt-3.5-turbo',
|
||||
// embeddings
|
||||
TextEmbedding3Large = 'text-embedding-3-large',
|
||||
TextEmbedding3Small = 'text-embedding-3-small',
|
||||
TextEmbeddingAda002 = 'text-embedding-ada-002',
|
||||
// moderation
|
||||
TextModerationLatest = 'text-moderation-latest',
|
||||
TextModerationStable = 'text-moderation-stable',
|
||||
}
|
||||
|
||||
export type AvailableModel = keyof typeof AvailableModels;
|
||||
|
||||
export function getTokenEncoder(model?: string | null): Tiktoken | undefined {
|
||||
if (!model) return undefined;
|
||||
const modelStr = AvailableModels[model as AvailableModel];
|
||||
if (!modelStr) return undefined;
|
||||
if (modelStr.startsWith('gpt')) {
|
||||
return encoding_for_model(modelStr as TiktokenModel);
|
||||
} else if (modelStr.startsWith('dall')) {
|
||||
// dalle don't need to calc the token
|
||||
return undefined;
|
||||
} else {
|
||||
return get_encoding('cl100k_base');
|
||||
}
|
||||
}
|
||||
|
||||
// ======== ChatMessage ========
|
||||
|
||||
export const ChatMessageRole = Object.values(AiPromptRole) as [
|
||||
'system',
|
||||
'assistant',
|
||||
'user',
|
||||
];
|
||||
|
||||
export const PromptMessageSchema = z.object({
|
||||
role: z.enum(ChatMessageRole),
|
||||
content: z.string(),
|
||||
attachments: z.array(z.string()).optional(),
|
||||
params: z
|
||||
.record(z.union([z.string(), z.array(z.string())]))
|
||||
.optional()
|
||||
.nullable(),
|
||||
});
|
||||
|
||||
export type PromptMessage = z.infer<typeof PromptMessageSchema>;
|
||||
|
||||
export type PromptParams = NonNullable<PromptMessage['params']>;
|
||||
|
||||
export const ChatMessageSchema = PromptMessageSchema.extend({
|
||||
createdAt: z.date(),
|
||||
}).strict();
|
||||
|
||||
export type ChatMessage = z.infer<typeof ChatMessageSchema>;
|
||||
|
||||
export const ChatHistorySchema = z
|
||||
.object({
|
||||
sessionId: z.string(),
|
||||
tokens: z.number(),
|
||||
messages: z.array(ChatMessageSchema),
|
||||
})
|
||||
.strict();
|
||||
|
||||
export type ChatHistory = z.infer<typeof ChatHistorySchema>;
|
||||
|
||||
// ======== Provider Interface ========
|
||||
|
||||
export enum CopilotProviderType {
|
||||
FAL = 'fal',
|
||||
OpenAI = 'openai',
|
||||
@@ -25,24 +101,26 @@ export interface CopilotProvider {
|
||||
getCapabilities(): CopilotProviderCapability[];
|
||||
}
|
||||
|
||||
export const ChatMessageSchema = z
|
||||
.object({
|
||||
role: z.enum(
|
||||
Array.from(Object.values(AiPromptRole)) as [
|
||||
'system' | 'assistant' | 'user',
|
||||
]
|
||||
),
|
||||
content: z.string(),
|
||||
})
|
||||
.strict();
|
||||
|
||||
export type ChatMessage = z.infer<typeof ChatMessageSchema>;
|
||||
|
||||
export interface CopilotTextToTextProvider extends CopilotProvider {
|
||||
generateText(messages: ChatMessage[], model: string): Promise<string>;
|
||||
generateText(
|
||||
messages: PromptMessage[],
|
||||
model: string,
|
||||
options: {
|
||||
temperature?: number;
|
||||
maxTokens?: number;
|
||||
signal?: AbortSignal;
|
||||
user?: string;
|
||||
}
|
||||
): Promise<string>;
|
||||
generateTextStream(
|
||||
messages: ChatMessage[],
|
||||
model: string
|
||||
messages: PromptMessage[],
|
||||
model: string,
|
||||
options: {
|
||||
temperature?: number;
|
||||
maxTokens?: number;
|
||||
signal?: AbortSignal;
|
||||
user?: string;
|
||||
}
|
||||
): AsyncIterable<string>;
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user