💬

Goで作るセキュリティ分析LLMエージェント(13): 会話コンテキストの圧縮

に公開

この記事はアドベントカレンダー「Goで作るセキュリティ分析LLMエージェント」の13日目です。

今回のコードは https://github.com/m-mizutani/leveretday13-context-compression ブランチに格納されていますので適宜参照してください。

トークン制限の課題

LLMエージェントを運用していくうえで避けて通れない問題が、コンテキストウィンドウ(トークン限界)の超過です。これまでの記事でも何度か触れてきましたが、会話が長くなったりツールの実行結果が蓄積されたりすると、必ずこの問題に直面します。本記事では、この問題にどう対処するかを解説します。

LLMにおけるトークンに関する基礎知識

まず、トークンについての基礎知識を整理しておきます。ここでいうトークンとは、生成AI(特にLarge Language Model)が処理するために文字列を分割した単位のことです。英語の場合は1単語が1トークンとなることが多いのですが、必ずしもそうとは限りません。記号や特殊文字なども1トークンとして表されますし、長い単語は複数のトークンに分割されることもあります。

トークン数は原則としてサービスに問い合わせないと正確な値がわかりません。Geminiの場合、トークン数の取得自体は無料で提供されています。トークンの概念や分割方法はモデルやサービスごとに異なり、BPE(Byte Pair Encoding)やWordPieceといった手法が用いられています。また、言語によってもトークン数の傾向が変わります。英語では1単語が1トークン前後になることが多いのに対し、日本語では同じ文字列長に対してトークン数が多くなる傾向があります。

Geminiでトークン数を取得するには、以下のようなコードを使います。

client, err := genai.NewClient(ctx, &genai.ClientConfig{
    Project:  os.Getenv("GEMINI_PROJECT"),
    Location: location,
    Backend:  genai.BackendVertexAI,
})
if err != nil {
    return goerr.Wrap(err, "failed to create genai client")
}

// Count tokens
contents := genai.Text(text)
resp, err := client.Models.CountTokens(ctx, model, contents, nil)
if err != nil {
    return goerr.Wrap(err, "failed to count tokens")
}

fmt.Printf("Token count: %d tokens\n", resp.TotalTokens)

実際に実行してみると、トークン分割の様子を確認できます。

$ cd examples/count-token
$ env GEMINI_PROJECT=your-project go run . "hello world"
Token count: 2 tokens
$ go run . '{"hello":"world"}'
Token count: 5 tokens
$ go run . 'こんにちは'
Token count: 1 tokens
$ go run . '寿限無寿限無五劫の擦り切れ'
Token count: 12 tokens

英語の "hello world" は2トークン、JSON形式の文字列は括弧やコロンなどの記号も含めて5トークンとなります。日本語の場合、「こんにちは」は1トークン、「寿限無寿限無五劫の擦り切れ」(15文字)は12トークンとなっています。英語は文字数(10文字)に対してトークン数(2トークン)が少ないのに対し、日本語は文字数に近いトークン数になる傾向があります。

LLMエージェント利用時のトークン制限問題

モデルごとのトークン限界は年々増加していますが、それでも有限です。例えばGeminiの場合、100万トークン(正確には1,048,576トークン)が上限となっています。これはテキスト換算でおよそ4MB程度の容量に相当し、一見すると非常に長いように思えますが、LLMエージェントを運用していると意外と簡単に到達してしまいます。

LLMエージェントでトークン制限が問題になるケースは、大きく分けて2つあります。

(1) 会話が続くことで履歴が長くなる

すでに実装した通り、LLMエージェントとの「会話」の実態は「都度会話の履歴を全部投げつける」ということでした。LLMは基本的にステートレスであるため、前回までの文脈を保持するには、毎回すべての履歴を含めてリクエストする必要があります。そのため会話が続くと自ずと履歴が長くなっていきます。個々のメッセージが短くても、やり取りの回数が増えれば全体のトークン数は増加し、最終的にはトークン限界を超過してしまいます。特にツール呼び出しを繰り返すような処理では、ツールの実行結果も履歴に含まれるため、トークン消費が加速します。

(2) 一度に長いメッセージが投入される場合

もう一つのケースは、単一のメッセージが非常に長い場合です。ユーザが大きなファイルやログデータを直接入力しようとすると発生しやすい問題です。また、ツール実行の結果が非常に大きい場合(例えばログ検索の結果が数万行に及ぶなど)にも同様の問題が起こります。

このケースは会話履歴の要約だけでは解決できません。入力データそのもののサイズが限界を超えている場合、データ自体を圧縮したりフィルタリングしたりする必要があります。本記事では主に(1)のケース、つまり会話履歴の蓄積によるトークン超過への対処方法を扱います。

会話履歴の要約と戦略

会話履歴が限界を超えた場合、古い履歴を単純に削除するという対処も可能です。しかしこれには重大な問題があります。過去に会話した内容を「忘れて」しまうことになるからです。

例えば、最初にユーザから重要な指示や制約条件が与えられていた場合、それを削除してしまうと以降の応答がユーザの意図から外れてしまう可能性があります。また、調査の途中で得られた重要な知見や判断材料を失ってしまうと、同じ調査を何度も繰り返すことになりかねません。

そこで採用するのが、会話を要約することでサイズを圧縮するという手法です。これは単なるデータの可逆圧縮ではなく、要約によって要点だけを残すという手法です。例えば、途中でツールを呼び出した際のリクエストやレスポンスの詳細、知見を得る前の生ログデータなどは、後続の処理にはあまり意味を持ちません。このような情報を省略し、重要度の高い情報のみを保持することで、情報の欠落を最小限に抑えながらサイズを削減できます。

もちろん、要約によって一定の情報欠落は避けられません。そのため、可能であれば要約を行わないに越したことはありません。しかし現実的には避けられない問題であるため、適切な要約戦略を立てることが重要になります。

要約の具体的方法

大容量のデータを要約する方法はいくつかあります。今回はシンプルに、要約自体をLLMにクエリするという方法を採用します。これは、もともとトークン限界に収まっていたデータのサイズを圧縮したいだけであり、必要な履歴データをそのまま投入できるためです。

より高度な要約方法としては、例えば履歴を複数の断片に分割して個別に要約を作成し、それらを統合するといった手法もあります。しかし今回はシンプルな実装を優先し、このような高度な手法は扱いません。興味のある方は Map-Reduce 型の要約手法などを調べてみてください。

なお、要約自体もLLMへのクエリとなるため、要約対象のデータがトークン制限内に収まっている必要があります。要約対象が既にトークン制限を超えている場合は、先述したMap-Reduce型の分割要約などの別の手法が必要になります。今回の実装では、もともとトークン限界内だった履歴が時間経過で膨れ上がったケースを想定しているため、この点は問題になりません。

要約の方針

要約で最も重要な原則は、「何が分かったか(結果・知見)」を残し、「どうやって調べたか(プロセス)」を捨てることです。この観点から、何を保持し何を省略すべきかを考えます。

単に「要約して」とだけ指示すると、漠然としたまとめになったり、重要な情報が欠如したりする可能性があります。そのため、どういう情報を残すべきかをプロンプトエンジニアリングで明確に指示する必要があります。この方針はタスクによって大きく変わるため、エージェントごとに調整すべき部分です。

セキュリティ分析を主な目的とする場合、以下のような情報を優先的に保持すべきです。まず最も重要なのは、ユーザからの指示や質問、そして何をゴールにしているかという情報です。これらが失われると、エージェントの行動指針そのものが失われてしまいます。次に、分析の過程で得られた重要な知見や攻撃に関する情報も必須です。IOC(Indicator of Compromise: 侵害指標)は特に重要で、IPアドレス、ドメイン名、ファイルハッシュ、URLなどは確実に保持する必要があります。さらに、重要度や影響度の判定に利用できそうな証跡や痕跡、調査作業の進行状況や文脈、次以降の調査に必要そうな手がかり、ネクストステップなども保持すべき情報です。これらがあることで、調査を中断せずに継続できます。

逆に、省略してもよい情報もあります。ツール呼び出しの詳細(関数名やパラメータなど)や、その結果の生データ全体、ツール利用が失敗した履歴などは、最終的な結果には影響しないため削除できます。調査の過程で得られた冗長な情報や、調査の過程そのもの(手続きの詳細)も同様です。

こういった指示をプロンプトで与えることで、目的に応じた適切な要約を生成できます。

要約の範囲

会話履歴を要約する際、現状の履歴全体を要約するという選択肢もありますが、必ずしもそうする必要はありません。より効果的な戦略として、新しい会話は残し、古い部分のみを要約する方法があります。

例えば、要約対象は古い方の70%程度にしておき、新しい30%はそのまま残すという方針です。これにより、直近の文脈が要約によって失われることを防げます。特に、最近のやり取りには現在進行中の調査に関する重要な情報が含まれていることが多いため、これをそのまま保持することで調査の継続性を保ちやすくなります。

このあたりの勘所はタスクによって変わります。対話が短期間で完結するタスクであれば直近の履歴の重要性が高いですし、長期にわたる調査であれば初期の情報も重要になります。運用しながら適切な比率を見つけていくとよいでしょう。

いつ要約するか

要約を実行するタイミングには、大きく分けて2つの方式があります。

予防的圧縮(事前監視方式)

一つ目は、トークンサイズを監視しながら一定値を超えたら圧縮するという方式です。例えば、トークン限界の80%に達したら要約を実行するといった具合です。この方式の利点は、トークン超過エラーが発生する前に対処できることです。

ただし、正確なトークン数を取得するにはAPIを呼び出す必要があります。Geminiの場合、トークン数の取得自体は無料ですが、毎回チェックしていると応答時間が劣化します。また、レート制限にも注意が必要です。また、モデルごとにトークン限界が異なるため、複数のモデルを使い分ける場合はそれぞれの制限を管理する必要があり、実装が複雑になります。バイト数による近似であれば、モデルに依存しない単純な実装で済みます。

リアクティブ圧縮(エラー駆動方式)

もう一つの方式は、トークン超過エラーが起きたら発火するという方式です。今回の実装ではこちらを採用しています。実装がシンプルになり、不要な監視コストがかからないためです。

なお、今回の実装では文脈を可能な限り維持するためにエラーが発生するまで圧縮を行わない方式を採用していますが、コスト削減を優先する場合はもっと早い段階で要約を実行してもよいでしょう。GenerateContentのレスポンスには現在のトークン使用量(UsageMetadata)が含まれているため、この情報を利用して一定の割合(例えば限界の60-70%)に達したら予防的に圧縮する実装も可能です。文脈の維持とコストのバランスを考慮して、自分のユースケースに合った戦略を選択してください。

両方式の比較は以下の通りです。

方式 メリット デメリット
予防的圧縮 エラー発生前に対処できる トークン数取得のコスト、閾値設定の難しさ
リアクティブ圧縮 実装がシンプル、不要な監視コストなし 一度エラーが発生する、リトライが必要

Geminiの場合、トークン超過エラーは以下のようなメッセージで返されます。

% cd examples/too-large-request
% env GEMINI_PROJECT=your-project go run .
APIError - Code:400 Status:INVALID_ARGUMENT Message:The input token count (2500030) exceeds the maximum number of tokens allowed (1048576). Details:[]

このエラーメッセージを検知して要約処理を実行します。examples/too-large-request/main.go を実行すると、実際にこのエラーを確認できます。

ただし注意点として、API側の挙動が変わる(エラーメッセージが変更される)場合などに影響を受ける可能性があります。AI関連のツールやサービスはこういった仕様の更新が頻繁に行われるため、実装が不安定になりやすい面があります。重要なシステムでは、エラーパターンのテストを定期的に実行することをお勧めします。例えば、isTokenLimitError 関数が実際のエラーメッセージを正しく検知できるかのユニットテストを用意し、API仕様変更の影響を早期に検出できるようにしておくとよいでしょう。

実装例

それでは、実際のコードを見ていきましょう。全体の流れとしては、トークン超過エラーを検知したら要約を実行し、圧縮された履歴でリトライするという仕組みです。

エラーハンドリングとリトライロジック

まず、セッション内でのエラーハンドリングとリトライのロジックです。pkg/usecase/chat/session.goSend メソッド内で、以下のような処理を行っています。

pkg/usecase/chat/session.go
resp, err := s.gemini.GenerateContent(ctx, s.history.Contents, config)
if err != nil {
    // Check if error is due to token limit exceeded
    if isTokenLimitError(err) {
        // Attempt compression
        fmt.Println("\n📦 Token limit exceeded. Compressing conversation history...")

        compressedContents, compressErr := compressHistory(ctx, s.gemini, s.history.Contents)
        if compressErr != nil {
            return nil, goerr.Wrap(compressErr, "failed to compress history")
        }

        // Update history with compressed contents
        s.history.Contents = compressedContents

        // Save compressed history immediately
        if saveErr := saveHistory(ctx, s.repo, s.storage, s.alertID, s.history); saveErr != nil {
            fmt.Printf("⚠️  Warning: failed to save compressed history: %v\n", saveErr)
        }

        fmt.Println("✅ Conversation history compressed successfully. Retrying...")
        continue // Retry with compressed history
    }
    return nil, goerr.Wrap(err, "failed to generate content")
}

このコードでは、先述した通りシンプルにエラーが起きたときに要約を試みています。isTokenLimitError 関数でトークン超過エラーかどうかを判定し、該当する場合は compressHistory 関数で履歴を圧縮します。圧縮された履歴で s.history.Contents を差し替え、continue でループの先頭に戻ることでリトライします。トークン超過エラーでなければ、そのままエラーを返します。

重要な点として、要約後は即座に saveHistory で履歴を永続化しています。これにより、次回以降のセッションでも圧縮後の状態から再開できます。圧縮処理はコストがかかるため、一度圧縮したものは保存しておくことで、無駄な再圧縮を避けられます。

履歴の範囲計算

次に、compressHistory 関数の中身を見ていきましょう。まず最初に行うのが、要約対象となる履歴部分の抽出です。pkg/usecase/chat/compress.go では以下のようなコードで範囲を計算しています。

pkg/usecase/chat/compress.go
	// Calculate byte size for each content
	totalBytes := 0
	byteSizes := make([]int, len(contents))
	for i, content := range contents {
		size := contentSize(content)
		byteSizes[i] = size
		totalBytes += size
	}

	// Calculate compression threshold (70% of total bytes)
	compressThreshold := int(float64(totalBytes) * compressionRatio)

	// Find the index where we cross the 70% threshold
	cumulativeBytes := 0
	compressIndex := 0
	for i, size := range byteSizes {
		cumulativeBytes += size
		if cumulativeBytes >= compressThreshold {
			compressIndex = i + 1 // Include this message in compression
			break
		}
	}

	// If compression index is 0 or at the end, nothing to compress
	if compressIndex == 0 || compressIndex >= len(contents) {
		return nil, goerr.New("insufficient content to compress")
	}

ここでは対象をバイト数で計算しています。contentSize 関数は、各 Content を JSON にマーシャルしてバイト数を取得する関数です。

pkg/usecase/chat/compress.go
// contentSize calculates the byte size of a content by JSON marshaling
func contentSize(content *genai.Content) int {
	data, err := json.Marshal(content)
	if err != nil {
		return 0
	}
	return len(data)
}

この関数は、テキストだけでなくツール呼び出しの情報なども含めた正確なサイズを測定できます。

本来であればトークン数で計算すべきですが、都度APIでトークン数を問い合わせていると遅延が発生しますし、レート制限に引っかかる恐れもあります。そのため、多少粗くてもローカルで計算できるバイト数を使っています。

バイト数とトークン数には相関があるため、完全に正確でなくても圧縮の目的は達成できます。コード内の compressionRatio は 0.7 に設定されており、全体の70%のバイト数に達するまでのメッセージを要約対象としています。これにより、新しい30%の履歴はそのまま保持されます。

要約の実行と履歴の再構築

範囲計算が終わったら、実際に要約を実行して新しい履歴を構築します。

pkg/usecase/chat/compress.go
	// Extract contents to compress and to keep
	toCompress := contents[:compressIndex]
	toKeep := contents[compressIndex:]

	// Generate summary of compressed contents
	summary, err := summarizeContents(ctx, gemini, toCompress)
	if err != nil {
		return nil, goerr.Wrap(err, "failed to summarize contents")
	}

	// Create summary content as user message
	summaryContent := &genai.Content{
		Role: genai.RoleUser,
		Parts: []*genai.Part{
			{Text: "=== Previous Conversation Summary ===\n\n" + summary},
		},
	}

	// Return new history: summary + kept contents
	newContents := append([]*genai.Content{summaryContent}, toKeep...)
	return newContents, nil
}

まず、計算した compressIndex を使って履歴を2つに分割します。toCompress が要約対象の古い履歴、toKeep がそのまま保持する新しい履歴です。

次に、summarizeContents 関数を呼び出して要約を生成します。この関数は、要約対象の履歴に要約用のプロンプト(後述)をユーザメッセージとして追加し、LLMに投げて要約テキストを取得します。

要約結果は新しいユーザメッセージとして履歴の先頭に挿入されます。Rolegenai.RoleUser に設定し、テキストの冒頭に === Previous Conversation Summary === というマーカーを付けることで、LLMが「これは過去の会話の要約である」と認識できるようにしています。最後に、要約メッセージと保持する履歴を結合して返します。これにより、元の履歴よりも大幅に小さい新しい履歴が構築され、トークン限界を回避できます。

要約用のプロンプト

最後に、要約を生成する際に使用するプロンプトです。pkg/usecase/chat/prompt/summarize.md に以下のような内容が格納されています。

pkg/usecase/chat/prompt/summarize.md
You are an assistant for security alert analysis.

## Context and Purpose

The conversation history has exceeded the token limit. Create a summary that will replace the older parts of the conversation, preserving all critical information needed to continue the security investigation.

This summary will be inserted at the beginning of the conversation history. Focus on what matters for ongoing analysis, not the investigation process itself.

## What to Preserve (Highest Priority)

**1. User's Intent and Goals (MOST CRITICAL)**
- User's questions and what they want to know
- Investigation goals and what conclusion the user seeks
- Explicit instructions or constraints the user has given
- User's concerns or areas of focus

**2. Attack and Security Intelligence**
- Key findings about the incident (malicious/benign/false positive)
- Attack patterns, techniques, TTPs identified
- IOCs: IP addresses, domains, file hashes, URLs, email addresses, usernames
- Evidence supporting severity/impact assessment
- Timeline of the attack or suspicious activities

**3. Investigation Progress and Context**
- Current state of the investigation
- Important insights or discoveries from the analysis
- Clues or leads for next steps
- What has been verified vs. what remains uncertain

**4. Next Steps and Actions**
- Recommended next steps in the investigation
- Decisions requiring user input
- Outstanding questions that need answers

## What to Deprioritize or Omit (Lowest Priority)

**Do NOT include:**
- Tool call details (function names, parameters, how they were invoked)
- Full tool output or raw data dumps
- Failed tool calls or error messages
- Exploratory queries that yielded no useful information
- The investigation process itself (step-by-step procedures)
- Redundant or repeated information
- Assistant's internal reasoning or thought process

**Remember:** Summarize RESULTS and FINDINGS, not the PROCESS of obtaining them.

## Output Format

Format the summary in markdown:

- **User's Goals**: What the user wants to achieve or understand
- **Investigation Status**: Current understanding of the incident
- **Key Findings**: Critical security conclusions and determinations
- **Attack Intelligence**: IOCs, TTPs, timeline, attack patterns
- **Evidence**: Important facts supporting severity/impact assessment
- **Next Steps**: What to investigate next or decisions needed
- **Open Questions**: Unresolved issues requiring attention

Be extremely concise. One sentence per point is ideal. Preserve facts, not explanations.

このプロンプトは、先ほど説明した要約の方針を具体的に指示したものです。あくまで一例であり、実際の運用に合わせて調整できます。

プロンプトでは「何を残すか」「何を捨てるか」を明確に指示することが重要です。特にセキュリティ分析では、IOCや攻撃パターンなどの重要情報を確実に保持する必要があります。一方で、ツール実行の詳細や調査プロセスは省略可能です。「RESULTS and FINDINGS, not the PROCESS」という原則を明記することで、LLMに適切な要約を生成させられます。

また、出力フォーマットを指定することで、要約結果の構造を一定に保つことができます。構造化された要約は、後続の処理でも扱いやすくなります。各項目を1文で簡潔に記述するよう指示しているのもポイントです。冗長な説明を避け、事実のみを保持することで、要約後のトークン数を最小限に抑えられます。

このあたりのプロンプトは、実際に動かしてみながら調整していく必要があります。最初から完璧なプロンプトを作ることは難しいため、実運用での要約結果を観察しながら改善していくとよいでしょう。

まとめ

本記事では、LLMエージェントにおける会話コンテキストの圧縮戦略について解説しました。

会話履歴の圧縮で最も重要なのは、「何を残し、何を捨てるか」を明確に定義することです。LLMによる要約で重要な情報のみを抽出し、エージェントの文脈理解を保ちながらトークン数を削減できます。本記事で紹介した要約プロンプトは、セキュリティ分析に特化したものですが、自分のユースケースに合わせてカスタマイズしてください。

要約による情報欠落は避けられないため、可能な限り要約を発生させない設計(ツール実行結果のフィルタリングなど)も併せて検討するとよいでしょう。トークン管理はエージェント設計全体で考えるべき課題であり、今後の記事でもこの観点を意識しながら機能拡張を進めていきます。

Discussion