feat: init @affine/copilot (#2511)

plugins/copilot/src/core/langchain/message-history.ts

@@ -0,0 +1,109 @@
+import type { DBSchema, IDBPDatabase } from 'idb';
+import { openDB } from 'idb';
+import {
+  AIChatMessage,
+  type BaseChatMessage,
+  BaseChatMessageHistory,
+  ChatMessage,
+  HumanChatMessage,
+  type StoredMessage,
+  SystemChatMessage,
+} from 'langchain/schema';
+
+interface ChatMessageDBV1 extends DBSchema {
+  chat: {
+    key: string;
+    value: {
+      /**
+       * ID of the chat
+       */
+      id: string;
+      messages: StoredMessage[];
+    };
+  };
+}
+
+export const conversationHistoryDBName = 'affine-copilot-chat';
+
+export class IndexedDBChatMessageHistory extends BaseChatMessageHistory {
+  public id: string;
+  private messages: BaseChatMessage[] = [];
+
+  private readonly dbPromise: Promise<IDBPDatabase<ChatMessageDBV1>>;
+  private readonly initPromise: Promise<void>;
+
+  constructor(id: string) {
+    super();
+    this.id = id;
+    this.messages = [];
+    this.dbPromise = openDB<ChatMessageDBV1>('affine-copilot-chat', 1, {
+      upgrade(database, oldVersion) {
+        if (oldVersion === 0) {
+          database.createObjectStore('chat', {
+            keyPath: 'id',
+          });
+        }
+      },
+    });
+    this.initPromise = this.dbPromise.then(async db => {
+      const objectStore = db
+        .transaction('chat', 'readonly')
+        .objectStore('chat');
+      const chat = await objectStore.get(id);
+      if (chat != null) {
+        this.messages = chat.messages.map(message => {
+          switch (message.type) {
+            case 'ai':
+              return new AIChatMessage(message.data.content);
+            case 'human':
+              return new HumanChatMessage(message.data.content);
+            case 'system':
+              return new SystemChatMessage(message.data.content);
+            default:
+              return new ChatMessage(
+                message.data.content,
+                message.data.role ?? 'never'
+              );
+          }
+        });
+      }
+    });
+  }
+
+  protected async addMessage(message: BaseChatMessage): Promise<void> {
+    await this.initPromise;
+    this.messages.push(message);
+    const db = await this.dbPromise;
+    const objectStore = db.transaction('chat', 'readwrite').objectStore('chat');
+    const chat = await objectStore.get(this.id);
+    if (chat != null) {
+      chat.messages.push(message.toJSON());
+      await objectStore.put(chat);
+    } else {
+      await objectStore.add({
+        id: this.id,
+        messages: [message.toJSON()],
+      });
+    }
+  }
+
+  async addAIChatMessage(message: string): Promise<void> {
+    await this.addMessage(new AIChatMessage(message));
+  }
+
+  async addUserMessage(message: string): Promise<void> {
+    await this.addMessage(new HumanChatMessage(message));
+  }
+
+  async clear(): Promise<void> {
+    await this.initPromise;
+    this.messages = [];
+    const db = await this.dbPromise;
+    const objectStore = db.transaction('chat', 'readwrite').objectStore('chat');
+    await objectStore.delete(this.id);
+  }
+
+  async getMessages(): Promise<BaseChatMessage[]> {
+    return await this.initPromise.then(() => this.messages);
+  }
+}

plugins/copilot/src/core/langchain/vector-store.ts

@@ -0,0 +1,118 @@
+// fixme: vector store has not finished
+import type { DBSchema } from 'idb';
+import { Document } from 'langchain/document';
+import type { Embeddings } from 'langchain/embeddings';
+import { VectorStore } from 'langchain/vectorstores';
+import { similarity as ml_distance_similarity } from 'ml-distance';
+
+// eslint-disable-next-line @typescript-eslint/no-unused-vars
+interface VectorDBV1 extends DBSchema {
+  vector: {
+    key: string;
+    value: Vector;
+  };
+}
+
+interface Vector {
+  id: string;
+
+  content: string;
+  embedding: number[];
+  metadata: Record<string, unknown>;
+}
+
+export interface MemoryVectorStoreArgs {
+  similarity?: typeof ml_distance_similarity.cosine;
+}
+
+export class IndexedDBVectorStore extends VectorStore {
+  memoryVectors: any[] = [];
+
+  similarity: typeof ml_distance_similarity.cosine;
+
+  constructor(
+    embeddings: Embeddings,
+    { similarity, ...rest }: MemoryVectorStoreArgs = {}
+  ) {
+    super(embeddings, rest);
+
+    this.similarity = similarity ?? ml_distance_similarity.cosine;
+  }
+
+  async addDocuments(documents: Document[]): Promise<void> {
+    const texts = documents.map(({ pageContent }) => pageContent);
+    return this.addVectors(
+      await this.embeddings.embedDocuments(texts),
+      documents
+    );
+  }
+
+  async addVectors(vectors: number[][], documents: Document[]): Promise<void> {
+    const memoryVectors = vectors.map((embedding, idx) => ({
+      content: documents[idx].pageContent,
+      embedding,
+      metadata: documents[idx].metadata,
+    }));
+
+    this.memoryVectors = this.memoryVectors.concat(memoryVectors);
+  }
+
+  async similaritySearchVectorWithScore(
+    query: number[],
+    k: number
+  ): Promise<[Document, number][]> {
+    const searches = this.memoryVectors
+      .map((vector, index) => ({
+        similarity: this.similarity(query, vector.embedding),
+        index,
+      }))
+      .sort((a, b) => (a.similarity > b.similarity ? -1 : 0))
+      .slice(0, k);
+
+    const result: [Document, number][] = searches.map(search => [
+      new Document({
+        metadata: this.memoryVectors[search.index].metadata,
+        pageContent: this.memoryVectors[search.index].content,
+      }),
+      search.similarity,
+    ]);
+
+    return result;
+  }
+
+  static async fromTexts(
+    texts: string[],
+    metadatas: object[] | object,
+    embeddings: Embeddings,
+    dbConfig?: MemoryVectorStoreArgs
+  ): Promise<IndexedDBVectorStore> {
+    const docs: Document[] = [];
+    for (let i = 0; i < texts.length; i += 1) {
+      const metadata = Array.isArray(metadatas) ? metadatas[i] : metadatas;
+      const newDoc = new Document({
+        pageContent: texts[i],
+        metadata,
+      });
+      docs.push(newDoc);
+    }
+    return IndexedDBVectorStore.fromDocuments(docs, embeddings, dbConfig);
+  }
+
+  static async fromDocuments(
+    docs: Document[],
+    embeddings: Embeddings,
+    dbConfig?: MemoryVectorStoreArgs
+  ): Promise<IndexedDBVectorStore> {
+    const instance = new this(embeddings, dbConfig);
+    await instance.addDocuments(docs);
+    return instance;
+  }
+
+  static async fromExistingIndex(
+    embeddings: Embeddings,
+    dbConfig?: MemoryVectorStoreArgs
+  ): Promise<IndexedDBVectorStore> {
+    const instance = new this(embeddings, dbConfig);
+    return instance;
+  }
+}