Fixed types in token handler (#365)

road-core · Feb 9, 2025 · ef63fca · ef63fca
1 parent 089520d
commit ef63fca
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/ols/utils/token_handler.py b/ols/utils/token_handler.py
@@ -179,10 +179,10 @@ def truncate_rag_context(
 
     def limit_conversation_history(
         self, history: list[BaseMessage], model: str, limit: int = 0
-    ) -> tuple[list[str], bool]:
+    ) -> tuple[list[BaseMessage], bool]:
         """Limit conversation history to specified number of tokens."""
         total_length = 0
-        formatted_history: list[str] = []
+        formatted_history: list[BaseMessage] = []
 
         for original_message in reversed(history):
             # Restructure messages as per model