Merge pull request #38 from plasma-umass/ignore-litellm-warnings

Ignore any warnings LiteLLM might emit on import
plasma-umass · Apr 8, 2024 · 646d31e · 646d31e
2 parents f47eb65 + 7b17579
commit 646d31e
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 2 deletions.
diff --git a/src/chatdbg/assistant/assistant.py b/src/chatdbg/assistant/assistant.py
@@ -4,7 +4,12 @@
 import time
 import pprint
 
-import litellm
+import warnings
+
+with warnings.catch_warnings():
+    warnings.simplefilter("ignore")
+    import litellm
+
 import openai
 
 from ..util.trim import sandwich_tokens, trim_messages

diff --git a/src/chatdbg/util/trim.py b/src/chatdbg/util/trim.py
@@ -1,6 +1,10 @@
 
 import copy
-import litellm
+import warnings
+
+with warnings.catch_warnings():
+    warnings.simplefilter("ignore")
+    import litellm
 
 def sandwich_tokens(
     text: str, model: str, max_tokens: int = 1024, top_proportion: float = 0.5