はじめに

こんにちは！今回は、PythonのRichライブラリを使って、AIモデルのファンクションコーリング結果を美しく表示する方法をご紹介します。初心者の方でも理解しやすいように、step by stepで解説していきますので、ぜひ最後までお付き合いください。

ファンクションコーリングの結果を美しく表示する関数を作成中。。。これはめっちゃいいぞ！！！！ pic.twitter.com/Y3MPIeteQp
— Maki@Sunwood AI Labs. (@hAru_mAki_ch) July 29, 2024

必要なライブラリのインストール

まずは、必要なライブラリをインストールしましょう。Google Colabを使用している場合は、以下のセルを実行してRichをインストールしてください。

!pip install rich

必要なモジュールのインポート

次に、必要なモジュールをインポートします。

# Richライブラリから必要なコンポーネントをインポート
from rich.console import Console
from rich.panel import Panel
from rich.tree import Tree
from rich.syntax import Syntax
from rich.table import Table

# 型ヒントのためのDictとAnyをインポート
from typing import Dict, Any

# Consoleオブジェクトを作成（これを使って出力します）
console = Console()

ファンクションコーリング結果表示関数の定義

ここからが本題です。ファンクションコーリングの結果を美しく表示する関数を定義していきます。

def display_rich_function_call_response(response: Dict[str, Any]) -> None:
    # モデル情報を表示するためのツリーを作成
    model_info = Tree("Function Call Response")
    model_info.add(f"ID: [cyan]{response['id']}[/cyan]")
    model_info.add(f"Created: [cyan]{response['created']}[/cyan]")
    model_info.add(f"Model: [cyan]{response['model']}[/cyan]")
    model_info.add(f"Object: [cyan]{response['object']}[/cyan]")
    model_info.add(f"System Fingerprint: [cyan]{response['system_fingerprint']}[/cyan]")

    # 使用量情報を表示するためのテーブルを作成
    usage_table = Table(title="Usage", show_header=True, header_style="bold magenta")
    usage_table.add_column("Metric", style="dim")
    usage_table.add_column("Value", justify="right")
    usage = response['usage']
    usage_table.add_row("Prompt Tokens", str(usage['prompt_tokens']))
    usage_table.add_row("Completion Tokens", str(usage['completion_tokens']))
    usage_table.add_row("Total Tokens", str(usage['total_tokens']))

    # 選択肢（Choices）情報を表示するためのツリーを作成
    choices_tree = Tree("Choices")
    for idx, choice in enumerate(response['choices']):
        choice_tree = choices_tree.add(f"Choice {idx}")
        choice_tree.add(f"Finish Reason: [green]{choice['finish_reason']}[/green]")
        choice_tree.add(f"Index: [green]{choice['index']}[/green]")

        message = choice['message']
        message_tree = choice_tree.add("Message")
        message_tree.add(f"Role: [blue]{message['role']}[/blue]")
        content_syntax = Syntax(message['content'], "markdown", theme="monokai", line_numbers=True)
        message_tree.add(Panel(content_syntax, title="Content", border_style="blue"))

        if message['tool_calls']:
            tool_calls_tree = choice_tree.add("Tool Calls")
            for tool_call in message['tool_calls']:
                tool_call_tree = tool_calls_tree.add(f"Tool Call {tool_call['index']}")
                tool_call_tree.add(f"ID: [yellow]{tool_call['id']}[/yellow]")
                tool_call_tree.add(f"Type: [yellow]{tool_call['type']}[/yellow]")
                function_tree = tool_call_tree.add("Function")
                function_tree.add(f"Name: [red]{tool_call['function']['name']}[/red]")
                args_syntax = Syntax(tool_call['function']['arguments'], "json", theme="monokai")
                function_tree.add(Panel(args_syntax, title="Arguments", border_style="red"))

    # 作成したツリーとテーブルを表示
    console.print(Panel(model_info, title="Function Call Response", border_style="bold"))
    console.print(usage_table)
    console.print(Panel(choices_tree, title="Choices", border_style="bold"))

この関数は少し長いので、部分ごとに詳しく解説していきます。

モデル情報の表示

まず、モデルの基本情報を表示するパートを見てみましょう。

# 使用例
response = {
    "id": "chatcmpl-123",
    "created": 1677858242,
    "model": "gpt-3.5-turbo-0301",
    "object": "chat.completion",
    "system_fingerprint": "fp_44709d6fcb",
    "usage": {
        "prompt_tokens": 13,
        "completion_tokens": 7,
        "total_tokens": 20
    },
    "choices": [
        {
            "finish_reason": "stop",
            "index": 0,
            "message": {
                "content": "Hello! How can I assist you today?",
                "role": "assistant",
                "tool_calls": [
                    {
                        "index": 0,
                        "id": "call_abc123",
                        "type": "function",
                        "function": {
                            "name": "get_current_weather",
                            "arguments": '{"location": "New York", "unit": "celsius"}'
                        }
                    }
                ]
            }
        }
    ]
}

# モデル情報を表示するためのツリーを作成
model_info = Tree("Function Call Response")
# IDを追加（シアン色で表示）
model_info.add(f"ID: [cyan]{response['id']}[/cyan]")
# 作成日時を追加（シアン色で表示）
model_info.add(f"Created: [cyan]{response['created']}[/cyan]")
# モデル名を追加（シアン色で表示）
model_info.add(f"Model: [cyan]{response['model']}[/cyan]")
# オブジェクトタイプを追加（シアン色で表示）
model_info.add(f"Object: [cyan]{response['object']}[/cyan]")
# システムフィンガープリントを追加（シアン色で表示）
model_info.add(f"System Fingerprint: [cyan]{response['system_fingerprint']}[/cyan]")

この部分では、Treeオブジェクトを使って階層構造を作成し、モデルの基本情報を追加しています。[cyan]のような記述は、Richの色付け機能を使用しています。

使用量情報の表示

次に、トークン使用量を表示するパートです。

# 使用量情報を表示するためのテーブルを作成
usage_table = Table(title="Usage", show_header=True, header_style="bold magenta")
# 'Metric'列を追加（薄い色で表示）
usage_table.add_column("Metric", style="dim")
# 'Value'列を追加（右寄せで表示）
usage_table.add_column("Value", justify="right")
usage = response['usage']
# プロンプトトークン数を行として追加
usage_table.add_row("Prompt Tokens", str(usage['prompt_tokens']))
# 完了トークン数を行として追加
usage_table.add_row("Completion Tokens", str(usage['completion_tokens']))
# 合計トークン数を行として追加
usage_table.add_row("Total Tokens", str(usage['total_tokens']))

ここでは、Tableオブジェクトを使ってトークン使用量を表形式で表示しています。列や行を追加して、見やすい形式にしています。

選択肢（Choices）情報の表示

最後に、モデルの応答内容を表示するパートです。

# 選択肢（Choices）情報を表示するためのツリーを作成
choices_tree = Tree("Choices")
for idx, choice in enumerate(response['choices']):
    # 各選択肢のサブツリーを作成
    choice_tree = choices_tree.add(f"Choice {idx}")
    # 終了理由を追加（緑色で表示）
    choice_tree.add(f"Finish Reason: [green]{choice['finish_reason']}[/green]")
    # インデックスを追加（緑色で表示）
    choice_tree.add(f"Index: [green]{choice['index']}[/green]")

    message = choice['message']
    # メッセージのサブツリーを作成
    message_tree = choice_tree.add("Message")
    # ロールを追加（青色で表示）
    message_tree.add(f"Role: [blue]{message['role']}[/blue]")
    # コンテンツをシンタックスハイライト付きで表示
    content_syntax = Syntax(message['content'], "markdown", theme="monokai", line_numbers=True)
    message_tree.add(Panel(content_syntax, title="Content", border_style="blue"))

    # ツールコールがある場合の処理
    if message['tool_calls']:
        tool_calls_tree = choice_tree.add("Tool Calls")
        for tool_call in message['tool_calls']:
            # 各ツールコールのサブツリーを作成
            tool_call_tree = tool_calls_tree.add(f"Tool Call {tool_call['index']}")
            # IDを追加（黄色で表示）
            tool_call_tree.add(f"ID: [yellow]{tool_call['id']}[/yellow]")
            # タイプを追加（黄色で表示）
            tool_call_tree.add(f"Type: [yellow]{tool_call['type']}[/yellow]")
            # 関数のサブツリーを作成
            function_tree = tool_call_tree.add("Function")
            # 関数名を追加（赤色で表示）
            function_tree.add(f"Name: [red]{tool_call['function']['name']}[/red]")
            # 引数をJSONとしてシンタックスハイライト付きで表示
            args_syntax = Syntax(tool_call['function']['arguments'], "json", theme="monokai")
            function_tree.add(Panel(args_syntax, title="Arguments", border_style="red"))

この部分では、モデルの応答内容を階層構造で表示しています。特に、ツールコール（関数呼び出し）がある場合は、その詳細も表示するようになっています。

表示の実行

最後に、作成したツリーとテーブルを表示します。

# モデル情報をパネルとして表示
console.print(Panel(model_info, title="Function Call Response", border_style="bold"))
# 使用量テーブルを表示
console.print(usage_table)
# 選択肢情報をパネルとして表示
console.print(Panel(choices_tree, title="Choices", border_style="bold"))

これらのconsole.print()呼び出しで、実際に画面に情報が表示されます。

関数の使用例

では、この関数を実際に使ってみましょう。以下は架空のレスポンスを使用した例です。

# 関数を呼び出してレスポンスを表示
display_rich_function_call_response(response)

このコードを実行すると、美しく整形されたファンクションコーリングの結果が表示されます。

全体コード

from rich.console import Console
from rich.panel import Panel
from rich.tree import Tree
from rich.syntax import Syntax
from rich.table import Table
from typing import Dict, Any

console = Console()

def display_rich_function_call_response(response: Dict[str, Any]) -> None:
    model_info = Tree("Function Call Response")
    model_info.add(f"ID: [cyan]{response['id']}[/cyan]")
    model_info.add(f"Created: [cyan]{response['created']}[/cyan]")
    model_info.add(f"Model: [cyan]{response['model']}[/cyan]")
    model_info.add(f"Object: [cyan]{response['object']}[/cyan]")
    model_info.add(f"System Fingerprint: [cyan]{response['system_fingerprint']}[/cyan]")

    usage_table = Table(title="Usage", show_header=True, header_style="bold magenta")
    usage_table.add_column("Metric", style="dim")
    usage_table.add_column("Value", justify="right")
    usage = response['usage']
    usage_table.add_row("Prompt Tokens", str(usage['prompt_tokens']))
    usage_table.add_row("Completion Tokens", str(usage['completion_tokens']))
    usage_table.add_row("Total Tokens", str(usage['total_tokens']))

    choices_tree = Tree("Choices")
    for idx, choice in enumerate(response['choices']):
        choice_tree = choices_tree.add(f"Choice {idx}")
        choice_tree.add(f"Finish Reason: [green]{choice['finish_reason']}[/green]")
        choice_tree.add(f"Index: [green]{choice['index']}[/green]")

        message = choice['message']
        message_tree = choice_tree.add("Message")
        message_tree.add(f"Role: [blue]{message['role']}[/blue]")
        content_syntax = Syntax(message['content'], "markdown", theme="monokai", line_numbers=True)
        message_tree.add(Panel(content_syntax, title="Content", border_style="blue"))

        if message['tool_calls']:
            tool_calls_tree = choice_tree.add("Tool Calls")
            for tool_call in message['tool_calls']:
                tool_call_tree = tool_calls_tree.add(f"Tool Call {tool_call['index']}")
                tool_call_tree.add(f"ID: [yellow]{tool_call['id']}[/yellow]")
                tool_call_tree.add(f"Type: [yellow]{tool_call['type']}[/yellow]")
                function_tree = tool_call_tree.add("Function")
                function_tree.add(f"Name: [red]{tool_call['function']['name']}[/red]")
                args_syntax = Syntax(tool_call['function']['arguments'], "json", theme="monokai")
                function_tree.add(Panel(args_syntax, title="Arguments", border_style="red"))

    console.print(Panel(model_info, title="Function Call Response", border_style="bold"))
    console.print(usage_table)
    console.print(Panel(choices_tree, title="Choices", border_style="bold"))

まとめ

以上で、Richライブラリを使ってファンクションコーリングの結果を美しく表示する方法の解説を終わります。この方法を使うことで、複雑なJSON形式のレスポンスを見やすく整理し、色付けされた形で表示することができます。

特に、大量のデータを扱う場合や、複数の関数呼び出しが含まれる複雑なレスポンスを扱う場合に、この表示方法は非常に役立ちます。ぜひ、皆さんのプロジェクトでも活用してみてください！

質問やコメントがありましたら、お気軽にお問い合わせください。Happy coding!

📒ノートブック