Generator Coroutines
Stateful generators
Created: 2025-05-25
A coroutine is a special type of generator that can receive values from the caller. This allows for more complex interactions between the generator and the caller, enabling the generator to act as a state machine or to maintain state across multiple calls.
In practice this is achieved by using the yield statement in a function, which allows the function to pause its execution and yield control back to the caller. The caller can then send values back to the generator using the send() method, which resumes the generator's execution at the point where it was paused.
Note that the generator must be 'primed' first. Once the function is called, call next() on the generator to advance it to the first yield statement. After that, you can use send() to send values back into the generator.
Here's a contrived example of a coroutine that processes characters from a string and yields them back to the caller
Something to keep in mind, is that nothing is executed inside stream_characters() until the first next() is called. Once the first next() is called, the function will run until it hits the first yield statement.
This is where we have to think differently about the yield statement. The first yield statement is where the function will pause and wait for a value to be sent back into it. This is the point where we can send a value back into the generator using the send() method.
def stream_characters():
print("Generator started")
text = yield
print("Generator resumed with text:", text)
yield
print("Generator acknowledged text:", text)
if text is not None:
for char in text:
yield char
gen = stream_characters()
next(gen)
gen.send("Hello World") # Send a string to the generator
for char in gen:
print("Received character:", char)
Generator started Generator resumed with text: Hello World Generator acknowledged text: Hello World Received character: H Received character: e Received character: l Received character: l Received character: o Received character: Received character: W Received character: o Received character: r Received character: l Received character: d
A More Practical Example
What's interesting about this is that we can treat these generators somewhat like a thread. We can send messages to them, and they can send messages back to us. They can share state, but for practicality, we should avoid that. In fact, it's best to keep the state inside the generator and not share it with the caller. This way, we can keep the generator as a pure function.
For this example, let's create a coroutine that simulates an LLM conversation. What we want to do is maintain one or more conversations and not need to worry about the state of each conversation in the caller. The generator will maintain the state of the conversation and return the conversation history when requested. We're not going to use any LLMs here, but we will simulate conversation history.
While we're at it, let's make a decorator that will automatically prime the coroutine for us. This will allow us to use the coroutine without having to call next() on it first.
from typing import Literal, Counter
from functools import wraps
from operator import methodcaller
def coroutine(func):
@wraps(func)
def wrapper(*args, **kwargs):
gen = func(*args, **kwargs)
next(gen) # Prime the coroutine
return gen
return wrapper
@coroutine
def conversation(mode: Literal["lower", "upper"]):
# The conversation history is maintained in the generator
history = []
token_counts = Counter()
user_message = yield # Received via .send()
transformer = methodcaller(mode)
while 1:
if user_message is None: # End of conversation, return the history and token counts
yield history, token_counts.most_common()
return
# Echo the user message in all caps
token_counts.update(user_message.split())
llm_message = transformer(user_message)
token_counts.update(llm_message.split())
history.extend([{
"role": "user",
"content": user_message,
},
{
"role": "llm",
"content": llm_message,
}])
user_message = yield llm_message
# Let's maintain 2 conversations
convo_lower = conversation("lower")
convo_upper = conversation("upper")
# Send a message to the lower case conversation
response = convo_lower.send("Hello World")
# If we don't care about the response, we can just call send() again
for i in range(5):
convo_lower.send(f"Hello World {i}")
# Send a message to the upper case conversation
response = convo_upper.send("Hello World")
# If we don't care about the response, we can just call send() again
for i in range(5):
convo_upper.send(f"Hello World {i}")
# Now we can end the conversations and get the history and token counts
history_lower, token_counts_lower = convo_lower.send(None)
print(history_lower)
print(token_counts_lower)
[{'role': 'user', 'content': 'Hello World'}, {'role': 'llm', 'content': 'hello world'}, {'role': 'user', 'content': 'Hello World 0'}, {'role': 'llm', 'content': 'hello world 0'}, {'role': 'user', 'content': 'Hello World 1'}, {'role': 'llm', 'content': 'hello world 1'}, {'role': 'user', 'content': 'Hello World 2'}, {'role': 'llm', 'content': 'hello world 2'}, {'role': 'user', 'content': 'Hello World 3'}, {'role': 'llm', 'content': 'hello world 3'}, {'role': 'user', 'content': 'Hello World 4'}, {'role': 'llm', 'content': 'hello world 4'}]
[('Hello', 6), ('World', 6), ('hello', 6), ('world', 6), ('0', 2), ('1', 2), ('2', 2), ('3', 2), ('4', 2)]
history_upper, token_counts_upper = convo_upper.send(None)
print(history_upper)
print(token_counts_upper)
[{'role': 'user', 'content': 'Hello World'}, {'role': 'llm', 'content': 'HELLO WORLD'}, {'role': 'user', 'content': 'Hello World 0'}, {'role': 'llm', 'content': 'HELLO WORLD 0'}, {'role': 'user', 'content': 'Hello World 1'}, {'role': 'llm', 'content': 'HELLO WORLD 1'}, {'role': 'user', 'content': 'Hello World 2'}, {'role': 'llm', 'content': 'HELLO WORLD 2'}, {'role': 'user', 'content': 'Hello World 3'}, {'role': 'llm', 'content': 'HELLO WORLD 3'}, {'role': 'user', 'content': 'Hello World 4'}, {'role': 'llm', 'content': 'HELLO WORLD 4'}]
[('Hello', 6), ('World', 6), ('HELLO', 6), ('WORLD', 6), ('0', 2), ('1', 2), ('2', 2), ('3', 2), ('4', 2)]