小ネタ：yaspinでターミナルにクルクルを出そう #ai #azure #openai #llm #python

簡単なスクリプト

次のような Azure OpenAI に問い合わせる簡単なPythonスクリプトがあるとします。

import os
from openai import AzureOpenAI
from dotenv import load_dotenv
load_dotenv()

def send_prompt(prompt):
    response = None
    try:
        client = AzureOpenAI(
            api_key = os.getenv("AZURE_OPENAI_API_KEY"),
            azure_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT"),
            api_version = os.getenv("AZURE_OPENAI_API_VERSION")
        )
        response = client.chat.completions.create(
            model = os.getenv("AZURE_OPENAI_MODEL_NAME"),
            messages = [{"role": "user", "content": prompt}],
            max_tokens = 1024,
            temperature = 0.95
        )
    except Exception as e:
        print(f"error: {e}")
        response = None
    return response

res = send_prompt("What is Docker?")
if res:
    model_dump = res.model_dump()
    print({"response": model_dump['choices'][0]['message']['content']})

これを実行すると、応答が返ってくるまで何も表示されないので「本当に大丈夫なのかな？」と心配になってきます。ウェブなら待ちのときに、次のようなクルクル（spinnerというらしい？）を出すこともできますが…。

https://commons.wikimedia.org/wiki/File:Loading_icon.gif

yaspin: Yet Another Terminal Spinner for Python

そこで yaspin の出番です。次のような処理待ちのクルクルを簡単に作ることができます。

次がサンプルスクリプトです。

import time
from yaspin import yaspin

with yaspin(text="Processing", color="yellow") as spinner:
    time.sleep(5)
    spinner.ok("✅ ")

これを実行すると、次のように5秒間クルクルしてから完了します。

では早速yaspinを先のAzure OpenAIスクリプトに組み込んでみましょう。

import os
from yaspin import yaspin
from openai import AzureOpenAI
from dotenv import load_dotenv
load_dotenv()

def send_prompt(prompt):
    response = None
    with yaspin(text="Processing", color="yellow") as spinner:
      try:
        client = AzureOpenAI(
            api_key = os.getenv("AZURE_OPENAI_API_KEY"),
            azure_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT"),
            api_version = os.getenv("AZURE_OPENAI_API_VERSION")
        )
        response = client.chat.completions.create(
            model = os.getenv("AZURE_OPENAI_MODEL_NAME"),
            messages = [{"role": "user", "content": prompt}],
            max_tokens = 1024,
            temperature = 0.95
        )
        spinner.ok("✅ ")
      except Exception as e:
        print(f"error: {e}")
        response = None
        spinner.fail("💥 ")
    return response

res = send_prompt("What is Docker?")
if res:
    model_dump = res.model_dump()
    print({"response": model_dump['choices'][0]['message']['content']})

差分は次のようになります。

--- 10__simple.py    2024-10-28 14:43:36.765549180 +0900
+++ 11__yaspin.py    2024-10-28 17:26:28.336770580 +0900
@@ -1,13 +1,15 @@
 #!python3

 import os
+from yaspin import yaspin
 from openai import AzureOpenAI
 from dotenv import load_dotenv
 load_dotenv()

 def send_prompt(prompt):
     response = None
-    try:
+    with yaspin(text="Processing", color="yellow") as spinner:
+      try:
         client = AzureOpenAI(
             api_key = os.getenv("AZURE_OPENAI_API_KEY"),
             azure_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT"),
@@ -19,9 +21,11 @@
             max_tokens = 1024,
             temperature = 0.95
         )
-    except Exception as e:
+        spinner.ok("✅ ")
+      except Exception as e:
         print(f"error: {e}")
         response = None
+        spinner.fail("💥 ")
     return response

 res = send_prompt("What is Docker?")

これを実行すると、処理待ち中はクルクルしています。

まとめ

yaspinを使うと、処理待ち中のターミナルにクルクルを出すことができます。単純なクルクルだけでなく、いくつか種類があるのでお好みで使い分けてみてください。