Fix failing test case

ag2ai · Jan 21, 2025 · 7a4232a · 7a4232a
1 parent a0f9e97
commit 7a4232a
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 4 deletions.
diff --git a/autogen/token_count_utils.py b/autogen/token_count_utils.py
@@ -11,16 +11,14 @@
 
 import tiktoken
 
-# from .agentchat.contrib.img_utils import num_tokens_from_gpt_image
+from .agentchat.contrib.img_utils import num_tokens_from_gpt_image
 from .import_utils import optional_import_block
 
 # if PIL is not imported, we will redefine num_tokens_from_gpt_image to return 0 tokens for images
 # Otherwise, it would raise an ImportError
 with optional_import_block() as result:
     import PIL  # noqa: F401
 
-    from .agentchat.contrib.img_utils import num_tokens_from_gpt_image
-
 pil_imported = result.is_successful
 if not pil_imported:
 

diff --git a/test/test_token_count.py b/test/test_token_count.py
@@ -11,7 +11,7 @@
 from autogen.import_utils import optional_import_block
 
 with optional_import_block() as result:
-    from autogen.agentchat.contrib.img_utils import num_tokens_from_gpt_image  # noqa: F401
+    from PIL import Image  # noqa: F401
 
 img_util_imported = result.is_successful