Cancel a message delta when we receive "overloaded" error

lib/chains/llm_chain.ex

@@ -454,7 +454,7 @@ defmodule LangChain.Chains.LLMChain do
   completes the message, the LLMChain is updated to clear the `delta` and the
   `last_message` and list of messages are updated.
   """
-  @spec apply_delta(t(), MessageDelta.t()) :: t()
+  @spec apply_delta(t(), MessageDelta.t() | {:error, LangChainError.t()}) :: t()
   def apply_delta(%LLMChain{delta: nil} = chain, %MessageDelta{} = new_delta) do
     %LLMChain{chain | delta: new_delta}
   end
@@ -464,6 +464,11 @@ defmodule LangChain.Chains.LLMChain do
     delta_to_message_when_complete(%LLMChain{chain | delta: merged})
   end
 
+  # Handle when the server is overloaded and cancelled the stream on the server side.
+  def apply_delta(%LLMChain{} = chain, {:error, %LangChainError{type: "overloaded"}}) do
+    cancel_delta(chain, :cancelled)
+  end
+
   @doc """
   Convert any hanging delta of the chain to a message and append to the chain.
 

test/chains/llm_chain_test.exs

@@ -472,6 +472,27 @@ defmodule LangChain.Chains.LLMChainTest do
       assert tool_call.arguments == %{"expression" => "100 + 300 - 200"}
       assert updated_chain.messages == [last]
     end
+
+    test "cancels the current delta when applying an overloaded error", %{chain: chain} do
+      assert chain.messages == []
+
+      updated_chain =
+        chain
+        |> LLMChain.apply_delta(
+          MessageDelta.new!(%{role: :assistant, content: "Greetings from "})
+        )
+        |> LLMChain.apply_delta(MessageDelta.new!(%{content: "your "}))
+        |> LLMChain.apply_delta(MessageDelta.new!(%{content: "favorite "}))
+        |> LLMChain.apply_delta({:error, LangChainError.exception(type: "overloaded", message: "Overloaded")})
+
+      # the delta is complete and removed from the chain
+      assert updated_chain.delta == nil
+      # the delta is converted to a message and applied to the messages
+      assert [%Message{} = new_message] = updated_chain.messages
+      assert new_message.role == :assistant
+      assert new_message.content == "Greetings from your favorite "
+      assert new_message.status == :cancelled
+    end
   end
 
   describe "apply_deltas/2" do