#!/usr/bin/env python3 """LiteLLMを使ってOpenAI/Geminiのモデル一覧取得と応答テストを行うモジュール。 このモジュールはLiteLLMライブラリを利用し、OpenAIおよびGoogle GeminiのAIモデルに対して 以下の操作を実行します。 - 利用可能なモデルの一覧を取得し、コンソールに表示する。 - 指定されたモデルに対して接続テストを行い、応答を表示する。 APIキーの管理には、`tkai_lib` または `tkai_lib_litellm` ライブラリ（もし利用可能であれば） または環境変数を使用します。特に、Google AI StudioのAPIキーは `GOOGLE_API_KEY` または `GEMINI_API_KEY` のどちらかに設定することで利用可能です。 使用例:: python list_litellm_models.py mode=list python list_litellm_models.py mode=list api=openai python list_litellm_models.py mode=test model=gpt-4o-mini python list_litellm_models.py mode=test model=gemini-2.0-flash python list_litellm_models.py --mode test --model openai/gpt-4o-mini .. seealso:: :doc:`list_models_usage` """ from __future__ import annotations import argparse import os import sys from collections.abc import Sequence from typing import Any from litellm import completion, get_valid_models try: from tkai_lib_litellm import read_ai_config except ImportError: try: from tkai_lib import read_ai_config except ImportError: read_ai_config = None # type: ignore[assignment] PROVIDERS = ("openai", "gemini") DEFAULT_TEST_PROMPT = ( "これはAPI接続テストです。日本語で、あなたのモデル名を名乗らずに、" "『接続テストに成功しました』と短く回答してください。" ) def parse_key_value_args(argv: Sequence[str]) -> list[str]: """mode=list のような形式のコマンドライン引数をargparse互換形式へ変換する。 特定のキーと値のペアを `--key value` の形式に変換します。 これにより、ユーザーは `mode=list` のような簡潔な記法と、 標準的な `--mode list` の両方を使用できます。 :param argv: コマンドライン引数のリスト。 :type argv: Sequence[str] :returns: 変換されたコマンドライン引数のリスト。 :rtype: list[str] """ supported = {"mode", "api", "model", "prompt", "config"} converted: list[str] = [] for arg in argv: if "=" in arg and not arg.startswith("--"): key, value = arg.split("=", 1) if key in supported: converted.extend((f"--{key}", value)) continue converted.append(arg) return converted def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: """コマンドライン引数をパースする。 スクリプトの実行モード（モデル一覧表示またはテスト）と、 それに伴うオプション（プロバイダー、モデル名、プロンプト、設定ファイル）を 定義し、パースします。 :param argv: コマンドライン引数のリスト。None の場合は sys.argv[1:] を使用。 :type argv: Sequence[str] | None :returns: パースされた引数を格納するNamespaceオブジェクト。 :rtype: argparse.Namespace """ parser = argparse.ArgumentParser( description="LiteLLMでOpenAI/Geminiのモデル一覧取得と応答テストを行います。" ) parser.add_argument( "--mode", choices=("list", "test"), default="list", help="list: モデル一覧、test: 指定モデルへテスト送信", ) parser.add_argument( "--api", choices=("all", *PROVIDERS), default="all", help="listモードで対象にするプロバイダー。既定値: all", ) parser.add_argument( "--model", help="testモードで使用するモデル名。例: gpt-4o-mini, gemini-2.0-flash", ) parser.add_argument( "--prompt", default=DEFAULT_TEST_PROMPT, help="testモードで送信するプロンプト", ) parser.add_argument( "--config", default="ai.env", help="read_ai_config() に渡す設定ファイル。既定値: ai.env", ) return parser.parse_args(parse_key_value_args(argv or sys.argv[1:])) def load_api_keys(config_path: str) -> None: """添付ライブラリを使ってAPIキー設定を読み込む。 `tkai_lib_litellm` または `tkai_lib` が利用可能な場合、指定された設定ファイルから APIキーを環境変数に読み込みます。これらのライブラリが見つからない場合は、 既存の環境変数がそのまま使用されます。 また、`GOOGLE_API_KEY` が設定されていれば、LiteLLMがGeminiモデルで使用できるよう `GEMINI_API_KEY` にも設定をコピーします。 :param config_path: APIキー設定ファイルのパス。 :type config_path: str :returns: なし :rtype: None """ if read_ai_config is None: print( "Warning: tkai_lib_litellm または tkai_lib が見つからないため、" "現在の環境変数だけを使用します。", file=sys.stderr, ) else: read_ai_config(config_path) # 既存環境の GOOGLE_API_KEY をLiteLLM用にも利用する。 if not os.getenv("GEMINI_API_KEY") and os.getenv("GOOGLE_API_KEY"): os.environ["GEMINI_API_KEY"] = os.environ["GOOGLE_API_KEY"] def required_key_exists(provider: str) -> bool: """指定されたプロバイダーに必要なAPIキーが環境変数に存在するかを確認する。 - 'openai' の場合: `OPENAI_API_KEY` の有無をチェックします。 - 'gemini' の場合: `GEMINI_API_KEY` または `GOOGLE_API_KEY` のいずれかの有無をチェックします。 :param provider: チェックするプロバイダー名 ('openai' または 'gemini')。 :type provider: str :returns: 必要なAPIキーが存在すれば True、そうでなければ False。 :rtype: bool :raises ValueError: 未知のプロバイダーが指定された場合。 """ if provider == "openai": return bool(os.getenv("OPENAI_API_KEY")) if provider == "gemini": return bool(os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY")) raise ValueError(f"Unknown provider: {provider}") def list_provider_models(provider: str) -> list[str]: """プロバイダーAPIへ問い合わせ、利用可能モデル名を返す。 LiteLLMの `get_valid_models` 関数を使用して、指定されたプロバイダーが提供する 利用可能なモデルのリストを取得します。エンドポイントの有効性も確認されます。 :param provider: モデル一覧を取得するプロバイダー名。 :type provider: str :returns: 利用可能なモデル名のソート済みリスト。 :rtype: list[str] """ models = get_valid_models( check_provider_endpoint=True, custom_llm_provider=provider, ) return sorted({str(model) for model in models}, key=str.casefold) def print_provider_models(provider: str) -> bool: """指定されたプロバイダーの利用可能モデル一覧をコンソールに表示する。 まず、必要なAPIキーが設定されているかを確認します。 キーが設定されていれば、`list_provider_models` を呼び出してモデル一覧を取得し、表示します。 エラーが発生した場合やキーが未設定の場合は、エラーメッセージを表示します。 :param provider: モデル一覧を表示するプロバイダー名。 :type provider: str :returns: モデル一覧の取得と表示が成功した場合は True、失敗した場合は False。 :rtype: bool """ title = "OpenAI" if provider == "openai" else "Gemini (Google AI)" print(f"\n--- {title}: 利用可能モデル一覧 ---") if not required_key_exists(provider): key_name = ( "OPENAI_API_KEY" if provider == "openai" else "GOOGLE_API_KEY / GEMINI_API_KEY" ) print(f"SKIP: {key_name} が設定されていません", file=sys.stderr) return False try: models = list_provider_models(provider) except Exception as exc: print(f"ERROR: {title} のモデル一覧取得に失敗しました", file=sys.stderr) print(f" [{type(exc).__name__}] {exc}", file=sys.stderr) return False for model in models: print(f"ID: {model}") print(f"--- 完了: {len(models)} 件 ---") return True def normalize_model_name(model: str) -> tuple[str, str]: """モデル名からプロバイダーを判定し、LiteLLM形式へ正規化する。 入力されたモデル名に基づいてプロバイダーを識別し、LiteLLMが認識する形式 (例: 'openai/gpt-4o-mini', 'gemini/gemini-2.0-flash') に変換します。 モデル名に `/` が含まれる場合はそれを使用してプロバイダーを判断し、 そうでなければ一般的なプレフィックス（'gemini' など）で判断します。 Gemini以外のモデルはデフォルトでOpenAIとして扱われます。 :param model: ユーザーが指定したモデル名。 :type model: str :returns: (プロバイダー名, LiteLLM形式のモデル名) のタプル。 :rtype: tuple[str, str] :raises ValueError: モデル名が空の場合、または対応していないプロバイダーが指定された場合。 """ normalized = model.strip() if not normalized: raise ValueError("model が空です") if "/" in normalized: provider = normalized.split("/", 1)[0].lower() if provider == "google": # 'google/gemini-pro' のようなケースを 'gemini/gemini-pro' に変換 provider = "gemini" normalized = "gemini/" + normalized.split("/", 1)[1] if provider not in PROVIDERS: raise ValueError( f"対応していないプロバイダーです: {provider} " f"(対応: {', '.join(PROVIDERS)})" ) return provider, normalized lower_name = normalized.lower() if lower_name.startswith("gemini"): return "gemini", f"gemini/{normalized}" # OpenAIモデルは gpt-, o1, o3, o4 など複数の接頭辞があるため、 # Gemini以外はOpenAIとして扱う。 return "openai", f"openai/{normalized}" def extract_response_text(response: Any) -> str: """LiteLLMの応答オブジェクトからメッセージ本文を取り出す。 LiteLLMの `completion` 関数が返す応答オブジェクトから、 AIモデルの生成したメッセージのテキストコンテンツを抽出します。 応答の構造が異なる可能性に対応するため、複数のアクセス方法を試行します。 :param response: LiteLLMの `completion` 関数が返した応答オブジェクト。 :type response: Any :returns: 抽出されたメッセージ本文の文字列。本文がNoneの場合は空文字列。 :rtype: str :raises ValueError: 応答オブジェクトから本文を抽出できなかった場合。 """ try: content = response.choices[0].message.content except (AttributeError, IndexError, TypeError): try: content = response["choices"][0]["message"]["content"] except (KeyError, IndexError, TypeError) as exc: raise ValueError("応答本文を取得できませんでした") from exc if content is None: return "" if isinstance(content, str): return content return str(content) def test_model(model: str, prompt: str) -> bool: """指定モデルへプロンプトを送信し、回答を表示する。 まず、モデル名を正規化し、必要なAPIキーが設定されているかを確認します。 その後、LiteLLMの `completion` 関数を使用して指定されたプロンプトをモデルに送信し、 得られた応答のテキストコンテンツをコンソールに表示します。 APIキーの不足、モデルの認識失敗、またはAPI呼び出し中のエラーが発生した場合は、 適切なエラーメッセージを表示します。 :param model: テスト対象のモデル名。 :type model: str :param prompt: モデルに送信するプロンプト。 :type prompt: str :returns: テストが成功した場合は True、失敗した場合は False。 :rtype: bool """ try: provider, litellm_model = normalize_model_name(model) except ValueError as exc: print(f"ERROR: {exc}", file=sys.stderr) return False if not required_key_exists(provider): key_name = ( "OPENAI_API_KEY" if provider == "openai" else "GOOGLE_API_KEY / GEMINI_API_KEY" ) print(f"ERROR: {key_name} が設定されていません", file=sys.stderr) return False print("--- モデル応答テスト ---") print(f"Provider : {provider}") print(f"Model : {litellm_model}") print(f"Prompt : {prompt}") print("--- 応答 ---") try: response = completion( model=litellm_model, messages=[{"role": "user", "content": prompt}], ) text = extract_response_text(response) except Exception as exc: print("ERROR: モデルからの回答取得に失敗しました", file=sys.stderr) print(f" [{type(exc).__name__}] {exc}", file=sys.stderr) print( " APIキー、モデル名、利用権限、クォータを確認してください。", file=sys.stderr, ) return False print(text if text else "(空の応答)") print("--- テスト完了 ---") return True def main(argv: Sequence[str] | None = None) -> int: """スクリプトのメイン処理を実行する。 コマンドライン引数をパースし、APIキーを読み込み、指定されたモードに応じて モデル一覧表示またはモデル応答テストを実行します。 - `mode=test` の場合、`--model` オプションが必須です。 - `mode=list` の場合、`--api` オプションで対象プロバイダーを選択できます。 :param argv: コマンドライン引数のリスト。None の場合は `sys.argv[1:]` を使用。 :type argv: Sequence[str] | None :returns: 終了コード (0: 成功、1: 失敗、2: 不適切な引数)。 :rtype: int """ args = parse_args(argv) load_api_keys(args.config) if args.mode == "test": if not args.model: print( "ERROR: mode=test では model=XXX の指定が必要です。\n" "例: python list_litellm_models.py mode=test model=gpt-4o-mini", file=sys.stderr, ) return 2 return 0 if test_model(args.model, args.prompt) else 1 selected = PROVIDERS if args.api == "all" else (args.api,) results = [print_provider_models(provider) for provider in selected] return 0 if all(results) else 1 if __name__ == "__main__": sys.exit(main())