1from typing import List
 2
 3import rich
 4from pydantic import BaseModel
 5
 6from openai import OpenAI
 7
 8
 9class Step(BaseModel):
10    explanation: str
11    output: str
12
13
14class MathResponse(BaseModel):
15    steps: List[Step]
16    final_answer: str
17
18
19client = OpenAI()
20
21with client.chat.completions.stream(
22    model="gpt-4o-2024-08-06",
23    messages=[
24        {"role": "system", "content": "You are a helpful math tutor."},
25        {"role": "user", "content": "solve 8x + 31 = 2"},
26    ],
27    response_format=MathResponse,
28) as stream:
29    for event in stream:
30        if event.type == "content.delta":
31            print(event.delta, end="", flush=True)
32        elif event.type == "content.done":
33            print("\n")
34            if event.parsed is not None:
35                print(f"answer: {event.parsed.final_answer}")
36        elif event.type == "refusal.delta":
37            print(event.delta, end="", flush=True)
38        elif event.type == "refusal.done":
39            print()
40
41print("---------------")
42rich.print(stream.get_final_completion())