Feature/streaming responses

[W3NDO]

• returns {:error, "streaming"} when evauating the rag triad of a streaming request.
• added keyword() in the @spec for generate_response and the @callback for generate_text
• Update the function clauses to consider the new opts argument.
• added 2 tests.
  - A test in rag/generation_test.exs that confirms a stream response is returned.
  - A test in rag/evaluation_test.exs that ensures an error is returned when attempting to evaluate a streaming request

[W3NDO]

@joelpaulkoch I think this is ready for a review so we can see if I am heading in the right direction. I am unsure about the evaluation test and what other behaviours I may need to update for the Rag.Ai.Provider

[joelpaulkoch]

@W3NDO just want to let you know that I'll have a look at it tomorrow, thank you!

[W3NDO]

@joelpaulkoch Sounds good. I should be available to make changes.

[joelpaulkoch]

This goes very much in the right direction, thank you!

I left some comments and we still have to implement the streaming in the providers, Rag.Ai.Nx, Rag.Ai.OpenAI, Rag.Ai.Cohere.

Just let me know if you have further questions or need help with something.

[W3NDO]

@joelpaulkoch
I wanted to ask if you could explain how we would go about doing this for Rag.Ai.Nx for example?

I left some comments and we still have to implement the streaming in the providers, Rag.Ai.Nx, Rag.Ai.OpenAI, Rag.Ai.Cohere.

What I was thinking is that in the Rag.Ai.Nx it would be an extension of generate_text/3. However, I got stuck on how to write a test for it. Should we also just raise Streaming. If this is the case, would it mean that we do something similar for Rag.Ai.Cohere and Rag.Ai.OpenAI.

[joelpaulkoch]

Sorry, took me a while to respond.
For Rag.Ai.Cohere and Rag.Ai.OpenAI we can add a stream: true option in generate_text/3 and then pass "stream" => true to the API.
In the case of Rag.Ai.Nx streaming must be configured when creating the serving.
So, there is nothing we can do with options, maybe we should raise and say "streaming must be configured when creating the serving" when someone passes a stream option?
We must still handle the stream inside Rag.Ai.Nx, at the moment this is the code in generate_text/3:

  def generate_text(%__MODULE__{} = provider, prompt, _opts) when is_binary(prompt) do
    %{results: [result]} =
      Nx.Serving.batched_run(provider.text_serving, prompt)

    {:ok, result.text}
  rescue
    error ->
      {:error, error}
  end

When we're streaming we would receive only the stream from Nx.Serving.batched_run/2.
So we would have to implement it somehow like this:

  def generate_text(%__MODULE__{} = provider, prompt, _opts) when is_binary(prompt) do
     case Nx.Serving.batched_run(provider.text_serving, prompt) do
            %{results: [result]} -> {:ok, result.text}
           stream -> {:ok, stream}
     end
  rescue
    error ->
      {:error, error}
  end

Regarding testing, at the moment we mock Nx.Serving.batched_run/2 in test/rag/generation/nx_test.exs:

      expect(Nx.Serving, :batched_run, fn _serving, prompt ->
        assert prompt == "a prompt"
        %{results: [%{text: "a response"}]}
      end)

We could do the same in a test for streaming but return a stream instead:

      expect(Nx.Serving, :batched_run, fn _serving, prompt ->
        assert prompt == "a prompt"
        Stream.take_while(["this", "is", "a", "stream"], &(&1))
      end)

[W3NDO]

@joelpaulkoch ,I will start working on this this week.

[W3NDO]

@joelpaulkoch sorry this took a while. I ran into some issues with testing the streaming options for Cohere and OpenAI. Further I could not figure out how to get the nx test working.

[W3NDO]

@joelpaulkoch we have 3 tests failing that I am unsure how to make pass. I have explained the issues I may be facing.

[joelpaulkoch]

Thank you a lot! I will have another look at the complete PR soon but I think it's basically ready to merge.