型駆動でAIパイプラインを組むためのPython高度テクニック

高度な型安全が必要になる場面

• LLMチェーンが複数モデルをまたぎ、入力/出力の型が動的に変わる。
• Feature StoreやVector Storeを共通のインターフェースで扱いたい。
• ルーターパターンで複数のハンドラに処理を振り分ける際、キーと関数の整合性をコンパイル時に確保したい。
• DataclassやPydanticのラッパーを自作し、型チェッカにカスタム動作を教えたい。

ここでは、標準ライブラリ typing と mypyプラグインAPIを使った高度な型安全テクニックを紹介します。

NewType と LiteralStringでプロンプトインジェクションを抑止

from typing import NewType, LiteralString

SystemPrompt = NewType("SystemPrompt", str)
UserPrompt = NewType("UserPrompt", str)

ALLOWED_TEMPLATE: LiteralString = "You are a helpful assistant"

def build_system_prompt(template: LiteralString) -> SystemPrompt:
    return SystemPrompt(template)

system_prompt = build_system_prompt(ALLOWED_TEMPLATE)
user_prompt = UserPrompt("Summarize the attached report")

# stringをそのまま代入するとmypyエラー
# user_prompt = "free form"  # error: incompatible type

NewTypeでプロンプトの役割を分離し、LiteralStringでテンプレートをコンパイル時に限定できます。外部入力をそのままSystem Promptに流し込むと型エラーになるため、レビュー漏れを防げます。

TypeVarとProtocolでモデルの組合せ爆発を制御

from typing import TypeVar, Protocol, Generic

Doc = TypeVar("Doc", contravariant=True)
Result = TypeVar("Result", covariant=True)

class Stage(Protocol[Doc, Result]):
    def __call__(self, item: Doc) -> Result: ...

class Pipeline(Generic[Doc, Result]):
    def __init__(self, stage: Stage[Doc, Result]):
        self._stage = stage

    def run(self, data: Doc) -> Result:
        return self._stage(data)

class JsonToEmbedding(Stage[str, list[float]]):
    def __call__(self, item: str) -> list[float]:
        obj = json.loads(item)
        return embed(obj["text"])  # 型が保証される

pipe = Pipeline(JsonToEmbedding())
vector = pipe.run('{"text": "hello"}')

Stageに逆変/共変TypeVarを指定し、パイプラインをGenericsで結ぶことで、Stage同士の型整合性をmypyが保証します。

Router + TypedDict + Literalで安全なハンドラ登録

from typing import Callable, TypedDict, Literal, overload

class ClassificationRequest(TypedDict):
    type: Literal['classification']
    text: str

class EmbeddingRequest(TypedDict):
    type: Literal['embedding']
    text: str

Request = ClassificationRequest | EmbeddingRequest

Handlers = dict[str, Callable[[Request], str | list[float]]]

handlers: Handlers = {}

@overload
def register(type_: Literal['classification']) -> Callable[[Callable[[ClassificationRequest], str]], None]: ...

@overload
def register(type_: Literal['embedding']) -> Callable[[Callable[[EmbeddingRequest], list[float]]], None]: ...

def register(type_):
    def decorator(func):
        handlers[type_] = func
        return func
    return decorator

@register('classification')
def classify(req: ClassificationRequest) -> str:
    return classify_text(req['text'])

LiteralキーとTypedDictハンドラの型を一致させるため、@overloadで登録関数の型を分岐させています。存在しないリクエストタイプを登録しようとするとmypyが検知します。

dataclass_transformで自作デコレータに型認識させる

from typing import dataclass_transform

@dataclass_transform()
def schema(cls: type[T]) -> type[T]:
    return dataclasses.dataclass(eq=False, frozen=True)(cls)

@schema
class PromptTemplate:
    name: str
    input_vars: tuple[str, ...]

dataclass_transformを付与すると、カスタムデコレータでもmypyがdataclassとして扱い、生成フィールドや補完を理解できるようになります。AI系DSLを自作するときに有用です。

mypy プラグインでPydantic v2を理解させる

pyproject.tomlで pydantic.mypy を有効にし、ConfigDictに応じた型推論をさせます。

[tool.mypy]
plugins = ["pydantic.mypy"]

class QueryModel(BaseModel):
    model_config = ConfigDict(extra='forbid')
    vector: list[float]
    metadata: dict[str, str]

これでextra禁止が型チェックでも反映され、未定義フィールドアクセスをコンパイル時に検知できます。

Typed DFSとExhaustive MatchでLLMステートを管理

from typing import Literal, assert_never

State = Literal['init', 'collect', 'generate', 'done']

def next_state(state: State, event: str) -> State:
    match state:
        case 'init' if event == 'input':
            return 'collect'
        case 'collect' if event == 'ready':
            return 'generate'
        case 'generate' if event == 'finish':
            return 'done'
        case 'done':
            return 'done'
        case other:
            assert_never(other)

assert_neverを使うと、新しい状態を追加した際にmatchの抜け漏れをmypyが指摘してくれるため、LLMステートマシンを安全に保守できます。

ストレージごとのNewTypeとFactoryでミスを防ぐ

VectorId = NewType('VectorId', str)
DocumentId = NewType('DocumentId', str)

class VectorStore(Protocol):
    def upsert(self, vid: VectorId, vector: list[float]) -> None: ...

class DocumentStore(Protocol):
    def save(self, did: DocumentId, text: str) -> None: ...

vector_store.upsert(VectorId('vec-1'), embedding)  # OK
vector_store.upsert(DocumentId('doc-1'), embedding)  # mypy error

IDの混在はLLMパイプラインで頻出するバグです。NewTypeで意味付きIDを作り、Factory関数で生成する習慣を付けます。

Typed Settingsで環境変数を型付け

from pydantic_settings import BaseSettings

class Settings(BaseSettings):
    openai_api_key: str
    redis_url: AnyUrl
    concurrency_limit: int = 8

settings = Settings()  # 型安全に参照できる

環境変数をBaseSettings経由にすると、型未満の設定を起動時に検知できます。Literalと組み合わせればフラグ値の制約も可能です。

テストコードでも型を守る

from typing import Callable

Handler = Callable[[str], str]

@pytest.fixture
def handler() -> Handler:
    return lambda text: text.upper()

@pytest.mark.parametrize('text', ['abc', 'xyz'])
def test_handler(handler: Handler, text: str) -> None:
    result = handler(text)
    assert result.isupper()

テストにも型ヒントを付けると、fixtureのスコープやParametrizeの値がずれた際にmypyが警告を出します。AIプロジェクトはテストコード量が多いため効果的です。

まとめ

• NewType, Literal, Protocol, TypeVar, dataclass_transformを組み合わせることで、Pythonでも高度な型安全が実現できる。
• LLMやデータパイプラインではID・ステート・フォーマットが乱れがちなので、型で契約を表現し、CIで守らせる。
• カスタムデコレータやDSLを使う場合は、mypyプラグインやdataclass_transformで型チェッカに意図を伝える。

「とりあえずdictとAnyで書いてから直す」のではなく、最初に型の契約を敷き、AIコードもその枠内で生成・レビューする運用に切り替えましょう。