From aba7dcea9ca13fdc18c867bcddf343d5f0291ab5 Mon Sep 17 00:00:00 2001
From: Sameer Kankute <sameer@berri.ai>
Date: Sat, 10 Jan 2026 01:03:50 +0530
Subject: [PATCH] Fix : litellm import error

---
 litellm/llms/bedrock/count_tokens/handler.py    | 17 ++++++++---------
 .../llm_passthrough_endpoints.py                |  7 +++++++
 2 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/litellm/llms/bedrock/count_tokens/handler.py b/litellm/llms/bedrock/count_tokens/handler.py
index 60ace7f3369..e8366165b65 100644
--- a/litellm/llms/bedrock/count_tokens/handler.py
+++ b/litellm/llms/bedrock/count_tokens/handler.py
@@ -6,10 +6,9 @@ Simplified handler leveraging existing LiteLLM Bedrock infrastructure.
 
 from typing import Any, Dict
 
-from fastapi import HTTPException
-
 import litellm
 from litellm._logging import verbose_logger
+from litellm.llms.bedrock.common_utils import BedrockError
 from litellm.llms.bedrock.count_tokens.transformation import BedrockCountTokensConfig
 from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
 
@@ -97,9 +96,9 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig):
             if response.status_code != 200:
                 error_text = response.text
                 verbose_logger.error(f"AWS Bedrock error: {error_text}")
-                raise HTTPException(
-                    status_code=400,
-                    detail={"error": f"AWS Bedrock error: {error_text}"},
+                raise BedrockError(
+                    status_code=response.status_code,
+                    message=f"AWS Bedrock error: {error_text}",
                 )
 
             bedrock_response = response.json()
@@ -115,12 +114,12 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig):
 
             return final_response
 
-        except HTTPException:
-            # Re-raise HTTP exceptions as-is
+        except BedrockError:
+            # Re-raise Bedrock exceptions as-is
             raise
         except Exception as e:
             verbose_logger.error(f"Error in CountTokens handler: {str(e)}")
-            raise HTTPException(
+            raise BedrockError(
                 status_code=500,
-                detail={"error": f"CountTokens processing error: {str(e)}"},
+                message=f"CountTokens processing error: {str(e)}",
             )
diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
index 84550092d2e..d9798dae690 100644
--- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
+++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
@@ -776,6 +776,7 @@ async def handle_bedrock_count_tokens(
     - /v1/messages/count_tokens
     - /v1/messages/count-tokens
     """
+    from litellm.llms.bedrock.common_utils import BedrockError
     from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler
     from litellm.proxy.proxy_server import llm_router
 
@@ -822,6 +823,12 @@ async def handle_bedrock_count_tokens(
 
         return result
 
+    except BedrockError as e:
+        # Convert BedrockError to HTTPException for FastAPI
+        verbose_proxy_logger.error(f"BedrockError in handle_bedrock_count_tokens: {str(e)}")
+        raise HTTPException(
+            status_code=e.status_code, detail={"error": e.message}
+        )
     except HTTPException:
         # Re-raise HTTP exceptions as-is
         raise