型安全PythonでAIコードベースを堅牢化する実装ガイド

型安全性を無視したAIコードが抱える典型的なリスク

• LLMチェーンの入力/出力が文字列頼みで、JSONやベクトルの型崩れが実行時まで気付けない。
• データスキーマの改修でフィールド名が変わっても、NotebookやETLの随所に散らばる辞書アクセスが静的検知されない。
• 非同期＋並列推論ジョブでFutureの戻り値がAny化し、エラー時のハンドリングが困難になる。

これらは「型ヒントを少し書く」程度では防げません。ここでは、Pythonの型システムと周辺ツールをフル活用し、AIプロジェクトを堅牢化するための設計パターンとコードを紹介します。

プロジェクトの型方針を宣言する

1. pyproject.toml で mypy, ruff を必須化。
2. ルートに mypy.ini を置き、disallow_untyped_defs = True など厳格モードにする。
3. 生成AIが吐くコードもCIで型チェックを通させ、例外を設けない。

[tool.mypy]
disallow_untyped_defs = true
disallow_any_generics = true
warn_return_any = true
implicit_optional = false
python_version = "3.12"
plugins = ["pydantic.mypy", "sqlmypy"]

[tool.ruff.lint]
select = ["F", "E", "I", "UP", "ANN", "B", "RUF"]

Pydantic + Protocolで入出力を明示する

LLM用のプロンプト/レスポンスをdictのまま扱うと即脆くなります。ProtocolとPydanticを組み合わせ、モデル間の契約を型で表現します。

from typing import Protocol, TypedDict
from pydantic import BaseModel, Field

class EmbeddingResult(BaseModel):
    vector: list[float] = Field(..., min_items=1)
    model_name: str

class EmbeddingClient(Protocol):
    def embed(self, text: str) -> EmbeddingResult: ...

class OpenAIEmbeddingClient:
    def __init__(self, api_key: str) -> None:
        self._api_key = api_key

    def embed(self, text: str) -> EmbeddingResult:
        response = call_openai(text=text, key=self._api_key)
        return EmbeddingResult.model_validate(response)

class VectorDoc(TypedDict):
    id: str
    embedding: list[float]
    metadata: dict[str, str]

Protocolによって「embedを実装するクライアント」を抽象化し、TypedDictでストレージ構造を文書化できます。PydanticがJSON整形ミスを即エラーに変換するため、LLM出力の検証にも有効です。

LLMレスポンスをStructured Outputで型付けする

OpenAIやAnthropicのSDKには構造化出力の仕組みがあります。手動でJSONをパースするのではなく、型定義を先に書くようにします。

class Answer(BaseModel):
    reasoning: str
    final_answer: str
    citations: list[str]

response = client.responses.parse(
    model="gpt-4.1",
    input="与えられたドキュメントからFAQを生成",
    output_schema=Answer,
)

answer: Answer = response.output[0].content[0].parsed
print(answer.final_answer)

Answer型が変われば、その場で型エラーが発生し、IDE補完も効くようになります。

データパイプラインにTypedDictとProtocolを流し込む

ETLやRAGで扱う中間データは「とりあえずdict」になりがちです。typing.TypedDictとProtocolを併用して、変換処理の契約を守らせます。

class RawDocument(TypedDict):
    id: str
    title: str
    body: str
    published_at: str

class CleanDocument(TypedDict):
    id: str
    tokens: list[str]
    embedding: list[float]

class Transformer(Protocol):
    def __call__(self, doc: RawDocument) -> CleanDocument: ...

class SimpleTransformer:
    def __init__(self, embedder: EmbeddingClient) -> None:
        self._embedder = embedder

    def __call__(self, doc: RawDocument) -> CleanDocument:
        vector = self._embedder.embed(doc["body"]).vector
        return {
            "id": doc["id"],
            "tokens": tokenize(doc["body"]),
            "embedding": vector,
        }

このように型を通すことで、タスクランナー側もTransformerを実装しているかコンパイル時に検知できます。

typing.Annotated と Literal でプロンプトを文書化

from typing import Annotated, Literal

Role = Literal["system", "user", "assistant"]

Message = Annotated[
    dict[str, str],
    "OpenAI ChatCompletion互換メッセージ",
]

def build_prompt(role: Role, content: str) -> Message:
    return {"role": role, "content": content}

Literalで許可される役割を限定し、Annotatedを使ってメッセージ形式を IDE/ドキュメント上で分かりやすく表現できます。

TypedDict と pandas を橋渡しする

import pandas as pd
from typing import TypedDict

class FeatureRow(TypedDict):
    customer_id: str
    feature_vector: list[float]

rows: list[FeatureRow] = fetch()
df = pd.json_normalize(rows)

def ensure_schema(df: pd.DataFrame) -> pd.DataFrame:
    assert set(df.columns) == {"customer_id", "feature_vector"}
    return df

TypedDictでAPIレスポンスを型付けし、pandasに渡す前にスキーマ検証を挟むだけでも、不可解な欠損を未然に防げます。

mypy + pytest をCIに組み込む

# .github/workflows/ci.yml
name: ci
on: [push]
jobs:
  lint-test:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4
      - uses: actions/setup-python@v5
        with: { python-version: '3.12' }
      - run: pip install -r requirements-dev.txt
      - run: mypy src tests
      - run: pytest --maxfail=1 --disable-warnings

生成AIが書いたコードも例外なくこのCIを通し、型崩れや暗黙のAnyを早期に検知します。

非同期推論ジョブをTypedDictで管理する

from typing import TypedDict, Awaitable

class TaskSpec(TypedDict):
    id: str
    prompt: str
    temperature: float

async def dispatch(tasks: list[TaskSpec], llm: AsyncLLMClient) -> list[str]:
    jobs: list[Awaitable[str]] = []
    for t in tasks:
        jobs.append(llm.generate(prompt=t["prompt"], temperature=t["temperature"]))
    return await asyncio.gather(*jobs)

TaskSpecが明示されることで、Jobスケジューラとの連携やメタデータ保管が楽になります。

まとめ

• 型方針をチーム標準として宣言し、CIで強制する。
• Pydantic/TypedDict/Protocolを組み合わせ、LLM入出力・ETL・API層などすべてのデータを型で表現する。
• 生成AIが生成するコードにも同じルールを適用し、レビューと自動テストで仕上げる。

Pythonは動的言語ですが、型システムとツールを活かせばAI時代の大規模コードベースを十分に安定させられます。型を先に設計し、実装とAI支援をその上に載せる習慣をつけましょう。