コード上の知識アシスタントの構築

When developers join a new project or need to work across an unfamiliar codebase, knowledge assistants like Databricks Knowledge Assistant help them get up to speed by answering natural-language questions about the code. But answer quality depends heavily on how the source code and surrounding context were prepared and added. One key factor is chunking: how you split source files into pieces for indexing and retrieval. Code makes this tricky. If you break a function mid-body or strip its class context, even a capable assistant will struggle to answer questions about it.

We built three Knowledge Assistants over our Casper’s Kitchens demo GitHub repository, each using a different chunking strategy, from a simple fixed-size baseline to a structure-aware approach that parses code into its syntactic components. The repository simulates a ghost kitchen business on Databricks, using a wide range of features including Lakeflow pipelines, DSPy agents, and Databricks Asset Bundles (DABs), with documentation in markdown files and notebook cells. The cross-file dependencies, mixed file formats, and domain-specific patterns make it the kind of project where a capable knowledge assistant would be a huge help.

This post walks through what makes working with code different from working with typical business documents, how we deployed each chunking strategy as a Databricks Knowledge Assistant, and how we used MLflow’s evaluation framework to compare them. You can find all the code here.

How Knowledge Assistants Works (and Why Code Is Different)

Under the hood, knowledge assistants use various forms of retrieval-augmented generation (RAG). They retrieve relevant chunks of the source data, often from a vector search index, and pass them to a large language model as context for generating an answer to a user query.

Databricks Knowledge Assistant builds on this foundation with sophisticated retrieval techniques including Instructed Retriever, which incorporates query decomposition, context-informed re-ranking, and reasoning over document metadata. These capabilities go a long way toward handling the complexity of real-world codebases, and they work best when the underlying chunks preserve meaningful semantic boundaries.

Knowledge assistants are most commonly constructed and evaluated over collections of business documents, which tend to flow linearly, with paragraphs and sections. Code has nested hierarchies: files contain classes, classes contain methods, methods contain logic blocks. The semantic unit in code is often a complete function, not a paragraph.

This creates specific challenges, including:

• Semantic boundaries: Splitting a function mid-body loses the context needed to understand what it does. A chunk containing deletion_order = ['experiments', 'jobs'... is less useful if it doesn’t show that this variable is inside UCState.clear_all().
• Cross-file dependencies: Code references other code. Understanding one function often requires context from its class, its imports, or related functions.
• Mixed file types: Our codebase has .py files, .ipynb notebooks (JSON with code/Markdown cells), .md documentation, and .yaml configuration, each requiring different parsing approaches.

Because Databricks Knowledge Assistant lets you use your own vector index, you can prepare chunks however you want and just point Knowledge Assistant at the result. This allowed us to compare different approaches to preparing our codebase for RAG and pick the best one.

Chunking Strategies

To see how chunking strategies differ in practice, consider what happens when you ask: “In what order does resource cleanup occur?” The answer lives in a utility class that tracks experiments, jobs, and pipelines. Its logic spans initialization, a deletion order list, and cleanup methods. Here’s how each method works and how it affects the context retrieved about the resource cleanup class, UCState.

Naive Baseline: Fixed-Size Character Chunks

The simplest approach is to split the source files at fixed character intervals with overlap, treating code as plain text. This is not what you would choose for a production-ready RAG system today. It ignores syntax and semantic boundaries, so it fails in exactly the ways code queries care about. But it’s also extremely easy to implement, often “good enough” for quick experiments or doc-heavy repos, and common as a first pass, so it’s a useful baseline.

Here's what naive chunking produces for a search of deletion_order in our codebase:

The variable name got cut in two (eletion instead of deletion), and the chunk doesn't include the method name. If someone searches for “UCState deletion order,” this chunk won’t match well. Furthermore, the deletion_order list in the method got cut off.

Language-Aware: LangChain Heuristic Splitters

LangChain’s RecursiveCharacterTextSplitter.from_language() uses language-specific separators (like \nclass and \ndef for Python) to prefer splitting at logical boundaries. It tries to keep functions intact but still enforces strict size limits. Conceptually, this improves on naive chunking by prioritizing splits at likely semantic boundaries (like def and class) instead of arbitrary character counts, so chunks are more likely to contain complete units of logic.

Here’s what this approach produced for the same search:

The chunk starts at a more natural boundary, but it still lacks context showing which file or function it belongs to, and it cuts off right after the start of a for loop.

AST-Based: Tree-Sitter with Metadata Headers

Abstract syntax tree (AST) を利用したチャンキングは、Tree-sitter のようなパーサーを使用して実際のコード構造を理解します。AST は、コードが言語の文法規則に従ってどのように構成されているか、その構文構造を捉えるコードのツリー表現です。文字境界で分割したり、ヒューリスティックなパターンを使用したりする代わりに、AST ベースのチャンキング戦略では、コードを構文ツリーに解析し、関数、クラス、またはステートメントブロックなどの意味論的な境界でチャンクを作成します。また、関数の中間で分割するのではなく、完全なユニットをまとめておく必要がある場合は、サイズ制限を超えることもできます。

AST ベースの分割を処理するために、ASTChunk Python ライブラリを使用しました。このライブラリには チャンク拡張 オプションが含まれており、各チャンクの先頭にファイルパスとクラス/関数階層を示すメタデータヘッダーが追加されます。このコンテキストは埋め込みの一部となり、クエリ用語がチャンク本体に含まれていない場合でも、検索がクエリと関連コードを一致させるのに役立ちます。

このアプローチでクエリ用に生成されたチャンクは次のとおりです。

ヘッダーは、このコードがどこにあるかを正確に示しています。utils/uc_state/state_manager.py → class UCState: → def clear_all(...)。埋め込まれると、このチャンクは「UCState」、「clear_all」、「deletion order」に関するクエリとの意味論的な関連性が強くなります。

この段階で、どのメソッドが Knowledge Assistant でうまく機能するかについて、いくつかの直感がありました。しかし、確実にするためには、体系的な評価を実行する必要がありました。

MLflow を使用した評価設定

MLflow の GenAI 評価フレームワークは、LLM、エージェント、および検索システムを比較するための完全なツールキットを提供します。評価データセット、予測関数、および LLM ジャッジを提供すると、各質問がパイプラインを通過し、結果がスコアリングされます。ここでは、3 つのチャンキング方法を比較するためにこれを使用しました。

評価データセット

広範な概念トピックからコードに関する詳細なクエリまで、さまざまなカテゴリにわたる 46 の質問を作成しました。

コードベースに異なるコンテキストで構造的に類似したコードがある場合（たとえば、重複する関数名を持つ 2 つのアプリ、並列データベーススキーマ、または微妙に異なる設定ファイルなど）のあいまいさ解消質問にデータセットを意図的に重み付けしました。これらは、チャンキングの弱点を最も明確に露呈するクエリです。チャンクにコードの場所に関するメタデータがない場合、検索システムは異なるコンテキストに存在する類似したクラスや関数を区別するのに苦労します。

LLM ジャッジ

品質の異なる側面を捉える 3 つの主要な LLM ジャッジを使用しました。

• RetrievalSufficiency (組み込み): 取得されたチャンクには、質問に答えるのに十分な情報が含まれていますか？これは、生成とは独立して検索品質を測定するため、チャンキング戦略を比較する上での主要な指標です。
• RetrievalGroundedness (組み込み): 回答は取得されたコンテキストに基づいていますか、それともチャンクに存在しない情報を導入していますか？
• answer_correctness (カスタム): このカスタムスコアラーは、各回答を正解、部分正解、または不正解としてランク付けし、厳密な正誤判定よりも少しニュアンスがあります。断片化または不完全なコンテキストの可能性を考えると、詳細が欠けている、または小さな不正確さがある回答に注意を払いたいと思います。

評価の実行

比較を公平にするために、すべての戦略で同じターゲットチャンクサイズ（1,000 文字）、オーバーラップ（200 文字）、および埋め込みモデル（databricks-gte-large-en）を使用しました。実際には、最終的なチャンクサイズは異なります（たとえば、AST ベースのチャンキングは完全な意味論的ユニットを保持するために拡張される場合がありますが、非常に小さいファイルは自然に小さいチャンクを生成します）。

各チャンキング戦略について、チャンクを Delta テーブルに書き込み、管理された埋め込みを持つ Vector Search インデックスを作成し（Databricks Knowledge Assistant で必要とされる databricks-gte-large-en 埋め込みモデルを使用）、そのインデックスを Knowledge Assistant エンドポイントにアタッチしました。完全なセットアップについては、ドキュメントを参照してください。

各チャンキング戦略を、Knowledge Assistant エンドポイントに直接クエリすることで評価しました。MLflow の to_predict_fn() は、サービングエンドポイントを予測関数としてラップします。Knowledge Assistant は完全な MLflow トレース（検索スパンを含む）を生成するため、組み込みジャッジは取得されたチャンクと最終的な応答の両方を検査できます。

LLM ジャッジは、Databricks Model Serving を介して LLM ジャッジを呼び出します。databricks-claude-opus-4-6 を使用しました。

評価の実行が完了すると、MLflow の実験 UI で 3 つの戦略すべての結果を並べて比較できます。

結果と学んだこと

46 の質問すべてを各 Knowledge Assistant に実行し、3 つのジャッジで結果をスコアリングしました。結果は次のとおりです。

3つの戦略すべてで85%以上のRetrieval Sufficiencyを達成しており、これはKnowledge Assistantの検索手法がコードのチャンク化方法に関わらず関連コンテキストを見つけられることを意味します。検索レベルでの差はわずかです。

カスタムの正解率の結果は、より興味深いストーリーを語っています。ASTベースのチャンク化は、NaiveやLanguage-Awareの59%、61%と比較して、70%の時間で完全に正しい回答を生成します。3つの戦略すべてで、ほとんどの場合、少なくとも部分的に正しい回答が生成されます。より良いチャンクは、ナレッジアシスタントが質問により完全に回答するのに役立ちます。

この利点は特定の質問タイプに集中しています。ASTベースのチャンク化は、構造的に類似したコードがモジュール間に存在する場合のあいまいさ解消の質問で優れていました。これは、ファイルパス、クラス、関数名などのメタデータが追加されたことで、必要なコンテキストが提供されたためです。3つの戦略すべてで、値の検索や完全な定義の取得においては同等でした。

MLflowトレースを使用すると、個々の質問を詳しく調べ、どのチャンクが取得され、回答がどこで分岐したかを正確に確認できます。

この調査では、いくつかの疑問が残りました。ASTベースのチャンク化を使用した改善は、主に平均チャンクサイズが大きいことによるものだったのでしょうか？結果は、LLMジャッジを駆動するモデルの選択にどの程度依存していましたか？評価質問では、実際のユーザーが尋ねる可能性のある主要なカテゴリを見逃していませんでしたか？

Lessons Learned

Databricks Knowledge Assistantは、すぐに高い能力を発揮します。3つの戦略すべてでRetrieval Sufficiencyは高く、ほぼすべての質問で少なくとも部分的に正しい回答が得られました。

データ準備は依然として重要です。ASTベースのチャンク化は、この評価でGroundednessとCorrectnessを向上させました。特に、類似コードのあいまいさ解消に関わる質問で効果がありました。検索と回答品質のわずかな改善でも、1日に数十件の質問をする開発者チーム全体で複利効果を生み出します。

カスタムLLMジャッジは、本当に重要なことを測定するのに役立ちます。MLflowのmake_judge() APIを使用すると、ユースケース固有のLLMジャッジを簡単に構築できます。カスタムのanswer_correctnessジャッジは、単純な合格/不合格の正解率ジャッジよりも、よりニュアンスのあるビューを提供できました。

MLflowトレースは、評価ループを簡素化します。個々の質問を調査して、どのチャンクが取得され、回答がどこで間違っていたかを正確に確認できます。トレースは永続化されるため、エンドポイントを再クエリせずに異なるジャッジで再スコアリングできます。

References

• Databricks Agent Bricks: Knowledge Assistant—カスタムベクトル検索インデックスでKnowledge Assistantを作成するためのセットアップガイド。
• MLflow GenAI evaluation framework—mlflow.genai.evaluate()、組み込みLLMジャッジ、およびカスタムスコアラーAPIのドキュメント。
• cAST: Enhancing Code Retrieval-Augmented Generation with Structural Chunking via Abstract Syntax Tree—ASTチャンクアプローチの動機となった論文。複数のコードRAGタスクにわたるベンチマークが含まれています。ASTChunk Pythonライブラリ実装を使用しました。
• LangChain RecursiveCharacterTextSplitter—比較で使用した言語認識テキストスプリッターのAPIリファレンス。

Try It Yourself

このデモは、Casper’s Kitchensリポジトリで実際に試すことができます。独自のコードベースのチャンク化戦略を評価する場合でも、他のRAG改善を探索する場合でも、この評価フレームワークはアプローチを比較するための再現可能な方法を提供します。

1. 質問と期待される回答を含む評価データセットを作成します。
2. チャンク化戦略を実装します（または、開始点として当社の戦略を使用します）。
3. MLflow LLMジャッジをセットアップします。組み込みオプションから始めて、ギャップが見つかったらカスタムジャッジを追加します。
4. 各戦略で新しいインデックスを使用して評価を実行します。

(このブログ記事はAI翻訳ツールを使用して翻訳されています) 原文記事