LoginSignup
0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

@poruruba

ChromaによるRAGを使った専用チャットエージェントを作る

Posted at

Mastra→Gemini→Chromaの組み合わせを使ってPDFをRAG化し、専用チャットエージェントを作ります。Notebook.lmのようなものです。

以下、速足で。
・DockerへのChromaのセットアップ
・PDFをGeminiを使ってベクトル化してChromaに登録
・Mastraのチャットに取り込み

DockerへのChromaのセットアップ

docker-compose.yaml
version: "3"

services:
  chroma:
    image: chromadb/chroma:latest
    ports:
      - "8001:8000"
    volumes:
      - /share/Container/chroma:/chroma/chroma_data
    environment:
      - PERSISTENT_DATA=True
      - CHROMA_SERVER_CORS_ALLOW_ORIGINS='["*"]'
      - CHROMA_SERVER_MODE=default
      - CHROMA_SERVER_ENABLE=rest
    restart: unless-stopped

これで、ポート8001に立ち上がります。
ブラウザで以下を開いてみてください。SwaggerViewerが開きます。
http://[立ち上げたサーバのドメイン名]jp:8001/docs/

PDFをGeminiを使ってベクトル化してChromaに登録

ベクトル化にはGemini、分割にはMastraを使いました。
PDFからテキストの抽出には、pdf-parseライブラリを使いました。

const GOOGLE_GENERATIVE_AI_API_KEY = process.env.GOOGLE_GENERATIVE_AI_API_KEY;
const CHROMA_ENDPOINT = process.env.CHROMA_ENDPOINT;
const COLLECTION_NAME = "second_collection";

const { GoogleGenAI } = require('@google/genai');
const { ChromaVector } = require('@mastra/chroma');
const { createVectorQueryTool, createDocumentChunkerTool, MDocument } = require('@mastra/rag');
const fs = require('fs').promises;
const pdfParse = require('pdf-parse');

Chromaにコレクションの作成

  var chromaVector = new ChromaVector(
    { path: CHROMA_ENDPOINT }
  );

  var collections = await chromaVector.listIndexes();
  if( !collections.includes(COLLECTION_NAME) ){
    var result = await chromaVector.createIndex({ 
      indexName: COLLECTION_NAME,
      dimension: 3072
    });
  }
  console.log(await chromaVector.listIndexes());

PDFファイルの読み込み、テキストの抽出

  const dataBuffer = await fs.readFile('sample.pdf');
  var data = await pdfParse(dataBuffer);
  console.log(data.text);

テキストの分割とベクトル化と登録

async function getGeminiEmbedding(text) {
  const response = await ai.models.embedContent({
    model: 'gemini-embedding-001',
    contents: text
  });

  return response.embeddings;
}

const ai = new GoogleGenAI({
  apiKey: GOOGLE_GENERATIVE_AI_API_KEY
});

  const doc = MDocument.fromText(data.text);

  const chunker = createDocumentChunkerTool({
    doc: doc,
    params: {
      strategy: "recursive",
      maxSize: 512,
      overlap: 50,
      separators: ["\n", "\n\n", " "],
    },
  });
  const { chunks } = await chunker.execute();

  console.log("Number of chunks:", chunks.length);

  for( let chunk of chunks ){
    console.log(chunk.text);
    const embeddings = await getGeminiEmbedding(chunk.text);

    await chromaVector.upsert({
      indexName: COLLECTION_NAME,
      documents: [chunk.text],
      vectors: [embeddings[0].values],
      metadata: [ { source: "PDF" } ],
    });
  }

Mastraのチャットに取り込み

MastraとGeminiと相性が悪いのか、MastraのRAG機能とGeminiを組み合わせた動作はエラーとなってしまいました。
あるべき姿と、回避した姿の両方を示しておきます。

あるべき姿

準備

const { google } = require("@ai-sdk/google");
const { Agent } = require('@mastra/core/agent');
const { Memory } = require('@mastra/memory');
const { LibSQLStore } = require('@mastra/libsql');
const { ChromaVector } = require('@mastra/chroma');
const { createVectorQueryTool } = require('@mastra/rag');
const { RuntimeContext } = require('@mastra/core/runtime-context');

const CHROMA_ENDPOINT = process.env.CHROMA_ENDPOINT;
const COLLECTION_NAME = "second_collection";

const runtimeContext = new RuntimeContext();
let agent;
let memory;
let vectorQueryTool;

    const chromaVector = new ChromaVector(
      { path: CHROMA_ENDPOINT }
    );
    vectorQueryTool = createVectorQueryTool({
      vectorStoreName: "chromaVector",
      vectorStore: chromaVector,
      description: "「M5StackとJavascriptではじめるIoTデバイス制御」の情報で、ESP32をJavascriptで制御するためのライブラリについての説明です。",
      indexName: COLLECTION_NAME,
      model: google.textEmbeddingModel("gemini-embedding-001"),
    });
    
    memory = new Memory({
      storage: new LibSQLStore({
        url: "file:" + process.env.THIS_BASE_PATH + "/data/mastra/mastra.db"
      }),
      options: {
        threads: {
          generateTitle: true
        }
      }
    });

    agent = new Agent({
      name: 'Chat Agent',
      instructions: `自由に会話をします。`,
      model: google('gemini-1.5-pro-latest'),
      tools: { vectorQueryTool },
      memory: memory
    });

    console.log("Mastra Agent started");

チャットとの組み合わせ
body.messageにRAGに登録したPDFの内容に対して調べたい言葉を入れます。

    const response = await agent.generate(
      [
        {
          role: "user", // "system" "assistant" "user"
          content: body.message,
        },
      ],
    );
    console.log(response);

回避版

準備

const { google } = require("@ai-sdk/google");
const { Agent } = require('@mastra/core/agent');
const { Memory } = require('@mastra/memory');
const { LibSQLStore } = require('@mastra/libsql');
const { ChromaVector } = require('@mastra/chroma');
const { createVectorQueryTool } = require('@mastra/rag');
const { RuntimeContext } = require('@mastra/core/runtime-context');

const CHROMA_ENDPOINT = process.env.CHROMA_ENDPOINT;
const COLLECTION_NAME = "second_collection";

const runtimeContext = new RuntimeContext();
let agent;
let memory;
let vectorQueryTool;

    const chromaVector = new ChromaVector(
      { path: CHROMA_ENDPOINT }
    );
    vectorQueryTool = createVectorQueryTool({
      vectorStoreName: "chromaVector",
      vectorStore: chromaVector,
      description: "「M5StackとJavascriptではじめるIoTデバイス制御」の情報で、ESP32をJavascriptで制御するためのライブラリについての説明です。",
      indexName: COLLECTION_NAME,
      model: google.textEmbeddingModel("gemini-embedding-001"),
    });
    
    memory = new Memory({
      storage: new LibSQLStore({
        url: "file:" + process.env.THIS_BASE_PATH + "/data/mastra/mastra.db"
      }),
      options: {
        threads: {
          generateTitle: true
        }
      }
    });

    agent = new Agent({
      name: 'Chat Agent',
      instructions: `自由に会話をします。`,
      model: google('gemini-1.5-pro-latest'),
      memory: memory
    });

    console.log("Mastra Agent started");

チャットとの組み合わせ

    const queryResult = await vectorQueryTool.execute({
      context: { queryText: body.message, topK: 10 },
      runtimeContext,
    });    

    const contextText = queryResult.sources
      .map((r, i) => `(${i + 1}) ${r.document}`)
      .join('\n');

    const response = await agent.generate(
      [
        {
          role: "user", // "system" "assistant" "user"
          content: body.message + " [参考情報]" + contextText,
        },
      ],
    );
    console.log(response);

参考

https://mastra.ai/ja/docs
https://ai.google.dev/gemini-api/docs/embeddings?hl=ja

Chromaライブラリで直接登録する場合
https://docs.trychroma.com/docs/overview/introduction

以上

0
0
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?