Skip to content

Commit

Permalink
feat: add session impl (#6254)
Browse files Browse the repository at this point in the history
  • Loading branch information
darkskygit committed Apr 10, 2024
1 parent 8a02c81 commit 46a368d
Show file tree
Hide file tree
Showing 16 changed files with 1,033 additions and 88 deletions.
Original file line number Diff line number Diff line change
@@ -1,11 +1,12 @@
-- CreateTable
CREATE TABLE "ai_sessions" (
"id" VARCHAR NOT NULL,
"id" VARCHAR(36) NOT NULL,
"user_id" VARCHAR NOT NULL,
"workspace_id" VARCHAR NOT NULL,
"doc_id" VARCHAR NOT NULL,
"prompt_name" VARCHAR NOT NULL,
"action" BOOLEAN NOT NULL,
"flavor" VARCHAR NOT NULL,
"model" VARCHAR NOT NULL,
"messages" JSON NOT NULL,
"created_at" TIMESTAMPTZ(6) NOT NULL DEFAULT CURRENT_TIMESTAMP,
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,90 @@
/*
Warnings:
- You are about to drop the `ai_prompts` table. If the table is not empty, all the data it contains will be lost.
- You are about to drop the `ai_sessions` table. If the table is not empty, all the data it contains will be lost.
*/
-- DropForeignKey
ALTER TABLE "ai_sessions" DROP CONSTRAINT "ai_sessions_doc_id_workspace_id_fkey";

-- DropForeignKey
ALTER TABLE "ai_sessions" DROP CONSTRAINT "ai_sessions_user_id_fkey";

-- DropForeignKey
ALTER TABLE "ai_sessions" DROP CONSTRAINT "ai_sessions_workspace_id_fkey";

-- DropTable
DROP TABLE "ai_prompts";

-- DropTable
DROP TABLE "ai_sessions";

-- CreateTable
CREATE TABLE "ai_prompts_messages" (
"prompt_id" INTEGER NOT NULL,
"idx" INTEGER NOT NULL,
"role" "AiPromptRole" NOT NULL,
"content" TEXT NOT NULL,
"params" JSON,
"created_at" TIMESTAMPTZ(6) NOT NULL DEFAULT CURRENT_TIMESTAMP
);

-- CreateTable
CREATE TABLE "ai_prompts_metadata" (
"id" SERIAL NOT NULL,
"name" VARCHAR(32) NOT NULL,
"action" VARCHAR,
"model" VARCHAR,
"created_at" TIMESTAMPTZ(6) NOT NULL DEFAULT CURRENT_TIMESTAMP,

CONSTRAINT "ai_prompts_metadata_pkey" PRIMARY KEY ("id")
);

-- CreateTable
CREATE TABLE "ai_sessions_messages" (
"id" VARCHAR(36) NOT NULL,
"session_id" VARCHAR(36) NOT NULL,
"role" "AiPromptRole" NOT NULL,
"content" TEXT NOT NULL,
"created_at" TIMESTAMPTZ(6) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"updated_at" TIMESTAMPTZ(6) NOT NULL,

CONSTRAINT "ai_sessions_messages_pkey" PRIMARY KEY ("id")
);

-- CreateTable
CREATE TABLE "ai_sessions_metadata" (
"id" VARCHAR(36) NOT NULL,
"user_id" VARCHAR(36) NOT NULL,
"workspace_id" VARCHAR(36) NOT NULL,
"doc_id" VARCHAR(36) NOT NULL,
"prompt_name" VARCHAR(32) NOT NULL,
"created_at" TIMESTAMPTZ(6) NOT NULL DEFAULT CURRENT_TIMESTAMP,

CONSTRAINT "ai_sessions_metadata_pkey" PRIMARY KEY ("id")
);

-- CreateIndex
CREATE UNIQUE INDEX "ai_prompts_messages_prompt_id_idx_key" ON "ai_prompts_messages"("prompt_id", "idx");

-- CreateIndex
CREATE UNIQUE INDEX "ai_prompts_metadata_name_key" ON "ai_prompts_metadata"("name");

-- AddForeignKey
ALTER TABLE "ai_prompts_messages" ADD CONSTRAINT "ai_prompts_messages_prompt_id_fkey" FOREIGN KEY ("prompt_id") REFERENCES "ai_prompts_metadata"("id") ON DELETE CASCADE ON UPDATE CASCADE;

-- AddForeignKey
ALTER TABLE "ai_sessions_messages" ADD CONSTRAINT "ai_sessions_messages_session_id_fkey" FOREIGN KEY ("session_id") REFERENCES "ai_sessions_metadata"("id") ON DELETE CASCADE ON UPDATE CASCADE;

-- AddForeignKey
ALTER TABLE "ai_sessions_metadata" ADD CONSTRAINT "ai_sessions_metadata_user_id_fkey" FOREIGN KEY ("user_id") REFERENCES "users"("id") ON DELETE CASCADE ON UPDATE CASCADE;

-- AddForeignKey
ALTER TABLE "ai_sessions_metadata" ADD CONSTRAINT "ai_sessions_metadata_workspace_id_fkey" FOREIGN KEY ("workspace_id") REFERENCES "workspaces"("id") ON DELETE CASCADE ON UPDATE CASCADE;

-- AddForeignKey
ALTER TABLE "ai_sessions_metadata" ADD CONSTRAINT "ai_sessions_metadata_doc_id_workspace_id_fkey" FOREIGN KEY ("doc_id", "workspace_id") REFERENCES "snapshots"("guid", "workspace_id") ON DELETE CASCADE ON UPDATE CASCADE;

-- AddForeignKey
ALTER TABLE "ai_sessions_metadata" ADD CONSTRAINT "ai_sessions_metadata_prompt_name_fkey" FOREIGN KEY ("prompt_name") REFERENCES "ai_prompts_metadata"("name") ON DELETE CASCADE ON UPDATE CASCADE;
3 changes: 3 additions & 0 deletions packages/backend/server/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,7 @@
"keyv": "^4.5.4",
"lodash-es": "^4.17.21",
"mixpanel": "^0.18.0",
"mustache": "^4.2.0",
"nanoid": "^5.0.6",
"nest-commander": "^3.12.5",
"nestjs-throttler-storage-redis": "^0.4.1",
Expand All @@ -87,6 +88,7 @@
"semver": "^7.6.0",
"socket.io": "^4.7.4",
"stripe": "^14.18.0",
"tiktoken": "^1.0.13",
"ts-node": "^10.9.2",
"typescript": "^5.3.3",
"ws": "^8.16.0",
Expand All @@ -105,6 +107,7 @@
"@types/keyv": "^4.2.0",
"@types/lodash-es": "^4.17.12",
"@types/mixpanel": "^2.14.8",
"@types/mustache": "^4",
"@types/node": "^20.11.20",
"@types/nodemailer": "^6.4.14",
"@types/on-headers": "^1.0.3",
Expand Down
71 changes: 49 additions & 22 deletions packages/backend/server/schema.prisma
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ model User {
pagePermissions WorkspacePageUserPermission[]
connectedAccounts ConnectedAccount[]
sessions UserSession[]
AiSession AiSession[]
aiSessions AiSession[]
@@map("users")
}
Expand Down Expand Up @@ -97,7 +97,7 @@ model Workspace {
permissions WorkspaceUserPermission[]
pagePermissions WorkspacePageUserPermission[]
features WorkspaceFeatures[]
AiSession AiSession[]
aiSessions AiSession[]
@@map("workspaces")
}
Expand Down Expand Up @@ -323,7 +323,7 @@ model Snapshot {
// but the created time of last seen update that has been merged into snapshot.
updatedAt DateTime @map("updated_at") @db.Timestamptz(6)
AiSession AiSession[]
aiSessions AiSession[]
@@id([id, workspaceId])
@@map("snapshots")
Expand Down Expand Up @@ -432,39 +432,66 @@ enum AiPromptRole {
user
}

model AiPrompt {
id String @id @default(uuid()) @db.VarChar
// prompt name
name String @db.VarChar(20)
model AiPromptMessage {
promptId Int @map("prompt_id") @db.Integer
// if a group of prompts contains multiple sentences, idx specifies the order of each sentence
idx Int @db.Integer
// system/assistant/user
role AiPromptRole
// prompt content
content String @db.Text
params Json? @db.Json
createdAt DateTime @default(now()) @map("created_at") @db.Timestamptz(6)
prompt AiPrompt @relation(fields: [promptId], references: [id], onDelete: Cascade)
@@unique([promptId, idx])
@@map("ai_prompts_messages")
}

model AiPrompt {
id Int @id @default(autoincrement()) @db.Integer
name String @unique @db.VarChar(32)
// an mark identifying which view to use to display the session
// it is only used in the frontend and does not affect the backend
action String? @db.VarChar
model String? @db.VarChar
createdAt DateTime @default(now()) @map("created_at") @db.Timestamptz(6)
messages AiPromptMessage[]
sessions AiSession[]
@@map("ai_prompts_metadata")
}

model AiSessionMessage {
id String @id @default(uuid()) @db.VarChar(36)
sessionId String @map("session_id") @db.VarChar(36)
role AiPromptRole
content String @db.Text
createdAt DateTime @default(now()) @map("created_at") @db.Timestamptz(6)
updatedAt DateTime @updatedAt @map("updated_at") @db.Timestamptz(6)
@@unique([name, idx])
@@map("ai_prompts")
session AiSession @relation(fields: [sessionId], references: [id], onDelete: Cascade)
@@map("ai_sessions_messages")
}

model AiSession {
id String @id @default(uuid()) @db.VarChar
userId String @map("user_id") @db.VarChar
workspaceId String @map("workspace_id") @db.VarChar
docId String @map("doc_id") @db.VarChar
promptName String @map("prompt_name") @db.VarChar
action Boolean @db.Boolean
model String @db.VarChar
messages Json @db.Json
id String @id @default(uuid()) @db.VarChar(36)
userId String @map("user_id") @db.VarChar(36)
workspaceId String @map("workspace_id") @db.VarChar(36)
docId String @map("doc_id") @db.VarChar(36)
promptName String @map("prompt_name") @db.VarChar(32)
createdAt DateTime @default(now()) @map("created_at") @db.Timestamptz(6)
updatedAt DateTime @updatedAt @map("updated_at") @db.Timestamptz(6)
user User @relation(fields: [userId], references: [id], onDelete: Cascade)
workspace Workspace @relation(fields: [workspaceId], references: [id], onDelete: Cascade)
doc Snapshot @relation(fields: [docId, workspaceId], references: [id, workspaceId], onDelete: Cascade)
user User @relation(fields: [userId], references: [id], onDelete: Cascade)
workspace Workspace @relation(fields: [workspaceId], references: [id], onDelete: Cascade)
doc Snapshot @relation(fields: [docId, workspaceId], references: [id, workspaceId], onDelete: Cascade)
prompt AiPrompt @relation(fields: [promptName], references: [name], onDelete: Cascade)
messages AiSessionMessage[]
@@map("ai_sessions")
@@map("ai_sessions_metadata")
}

model DataMigration {
Expand Down
1 change: 1 addition & 0 deletions packages/backend/server/src/config/affine.env.ts
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ AFFiNE.ENV_MAP = {
MAILER_SECURE: ['mailer.secure', 'boolean'],
THROTTLE_TTL: ['rateLimiter.ttl', 'int'],
THROTTLE_LIMIT: ['rateLimiter.limit', 'int'],
COPILOT_OPENAI_API_KEY: 'plugins.copilot.openai.apiKey',
REDIS_SERVER_HOST: 'plugins.redis.host',
REDIS_SERVER_PORT: ['plugins.redis.port', 'int'],
REDIS_SERVER_USER: 'plugins.redis.username',
Expand Down
4 changes: 1 addition & 3 deletions packages/backend/server/src/config/affine.self.ts
Original file line number Diff line number Diff line change
Expand Up @@ -39,9 +39,7 @@ if (env.R2_OBJECT_STORAGE_ACCOUNT_ID) {
}

AFFiNE.plugins.use('copilot', {
openai: {
apiKey: 'test',
},
openai: {},
});
AFFiNE.plugins.use('redis');
AFFiNE.plugins.use('payment', {
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
import { PrismaClient } from '@prisma/client';

import { prompts } from './utils/prompts';

export class Prompts1712068777394 {
// do the migration
static async up(db: PrismaClient) {
await db.$transaction(async tx => {
await Promise.all(
prompts.map(prompt =>
tx.aiPrompt.create({
data: {
name: prompt.name,
action: prompt.action,
model: prompt.model,
messages: {
create: prompt.messages.map((message, idx) => ({
idx,
role: message.role,
content: message.content,
params: message.params,
})),
},
},
})
)
);
});
}

// revert the migration
static async down(_db: PrismaClient) {}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@
import { PrismaClient } from '@prisma/client';

import { Quotas } from '../../core/quota';
import { upgradeQuotaVersion } from './utils/user-quotas';

export class RefreshFreePlan1712224382221 {
// do the migration
static async up(db: PrismaClient) {
// free plan 1.0
const quota = Quotas[4];
await upgradeQuotaVersion(db, quota, 'free plan 1.1 migration');
}

// revert the migration
static async down(_db: PrismaClient) {}
}

0 comments on commit 46a368d

Please sign in to comment.