2025年9月5日

データサイエンスエージェントでArtifactを使えるようにしてみた!(ローカル開発編)


Content

こんにちは!みっちーです!


ADK(Agent Development Kit)で構築したデータサイエンスエージェントは、SQLでのデータ抽出からPythonでの高度な分析・可視化まで、一連のワークフローを自動で実行できる非常に強力なツールです。しかし、その分析プロセスがブラックボックスでは、結果の信頼性や再現性を確保できません。


そこで今回は、このデータサイエンスエージェントの再現性を担保するというテーマに挑みます。エージェントが内部で生成したSQLクエリと、その結果から作成したグラフ画像の両方を、Artifact(アーティファクト)として自動的に保存する機能を実装します。


この記事では、「このグラフは、このSQLで抽出したデータに基づいています」と、誰でも検証可能な、透明性の高い分析プロセスをエージェントに組み込む本格的な手法をご紹介します!

前提条件

ADKのArtifactとは?

今回の実装の鍵となる「Artifact(アーティファクト)」について、少しだけ触れておきましょう。ADKにおけるArtifactとは、エージェントがセッション中またはセッションを越えてファイルを扱うための仕組みです。テキストだけでなく、画像、音声、PDFといったバイナリデータを効率的に管理できます。

なぜこれが重要かというと、エージェントが生成したグラフ画像のような「成果物」を永続化し、後からいつでも参照・検証できるようにするためです。ローカル開発ではメモリ上に保存されますが、本番環境ではGoogle Cloud Storageなどと連携させることで、堅牢なデータ管理が実現します。詳しくは、公式ドキュメントも参考にしてください。

実装手順と動作確認

ゴールとアーキテクチャ

今回のゴールは、データサイエンスエージェントがデータ分析と可視化を実行する際に、裏側で生成されたSQLクエリ(.sqlファイル)とグラフ画像(.pngファイル)を、自動的にアーティファクトとして保存することです。

このエージェントの挙動を理解する鍵は、その階層構造とToolContextの状態(state)管理にあります。

  • 司令塔 (root_agent): ユーザーからの指示を解釈し、専門家であるサブエージェントをツール経由で呼び出します。
  • SQL担当 (db_agent): 自然言語からSQLを生成し、実行します。重要なのは、このエージェントが内部で生成したSQLクエリをToolContextのstateに保存していることです。
  • Python担当 (ds_agent): 受け取ったデータをVertex AI Code Executorを使って分析・可視化し、結果のテキストとグラフ画像を返します。

ここで登場する Vertex AI Code Executor とは、Google Cloudが提供する安全なサンドボックス環境でPythonコードを実行するためのツールです。これにより、エージェントは単なるテキスト応答だけでなく、データ分析、グラフ作成、機械学習モデルのトレーニングといった高度なタスクを動的に実行できます。今回のケースでは、ds_agent がこのCode Executorを利用して、SQLの実行結果からグラフ画像を生成しています。

この仕組みを利用し、それぞれのサブエージェントを呼び出すツール(call_db_agent, call_ds_agent)の中で、適切な情報源からアーティファクトを保存するのが、最も確実な実装戦略です。

コードの修正

それでは、python/agents/data-science/data_science/tools.pyを修正していきましょう。

1. call_db_agentの修正: stateからSQLクエリを保存する

まず、call_db_agent関数を修正します。db_agentを呼び出した後、tool_context.state[“sql_query”]にアクセスして、db_agent自身が保存したSQLクエリを直接取得します。取得したクエリをgenerated_query.sqlとしてアーティファクトに保存します。

async def call_db_agent(
question: str,
tool_context: ToolContext,
):
"""
Tool to call the database (nl2sql) agent.

This tool invokes the db_agent to generate and execute a SQL query.
After execution, it retrieves the generated SQL from the tool_context's state
and saves it as an artifact for reproducibility.
"""
agent_tool = AgentTool(agent=db_agent)

db_agent_output = await agent_tool.run_async(
args={"request": question}, tool_context=tool_context
)
tool_context.state["db_agent_output"] = db_agent_output

# --- stateからSQLクエリを取得して保存するロジック ---
if "sql_query" in tool_context.state:
extracted_sql = tool_context.state["sql_query"]
file_name = "generated_query.sql"

part = types.Part.from_bytes(
data=extracted_sql.encode("utf-8"), mime_type="text/x-sql"
)

await tool_context.save_artifact(filename=file_name, artifact=part)
print(f"Successfully saved SQL query from state to artifact: {file_name}")

return db_agent_output

2. call_ds_agentの修正: 出力からグラフ画像を保存する

次に、call_ds_agent関数です。ds_agentCodeExecutorと連携しているため、その出力(ds_agent_output)はテキストとファイルの複合オブジェクトになっています。生成されたグラフ画像はds_agent_output.filesリストに含まれているので、これをループしてMIMEタイプがimage/で始まるファイルを見つけ、analysis_graph.pngとして保存します。

async def call_ds_agent(
question: str,
tool_context: ToolContext,
):
"""
Tool to call the data science (nl2py) agent.

After the ds_agent runs and potentially generates a graph, this tool
extracts the graph file from the agent's output and saves it as an artifact.
"""

if question == "N/A":
return tool_context.state["db_agent_output"]

# The result of the SQL query is expected to be in the state.
input_data = tool_context.state.get("query_result", "")

question_with_data = f"""
Question to answer: {question}

The data to analyze, which resulted from a previous SQL query, is as follows:
{input_data}

"""

agent_tool = AgentTool(agent=ds_agent)

ds_agent_output = await agent_tool.run_async(
args={"request": question_with_data}, tool_context=tool_context
)
tool_context.state["ds_agent_output"] = ds_agent_output

# --- Artifact Saving Logic for Graph ---
# The ds_agent with a CodeExecutor returns an AgentToolOutput object.
# The `files` attribute contains any files generated by the code execution.
if hasattr(ds_agent_output, 'files') and ds_agent_output.files:
for file in ds_agent_output.files:
# Assuming the first image file is the graph we want to save.
if file.mime_type.startswith("image/"):
graph_file_name = "analysis_graph.png"
await tool_context.save_artifact(filename=graph_file_name, artifact=file)
print(f"Successfully saved graph to artifact: {graph_file_name}")
# Stop after saving the first graph.
break

return ds_agent_output

動作確認

実装が正しく機能するか、確認しましょう。ターミナルでpoetry run adk webを実行し、エージェントにSQLでの集計とPythonでの可視化の両方を必要とする、以下のプロンプトを試してみてください。

trainテーブルの日ごとの売上合計を取得し、その結果を可視化してください

エージェントが処理を終えた後、ADKのUI上のArtifactタブ上に以下のように表示されれば大成功です!


Downloadボタンを押せば、ダウンロードすることも可能です!ぜひ、ダウンロードして生成されたSQLやグラフを確認してみてください!

まとめと次のステップ

今回は、データサイエンスエージェントのアーキテクチャの核心であるToolContextの状態(state)と、CodeExecutorの出力構造を理解し、それぞれに最適な方法でSQLクエリとグラフ画像をアーティファクトとして保存しました。「何が実行されたのかわからない」というAIエージェントの課題を克服し、分析プロセスの透明性と再現性を飛躍的に向上させることができました!

ローカル環境では、これらのアーティファクトはメモリ上に保存されます。次回の記事では、この仕組みをGoogle Cloud Storage(GCS)に接続し、本番環境でもファイルが永続化されるようにする「GCS活用編」をお届けします。お楽しみに!

2025年9月5日 データサイエンスエージェントでArtifactを使えるようにしてみた!(ローカル開発編)

Category Google Cloud

ご意見・ご相談・料金のお見積もりなど、
お気軽にお問い合わせください。

お問い合わせはこちら