コードベースをチャットコンテキストとして使用する

概要

私たちは Duo Chat リクエストへの追加コンテキストとしてコードベースを含める機能を導入しています。これはリポジトリ全体または リポジトリ内のサブディレクトリを指します。

これを実現するために、コードベースをコードエンベディングと呼ばれるベクトルエンベディングとしてインデックス化します。

ユーザーが Duo Chat で質問すると、システムはコードエンベディングに対してセマンティック検索を実行してリポジトリから関連するコンテキストを取得し、それを大規模言語モデルで処理して有用なレスポンスを生成します。

この新機能は、Duo Pro または Duo Enterprise アドオンを持つ GitLab Premium または Ultimate ユーザーが利用できます。

エピック： この機能の作業はこのエピックで追跡されています。

動機

現在、私たちはお客様がリポジトリとコードベースを理解するための支援が十分ではありません。競合他社はより広い視野をサポートしており、ユーザーはリポジトリ全体について質問したり、複数のフォルダ、複数のファイル、コードの一部にコンテキストをスコープしたりできます。この機能的なギャップはお客様からよく言及されており、こちらはこのスペースにおける最近の調査まとめです。LLM は私たちが提供するコンテキストと同等の性能しか発揮できないため、この分野での競合との同等性を達成することが重要です。

このイニシアチブは、自然言語クエリを通じてユーザーがコードベース全体と対話できるようにすることで、GitLab の Duo Chat 提供における重要な機能的ギャップを埋めることを目指しています。この機能により、ユーザーはリポジトリの理解、ナビゲーション、変更計画をより効果的に行えるようになります。これは競合製品ですでに提供されている機能です。

ゴール

主なゴールは、Duo Chat への追加コンテキストとしてコードベースを追加することです。このイニシアチブでは、これはリポジトリとディレクトリにスコープされています。その後、リポジトリまたはディレクトリに対してセマンティック検索が行われ、その結果が AI モデルに送信される Chat プロンプトを強化するために使用されます。

セマンティック検索をサポートするために、コードエンベディングの作成がこの作業の初期スコープに含まれています。デフォルトブランチのインデックス化は作業の第 1 フェーズで行われ、フィーチャーブランチは第 2 フェーズでインデックス化されます。

Duo が有効なプロジェクトまたは名前空間に対してのみエンベディングを生成します。

非ゴール

以下はこのイニシアチブのスコープ外ですが、理論的にはその上に構築できます：

• エージェンティック Duo Chat への追加コンテキストとしてのコードベース。
• Duo Chat スラッシュコマンドへの追加コンテキストとしてのコードベース。
• コードサジェスションへの追加コンテキストとしてのコードベース。
• ローカルで変更されたファイルのベクトルエンベディングとしてのインデックス化とクエリのサポート。
• Duo Chat への追加コンテキストとしてのコードベースのナレッジグラフ表現。

上記のトピックに関する提案された計画については、次のステップと将来への備えを参照してください。

提案

コードベースをチャットコンテキストとしてサポートするために、以下が必要です：

1. コードエンベディングを導入する

   • これはコードベース内のファイルのベクトル表現です。
   • コードベースをベクトルエンベディングとしてインデックス化するパイプラインを含みます
   • エンベディングに対してセマンティック検索を実行する機能を提供します
   • 2 つのフェーズで開発されます：
     • フェーズ 1：メインブランチでのコードエンベディングのサポート
     • フェーズ 2：フィーチャーブランチでのコードエンベディングのサポート

2. Duo Chat を更新して追加コンテキストとしてのコードベースをサポートする。

   • Chat で質問する際、ユーザーは以下を追加コンテキストとして含めるオプションを持ちます：
     • repository - プロジェクトのコードベース全体を指します
     • directory - repository の「サブセット」であり、プロジェクト内のサブフォルダを指します
   • repository が追加コンテキストとして選択された場合：
     • リポジトリ内のファイルのコードエンベディング表現に対してセマンティック検索が行われます。検索結果はその後、AI モデルに送信される Chat プロンプトを強化するために使用されます。
   • directory が追加コンテキストとして選択された場合：
     • ディレクトリ内のファイルのコードエンベディング表現に対してセマンティック検索が行われます。検索結果はその後、AI モデルに送信される Chat プロンプトを強化するために使用されます。
     • 追加の考慮事項：ディレクトリにファイルが少ない場合、ファイルの内容が AI モデルに送信される Chat プロンプトを強化するための追加コンテキストとして直接含まれます。

設計と実装の詳細

コンポーネント

このイニシアチブは以下のコンポーネントを導入または更新します：

コードエンベディング

このコンポーネントの詳細な設計については、**コードエンベディングブループリント**を参照してください。

これはコードベース内のファイルのベクトル表現です。ブランチに変更がプッシュされたときにエンベディングを生成するインデックス化パイプラインを導入します。また、これらのエンベディングを検索する機能も導入します。

コードエンベディングのインデックス化

• 変更は GitLab Elasticsearch Indexer と GitLab Rails の両方で行われます。
• Rails では、AI Context Abstraction Layer を使用します。
• エンベディングを生成する前にコードファイルを論理チャンクに解析するための新しいコードパーサーライブラリを使用します。
  • コードパーサーは異なるプロジェクトで使用できるように独自のリポジトリに置かれます。
  • 詳細な設計と実装については、One Parser 提案を参照してください。

コードエンベディングの検索

• 変更は GitLab Rails で行われます
• エンベディングの検索を実行するために AI Context Abstraction Layer を使用します。

Duo Chat

Duo Chat は GitLab の既存の AI 機能です。

現在の Duo Chat のワークフローとアーキテクチャの詳細については、以下のドキュメントを参照してください：

• Chat プロンプトの構築方法
• GraphQL API（aiAction）フロー
• Duo Chat プロセスフロー
• Duo Chat ツール

このイニシアチブでは、GitLab Language Server、GitLab Rails、および GitLab AI Gateway の変更を導入します。詳細な実装については、以下の Duo Chat へのコードベースのコンテキストとしての追加セクションを参照してください。

Duo Chat へのコードベースのコンテキストとしての追加

以下の図は、コードベース（リポジトリまたはディレクトリ）のセマンティック検索結果を追加コンテキストとして含めるワークフローを示しています。青色でハイライトされた領域が新しい変更を導入する場所です。

sequenceDiagram
    actor USR as User
    participant FE as IDE/Language Server
    box GitLab Rails
      participant GLRGQL as GraphQL API
      participant GLRLLMS as LLM Services
      participant GLRLLMCHAT as LLM Chat Module
      participant GLRLLMSEM as LLM Semantic Search Tool
      participant CES as Code Embeddings Search Service
    end
    participant CE as Code Embeddings Storage
    participant AIGW as AI Gateway
    participant LLM as LLM

    USR->>FE: Asks a question, indicating<br /> repository or directory<br /> as additional context

    FE->>GLRGQL: Sends chat request to `aiAction` mutation<br /> with repository or directory<br /> as additional context
    Note over FE: The Language Server will send the ID and category<br /> of the repository and directory additional contexts,<br /> but the content will still be empty at this point.
    GLRGQL->>GLRLLMS: Sends chat request, with<br /> repository or directory<br /> as additional context
    GLRLLMS->>GLRLLMCHAT: Sends chat request, with<br /> repository or directory<br /> as additional context

    rect rgb(240, 248, 255)
      GLRLLMCHAT->>GLRLLMSEM: Executes the Semantic Search Tool,<br /> with a repository or directory filter
      GLRLLMSEM->>CES: Performs the semantic search

      CES->>AIGW: Generates embeddings for the question
      AIGW-->>CES: Returns embeddings for the question
      CES->>CE: Queries Code Embeddings with<br /> the question embeddings as target<br /> and filtered by the given repository or directory
      CE-->>CES: Returns Code Embeddings
      CES-->>GLRLLMSEM: Returns semantic search result

      GLRLLMSEM-->>GLRLLMCHAT: Returns semantic search result
      GLRLLMCHAT->>GLRLLMCHAT: Adds the semantic search result as the content<br /> to the repository and directory additional context
      Note over GLRLLMCHAT: The content of the repository and directory<br /> additional contexts will be set here.

      GLRLLMCHAT->>AIGW: Request to /v2/chat/agent<br /> with repository or directory +<br /> semantic search result as additional context
      AIGW->>AIGW: Builds a prompt with the<br /> repository or directory +<br /> semantic search result<br /> as additional context
    end

    AIGW->>LLM: Sends prompt with the additional contexts

    LLM-->>AIGW: Returns the final answer
    AIGW-->>GLRLLMSEM: Returns the final answer
    GLRLLMSEM-->>GLRLLMCHAT: Returns the final answer
    GLRLLMCHAT-->>GLRLLMS: Returns the final answer
    GLRLLMS-->>GLRGQL: Returns the final answer
    GLRGQL-->>FE: Returns the final answer
    FE-->>USR: Shows the answer

追加コンテキストカテゴリ

1 つの追加コンテキストカテゴリを追加します：repository。ディレクトリが指定された場合は、metadata にディレクトリの相対パスを指定した repository 追加コンテキストとして扱われます。

ユニットプリミティブ

このイニシアチブの一部として以下のユニットプリミティブを追加します：

コンテキストの包含

• include_repository_context - Language Server と AIGW の両方で使用されます

ツール

• codebase_search - Rails で使用されます。セマンティック検索ツールのために使用されます

エンベディング生成

• generate_embeddings_codebase - エンベディング生成エンドポイント（/v1/proxy/vertex-ai）に使用されるユニットプリミティブ

コードエンベディング検索サービス

これはコードエンベディングへの呼び出しを処理するサービスクラスです。これは AI Context Abstraction Layer を使用します。これがどのように行われるかの詳細については、コードエンベディングブループリントの検索セクションを参照してください。

Duo セマンティック検索ツール

これはこのイニシアチブで導入される新しい Duo Chat ツールです。

スラッシュコマンドツールと同様に、ツールが必要かどうかを LLM が推論する必要はありません。このツールは repository 追加コンテキストが存在する限り呼び出されます。

Rails では、Chat の質問のマッチするエンベディングを取得するためにコードエンベディング検索サービスを呼び出す LLM ツールクラスを導入します。この新しいセマンティック検索ツールは LLM Chat モジュールから呼び出され、検索結果は repository 追加コンテキストの_コンテンツ_として含まれます。Chat リクエストは新しい追加コンテキストとともに AIGW に送信されます。

API 変更 - Duo Chat 利用可能な機能

Language Server は GraphQL クエリ {currentUser { duoChatAvailableFeatures } } を呼び出して利用可能な Duo Chat 機能のリストを取得します。

このイニシアチブの一部として、include_repository_context ユニットプリミティブをこの機能リストに追加します。

API 変更 - Chat リクエスト

Duo Chat で使用される GraphQL ミューテーション（aiAction） はすでに chat 入力のパラメーターとして additionalContext を受け付けています。

repository 追加コンテキストカテゴリを使用すると、aiAction ミューテーションへの chat 入力は次のようになります：

repository を追加コンテキストとして使用する場合

mutation newChatMessage {
  aiAction(
    input: {
      chat: {
        content: "the user question"
        additionalContext: [{
          category: "repository",
          id: "the-project-id",
          content: "", # should be empty
          metadata: {
            directory: "some/dir" # this is optional, specified when the user selects directory as an additional context
            branch: "some-branch" # including the branch is a second phase iteration
          }
        }]
      }
    }
  ) {
    requestId
  }
}

Duo Chat 変更 - フロントエンド

UI デザインを参照してください。

評価

コードベースのコンテキスト強化は、エンベディングの粒度や使用するエンベディングモデルなど、さまざまな要因によって異なる結果が出る可能性があります。この機能の MVC イテレーションを超えて、さまざまなエンベディングモデル、チャンキングの粒度、その他のエンベディングへのアプローチの有効性を評価すべきです。

評価のための考えられるアプローチ

次のステップと将来への備え

エージェンティック Chat アーキテクチャへの移植の提案された手順

エージェンティック Chat アーキテクチャを導入したら、AI Gateway 上の Duo ワークフローサービスまたは Language Server 上の Duo ワークフローエグゼキューターのどちらかがベクトルエンベディングをクエリする必要があります。

これをサポートするために、コードエンベディング検索サービスの上にDuo ワークフローサービスまたはDuo ワークフローエグゼキューターのどちらかから呼び出される API を導入します。

Duo Chat スラッシュコマンドへの追加コンテキストとしてのコードベース

スラッシュコマンドには /refactor、/fix、/test が含まれます。

スラッシュコマンドは**セマンティック検索ツールを使用するか、コードエンベディング検索サービス**を直接呼び出してコードエンベディングを検索することができます。

あるいは、エージェンティック Chat でのみスラッシュコマンドへの追加コンテキストとしてコードベースをサポートすることもできます。

コードサジェスションへの追加コンテキストとしてのコードベース

コード補完またはコード生成は**コードエンベディング検索サービス**を使用でき、これはコードエンベディングの検索に必要なすべてのロジックを抽象化します。

ローカルファイルインデックス化のサポートの提案された手順

TBA

ナレッジグラフとしてのコードベースのインデックス化

TBA

代替ソリューション

LLM がコードベースセマンティック検索の必要性を推論できるようにする

このソリューションでは、AIGW 上のコードベースセマンティック検索のツール定義を導入します。このツールは LLM に提供され、LLM は質問に基づいてツールが必要かどうかを推論します。

このソリューションを採用しない理由は以下のとおりです：

• コードベース検索ツールが必要かどうかを LLM が推論する必要はありません。repository または directory の追加コンテキストが含まれている場合は、すぐにコードベース検索を実行できます。
• MVC の提案とプロダクトの要件は、コードベースセマンティック検索が必要かどうかをユーザーが明示的に決定できるべきというものです。このソリューションでは、ユーザーが repository または directory を追加コンテキストとして指定できますが、コードベース検索が実行されるかどうかの最終決定は LLM に委ねられます。
• このソリューションは Rails と AIGW の間で追加のラウンドトリップが必要となり、レイテンシが高くなります。

ツール定義の PoC： MR: POC: コードベース検索ツール。

ワークフロー図：

sequenceDiagram
    actor USR as User
    participant FE as IDE/Language Server
    box GitLab Rails
      participant GLRGQL as GraphQL API
      participant GLRLLMS as LLM Services
      participant GLRLLMCHAT as LLM Chat Module
      participant GLRLLMSEM as LLM Semantic Search Tool
      participant CES as Code Embeddings Search Service
    end
    participant CE as Code Embeddings Storage
    participant AIGW as AI Gateway
    participant LLM as LLM

    USR->>FE: Asks a question, indicating<br /> codebase as additional context

    FE->>GLRGQL: Sends chat request to `aiAction` mutation<br /> with codebase as additional context
    GLRGQL->>GLRLLMS: Sends chat request, with<br /> codebase as additional context
    GLRLLMS->>GLRLLMCHAT: Sends chat request, with<br /> codebase as additional context
    GLRLLMCHAT->>AIGW: Request to /v2/chat/agent<br /> with codebase as additional context

    rect rgb(240, 248, 255)
      AIGW->>LLM: Sends chat request, indicating<br /> that codebase context is present
      LLM->>LLM: Determines that the<br /> semantic search tool is needed
      Note over LLM: The LLM will have a new available<br /> tool for semantic search.<br /> This tool will have instructions<br /> to perform a semantic search if<br /> the codebase context is present.
      LLM-->>AIGW: Returns response, indicating that<br /> the semantic search tool is needed
    end

    AIGW-->>GLRLLMCHAT: Returns response, indicating that<br /> the semantic search tool is needed

    GLRLLMCHAT->>GLRLLMSEM: Executes the Semantic Search Tool

    rect rgb(240, 248, 255)
      GLRLLMSEM->>CES: Performs the semantic search
      CES->>AIGW: Generates embeddings for the question
      AIGW-->>CES: Returns embeddings for the question
      CES->>CE: Queries Code Embeddings with<br /> the question embeddings as target
      CE-->>CES: Returns Code Embeddings
      CES-->>GLRLLMSEM: Returns semantic search result
    end

    GLRLLMSEM->>AIGW: Request to /v1/prompts/chat<br />with search results as additional context

    AIGW->>AIGW: Builds a prompt with the<br /> search results as additional context
    AIGW->>LLM: Sends prompt with the<br /> search results as additional context

    LLM-->>AIGW: Returns the final answer
    AIGW-->>GLRLLMSEM: Returns the final answer
    GLRLLMSEM-->>GLRLLMCHAT: Returns the final answer
    GLRLLMCHAT-->>GLRLLMS: Returns the final answer
    GLRLLMS-->>GLRGQL: Returns the final answer
    GLRGQL-->>FE: Returns the final answer
    FE-->>USR: Shows the answer

「コードベースに聞く」ツールとプロンプトを導入する

このソリューションでは、「コードベースに聞く」ツールを導入します：

• Rails では、repository 追加コンテキストがある場合に LLM Chat モジュールがツールが必要であることを判断します
• 「コードベースに聞く」ツールはセマンティック検索結果を取得するために「コードベースエンベディング検索サービス」を使用します
• 「コードベースに聞く」ツールは /v1/prompts/chat エンドポイントへのリクエストを送信します
• AIGW では、/v1/prompts/chat エンドポイントを通じて利用可能な ask_codebase の新しいプロンプトがあります
• ワークフローは本質的にスラッシュコマンドツールのワークフローと同様です

sequenceDiagram
    actor USR as User
    participant FE as IDE/Language Server
    box GitLab Rails
      participant GLRGQL as GraphQL API
      participant GLRLLMS as LLM Services
      participant GLRLLMCHAT as LLM Chat Module
      participant GLRLLMAC as LLM Ask Codebase Tool
      participant CES as Code Embeddings Search Service
    end
    participant CE as Code Embeddings Storage
    participant AIGW as AI Gateway
    participant LLM as LLM

    USR->>FE: Asks a question, indicating<br /> repository or directory<br /> as additional context

    FE->>GLRGQL: Sends chat request to `aiAction` mutation<br /> with repository or directory<br /> as additional context
    Note over FE: The Language Server will send the ID and category<br /> of the repository and directory additional contexts,<br /> but the content will still be empty at this point.
    GLRGQL->>GLRLLMS: Sends chat request, with<br /> repository or directory<br /> as additional context
    GLRLLMS->>GLRLLMCHAT: Sends chat request, with<br /> repository or directory<br /> as additional context

    rect rgb(240, 248, 255)
      GLRLLMCHAT->>GLRLLMAC: Executes the Ask Codebase Tool,<br /> with a repository or directory filter
      GLRLLMAC->>CES: Performs the semantic search

      CES->>AIGW: Generates embeddings for the question
      AIGW-->>CES: Returns embeddings for the question
      CES->>CE: Queries Code Embeddings with<br /> the question embeddings as target<br /> and filtered by the given repository or directory
      CE-->>CES: Returns Code Embeddings
      CES-->>GLRLLMAC: Returns semantic search result

      GLRLLMAC-->>GLRLLMCHAT: Returns semantic search result
      GLRLLMCHAT->>GLRLLMCHAT: Adds the semantic search result as additional context

      GLRLLMCHAT->>AIGW: Request to /v1/prompts/chat<br /> with repository or directory +<br /> semantic search result as additional context
      AIGW->>AIGW: Builds a prompt with the<br /> repository or directory +<br /> semantic search result<br /> as additional context
      Note over AIGW: This will be a new `ask_codebase` prompt
    end

    AIGW->>LLM: Sends prompt with the additional contexts

    LLM-->>AIGW: Returns the final answer
    AIGW-->>GLRLLMAC: Returns the final answer
    GLRLLMAC-->>GLRLLMCHAT: Returns the final answer
    GLRLLMCHAT-->>GLRLLMS: Returns the final answer
    GLRLLMS-->>GLRGQL: Returns the final answer
    GLRGQL-->>FE: Returns the final answer
    FE-->>USR: Shows the answer

repository と directory に異なるカテゴリとユニットプリミティブを使用する

リポジトリとディレクトリの追加コンテキストに単一の repository カテゴリを使用する代わりに、repository と directory のカテゴリを使用します。

注：カテゴリはユニットプリミティブと 1 対 1 のマッピングを持っており、概念的には「コードベースを含める」コンテキストに対して 1 つのユニットプリミティブしか持つべきではないため、このオプションを採用しないことに決めました。

aiAction ミューテーションへの chat 入力は次のようになります：

repository を追加コンテキストとして使用する場合

mutation newChatMessage {
  aiAction(
    input: {
      chat: {
        content: "the user question"
        additionalContext: [{
          category: "repository",
          id: "the-project-id",
          content: "", // should be empty
          metadata: {'branch': 'some-branch'} // including the branch is a second phase iteration
        }]
      }
    }
  ) {
    requestId
  }
}

directory を追加コンテキストとして使用する場合

mutation newChatMessage {
  aiAction(
    input: {
      chat: {
        content: "the user question"
        additionalContext: [{
          category: "directory",
          id: "file:///home/user/workspace/src/dir",
          content: "", // should be empty
          metadata: {
            'relativePath': 'src/dir',
            'branch': 'some-branch' // including the branch is a second phase iteration
          }
        }]
      }
    }
  ) {
    requestId
  }
}

ページのソースコードを見る - ページの編集 - please contribute.