Fix : litellm import error

This commit is contained in:
Sameer Kankute 2026-01-10 01:03:50 +05:30
parent 777ae4f530
commit aba7dcea9c
2 changed files with 15 additions and 9 deletions

View file

@ -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)}",
)

View file

@ -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