LlamaIndexのOpenAI Pydantic Program：構造化データ生成方法

## はじめに

OpenAIの新しいAPIを使用して、LlamaIndexを介して構造化データを生成する方法を紹介します。このガイドでは、ユーザーがPydanticオブジェクトを指定するだけで、簡単に構造化データを抽出できることを示します。

主に以下の設定を紹介します：
- `Album`オブジェクトへの抽出（複数の`Song`オブジェクトを含むことができます）


## 環境設定

まず、必要なライブラリをインストールしましょう。Google Colabで実行している場合は、以下のコマンドを実行してください。
<!-- #endregion -->

```python
# LlamaIndexとOpenAI関連のライブラリをインストール
%pip install llama-index-llms-openai
%pip install llama-index-program-openai
%pip install llama-index
```

```python
from google.colab import userdata
import os
os.environ["OPENAI_API_KEY"] = userdata.get('OPENAI_API_KEY')
```

## 1. `Album`オブジェクトへの抽出

### 1.1 基本的な実装（docstringなし）

まず、出力スキーマを定義します。

```python
# 必要なライブラリをインポート
from pydantic import BaseModel
from typing import List

# Songクラスを定義
class Song(BaseModel):
    title: str
    length_seconds: int

# Albumクラスを定義
class Album(BaseModel):
    name: str
    artist: str
    songs: List[Song]
```

次に、OpenAI Pydantic Programを定義します。

```python
# OpenAIPydanticProgramをインポート
from llama_index.program.openai import OpenAIPydanticProgram

# プロンプトテンプレートを定義
prompt_template_str = """\
Generate an example album, with an artist and a list of songs. \
Using the movie {movie_name} as inspiration.\
"""

# OpenAIPydanticProgramを初期化
program = OpenAIPydanticProgram.from_defaults(
    output_cls=Album,
    prompt_template_str=prompt_template_str,
    verbose=True
)
```

プログラムを実行して構造化出力を取得します。

```python
# プログラムを実行
output = program(
    movie_name="The Shining", description="Data model for an album."
)

# 結果を表示
print(output)
```

### 1.2 docstringを含む実装

モデルにdocstringを追加することで、より詳細な情報を提供できます。

```python
# docstringを含むSongクラスを定義
class Song(BaseModel):
    """Data model for a song."""
    title: str
    length_seconds: int

# docstringを含むAlbumクラスを定義
class Album(BaseModel):
    """Data model for an album."""
    name: str
    artist: str
    songs: List[Song]
```

OpenAI Pydantic Programの定義と実行は前述と同様です。

```python
# プロンプトテンプレートを定義
prompt_template_str = """\
Generate an example album, with an artist and a list of songs. \
Using the movie {movie_name} as inspiration.\
"""

# OpenAIPydanticProgramを初期化
program = OpenAIPydanticProgram.from_defaults(
    output_cls=Album,
    prompt_template_str=prompt_template_str,
    verbose=True
)

# プログラムを実行
output = program(movie_name="The Shining")

# 結果を表示
print(output)
```

## 2. 部分的な中間Pydanticオブジェクトのストリーミング

完全なJSONが生成されるまで待つ代わりに、利用可能になった時点で中間的なPydanticオブジェクトをストリームすることができます。

```python
# 必要なライブラリをインポート
from pydantic import BaseModel, Field

# CharacterInfoクラスを定義
class CharacterInfo(BaseModel):
    """Information about a character."""
    character_name: str
    name: str = Field(..., description="Name of the actor/actress")
    hometown: str

# Charactersクラスを定義
class Characters(BaseModel):
    """List of characters."""
    characters: list[CharacterInfo] = Field(default_factory=list)

# OpenAIPydanticProgramを初期化
from llama_index.program.openai import OpenAIPydanticProgram

prompt_template_str = "Information about 3 characters from the movie: {movie}"

program = OpenAIPydanticProgram.from_defaults(
    output_cls=Characters,
    prompt_template_str=prompt_template_str
)

# 部分的なオブジェクトをストリーミング
for partial_object in program.stream_partial_objects(movie="Harry Potter"):
    print(partial_object)
```

## 3. `Album`のリスト抽出（並列関数呼び出し）

OpenAIの最新の並列関数呼び出し機能を使用して、単一のプロンプトから複数の構造化データを同時に抽出できます。

```python
# OpenAIをインポート
from llama_index.llms.openai import OpenAI

# プロンプトテンプレートを定義
prompt_template_str = """\
Generate 4 albums about spring, summer, fall, and winter.
"""

# OpenAIPydanticProgramを初期化（並列関数呼び出しを有効化）
program = OpenAIPydanticProgram.from_defaults(
    output_cls=Album,
    llm=OpenAI(model="gpt-3.5-turbo-1106"),
    prompt_template_str=prompt_template_str,
    allow_multiple=True,
    verbose=True,
)

# プログラムを実行
output = program()

# 結果を表示
print(output)
```

## まとめ

このガイドでは、OpenAI Pydantic Programを使用して構造化データを生成する方法を詳しく紹介しました。単純なオブジェクトから再帰的な構造まで、さまざまなユースケースに対応できることがわかりました。この手法を活用することで、自然言語からの構造化データ抽出がより簡単かつ効率的になります。

OpenAI Pydantic Programは、データ抽出、情報整理、そして複雑な構造のモデリングに非常に有用なツールです。特に、並列関数呼び出しやストリーミング機能を使用することで、より柔軟で効率的なデータ生成が可能になります。

この技術を使用することで、以下のような利点があります：

1. 構造化データの迅速な生成
2. 複雑なデータモデルの簡単な定義と抽出
3. リアルタイムでの部分的なデータ取得
4. 並列処理による効率的なデータ生成

ぜひ、あなたのプロジェクトでOpenAI Pydantic Programを活用し、データ処理のワークフローを改善してみてください。

## 📒ノートブック

https://colab.research.google.com/drive/1jzXgjqbqjj0tJEw0xJ8spJZ04c6MRDvH?usp=sharing

## 参考サイト

https://docs.llamaindex.ai/en/stable/examples/output_parsing/openai_pydantic_program/