GitHub Copilot RAG-Extension

こちらの GitHub Copilot RAG-Extensionを試してみたので備忘録です。

GitHub Copilot Extensions とは

GitHub Copilot Extensionsとは端的に言うと、GitHub Copilot Chatのエージェントやスキルセット (@で呼び出せる拡張機能) を自作できるものです。

今回はこの Copilot ExtensionsでRAGを構築するサンプルがあるので動かしてみました。

GitHub アプリの作成

まずは上記のドキュメント通りにGitHub アプリを作成します。

1. あとで設定変更するので[Homepage URL]は一旦仮で入れておきます
2. Webhook [Active] のチェックを外しておきます
3. Permissionsの [Account permissions] > [Copilot Chat] を Read-onlyにしておきます

作成後に、App IDとClient Secretsをコピーしておきましょう。

次に、こちらのリポジトリをフォーク & クローンします。

ngrokを使用してlocalhost:8080をポートフォワーディングするように起動しておきます。

ngrok http 8080

...
Session Status online 
Account   ××××××××
Version   3.19.1
Region    Japan (jp) 
Latency   ××ms
Forwarding  https://××××××××××.ngrok-free.app -> http://localhost:8080

ngrok が生成した https://××××××××××.ngrok-free.app という Forwarding URL をコピーしておきます。次に、これまでにコピーした値を使って環境変数を設定します。

export PORT=8080
export CLIENT_ID=<コピーしたApplication id>
export CLIENT_SECRET=<コピーしたシークレット> 
export FQDN=<コピーしたngrokのurl>

依存関係をインストールし、アプリケーションを起動します。

go mod tidy
go run .

ここで、GitHub Appsの設定画面に戻ります。

1. [General] の設定から [Homepage URL] を<ngrokのURL>/auth/callback に修正します
2. [Callback URL] を <ngrokのURL>/auth/callback に修正します
3. [Copilot] の設定から [App Type] を [Agent] にします
4. [Pre-authorization URL] を <ngrokのURL>/auth/authorizationに修正します
5. [Agent Definition] のURLを <ngrokのURL>/agent に修正します

アプリケーションに戻り、data ディレクトリ内の既存のMarkdownファイルを削除し、新しいMarkdownファイルを追加します。例として、以下のようなドキュメントを使用します。

sample.md

# ドラゴンボールの物語概要

『ドラゴンボール』は、鳥山明による日本の漫画作品で、1984 年から 1995 年まで『週刊少年ジャンプ』で連載されました。物語は、主人公・孫悟空と、七つ集めるとどんな願いも叶う秘宝「ドラゴンボール」を巡る冒険を描いています。連載終了後も、テレビアニメ、映画、ゲームなど多岐にわたるメディア展開が続いており、2024 年秋には新作アニメ『ドラゴンボール DAIMA』が放送予定です。

## 主要なストーリーライン

1. **少年編**：幼少期の孫悟空が、ブルマや亀仙人、クリリンなどと出会い、ドラゴンボールを探す冒険に出ます。天下一武道会への参加や、ピッコロ大魔王との戦いが描かれます。

2. **サイヤ人編**：成長した悟空の前に、同じサイヤ人であるベジータやナッパが地球に襲来。激闘の末、ベジータは撤退します。

3. **フリーザ編**：悟空たちはナメック星で、宇宙の帝王フリーザと対決。悟空は伝説のスーパーサイヤ人に覚醒し、フリーザを打ち破ります。

4. **人造人間・セル編**：未来から来たトランクスの警告により、人造人間やセルとの戦いが始まります。悟飯がスーパーサイヤ人 2 に覚醒し、セルを撃破します。

5. **魔人ブウ編**：復活した魔人ブウとの戦いが描かれます。悟空はスーパーサイヤ人 3 やフュージョンなど新たな力を駆使し、最終的に元気玉でブウを倒します。

## メディア展開

- **テレビアニメ**：『ドラゴンボール』『ドラゴンボール Z』『ドラゴンボール GT』『ドラゴンボール改』『ドラゴンボール超』など、多くのシリーズが放送されています。

- **映画**：多数の劇場版アニメが公開されており、最新作は『ドラゴンボール超 ブロリー』です。

- **ゲーム**：多くのプラットフォームでゲームが発売されており、シリーズ累計販売本数は全世界で 5000 万本に達します。

## 世界的な人気

『ドラゴンボール』は、全世界で 2 億 6000 万部以上の単行本売上を記録しており、80 か国以上でアニメが放送されるなど、世界中で絶大な人気を誇る作品です。

VSCodeを開き、作成したCopilotエージェントをテストしてみましょう。@マークの後に作成したGitHub アプリ名を入力して質問してみます。（初回実行時は認証フローが発生します）

RAGが正常に動作していることが確認できました！

コードを読んでみると、datasets.go では埋め込み処理と類似度検索が実装されており、ユーザーの質問に類似したMarkdownドキュメントを取得しています。

service.go のgenerateCompletionメソッドでは、これらのドキュメントをコンテキストとしてLLMに渡して回答を生成しています。なお、モデルをGPT4に変更し、システムプロンプトを日本語にして動作確認しました。

しかし、Markdownドキュメントの数が増えると、埋め込み処理で error fetching embeddings: unexpected status code: 429 というエラーが発生するようになりました。
数個のMarkdownファイルなら問題ありませんが、複数のファイルを処理する場合、このままだと支障が出そうです。

ベクトルデータストアを使用する

そこで専用のベクトルデータストアを使用して、類似検索をそちらで行うことにします。

今回はベクトルデータストアとしてAzure CosmosDB for No SQLのベクトル検索機能を利用します。

なお、CosmosDBのベクトル検索機能は現在プレビュー段階です。ベクトル インデックス作成と検索の機能を有効にするの手順に従って、機能を有効化してからCosmosDBを作成します。

以下のコードは、Lang Chain と Azure OpenAI を使用して、Markdownファイルをチャンクに分割し、埋め込み処理を行ってCosmosDBに格納するものです。また、パーティションキーとして使用するメタデータを付与しています（詳細は後述）。

なお、ハイブリッド検索（全文検索とベクトル検索の組み合わせ）も利用可能ですが、現時点では英語のみ対応であるため、今回は有効化していません。

サンプルコード

https://python.langchain.com/docs/integrations/vectorstores/azure_cosmos_db_no_sql/

cosmos_upload.py

from langchain_community.vectorstores.azure_cosmos_db_no_sql import (
    AzureCosmosDBNoSqlVectorSearch,
)
from azure.cosmos import CosmosClient, PartitionKey
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import DirectoryLoader
import os
from dotenv import load_dotenv
from langchain_openai import AzureOpenAIEmbeddings
from langchain_community.vectorstores.azuresearch import AzureSearch

load_dotenv()

os.environ["OPENAI_API_VERSION"] = "2024-08-01-preview"
os.environ["AZURE_OPENAI_ENDPOINT"] = os.environ.get("AZURE_OPENAI_ENDPOINT")

# コスモスDBにマークダウンを保存する
loader = DirectoryLoader("./knowledge/", glob="*.md")
data = loader.load()
len(data)

text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000, chunk_overlap=150)
docs = text_splitter.split_documents(data)
print(docs[0])

indexing_policy = {
    "indexingMode": "consistent",
    "includedPaths": [{"path": "/*"}],
    "excludedPaths": [{"path": '/"_etag"/?'}],
    "vectorIndexes": [{"path": "/embedding", "type": "diskANN"}],
    # "fullTextIndexes": [{"path": "/text"}],
}

vector_embedding_policy = {
    "vectorEmbeddings": [
        {
            "path": "/embedding",
            "dataType": "float32",
            "distanceFunction": "cosine",
            "dimensions": 1536,
        }
    ]
}

# フルテキスト検索は現在en-USのみサポートしている状況なので利用しない
# https://learn.microsoft.com/en-us/azure/cosmos-db/gen-ai/full-text-search#full-text-policy
# full_text_policy = {
#     "defaultLanguage": "en-US",
#     "fullTextPaths": [{"path": "/text", "language": "en-US"}],
# }

HOST = os.environ.get("AZURE_COSMOS_DB_HOST")
KEY = os.environ.get("AZURE_COSMOS_DB_KEY")

cosmos_client = CosmosClient(HOST, KEY)
database_name = "knowledge_db"
container_name = "knowledge_container"
# pkはメタデータに付与
partition_key = PartitionKey(path="/metadata/pk")
cosmos_container_properties = {"partition_key": partition_key}

openai_embeddings = AzureOpenAIEmbeddings(
    deployment="text-embedding-ada-002",
    model="text-embedding-ada-002",
    chunk_size=1,
    openai_api_key=os.environ.get("AZURE_OPENAI_API_KEY"),
    openai_api_type="azure"
)

# 各ドキュメントにメタデータを付与する
for doc in docs:
    if not hasattr(doc, "metadata") or doc.metadata is None:
        doc.metadata = {}
    # パーティションキーの設定
    doc.metadata["pk"] = "knowledge"

AzureCosmosDBNoSqlVectorSearch.from_documents(
    documents=docs,
    embedding=openai_embeddings,
    cosmos_client=cosmos_client,
    database_name=database_name,
    container_name=container_name,
    vector_embedding_policy=vector_embedding_policy,
    # full_text_policy=full_text_policy,
    indexing_policy=indexing_policy,
    cosmos_container_properties=cosmos_container_properties,
    cosmos_database_properties={},
    full_text_search_enabled=False,
)

格納後⇩

Copilot Extensionのコードに戻り一部修正します。

datasets.goに Cosmos DB へのベクトル検索処理を追加します。

embedding/datasets.go

package embedding

...

type CosmosClient struct {
	Client *azcosmos.Client
}

type Doc struct {
	Id        string    `json:"id"`
	Text      string    `json:"text"`
	Embedding []float32 `json:"embedding"`
}

func NewCosmosClient() (*CosmosClient, error) {
	key := os.Getenv("AZURE_COSMOS_ACCOUNT_KEY")
	endPoint := os.Getenv("AZURE_COSMOS_ENDPOINT")
	cred, err := azcosmos.NewKeyCredential(key)
	if err != nil {
		return nil, fmt.Errorf("failed to get default credentials: %w", err)
	}
	client, err := azcosmos.NewClientWithKey(endPoint, cred, nil)
	if err != nil {
		return nil, fmt.Errorf("faild to create client: %w", err)
	}
	return &CosmosClient{Client: client}, nil
}

func (c *CosmosClient) GetDocs(vector []float32) ([]Doc, error) {
	ctx := context.Background()
	dbName := os.Getenv("AZURE_COSMOS_DB_NAME")
	containerName := os.Getenv("AZURE_COSMOS_CONTAINER_NAME")
	pK := azcosmos.NewPartitionKeyString("knowledge") //PK指定
	db, err := c.Client.NewDatabase(dbName)
	if err != nil {
		return nil, fmt.Errorf("failed to get database: %w", err)
	}
	container, err := db.NewContainer(containerName)
	if err != nil {
		return nil, fmt.Errorf("failed to get container: %w", err)
	}

	query := "SELECT TOP 10 c.text, VectorDistance(c.embedding,@embedding) AS SimilarityScore FROM c ORDER BY VectorDistance(c.embedding,@embedding)"
	queryOptions := azcosmos.QueryOptions{
		QueryParameters: []azcosmos.QueryParameter{
			{Name: "@embedding", Value: vector},
		},
	}

	pager := container.NewQueryItemsPager(query, pK, &queryOptions)
	docs := []Doc{}

	requestCharge := float32(0)

	for pager.More() {
		response, err := pager.NextPage(ctx)
		if err != nil {
			return nil, err
		}

		requestCharge += response.RequestCharge
		for _, bytes := range response.Items {
			item := Doc{}
			err := json.Unmarshal(bytes, &item)
			if err != nil {
				return nil, err
			}
			docs = append(docs, item)
		}
	}
	return docs, nil
}

どうも Go言語のSDKではパーティションキーを指定せずにベクトル検索クエリをかけることができなそうでした。苦肉の策としてLangChainで格納する際にメタデータへ追加したknowledgeという値をパーティションキーに指定しています。

service.go のベクトル検索処理部分を修正しておきます。

agent/service.go

...
func (s *Service) generateCompletion(ctx context.Context, integrationID, apiToken string, req *copilot.ChatRequest, w io.Writer) error {
	// Initialize the datasets.  In a real application, these would be generated
	// ahead of time and stored in a database
	var err error
	var messages []copilot.ChatMessage

	// Create embeddings from user messages
	for i := len(req.Messages) - 1; i >= 0; i-- {
		msg := req.Messages[i]
		if msg.Role != "user" {
			continue
		}

		// Filter empty messages
		if msg.Content == "" {
			continue
		}

		emb, err := embedding.Create(ctx, integrationID, apiToken, msg.Content)
		if err != nil {
			return fmt.Errorf("error creating embedding for user message: %w", err)
		}
		// ここをCosmosDBを使ったベクトル検索に変更
		cosmos, _ := embedding.NewCosmosClient()
		docs, err := cosmos.GetDocs(emb)
		if err != nil {
			return fmt.Errorf("error computing best dataset")
		}
		if len(docs) == 0 {
			break
		}

		fmt.Printf("loading docs: %s\n", docs[0].Text)
		fileContents := docs[0].Text

		messages = append(messages, copilot.ChatMessage{
			Role: "system",
			Content: "あなたは、ユーザーのメッセージに返信する親切なアシスタントです。メッセージに返信するときは、以下のコンテキストを確実に使用してください。\n" +
				"Context: " + fileContents,
		})

		break
	}

	messages = append(messages, req.Messages...)

	chatReq := &copilot.ChatCompletionsRequest{
		Model:    copilot.ModelGPT4,
		Messages: messages,
		Stream:   true,
	}
 ...

コードの修正後に動作確認を行ったところ、期待通りの精度で適切な回答が得られるようになりました。

さいごに

普段使っているGitHub Copilotで自作Extensionを呼び出せるのいいですね。

自分なりのナレッジ集だったり、特定のドキュメントだったりを参照できるようにしておくと、かなり使い勝手良さそうです。AIエージェントブームではありますが、作りたいものが無くて困っていたので、今後はGitHub Copilot の自作Extensionsでよりエージェントっぽいものを何か作ってみようかな～と思います😊