[editor][server endpoints][3/n]: Proof-of-Concept for Streaming

[rossdanlm]

[editor][server endpoints][3/n]: Proof-of-Concept for Streaming

The actual call in the run command will be a bit more complex since we'll have to attach a stream callback and probably create a text queue iterator like what we did for Gradio, but this shows that it's possible and shouldn't be too complicated.

Sorry Jonathan, I didn't use any of your super useful generic functions, but we can do that later.

A big thing I wanted to point out is that in the yield generator, I need explicitly return a string (not a json object), otherwise this will not work, so the frontend needs to be able to parse that result somehow. We should sync on this but I feel jsons are strings anyways and as long as we follow the same output "formula" (ex: in the list_models we out data key vs other commands we directly return the AIConfig json), it should be fine.

Test plan

a9d769e6-d8d6-4aad-b450-6b27a75e509a.mp4

---

Stack created with Sapling. Best reviewed with ReviewStack.

• -> [editor][server endpoints][3/n]: Proof-of-Concept for Streaming #615
• [editor][server endpoints][2/n]: Add params arg to run command and check for undefined prompt_name #613
• [editor][server endpoints][1/n]: Add delete_prompt() event #612

[jonathanlastmileai]

Add it! :D

[jonathanlastmileai]

You may as well make the real endpoint. It's just us using it anyway; we're not committing to anything big here.

[jonathanlastmileai]

Isn't there a lot of duplication with HttpResponseWithAIConfig.to_flask_format() ?

[rossdanlm]

Yea you're right, I was mainly just playing around and trying to get things tested quickly

[jonathanlastmileai]

why are we sleeping?

[rossdanlm]

Just to prove that the streaming is going as planned and I can see outputs bit-by-bit

[jonathanlastmileai]

Nit: remove space before :. You can autoformat in-place with the linter.

[jonathanlastmileai]

Should we be explicitly constructing this? This feels like unnecessary duplication

[rossdanlm]

You are right lol, it's kind of not the greatest. It would be better to just overwrite output.data each time, but we'd still have to initialize it first to an ExecuteResult object (not too hard, just check if output is empty array first) but yea, just a proof of concept, nothing fancy this diff yet

[jonathanlastmileai]

How does this iterate over the tokens returned? Doesn't it just yield the same fixed contents of get_activities() every iteration ?

[rossdanlm]

The aiconfig_json is updated with different output each time. Sorry I forgot to include test video but you can see it now :)

[rholinshead]

So the response itself would be a literal string instead of an object with {data: string}?

A big thing I wanted to point out is that in the yield generator, I need explicitly return a string (not a json object)

Can you explain the reasoning for this? Couldn't you alternatively just yield an object instead of the string? I'm not super familiar with generators but I didn't think they were limited to primitive types?

[jonathanlastmileai]

So the response itself would be a literal string instead of an object with {data: string}?

A big thing I wanted to point out is that in the yield generator, I need explicitly return a string (not a json object)

Can you explain the reasoning for this? Couldn't you alternatively just yield an object instead of the string? I'm not super familiar with generators but I didn't think they were limited to primitive types?

This might be a flask limitation. @rossdanlm should help clarify. I suggest that if this works and there's any question whether flask can do it another way, and it's easy to deal with on frontend, let's just do that. For non-streaming endpoints I have a pretty good understanding what's possible and have returned JSON objects.

[rossdanlm]

Can you explain the reasoning for this? Couldn't you alternatively just yield an object instead of the string? I'm not super familiar with generators but I didn't think they were limited to primitive types?

This might be a flask limitation. @rossdanlm should help clarify.

Yea I wasn't able to get regular object to be passed, so had to convert the json to a string and it worked

I suggest that if this works and there's any question whether flask can do it another way, and it's easy to deal with on frontend, let's just do that.

Yea I agree, unless we're able to get this unblocked in 20-30 mins of investigation

[rholinshead]

Yielding string is fine (and probably expected) since we can stream the json chunks, something like:

    def generate(num_stream_steps: int):
        aiconfig_json: str | None = None
        prompt: Prompt = state.aiconfig.get_prompt(prompt_name)
        for i in range(num_stream_steps):
            time.sleep(1)
            output = ExecuteResult(
                output_type="execute_result",
                execution_count=0,
                data="Rossdan" + str(i + 1),
                metadata={},
            )
            prompt.outputs = [output]
            print(f"Done step {i+1}/{num_stream_steps}...")

            aiconfig_json = state.aiconfig.model_dump(exclude=EXCLUDE_OPTIONS)
            # print(f"{aiconfig_json=}\n")
            # yield_output = core_utils.JSONObject({"data": aiconfig_json})
            # print(f"{yield_output=}")
            print(f"{str(aiconfig_json)=}\n")
            yield json.dumps({"output_chunk": output.model_dump()})
            # yield aiconfig_json

            # HttpResponseWithAIConfig(
            #     message=f"Done step {i+1}/{num_stream_steps}...",
            #     aiconfig=state.aiconfig,
            # ).to_flask_format()

        if aiconfig_json is None:
            aiconfig_json = state.aiconfig.model_dump(exclude=EXCLUDE_OPTIONS)
        yield json.dumps({"aiconfig": aiconfig_json})

Where each intermediate chunk is the output chunk with accumulated content and then a final output of the full aiconfig

[jonathanlastmileai]

Yielding string is fine (and probably expected) since we can stream the json chunks, something like:

    def generate(num_stream_steps: int):
        aiconfig_json: str | None = None
        prompt: Prompt = state.aiconfig.get_prompt(prompt_name)
        for i in range(num_stream_steps):
            time.sleep(1)
            output = ExecuteResult(
                output_type="execute_result",
                execution_count=0,
                data="Rossdan" + str(i + 1),
                metadata={},
            )
            prompt.outputs = [output]
            print(f"Done step {i+1}/{num_stream_steps}...")

            aiconfig_json = state.aiconfig.model_dump(exclude=EXCLUDE_OPTIONS)
            # print(f"{aiconfig_json=}\n")
            # yield_output = core_utils.JSONObject({"data": aiconfig_json})
            # print(f"{yield_output=}")
            print(f"{str(aiconfig_json)=}\n")
            yield json.dumps({"output_chunk": output.model_dump()})
            # yield aiconfig_json

            # HttpResponseWithAIConfig(
            #     message=f"Done step {i+1}/{num_stream_steps}...",
            #     aiconfig=state.aiconfig,
            # ).to_flask_format()

        if aiconfig_json is None:
            aiconfig_json = state.aiconfig.model_dump(exclude=EXCLUDE_OPTIONS)
        yield json.dumps({"aiconfig": aiconfig_json})

Where each intermediate chunk is the output chunk with accumulated content and then a final output of the full aiconfig

Makes sense to me!

[rholinshead]

Ok, so for the streaming to work with the client side api we'll use (oboe), we need to wrap the entire response in an array, so:

• first chunk should sent "[{output_chunk: output}, \n" with opening "["
• next chunks send just "{output_chunk: output}, \n"
• end sends "{aiconfig: aiconfig_json}]" with closing "]"

See #651 for example, just need some improvements there:

• shouldn't need to yield [ or ] separately (needs some python string manipulation to add to the first chunk / end config text
• can we do this for all run_prompt calls? i.e. even for non-streaming run_prompt, immediately just send [{output: output}, {aiconfig: aiconfig}]" and ideally we can use one endpoint/handling for all run prompt calls

[rossdanlm]

Closing because we have updated streaming PR in #683