stream.py

import json
import math
import shutil
import sys
import time

from langchain_core.output_parsers import JsonOutputParser


def simulate_stream():
    data = {
        "name": "pokutuna",
        "age": 16,
        "okodukai": 3000,
        "food": [
            "tonkatu",
            "chikuzenni",
            "torotaku",
        ]
    }
    json_str = json.dumps(data)
    for char in json_str:
        yield char


parser = JsonOutputParser()

h1 = "stream: "
h2 = "parsed: "
width = shutil.get_terminal_size().columns * 0.8

print(h1)
print(h2)

buf = ""
for char in simulate_stream():
    buf += char
    parsed = json.dumps(parser.parse(buf))
    head = max(0, math.floor(len(buf) - width))

    sys.stdout.write("\033[2F")
    sys.stdout.write("\033[J")
    print(f"{h1}{buf[head:]}")
    print(f"{h2}{parsed[head:]}")

    time.sleep(0.1)