def main()

in projects/dynasor/openai_client.py [0:0]

• 38 lines of code
• 4 McCabe index (conditional complexity)

def main():
    args = parse_args()
    logger.debug("Args: %s", args)

    # stream = not args.no_stream
    # assert stream, "No streaming implementation is not supported yet."
    stream = True
    disable_dynasor = args.disable_dynasor

    client = OpenAI(
        api_key=args.api_key,
        base_url=args.base_url,
        max_retries=1,
    )

    logger.debug("Grab models")
    models = client.models.list()
    model = models.data[0].id
    logger.debug("Model: %s", model)

    dynasor_config = dict()
    if not disable_dynasor:
        dynasor_config = dict(
            probe_interval=args.probe_interval,
            certainty_window=args.certainty_window,
        )

    response = client.chat.completions.create(
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": args.prompt},
        ],
        model=model,
        max_tokens=args.max_tokens,
        extra_body={
            "dynasor": dynasor_config
        },
        stream=stream,
    )

    print("Prompt: ", args.prompt)
    print("-" * 10)
    print("Response: \n")
    for chunk in response:
        if chunk.choices[0].delta.content is not None:
            print(chunk.choices[0].delta.content, end="", flush=True)