Support token streaming with json adapter

[chenmoneygithub]

With this PR we are support token streaming when using JSONAdapter. Sample usage:

import asyncio
from typing import List

from pydantic import BaseModel

import dspy
from dspy.streaming import StatusMessage

lm = dspy.LM("openai/gpt-4o-mini", cache=False)
dspy.settings.configure(lm=lm)


class EntitiesModel(BaseModel):
    attributes: List[str]
    colors: List[str]
    animals: List[str]


class MySignature(dspy.Signature):
    """Extract entities from a text"""

    text: str = dspy.InputField(description="The text to extract entities from")
    entities: EntitiesModel = dspy.OutputField(description="The entities in the text")


# predict = dspy.ChainOfThought(MySignature)

predict = dspy.Predict("question->answer")

stream_listeners = [
    dspy.streaming.StreamListener(signature_field_name="reasoning"),
]
stream_predict = dspy.streamify(
    predict, stream_listeners=stream_listeners, include_final_prediction_in_output_stream=False
)


async def main():
    with dspy.settings.context(adapter=dspy.JSONAdapter()):
        output = stream_predict(text="The quick brown fox jumps over the lazy dog with piercing blue eyes")

        return_value = None
        async for value in output:
            if isinstance(value, StatusMessage):
                print(value.message)
            elif isinstance(value, dspy.Prediction):
                print(value)
            elif isinstance(value, dspy.streaming.StreamResponse):
                print(value)
    return return_value


import time

start = time.time()
asyncio.run(main())
end = time.time()
print(f"Time taken: {end - start} seconds")

Sample output:

StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk='The')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' text')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' describes')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' a')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' scene')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' involving')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' a')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' fox')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' and')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' a')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' dog')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=',')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' both')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' of')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' which')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' are')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' animals')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk='.')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' The')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' fox')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' is')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' characterized')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' by')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' its')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' quick')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk='ness')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' and')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' brown')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' color')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=',')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' while')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' the')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' dog')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' is')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' described')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' as')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' lazy')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' and')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' has')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' blue')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' eyes')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=',')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' which')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' are')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' an')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' attribute')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' of')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' its')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' appearance')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk='.')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' Therefore')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=',')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' the')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' entities')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' extracted')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' include')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' the')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' animals')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' (')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk='fox')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' and')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' dog')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk='),')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' their')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' attributes')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' (')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk='quick')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=',')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' brown')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=',')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=' lazy')
StreamResponse(predict_name='predict', signature_field_name='reasoning', chunk=', blue), and colors (brown, blue).')
Time taken: 2.8075790405273438 seconds