From a3284c9d7e66a4b3f42997c7e40f5fa55eee8d28 Mon Sep 17 00:00:00 2001
From: Ariel Frischer <arielfrischer@gmail.com>
Date: Tue, 11 Mar 2025 23:37:20 -0700
Subject: [PATCH] feat(anthropic_token_limiter.py): add dataclass import for
 future use and improve code readability by restructuring import statements

---
 ra_aid/anthropic_token_limiter.py | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/ra_aid/anthropic_token_limiter.py b/ra_aid/anthropic_token_limiter.py
index c46cbc9..ae82ab2 100644
--- a/ra_aid/anthropic_token_limiter.py
+++ b/ra_aid/anthropic_token_limiter.py
@@ -2,9 +2,16 @@
 
 from functools import partial
 from typing import Any, Dict, List, Optional, Sequence
+from dataclasses import dataclass
 
 from langchain_anthropic import ChatAnthropic
-from langchain_core.messages import AIMessage, BaseMessage, ToolMessage, trim_messages
+from langchain_core.messages import (
+    AIMessage,
+    BaseMessage,
+    RemoveMessage,
+    ToolMessage,
+    trim_messages,
+)
 from langchain_core.messages.base import message_to_dict
 
 from ra_aid.anthropic_message_utils import (
@@ -143,6 +150,7 @@ def state_modifier(
         print_messages_compact(messages)
         print("AFTER TRIMMING")
         print_messages_compact(result)
+
     return result