Is it possible to stream async using with_structured_output

[GaryFail]

I want to stream the response from an OpenAI call back to the user in a FastAPI endpoint. I want the output to be a structured output however when I do this I can't stream. If I remove the with_structured_output then everything is fine. My question is, how do I stream a response if I'm using a structured output?

The code to replicate this is:

import json
import logging
import os

from langchain_core.prompts import PromptTemplate
from langchain_openai import AzureChatOpenAI
from langgraph.graph import START, END, StateGraph

from infrastructure.persistence.state_persistence import get_checkpointer
from starlette.responses import ContentStream


llm = AzureChatOpenAI(
        model=os.environ.get("MODEL_NAME"), temperature=0.6, streaming=True
    )

class Joke:
    question: str
    answer: str

class DebugState:
    joke: Joke

prompt_template = PromptTemplate(
    template="Tell me a joke about a {subject}",
    input_variables=[
        "subject",
    ],
    template_format="f-string",
)

# If I remove with_structured_output, it will stream back fine
runnable = prompt_template | llm.with_structured_output(Joke)

async def joke_creator(state: DebugState):

    subject = "Python"

    response = await runnable.ainvoke({"subject": subject})

    return {"joke": response}


builder = StateGraph(DebugState)
builder.add_node("joke_creator", joke_creator)

builder.add_edge(START, "joke_creator")
builder.add_edge("joke_creator", END)


saver = get_checkpointer()

assistant = builder.compile(checkpointer=saver)


async def create_joke() -> ContentStream:

    config = {
        "configurable": {
            "thread_id": "debug_thread",
        },
    }

    event_stream = assistant.astream(
        input={},
        config=config,
        stream_mode="updates",
        subgraphs=True,
    )

    try:
        async for event in event_stream:
            for item in event:
                if isinstance(item, dict):
                    if "joke_creator" in item:
                        message = item["joke_creator"]["joke"]
                        payload = {
                            "type": "AIMessage",
                            "content": message["question"] + " " + message["answer"],
                        }
                        yield f"{json.dumps(payload)}\n\n"

                    


    except Exception as e:
        logging.error(f"Error: {str(e)}")
        payload = {
            "type": "AIMessage",
            "content": "An unexpected error occurred. Please try again later.",
        }
        yield f"{json.dumps(payload)}\n\n"

[GaryFail]

The error I receive is as follows

AttributeError("'AsyncStream' object has no attribute 'aclose'")

I've also tried setting the class to be a TypedDict as described here but still get the same issue.

class Joke(TypedDict):
question: Annotated[str, "The question part of the joke"]
answer: Annotated[str, "The answer part of the joke"]

class DebugState(TypedDict):
joke: Joke

It looks like the error may be on the openai.AsyncStream object but I'm no Python expert so am clutching at straws. This is on the error object: <openai.AsyncStream object at 0x000001745D077920>

I can use a PydanticOutputParser but I'm unsure of their reliability in comparison to using with_structured_output.

[steven-hammal]

Is this possibly a bug?