Add object streaming (#8)

README.md

@@ -179,6 +179,29 @@ for country, streamed_str in zip(countries, streamed_strs):
 # 24.67s : Chile - 2186 chars
 ```
 
+#### Object Streaming
+
+Structured outputs can also be streamed from the LLM by using the return type annotation `Iterable` (or `AsyncIterable`). This allows each item to be processed while the next one is being generated. See the example in [examples/quiz](examples/quiz/) for how this can be used to improve user experience by quickly displaying/using the first item returned.
+
+```python
+from collections.abc import Iterable
+from time import time
+
+
+@prompt("Create a Superhero team named {name}.")
+def create_superhero_team(name: str) -> Iterable[Superhero]:
+    ...
+
+
+start_time = time()
+for hero in create_superhero_team("The Food Dudes"):
+    print(f"{time() - start_time:.2f}s : {hero}")
+
+# 2.23s : name='Pizza Man' age=30 power='Can shoot pizza slices from his hands' enemies=['The Hungry Horde', 'The Junk Food Gang']
+# 4.03s : name='Captain Carrot' age=35 power='Super strength and agility from eating carrots' enemies=['The Sugar Squad', 'The Greasy Gang']
+# 6.05s : name='Ice Cream Girl' age=25 power='Can create ice cream out of thin air' enemies=['The Hot Sauce Squad', 'The Healthy Eaters']
+```
+
 ### Additional Features
 
 - The `@prompt` decorator can also be used with `async` function definitions, which enables making concurrent queries to the LLM.

examples/quiz/quiz.py → examples/quiz/0_quiz.py

@@ -8,21 +8,20 @@
 Run this example within this directory with:
 
 ```sh
-poetry run python quiz.py
+poetry run python 0_quiz.py
 ```
 
 or if you have installed magentic with pip:
 
 ```sh
-python quiz.py
+python 0_quiz.py
 ```
 
 ---
 
 Example run:
 
 ```
-% poetry run python quiz.py
 Enter a topic for a quiz: pizza
 Enter the number of questions: 3
 
@@ -43,9 +42,9 @@
 
 Quiz complete! You scored: 66%
 
-"Hey pizza enthusiast! Congrats on scoring 66/100 on the pizza quiz! You may not have
-aced it, but hey, you've still got a slice of the pie! Keep up the cheesy spirit and
-remember, there's always room for improvement... and extra toppings! 🍕🎉"
+Hey pizza enthusiast! Congrats on scoring 66/100 on the pizza quiz! You may not have
+ aced it, but hey, you've still got a slice of the pie! Keep up the cheesy spirit and
+ remember, there's always room for improvement... and extra toppings! 🍕🎉
 ```
 
 """

examples/quiz/quiz_async.py → examples/quiz/1_quiz_async.py

@@ -1,6 +1,6 @@
 """A simple quiz game.
 
-This example builds on the `quiz.py` example to demonstrate how using asyncio can greatly speed up queries. The quiz
+This example builds on the `0_quiz.py` example to demonstrate how using asyncio can greatly speed up queries. The quiz
 questions are now generated concurrently which means the quiz starts much more quickly after the user has entered the
 topic and number of questions. However since the questions are generated independently there is more likelihood of
 duplicates - increasing the model temperature can help with this.
@@ -10,21 +10,20 @@
 Run this example within this directory with:
 
 ```sh
-poetry run python quiz.py
+poetry run python 1_quiz_async.py
 ```
 
 or if you have installed magentic with pip:
 
 ```sh
-python quiz.py
+python 1_quiz_async.py
 ```
 
 ---
 
 Example run:
 
 ```
-% poetry run python examples/quiz/quiz_async.py
 Enter a topic for a quiz: France
 Enter the number of questions: 3
 

examples/quiz/2_quiz_streamed.py

@@ -0,0 +1,110 @@
+"""A simple quiz game.
+
+This example improves on the `1_quiz.py` example by using streaming to generate the questions. In `1_quiz.py` the
+questions were generated concurrently which allowed the quiz to start quickly but meant there was a chance of duplicate
+questions being generated. In this example the questions are streamed which allows us to show the first question to the
+user as soon as it is ready, while still making a single query to the LLM which avoids generating duplicate questions.
+
+The only change from `0_quiz.py` is the return type annotations of the `generate_questions` function changing from
+`list[Question]` to `Iterable[Question]`. This allows us to iterate through the questions as they are generated.
+
+---
+
+Run this example within this directory with:
+
+```sh
+poetry run python 2_quiz_streamed.py
+```
+
+or if you have installed magentic with pip:
+
+```sh
+python 2_quiz_streamed.py
+```
+
+---
+
+Example run:
+
+```
+Enter a topic for a quiz: NASA
+Enter the number of questions: 3
+
+1 / 3
+Q: When was NASA founded?
+A: 1958
+Correct! The answer is: 1958
+
+2 / 3
+Q: Who was the first person to walk on the moon?
+A: Neil Armstrong
+Correct! The answer is: Neil Armstrong
+
+3 / 3
+Q: What is the largest planet in our solar system?
+A: Jupyter
+Incorrect! The correct answer is: Jupiter
+
+Quiz complete! You scored: 66%
+
+Congratulations on your stellar performance! You may not have reached the moon,
+ but you definitely rocked that NASA quiz with a score of 66/100! Remember,
+ even astronauts have their off days. Keep reaching for the stars, and who knows,
+ maybe next time you'll be the one discovering a new galaxy!
+ Keep up the astronomical work! 🚀🌟
+```
+
+"""
+
+from collections.abc import Iterable
+
+from pydantic import BaseModel
+
+from magentic import prompt
+
+
+class Question(BaseModel):
+    question: str
+    answer: str
+
+
+@prompt("Generate {num} quiz questions about {topic}")
+def generate_questions(topic: str, num: int) -> Iterable[Question]:
+    ...
+
+
+@prompt("""Return true if the user's answer is correct.
+Question: {question.question}
+Answer: {question.answer}
+User Answer: {user_answer}""")
+def is_answer_correct(question: Question, user_answer: str) -> bool:
+    ...
+
+
+@prompt(
+    "Create a short and funny message of celebration or encouragment for someone who"
+    " scored {score}/100 on a quiz about {topic}."
+)
+def create_encouragement_message(score: int, topic: str) -> str:
+    ...
+
+
+topic = input("Enter a topic for a quiz: ")
+num_questions = int(input("Enter the number of questions: "))
+questions = generate_questions(topic, num_questions)
+
+user_points = 0
+for num, question in enumerate(questions, start=1):
+    print(f"\n{num} / {num_questions}")
+    print(f"Q: {question.question}")
+    user_answer = input("A: ")
+
+    if is_answer_correct(question, user_answer):
+        print(f"Correct! The answer is: {question.answer}")
+        user_points += 1
+    else:
+        print(f"Incorrect! The correct answer is: {question.answer}")
+
+score = 100 * user_points // num_questions
+print(f"\nQuiz complete! You scored: {score}%\n")
+print(create_encouragement_message(score, topic))

src/magentic/__init__.py

@@ -1,7 +1,7 @@
 from magentic.function_call import FunctionCall
 from magentic.prompt_chain import prompt_chain
 from magentic.prompt_function import prompt
-from magentic.streamed_str import AsyncStreamedStr, StreamedStr
+from magentic.streaming import AsyncStreamedStr, StreamedStr
 
 __all__ = [
     "FunctionCall",