From c453493ad35840a5dd36c2a4bf928e32b75671db Mon Sep 17 00:00:00 2001
From: Mark Ericksen <brainlid@gmail.com>
Date: Sun, 15 Dec 2024 15:23:08 -0700
Subject: [PATCH] summarized chain's last_message get updated - changes when
 the keep_count is 0

---
 lib/chains/summarize_conversation_chain.ex    |  8 ++++-
 .../summarize_conversation_chain_test.exs     | 30 ++++++++++++++-----
 2 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/lib/chains/summarize_conversation_chain.ex b/lib/chains/summarize_conversation_chain.ex
index 488ea3c..ab634db 100644
--- a/lib/chains/summarize_conversation_chain.ex
+++ b/lib/chains/summarize_conversation_chain.ex
@@ -322,7 +322,13 @@ defmodule LangChain.Chains.SummarizeConversationChain do
           [system | summary_messages]
       end
 
-    %LLMChain{to_summarize | messages: starting_messages ++ keeping_items}
+    new_messages = starting_messages ++ keeping_items
+
+    %LLMChain{
+      to_summarize
+      | messages: new_messages,
+        last_message: List.last(new_messages)
+    }
   end
 
   # Convert each `%Message{}` into a text message like `<user>The user message
diff --git a/test/chains/summarize_conversation_chain_test.exs b/test/chains/summarize_conversation_chain_test.exs
index 8d6c546..1ae1526 100644
--- a/test/chains/summarize_conversation_chain_test.exs
+++ b/test/chains/summarize_conversation_chain_test.exs
@@ -296,7 +296,24 @@ defmodule LangChain.Chains.SummarizeConversationChainTest do
       assert keep_2.role == :assistant
     end
 
-    @tag live_call: true, live_anthropic: true
+    test "set last_message correctly when keep_count is 0", %{llm_anthropic: llm, chain: chain} do
+      # Made NOT LIVE here
+      expect(ChatAnthropic, :call, fn _model, _messages, _tools ->
+        {:ok, Message.new_assistant!("- Fake OpenAI summary")}
+      end)
+
+      summarizer = SummarizeConversationChain.new!(%{llm: llm, threshold_count: 6, keep_count: 0})
+      chain = LLMChain.add_messages(chain, get_full_conversation())
+
+      summarized_chain = SummarizeConversationChain.summarize(summarizer, chain)
+
+      [_system, _summary_1, summary_2] = summarized_chain.messages
+      assert summary_2.role == :assistant
+      assert summary_2.content == "- Fake OpenAI summary"
+
+      assert summarized_chain.last_message == summary_2
+    end
+
     test "returns unmodified chain when LLM operation fails", %{
       llm_anthropic: llm,
       chain: chain
@@ -366,12 +383,11 @@ defmodule LangChain.Chains.SummarizeConversationChainTest do
 
       summarized_chain = SummarizeConversationChain.summarize(summarizer, original_chain)
 
-      IO.inspect(summarized_chain.messages)
-
-
-      # TODO:
+      # IO.inspect(summarized_chain.messages)
+      # IO.inspect(summarized_chain.last_message)
 
-      assert false
+      assert summarized_chain.last_message.role == :assistant
+      assert String.starts_with?(summarized_chain.last_message.content, "- ")
     end
   end
 
@@ -541,7 +557,7 @@ Florence: Hotel Davanzati - near the Ponte Vecchio
 Amalfi: Hotel Marina Riviera - beautiful sea views
 
 Would you like me to check current availability and rates for your September dates?|),
-Message.new_user!(
+      Message.new_user!(
         "That would be helpful! My dates are September 10-24. Could you also suggest some must-see attractions in each city?"
       ),
       Message.new_assistant!(