openai_test.py

"""
uv pip install openai
"""

import os

import logging

# logging.basicConfig(
#     level=logging.DEBUG,
# )

os.environ["NO_PROXY"] = "127.0.0.1"

from openai import OpenAI

client = OpenAI(base_url="http://127.0.0.1:8000/api/v1", api_key="sk-test")


def completionStreamTest():
    print("[*] Stream completion: ")

    completion = client.chat.completions.create(
        model="rwkv-latest",
        messages=[
            {
                "role": "User",
                "content": "请讲个关于一只灰猫和一个小女孩之间的简短故事。",
            },
        ],
        stream=True,
        max_tokens=2048,
    )

    isReasoning = False

    for chunk in completion:
        if chunk.choices[0].delta.reasoning_content and not isReasoning:
            print("<- Reasoning ->")
            isReasoning = True
        elif chunk.choices[0].delta.content and isReasoning:
            isReasoning = False
            print("<- Stop Reasoning ->")

        if chunk.choices[0].delta.reasoning_content:
            print(chunk.choices[0].delta.reasoning_content, end="", flush=True)
        if chunk.choices[0].delta.content:
            print(chunk.choices[0].delta.content, end="", flush=True)

    print("")


def completionTest():
    completion = client.chat.completions.create(
        model="rwkv-latest:thinking",
        messages=[
            {
                "role": "User",
                "content": "How many planets are there in our solar system?",
            },
        ],
        max_tokens=2048,
    )

    print("[*] Completion: ", completion)


if __name__ == "__main__":
    try:
        # completionTest()

        testRounds = input("Test rounds (Default: 10) :")

        for i in range(int(testRounds) if testRounds != "" else 10):
            print("\n", "=" * 10, i + 1, "/", testRounds, "=" * 10)
            completionStreamTest()
    except KeyboardInterrupt:
        pass