LangChainを使ってリポジトリ内のソースコードについて答えてもらう

Posted at 2023-12-27

はじめに

社内のメンバー数人でチーム開発を行うイベントがあり、「GitHubリポジトリを食わせて質問できるアプリ」を作ってみました。

使用技術

LangChain (TypeScript)
OpenAI API
Prisma

仕組み

RAGパターンを用いて、リポジトリのソースコードを元にVector Storeを作成し、回答を生成していきます。
RAGについてはこちらの記事がわかりやすかったです。

大きく分けて2つのフェーズがあるのでそれぞれ見ていきます。

リポジトリからソースコードを取得し解析する。

1. ソースコードの取得

まずはLangChainが用意してくれているLoaderを使ってリポジトリからソースコードを取得します。

const getSourceCodes = async (
  repository: Repository,
): Promise<SourceCode[]> => {
  const loader = new GithubRepoLoader(repository.githubUrl, {
    accessToken: !!process.env.GITHUB_ACCESS_TOKEN
      ? process.env.GITHUB_ACCESS_TOKEN
      : undefined,
    branch: 'main',
    recursive: true,
    unknown: 'warn',
    maxConcurrency: 5,
  });

  const documents = await loader.load();

  return documents.map((document) => ({
    fileName: String(document.metadata.source),
    content: document.pageContent,
  }));
};

2. 取得したソースコードを分割する

OpenAIのToken制限対策として取得したソースコードを適当な長さに分割します。
RecursiveCharacterTextSplitterは特定の区切り文字等を元に文章を分割する機能を提供してくれます。
各プログラミング言語に最適化して分割する機能もありますが、今回は使用していません。

const chunkSourceCodes = async (
  sourceCodes: SourceCode[],
  chunkSize: number = 10000,
  chunkOverlap: number = 100,
) => {
  const splitter = new RecursiveCharacterTextSplitter({
    chunkSize,
    chunkOverlap,
  });

  const chunkedSourceCodes = [] as SourceCode[];

  for (const sourceCode of sourceCodes) {
    const docs = await splitter.createDocuments([sourceCode.content]);
    const chunkedSourceCode = docs.map((doc) => ({
      fileName: sourceCode.fileName,
      content: doc.pageContent,
    }));
    chunkedSourceCodes.push(...chunkedSourceCode);
  }

  return chunkedSourceCodes;
};

3. LLMにソースコードを解説してもらう

質問フェーズでは質問のvectorとソースコードのvectorを類似度検索し、質問と類似度の高いソースコードを元に回答を生成します。
この時、質問＝自然言語、ソースコード＝プログラミング言語の状態だと類似度検索が上手く行かないため、LLMにソースコードを解説してもらい自然言語に変換します。

const analyzeSourceCode = async (
  sourceCode: SourceCode,
  model: OpenAI,
): Promise<AnalyzedSourceCode> => {
  const prompt = `
  あなたは優秀なシステムエンジニアです。
  {#ファイル名}と{#ファイルコンテンツ}を元に、ファイルの内容を解説してください。
  
  #ファイル名
  ${sourceCode.fileName}
  
  #ファイルコンテンツ
  ${sourceCode.content}
  `;
  
  const analyzeResult = await model.call(prompt);

  return {
    fileName: sourceCode.fileName,
    content: sourceCode.content,
    description: analyzeResult,
  };
};

4.ソースコードの解説文をvector化してDBに保存する

3でできたソースコードの解説文をOpenAI Embedding APIを使ってvector化します。
その後ソースコードの情報や解説文、解説文のvector値等をDBに保存します。
この処理は、LangChainのPrismaVectorStoreを使えばもっとシンプルに書けますが、
「vector以外をInsert→Embedding→vectorをUpdate」という処理順になり、時間がかかってしまうため、自前でEmbedding→Bulk Insertにしています。

const saveVector = async (
  repository: Repository,
  analyzedSourceCodes: AnalyzedSourceCode[],
) => {
  const prisma = new PrismaClient();

  const openAIEmbedding = new OpenAIEmbeddings({
    verbose: true,
  });

  const embedSourceCodes = await Promise.all(
    analyzedSourceCodes.map(async (analyzedSourceCode) => {
      const vector = await openAIEmbedding.embedQuery(
        analyzedSourceCode.description,
      );

      return {
        ...analyzedSourceCode,
        vector,
      };
    }),
  );

  const insertValues = embedSourceCodes.map((embedSourceCode) => {
    return Prisma.sql`(
    ${Prisma.join([
      repository.id,
      embedSourceCode.fileName,
      embedSourceCode.content,
      embedSourceCode.description,
      embedSourceCode.vector,
    ])}
    )`;
  });

  await prisma.$transaction(async (transaction) => {
    await transaction.sourceCodeEmbedding.deleteMany({
      where: {
        repositoryId: repository.id,
      },
    });

    await transaction.repository.update({
      where: {
        id: repository.id,
      },
      data: {
        lastAnalysisAt: new Date(),
      },
    });

    await transaction.$executeRaw`
       INSERT INTO "SourceCodeEmbedding" (
       "repositoryId",
       "fileName",
       "sourceCode",
       "description",
       "descriptionEmbedding"
       )
        VALUES ${Prisma.join(insertValues)}
    `;
  });
}

質問に対して解析結果を元に回答を生成する

質問した際の処理の流れは以下のようになります。

質問文をvector化する。
1のvectorとVectorStore内のvector(ソースコードの解説文をembbedingしたもの)で類似度検索し類似度上位数件のレコードを取得する。
2で取得したレコードからソースコードの解説文を取得し、質問文と合わせてプロンプトを組む。
3のプロンプトをOpenAIに投げて回答を生成する。

これらはLangChainのRetrievalQAとVectorStoreを組み合わせることで実装することができます。
今回は、日本語で回答してほしかったのでプロンプトのテンプレートだけ差し替えて、他は基本的な使い方をしています。

const retrievalQA = async (repositoryId: string, question: string) => {
  const vectorStore = PrismaVectorStore.withModel<SourceCodeEmbedding>(
    new PrismaClient(),
  ).create(new OpenAIEmbeddings(), {
    prisma: Prisma,
    tableName: 'SourceCodeEmbedding',
    vectorColumnName: 'descriptionEmbedding',
    columns: {
      id: PrismaVectorStore.IdColumn,
      description: PrismaVectorStore.ContentColumn,
    },
    filter: {
      repositoryId: {
        equals: repository.id,
      },
    },
  });
  
  const template = `あなたは優秀なシステムエンジニアです。
    以下の{#コンテキスト}を元に、{#質問}に答えてください。
    
    #コンテキスト
    {context}
    
    #質問
    {question}`;

  const prompt = new PromptTemplate({
    inputVariables: ['context', 'question'],
    template: template,
  });

  const chain = RetrievalQAChain.fromLLM(
    new OpenAI({
      modelName: 'gpt-3.5-turbo-1106',
    }),
    vectorStore.asRetriever(),
    { prompt, inputKey: 'question' },
  );

  const qaResult = await chain.call({
    question,
  });

  return qaResult.text;
};

動かしてみる

$ npm run qa:repository githubからソースコードを取得しているファイルはありますか？ codemon

> co-demon@0.0.1 qa:repository
> ./node_modules/.bin/ts-node scripts/repository-qa.ts githubからソースコードを取得しているファイルはありますか？ codemon

質問: githubからソースコードを取得しているファイルはありますか？
リポジトリ名: codemon
はい、このスクリプト `analyze-repository.ts` は、GitHubからソースコードを取得するための処理を含んでいます。具体的には、「`getSourceCodes`」関数がその役割を担っており、GitHub APIを使用してリポジトリからソースコードを取得しています。

それっぽい答えを返してくれました！

最後に

とりあえず動くところまで実装することができました。
今回は時間が足りずCLIで動作させるだけでしたが、チャット形式のGUIでサクサク質問できるようにしたりしたいですね。