レッスン 10: 埋め込みとベクター ストア

AIエージェントプラットフォームをゼロから構築 — xClawとの実戦

パート 3: RAG パイプライン

xdev.asia

はじめに

ドキュメントをチャンク化した後、類似性検索のためにテキストをベクトル (埋め込み) に変換し、ベクトル ストアに保存する必要があります。この記事では、埋め込みパイプラインとベクター ストアの抽象化を実装します。

---

1. インターフェースの埋め込み

// packages/core/src/rag/embedding.ts
export interface EmbeddingProvider {
  embed(texts: string[]): Promise<number[][]>;
  readonly dimensions: number;
  readonly model: string;
}

export class OpenAIEmbedding implements EmbeddingProvider {
  readonly dimensions = 1536;
  readonly model = 'text-embedding-3-small';
  private client: OpenAI;

  constructor(apiKey: string) {
    this.client = new OpenAI({ apiKey });
  }

  async embed(texts: string[]): Promise<number[][]> {
    // Batch processing — max 2048 texts per request
    const batches: string[][] = [];
    for (let i = 0; i < texts.length; i += 2048) {
      batches.push(texts.slice(i, i + 2048));
    }

    const allEmbeddings: number[][] = [];
    for (const batch of batches) {
      const response = await this.client.embeddings.create({
        model: this.model,
        input: batch,
      });
      allEmbeddings.push(...response.data.map(d => d.embedding));
    }

    return allEmbeddings;
  }
}

---

2. ベクトル ストアの抽象化

// packages/core/src/rag/vector-store.ts
export interface VectorStore {
  upsert(documents: VectorDocument[]): Promise<void>;
  search(query: number[], options: SearchOptions): Promise<SearchResult[]>;
  delete(ids: string[]): Promise<void>;
  createCollection(name: string): Promise<void>;
}

export interface VectorDocument {
  id: string;
  vector: number[];
  content: string;
  metadata: Record<string, unknown>;
}

export interface SearchResult {
  id: string;
  content: string;
  score: number;
  metadata: Record<string, unknown>;
}

export interface SearchOptions {
  collection: string;
  topK: number;
  filter?: Record<string, unknown>;
  minScore?: number;
}

---

3. MongoDB Atlas ベクトル検索

// packages/core/src/rag/stores/mongodb-vector-store.ts
export class MongoDBVectorStore implements VectorStore {
  private db: Db;

  constructor(db: Db) {
    this.db = db;
  }

  async createCollection(name: string): Promise<void> {
    const collection = this.db.collection(name);

    // Create vector search index
    await collection.createSearchIndex({
      name: 'vector_index',
      definition: {
        mappings: {
          dynamic: true,
          fields: {
            vector: {
              type: 'knnVector',
              dimensions: 1536,
              similarity: 'cosine',
            },
          },
        },
      },
    });
  }

  async upsert(documents: VectorDocument[]): Promise<void> {
    const ops = documents.map(doc => ({
      updateOne: {
        filter: { _id: doc.id },
        update: {
          $set: {
            vector: doc.vector,
            content: doc.content,
            metadata: doc.metadata,
            updatedAt: new Date(),
          },
        },
        upsert: true,
      },
    }));

    await this.db.collection('vectors').bulkWrite(ops);
  }

  async search(query: number[], options: SearchOptions): Promise<SearchResult[]> {
    const results = await this.db.collection(options.collection).aggregate([
      {
        $vectorSearch: {
          index: 'vector_index',
          path: 'vector',
          queryVector: query,
          numCandidates: options.topK * 10,
          limit: options.topK,
          filter: options.filter,
        },
      },
      {
        $project: {
          content: 1,
          metadata: 1,
          score: { $meta: 'vectorSearchScore' },
        },
      },
    ]).toArray();

    return results
      .filter(r => !options.minScore || r.score >= options.minScore)
      .map(r => ({
        id: r._id.toString(),
        content: r.content,
        score: r.score,
        metadata: r.metadata,
      }));
  }

  async delete(ids: string[]): Promise<void> {
    await this.db.collection('vectors').deleteMany({
      _id: { $in: ids },
    });
  }
}

---

4. バッチ埋め込みパイプライン

export async function embedAndStore(
  chunks: DocumentChunk[],
  embedding: EmbeddingProvider,
  vectorStore: VectorStore,
  collection: string,
): Promise<void> {
  // Embed in batches of 100
  const BATCH_SIZE = 100;

  for (let i = 0; i < chunks.length; i += BATCH_SIZE) {
    const batch = chunks.slice(i, i + BATCH_SIZE);
    const texts = batch.map(c => c.content);

    console.log(`Embedding batch ${i / BATCH_SIZE + 1}/${Math.ceil(chunks.length / BATCH_SIZE)}`);

    const vectors = await embedding.embed(texts);

    const documents: VectorDocument[] = batch.map((chunk, j) => ({
      id: chunk.id,
      vector: vectors[j],
      content: chunk.content,
      metadata: chunk.metadata,
    }));

    await vectorStore.upsert(documents);
  }
}

---

5. まとめ

• Embedding Provider — OpenAI、Cohere、ローカル モデルの抽象化
• Vector Store — MongoDB Atlas Search または ChromaDB
• バッチ処理 — レート制限を回避するためにチャンクをバッチで処理します。
• コサイン類似度 — テキスト埋め込みの最も一般的な距離メトリック

次の記事: RAG エンジン — 取得をエージェント フローに統合します。