From c27f0a926b8363bdd70c84bb7c1b341ac8f00cba Mon Sep 17 00:00:00 2001
From: Beibin Li <BeibinLi@users.noreply.github.com>
Date: Mon, 29 Jul 2024 11:26:09 -0700
Subject: [PATCH] Only catch not implemented error

---
 autogen/oai/client.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/autogen/oai/client.py b/autogen/oai/client.py
index 63a0923afe61..38af18eb0afc 100644
--- a/autogen/oai/client.py
+++ b/autogen/oai/client.py
@@ -274,7 +274,7 @@ def create(self, params: Dict[str, Any]) -> ChatCompletion:
             model = chunk.model.replace("gpt-35", "gpt-3.5")  # hack for Azure API
             try:
                 prompt_tokens = count_token(params["messages"], model)
-            except Exception as e:
+            except NotImplementedError as e:
                 # Catch token calculation error if streaming with customized models.
                 logger.warning(str(e))
                 prompt_tokens = 0