from litellm import completion

def stream_response(model, messages):
    stream = completion(
        model=model,
        messages=messages,
        stream=True
    )

    for chunk in stream:
        if chunk.choices[0].delta.content:
            yield chunk.choices[0].delta.content

# Usage with any provider
for token in stream_response("gpt-3.5-turbo", messages):
    print(token, end="", flush=True)

for token in stream_response("claude-3-sonnet-20240229", messages):
    print(token, end="", flush=True)

        __
        __