ADKでoutput_schemaを使ってAIエージェントの出力形式を指定する

初めに

AIエージェントの出力形式をプロンプトで指定しようとしていたのですが、ADKにはoutput_schemaという機能が存在するようです。実際にどのように使えるのか試してみました。

output_schema とは

Agentを作成する際、output_schemaを定義することで、出力のフォーマットを定義できるようです。
フォーマットといっても、現状はJSON形式のみ対応で、指定できるのはキーや値のみのようです（十分ありがたすぎる）。

試してみる

というわけで試してみましょう。

ただoutput_schemaに指定をするだけではなく、instructionで明示的に指定してあげる必要があります。

• instructionに出力形式を明示的に指定する
• 出力形式のクラスを作成する
• LlmAgentでoutput_schemaを指定する

# ポイント１：instructionに出力形式を明示的に指定する
instruction = """
    あなたはPHPフレームワークFlowに関する質問に答えるエージェントです。
    必要に応じてflow_rag_documentationを見て、質問に答えてください。
    
    [出力フォーマット]
    JSON オブジェクト
    
    [出力スキーマ]
    can_answer: 回答可能かどうか 
    answer: 回答内容
    例: {"can_answer": "回答可能", "answer": "FlowはPHPのフレームワークです。"}
"""

# ポイント２：出力形式のクラスを作成する
class TestOutputFormat(BaseModel):
    """
    Model to test output format of the agent.
    """
    can_answer: str = Field(description="回答できるかどうか")
    answer: str = Field(description="回答内容")

# ポイント３：LlmAgentでoutput_schemaを指定する
root_agent = Agent(
    name="flow_rag_agent",
    model=os.environ.get("MODEL"),
    instruction=instruction,
    tools=[ask_vertex_retrieval],
    output_schema=TestOutputFormat,
)

実際の実装

import os

from google.adk.agents.llm_agent import Agent
from google.adk.tools.retrieval.vertex_ai_rag_retrieval import VertexAiRagRetrieval
from vertexai.preview import rag
from pydantic import BaseModel, Field

from dotenv import load_dotenv

load_dotenv()

ask_vertex_retrieval = VertexAiRagRetrieval(
    name="flow_rag_documentation",
    description=(
        'This is documentation od PHP web Framework Flow'
    ),
    rag_resources=[
        rag.RagResource(
            rag_corpus= os.environ.get("RAG_CORPUS"),
        )
    ],
    similarity_top_k=1,
    vector_distance_threshold=0.6,
)

instruction = """
    あなたはPHPフレームワークFlowに関する質問に答えるエージェントです。
    必要に応じてflow_rag_documentationを見て、質問に答えてください。
    
    [出力フォーマット]
    JSON オブジェクト
    
    [出力スキーマ]
    can_answer: 回答可能かどうか 
    answer: 回答内容
    例: {"can_answer": "回答可能", "answer": "FlowはPHPのフレームワークです。"}
"""

class TestOutputFormat(BaseModel):
    """
    Model to test output format of the agent.
    """
    can_answer: str = Field(description="回答できるかどうか")
    answer: str = Field(description="回答内容")

root_agent = Agent(
    name="flow_rag_agent",
    model=os.environ.get("MODEL"),
    instruction=instruction,
    tools=[ask_vertex_retrieval],
    output_schema=TestOutputFormat,
)

adk webコマンドで立ち上げ、動作確認してみたところ、指定した形式で返ってくることが確認できました。
（Unicodeになってますが今回の趣旨とは別なので一旦無視します）

回答の上二つの'set_model_response`がありますね。
WebUI上でEventを確認してみると、一つ目はRAGへの問い合わせのようでした。
そして、二つ目がレスポンスの整形処理のようです

終わりに

今回はAIエージェントの出力形式を指定する方法を調べました。プロンプトで指定するよりも精度が高いものなのでしょうか。もしそうであれば、呼び出し側での分岐なども比較的考えやすくなりそうです。この辺はもう少し調べてみようと思います。
ここまでご覧いただきありがとうございました！

参考