Support streaming responses

[joelpaulkoch]

Streaming responses should be supported via http and nx.

[W3NDO]

Hi @joelpaulkoch is this open for public contribution?

If so, what are the requirements for this task?

[joelpaulkoch]

Hi, thanks for reaching out, contributions are very much welcome.

We want to support streaming responses, so we want a function in Rag.Generation that takes a Rag.Generation struct and a Rag.Ai.Provider implementation and returns a Rag.Generation struct with a stream as response.

1. We must update the Rag.Generation.t type, so that the response field can be a stream/enumerable.
2. We must provide a function for streaming the response. I think there are two options: either a new function Rag.Generation.stream_response/2, or we extend the existing Rag.Generation.generate_response/2 to take opts as third argument and pass in stream: true, what do you think?
3. We must update the Rag.Ai.Provider behaviour and it's implementations to support streaming. Here it's basically the same choice, either we create a new function stream_text/2 or pass in stream: true as opts to generate_text/3, although I think I would prefer passing stream: true
4. Last, after generating a response, we want to be able to evaluate it. Now that it can be a stream we must think about how we should handle it. I think for now we could just skip it and raise if we try to evaluate a streamed response.
   In the long term it would be nice to have evaluations with streams.
   Two options would be: 1. just pass the stream through and only evaluate when the response is complete, 2. evaluate each sentence/otherwise meaningful paragraph

In any case, please reach out if you have further questions, get stuck, or just want to tackle something else.

[W3NDO]

Awesome, I will get started right away.

For this task

We must provide a function for streaming the response. I think there are two options: either a new function Rag.Generation.stream_response/2, or we extend the existing Rag.Generation.generate_response/2 to take opts as third argument and pass in stream: true, what do you think?

and this one

We must update the Rag.Ai.Provider behaviour and it's implementations to support streaming. Here it's basically the same choice, either we create a new function stream_text/2 or pass in stream: true as opts to generate_text/3, although I think I would prefer passing stream: true

I think passing stream true as opts would be the better way to do it too. I think that's how most LLM libraries handle it as well.

I will have to think more about the evaluation, then I will get back to you about it.

In the meantime, let me open a new branch and get started. :)

[W3NDO]

@joelpaulkoch I am having some issues designing the tests for a stream. Specifically, I am thinking of how we can know that the stream part received is a partial part of the stream and how we can know that another is the final partial of the stream. The test I had thus far looked something like this

test "returns a stream response" do
      generation = %Generation{query: "query", prompt: "get a streaming response", stream?: true}

      stream_response_fn = fn "get a streaming response", _opts ->
        result = Stream.repeatedly(fn -> "This is a streamed response" end) |> Enum.take(3) |> Enum.join(" ")
        {:ok, result}
      end

      assert assert %{
                      response: "This is a streamed response This is a streamed response This is a streamed response"
                    } =
                      Generation.generate_response(generation, stream_response_fn, true)
    end

The problem here is that as far as I can tell this is no different that just sending the string "This is a streamed response This is a streamed response This is a streamed response" across.

How would you go about this?

Further, I have a question on whether we would need to update put_response/2 to cater for streaming responses.

I could open a PR with what I have so far

[joelpaulkoch]

Cool that you're working on this.

So, I think we shouldn't have a stream? field as part of the Generation struct.
Instead it should be possible to get a stream in the response field of the Generation struct.

The problem here is that as far as I can tell this is no different that just sending the string "This is a streamed response This is a streamed response This is a streamed response" across.

I agree with this point, as a consumer of the library you would want to be able to get a stream from the response field and continue from there.

With that, you're snippet would become:

test "returns a stream response" do
      generation = %Generation{query: "query", prompt: "get a streaming response"}

      stream_response_fn = fn "get a streaming response", _opts ->
        result = Stream.repeatedly(fn -> "This is a streamed response" end)
        {:ok, result}
      end

      %{response: response} = Generation.generate_response(generation, stream_response_fn, stream: true)

     assert Enum.take(response, 3) == "This is a streamed response This is a streamed response This is a streamed response"
    end

Do you agree? Is there something else to consider?

Further, I have a question on whether we would need to update put_response/2 to cater for streaming responses.

Yes, although we only have to update the typespec to allow Enumerable.t() and remove the is_binary guard.

  @doc """
  Puts `response` in `generation.response`.
  """
  @spec put_response(t(), response :: String.t() | Enumerable.t()) :: t()
  def put_response(%Generation{} = generation, response),
    do: %{generation | response: response}

Feel free to open a PR whenever you want to.

[W3NDO]

Do you agree? Is there something else to consider?

I agree, the other part I was considering however was if we are streaming a response, how do we indicate that the response is finalised or if there are still more chunks to come?

[joelpaulkoch]

I agree, the other part I was considering however was if we are streaming a response, how do we indicate that the response is finalised or if there are still more chunks to come?

About this point: I think we don't have to indicate that the response is finalised.
We receive a stream either from a streaming Nx.Serving, or via req and http.
We can just transform the stream, map to the actual text chunks and return this stream in the response field.
I believe the stream itself will indicate when it's done, see this example from req

[W3NDO]

OK, that makes sense. I had not thought about it like that.