Claude API×pgvectorでRAG実装

flowchart LR
  subgraph インデックス作成
    D[文書] --> C[チャンク分割] --> E[Embedding生成] --> V[(pgvector)]
  end
  subgraph クエリ処理
    Q[ユーザー質問] --> QE[Query Embedding] --> S[類似検索] --> V
    S --> CTX[コンテキスト構築] --> G[Claude生成] --> A[回答]
  end

flowchart TD
  SRC[ソース文書\nPDF/Markdown/HTML] -->|テキスト抽出| RAW[生テキスト]
  RAW -->|チャンク分割| CHK[チャンク群]
  CHK -->|Voyage AI Embeddings| VEC[ベクトル配列 1024次元]
  VEC -->|INSERT| PG[(PostgreSQL + pgvector)]
  PG -->|HNSWインデックス| IDX[高速検索準備完了]

flowchart LR
  U[ユーザー質問] --> QV[Query Embedding]
  QV --> SR[pgvector類似検索 top-k]
  SR --> CTX[コンテキスト組み立て]
  CTX --> CL[Claude API\nメッセージ生成]
  CL --> ANS[最終回答]

docker run -d \
  --name pgvector-dev \
  -e POSTGRES_PASSWORD=password \
  -e POSTGRES_DB=ragdb \
  -p 5432:5432 \
  pgvector/pgvector:pg16

# docker-compose.yml
services:
  db:
    image: pgvector/pgvector:pg16
    environment:
      POSTGRES_PASSWORD: password
      POSTGRES_DB: ragdb
    ports:
      - "5432:5432"
    volumes:
      - pgdata:/var/lib/postgresql/data

volumes:
  pgdata:

CREATE EXTENSION IF NOT EXISTS vector;

-- 動作確認
SELECT '[1,2,3]'::vector;

sudo apt install postgresql-16-pgvector

-- PostgreSQLにログイン後
CREATE EXTENSION IF NOT EXISTS vector;

CREATE TABLE documents (
  id          SERIAL PRIMARY KEY,
  content     TEXT        NOT NULL,
  embedding   VECTOR(1024),
  metadata    JSONB,
  created_at  TIMESTAMPTZ DEFAULT NOW(),
  updated_at  TIMESTAMPTZ DEFAULT NOW()
);

-- HNSWインデックス（pgvector 0.5.0以降の推奨）
-- m: グラフの最大接続数、ef_construction: 構築時の探索幅
CREATE INDEX ON documents
  USING hnsw (embedding vector_cosine_ops)
  WITH (m = 16, ef_construction = 64);

-- source_url重複排除用のUNIQUEインデックス（UPSERTで使用）
CREATE UNIQUE INDEX ON documents ((metadata->>'source_url'))
  WHERE metadata->>'source_url' IS NOT NULL;

pnpm add voyageai

// src/lib/embeddings.ts
import { VoyageAIClient } from "voyageai";

const voyageClient = new VoyageAIClient({ apiKey: process.env.VOYAGE_API_KEY });

/** ドキュメント格納時のEmbedding生成 */
export async function generateEmbedding(text: string): Promise<number[]> {
  const response = await voyageClient.embed({
    input: text,
    model: "voyage-3",
    inputType: "document",
  });
  return response.data[0].embedding;
}

/** クエリ検索時のEmbedding生成（input_typeがqueryに変わる） */
export async function generateQueryEmbedding(query: string): Promise<number[]> {
  const response = await voyageClient.embed({
    input: query,
    model: "voyage-3",
    inputType: "query",
  });
  return response.data[0].embedding;
}

// src/lib/batch-embed.ts
import { VoyageAIClient } from "voyageai";

const voyageClient = new VoyageAIClient({ apiKey: process.env.VOYAGE_API_KEY });
const BATCH_SIZE = 128; // Voyage AIの最大バッチサイズ

export async function batchEmbed(texts: string[]): Promise<number[][]> {
  const results: number[][] = [];

  for (let i = 0; i < texts.length; i += BATCH_SIZE) {
    const batch = texts.slice(i, i + BATCH_SIZE);
    let attempts = 0;

    while (attempts < 3) {
      try {
        const response = await voyageClient.embed({
          input: batch,
          model: "voyage-3",
          inputType: "document",
        });
        results.push(...response.data.map((d) => d.embedding));
        break;
      } catch (err: unknown) {
        attempts++;
        if (attempts >= 3) throw err;
        // 指数バックオフ: 1s, 2s, 4s
        await new Promise((r) => setTimeout(r, 1000 * 2 ** (attempts - 1)));
      }
    }
  }

  return results;
}

// src/lib/chunker.ts

interface ChunkOptions {
  maxChars?: number;  // tokenではなく文字数で管理するシンプル版
  overlap?: number;   // オーバーラップ文字数
}

/** maxCharsを超える単一段落を、上限以下の断片へ強制分割する */
function splitLongParagraph(para: string, maxChars: number): string[] {
  if (para.length <= maxChars) return [para];
  const pieces: string[] = [];
  for (let i = 0; i < para.length; i += maxChars) {
    pieces.push(para.slice(i, i + maxChars));
  }
  return pieces;
}

/**
 * 段落境界でチャンク分割する。
 * maxCharsを超えた時点で現在の蓄積を確定し、次のチャンクへオーバーラップする。
 * maxCharsを超える長さの段落は、事前にsplitLongParagraphで分割しておく。
 */
export function chunkByParagraph(
  text: string,
  { maxChars = 1200, overlap = 150 }: ChunkOptions = {}
): string[] {
  const paragraphs = text
    .split(/\n\n+/)
    .flatMap((para) => splitLongParagraph(para, maxChars));
  const chunks: string[] = [];
  let current = "";

  for (const para of paragraphs) {
    if (current.length + para.length > maxChars) {
      if (current) {
        chunks.push(current.trim());
        // オーバーラップ: 直前チャンクの末尾を引き継ぐ
        current = current.slice(-overlap) + "\n\n" + para;
      } else {
        current = para;
      }
    } else {
      current += (current ? "\n\n" : "") + para;
    }
  }

  if (current.trim()) chunks.push(current.trim());
  return chunks;
}

/** タイトルやメタデータをチャンク先頭に付加してEmbedding精度を向上させる */
export function enrichChunk(chunk: string, title: string, category: string): string {
  return `[${category}] ${title}\n\n${chunk}`;
}

-- $1: クエリEmbeddingベクトル（文字列形式）、$2: カテゴリフィルタ、$3: LIMIT数
SELECT
  id,
  content,
  metadata,
  1 - (embedding <=> $1::vector) AS similarity
FROM documents
WHERE
  ($2::text IS NULL OR metadata->>'category' = $2)
  AND 1 - (embedding <=> $1::vector) > 0.7  -- 低精度結果を除外
ORDER BY embedding <=> $1::vector
LIMIT $3;

// src/lib/search.ts
import { Pool } from "pg";

const pool = new Pool({ connectionString: process.env.DATABASE_URL });

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

export async function searchSimilarDocs(
  queryEmbedding: number[],
  limit = 5,
  category?: string
): Promise<SearchResult[]> {
  const vectorStr = `[${queryEmbedding.join(",")}]`;

  const { rows } = await pool.query<SearchResult>(
    `SELECT
       id,
       content,
       metadata,
       1 - (embedding <=> $1::vector) AS similarity
     FROM documents
     WHERE
       ($2::text IS NULL OR metadata->>'category' = $2)
       AND 1 - (embedding <=> $1::vector) > 0.7
     ORDER BY embedding <=> $1::vector
     LIMIT $3`,
    [vectorStr, category ?? null, limit]
  );

  return rows;
}

// src/lib/hyde.ts
import Anthropic from "@anthropic-ai/sdk";
import { generateQueryEmbedding } from "./embeddings";
import { searchSimilarDocs } from "./search";

const client = new Anthropic();

export async function hydeSearch(userQuery: string, limit = 5) {
  // Step1: 仮説的ドキュメントをClaudeで生成
  const hypoResponse = await client.messages.create({
    model: "claude-3-5-haiku-20241022",
    max_tokens: 256,
    messages: [
      {
        role: "user",
        content: `以下の質問に対して、完璧な答えを含むドキュメントの一節を書いてください。質問: ${userQuery}`,
      },
    ],
  });

  const firstBlock = hypoResponse.content[0];
  const hypoDoc =
    firstBlock?.type === "text" ? firstBlock.text : userQuery;

  // Step2: 仮説的ドキュメントのEmbeddingで検索
  const embedding = await generateQueryEmbedding(hypoDoc);
  return searchSimilarDocs(embedding, limit);
}

// src/lib/rag-tool-use.ts
import Anthropic from "@anthropic-ai/sdk";
import { generateQueryEmbedding } from "./embeddings";
import { searchSimilarDocs } from "./search";

const client = new Anthropic();

const searchTool: Anthropic.Tool = {
  name: "search_documents",
  description:
    "社内ドキュメントをベクトル検索で取得する。ユーザーの質問に関連する情報が必要な場合に使用する。",
  input_schema: {
    type: "object" as const,
    properties: {
      query: {
        type: "string",
        description: "検索クエリ文字列",
      },
      category: {
        type: "string",
        description: "絞り込むカテゴリ（省略可）",
      },
    },
    required: ["query"],
  },
};

// 初回・再生成の両方で同じ方針を適用するため定数として共有する
const SYSTEM_PROMPT =
  "あなたは社内ドキュメントに基づいて回答するアシスタントです。" +
  "必要な情報が不明な場合はsearch_documentsツールを使って検索してから回答してください。" +
  "検索結果が見つからなかった場合は「該当する情報が見つかりませんでした」と答えてください。";

export async function ragWithToolUse(userMessage: string): Promise<string> {
  const messages: Anthropic.MessageParam[] = [
    { role: "user", content: userMessage },
  ];

  let response = await client.messages.create({
    model: "claude-3-5-haiku-20241022",
    max_tokens: 1024,
    system: SYSTEM_PROMPT,
    tools: [searchTool],
    messages,
  });

  // Tool Useループ: stop_reasonが"tool_use"の間は検索→再生成を繰り返す
  while (response.stop_reason === "tool_use") {
    const toolUseBlock = response.content.find(
      (b): b is Anthropic.ToolUseBlock => b.type === "tool_use"
    );
    if (!toolUseBlock) break;

    const input = toolUseBlock.input as { query?: string; category?: string };
    if (typeof input.query !== "string" || input.query.length === 0) break;

    // 実際の検索実行
    const queryEmbedding = await generateQueryEmbedding(input.query);
    const docs = await searchSimilarDocs(queryEmbedding, 5, input.category);

    // アシスタントのtool_useと検索結果をメッセージ履歴に追加
    messages.push({ role: "assistant", content: response.content });
    messages.push({
      role: "user",
      content: [
        {
          type: "tool_result",
          tool_use_id: toolUseBlock.id,
          content:
            docs.length > 0
              ? JSON.stringify(docs)
              : "No relevant documents found.",
        },
      ],
    });

    // 検索結果を踏まえてClaudeに再度回答生成させる
    response = await client.messages.create({
      model: "claude-3-5-haiku-20241022",
      max_tokens: 2048,
      system: SYSTEM_PROMPT, // 再生成でも方針を維持する
      tools: [searchTool],
      messages,
    });
  }

  const textBlock = response.content.find(
    (b): b is Anthropic.TextBlock => b.type === "text"
  );
  return textBlock?.text ?? "回答を生成できませんでした。";
}

// src/lib/rerank.ts
import { VoyageAIClient } from "voyageai";
import { SearchResult } from "./search";

const voyageClient = new VoyageAIClient({ apiKey: process.env.VOYAGE_API_KEY });

export async function rerankDocs(
  query: string,
  docs: SearchResult[],
  topK = 5
): Promise<SearchResult[]> {
  // 空配列を渡すとVoyage AI側でエラーになるため事前に返す
  if (docs.length === 0) return [];

  const response = await voyageClient.rerank({
    query,
    documents: docs.map((d) => d.content),
    model: "rerank-2",
    topK,
  });

  return response.data.map((r) => docs[r.index]);
}

-- pgvector（コサイン類似度）+ PostgreSQL標準全文検索（ts_rank）の組み合わせ
-- 重み: ベクトル70%、全文検索30%
SELECT
  id,
  content,
  (
    0.7 * (1 - (embedding <=> $1::vector)) +
    0.3 * ts_rank(
      to_tsvector('simple', content),
      plainto_tsquery('simple', $2)
    )
  ) AS hybrid_score
FROM documents
ORDER BY hybrid_score DESC
LIMIT 5;

-- メタデータ付きUPSERT例
INSERT INTO documents (content, embedding, metadata)
VALUES ($1, $2::vector, $3::jsonb)
ON CONFLICT ((metadata->>'source_url')) WHERE metadata->>'source_url' IS NOT NULL
DO UPDATE SET
  content    = EXCLUDED.content,
  embedding  = EXCLUDED.embedding,
  metadata   = EXCLUDED.metadata,
  updated_at = NOW();