chore: initial public snapshot for github upload
This commit is contained in:
Your Name
@@ -0,0 +1,15 @@
|
||||
"""
|
||||
Anthropic CountTokens API implementation.
|
||||
"""
|
||||
|
||||
from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler
|
||||
from litellm.llms.anthropic.count_tokens.token_counter import AnthropicTokenCounter
|
||||
from litellm.llms.anthropic.count_tokens.transformation import (
|
||||
AnthropicCountTokensConfig,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"AnthropicCountTokensHandler",
|
||||
"AnthropicCountTokensConfig",
|
||||
"AnthropicTokenCounter",
|
||||
]
|
||||
@@ -0,0 +1,128 @@
|
||||
"""
|
||||
Anthropic CountTokens API handler.
|
||||
|
||||
Uses httpx for HTTP requests instead of the Anthropic SDK.
|
||||
"""
|
||||
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
|
||||
import httpx
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.llms.anthropic.common_utils import AnthropicError
|
||||
from litellm.llms.anthropic.count_tokens.transformation import (
|
||||
AnthropicCountTokensConfig,
|
||||
)
|
||||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
|
||||
|
||||
class AnthropicCountTokensHandler(AnthropicCountTokensConfig):
|
||||
"""
|
||||
Handler for Anthropic CountTokens API requests.
|
||||
|
||||
Uses httpx for HTTP requests, following the same pattern as BedrockCountTokensHandler.
|
||||
"""
|
||||
|
||||
async def handle_count_tokens_request(
|
||||
self,
|
||||
model: str,
|
||||
messages: List[Dict[str, Any]],
|
||||
api_key: str,
|
||||
api_base: Optional[str] = None,
|
||||
timeout: Optional[Union[float, httpx.Timeout]] = None,
|
||||
tools: Optional[List[Dict[str, Any]]] = None,
|
||||
system: Optional[Any] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Handle a CountTokens request using httpx.
|
||||
|
||||
Args:
|
||||
model: The model identifier (e.g., "claude-3-5-sonnet-20241022")
|
||||
messages: The messages to count tokens for
|
||||
api_key: The Anthropic API key
|
||||
api_base: Optional custom API base URL
|
||||
timeout: Optional timeout for the request (defaults to litellm.request_timeout)
|
||||
|
||||
Returns:
|
||||
Dictionary containing token count response
|
||||
|
||||
Raises:
|
||||
AnthropicError: If the API request fails
|
||||
"""
|
||||
try:
|
||||
# Validate the request
|
||||
self.validate_request(model, messages)
|
||||
|
||||
verbose_logger.debug(
|
||||
f"Processing Anthropic CountTokens request for model: {model}"
|
||||
)
|
||||
|
||||
# Transform request to Anthropic format
|
||||
request_body = self.transform_request_to_count_tokens(
|
||||
model=model,
|
||||
messages=messages,
|
||||
tools=tools,
|
||||
system=system,
|
||||
)
|
||||
|
||||
verbose_logger.debug(f"Transformed request: {request_body}")
|
||||
|
||||
# Get endpoint URL
|
||||
endpoint_url = api_base or self.get_anthropic_count_tokens_endpoint()
|
||||
|
||||
verbose_logger.debug(f"Making request to: {endpoint_url}")
|
||||
|
||||
# Get required headers
|
||||
headers = self.get_required_headers(api_key)
|
||||
|
||||
# Use LiteLLM's async httpx client
|
||||
async_client = get_async_httpx_client(
|
||||
llm_provider=litellm.LlmProviders.ANTHROPIC
|
||||
)
|
||||
|
||||
# Use provided timeout or fall back to litellm.request_timeout
|
||||
request_timeout = (
|
||||
timeout if timeout is not None else litellm.request_timeout
|
||||
)
|
||||
|
||||
response = await async_client.post(
|
||||
endpoint_url,
|
||||
headers=headers,
|
||||
json=request_body,
|
||||
timeout=request_timeout,
|
||||
)
|
||||
|
||||
verbose_logger.debug(f"Response status: {response.status_code}")
|
||||
|
||||
if response.status_code != 200:
|
||||
error_text = response.text
|
||||
verbose_logger.error(f"Anthropic API error: {error_text}")
|
||||
raise AnthropicError(
|
||||
status_code=response.status_code,
|
||||
message=error_text,
|
||||
)
|
||||
|
||||
anthropic_response = response.json()
|
||||
|
||||
verbose_logger.debug(f"Anthropic response: {anthropic_response}")
|
||||
|
||||
# Return Anthropic response directly - no transformation needed
|
||||
return anthropic_response
|
||||
|
||||
except AnthropicError:
|
||||
# Re-raise Anthropic exceptions as-is
|
||||
raise
|
||||
except httpx.HTTPStatusError as e:
|
||||
# HTTP errors - preserve the actual status code
|
||||
verbose_logger.error(f"HTTP error in CountTokens handler: {str(e)}")
|
||||
raise AnthropicError(
|
||||
status_code=e.response.status_code,
|
||||
message=e.response.text,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.error(f"Error in CountTokens handler: {str(e)}")
|
||||
raise AnthropicError(
|
||||
status_code=500,
|
||||
message=f"CountTokens processing error: {str(e)}",
|
||||
)
|
||||
@@ -0,0 +1,108 @@
|
||||
"""
|
||||
Anthropic Token Counter implementation using the CountTokens API.
|
||||
"""
|
||||
|
||||
import os
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler
|
||||
from litellm.llms.base_llm.base_utils import BaseTokenCounter
|
||||
from litellm.types.utils import LlmProviders, TokenCountResponse
|
||||
|
||||
# Global handler instance - reuse across all token counting requests
|
||||
anthropic_count_tokens_handler = AnthropicCountTokensHandler()
|
||||
|
||||
|
||||
class AnthropicTokenCounter(BaseTokenCounter):
|
||||
"""Token counter implementation for Anthropic provider using the CountTokens API."""
|
||||
|
||||
def should_use_token_counting_api(
|
||||
self,
|
||||
custom_llm_provider: Optional[str] = None,
|
||||
) -> bool:
|
||||
return custom_llm_provider == LlmProviders.ANTHROPIC.value
|
||||
|
||||
async def count_tokens(
|
||||
self,
|
||||
model_to_use: str,
|
||||
messages: Optional[List[Dict[str, Any]]],
|
||||
contents: Optional[List[Dict[str, Any]]],
|
||||
deployment: Optional[Dict[str, Any]] = None,
|
||||
request_model: str = "",
|
||||
tools: Optional[List[Dict[str, Any]]] = None,
|
||||
system: Optional[Any] = None,
|
||||
) -> Optional[TokenCountResponse]:
|
||||
"""
|
||||
Count tokens using Anthropic's CountTokens API.
|
||||
|
||||
Args:
|
||||
model_to_use: The model identifier
|
||||
messages: The messages to count tokens for
|
||||
contents: Alternative content format (not used for Anthropic)
|
||||
deployment: Deployment configuration containing litellm_params
|
||||
request_model: The original request model name
|
||||
|
||||
Returns:
|
||||
TokenCountResponse with token count, or None if counting fails
|
||||
"""
|
||||
from litellm.llms.anthropic.common_utils import AnthropicError
|
||||
|
||||
if not messages:
|
||||
return None
|
||||
|
||||
deployment = deployment or {}
|
||||
litellm_params = deployment.get("litellm_params", {})
|
||||
|
||||
# Get Anthropic API key from deployment config or environment
|
||||
api_key = litellm_params.get("api_key")
|
||||
if not api_key:
|
||||
api_key = os.getenv("ANTHROPIC_API_KEY")
|
||||
|
||||
if not api_key:
|
||||
verbose_logger.warning("No Anthropic API key found for token counting")
|
||||
return None
|
||||
|
||||
try:
|
||||
result = await anthropic_count_tokens_handler.handle_count_tokens_request(
|
||||
model=model_to_use,
|
||||
messages=messages,
|
||||
api_key=api_key,
|
||||
tools=tools,
|
||||
system=system,
|
||||
)
|
||||
|
||||
if result is not None:
|
||||
return TokenCountResponse(
|
||||
total_tokens=result.get("input_tokens", 0),
|
||||
request_model=request_model,
|
||||
model_used=model_to_use,
|
||||
tokenizer_type="anthropic_api",
|
||||
original_response=result,
|
||||
)
|
||||
except AnthropicError as e:
|
||||
verbose_logger.warning(
|
||||
f"Anthropic CountTokens API error: status={e.status_code}, message={e.message}"
|
||||
)
|
||||
return TokenCountResponse(
|
||||
total_tokens=0,
|
||||
request_model=request_model,
|
||||
model_used=model_to_use,
|
||||
tokenizer_type="anthropic_api",
|
||||
error=True,
|
||||
error_message=e.message,
|
||||
status_code=e.status_code,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.warning(f"Error calling Anthropic CountTokens API: {e}")
|
||||
return TokenCountResponse(
|
||||
total_tokens=0,
|
||||
request_model=request_model,
|
||||
model_used=model_to_use,
|
||||
tokenizer_type="anthropic_api",
|
||||
error=True,
|
||||
error_message=str(e),
|
||||
status_code=500,
|
||||
)
|
||||
|
||||
return None
|
||||
@@ -0,0 +1,107 @@
|
||||
"""
|
||||
Anthropic CountTokens API transformation logic.
|
||||
|
||||
This module handles the transformation of requests to Anthropic's CountTokens API format.
|
||||
"""
|
||||
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from litellm.constants import ANTHROPIC_TOKEN_COUNTING_BETA_VERSION
|
||||
|
||||
|
||||
class AnthropicCountTokensConfig:
|
||||
"""
|
||||
Configuration and transformation logic for Anthropic CountTokens API.
|
||||
|
||||
Anthropic CountTokens API Specification:
|
||||
- Endpoint: POST https://api.anthropic.com/v1/messages/count_tokens
|
||||
- Beta header required: anthropic-beta: token-counting-2024-11-01
|
||||
- Response: {"input_tokens": <number>}
|
||||
"""
|
||||
|
||||
def get_anthropic_count_tokens_endpoint(self) -> str:
|
||||
"""
|
||||
Get the Anthropic CountTokens API endpoint.
|
||||
|
||||
Returns:
|
||||
The endpoint URL for the CountTokens API
|
||||
"""
|
||||
return "https://api.anthropic.com/v1/messages/count_tokens"
|
||||
|
||||
def transform_request_to_count_tokens(
|
||||
self,
|
||||
model: str,
|
||||
messages: List[Dict[str, Any]],
|
||||
tools: Optional[List[Dict[str, Any]]] = None,
|
||||
system: Optional[Any] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Transform request to Anthropic CountTokens format.
|
||||
|
||||
Includes optional system and tools fields for accurate token counting.
|
||||
"""
|
||||
request: Dict[str, Any] = {
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
}
|
||||
|
||||
if system is not None:
|
||||
request["system"] = system
|
||||
|
||||
if tools is not None:
|
||||
request["tools"] = tools
|
||||
|
||||
return request
|
||||
|
||||
def get_required_headers(self, api_key: str) -> Dict[str, str]:
|
||||
"""
|
||||
Get the required headers for the CountTokens API.
|
||||
|
||||
Args:
|
||||
api_key: The Anthropic API key
|
||||
|
||||
Returns:
|
||||
Dictionary of required headers
|
||||
"""
|
||||
from litellm.llms.anthropic.common_utils import (
|
||||
optionally_handle_anthropic_oauth,
|
||||
)
|
||||
|
||||
headers: Dict[str, str] = {
|
||||
"Content-Type": "application/json",
|
||||
"x-api-key": api_key,
|
||||
"anthropic-version": "2023-06-01",
|
||||
"anthropic-beta": ANTHROPIC_TOKEN_COUNTING_BETA_VERSION,
|
||||
}
|
||||
headers, _ = optionally_handle_anthropic_oauth(headers=headers, api_key=api_key)
|
||||
return headers
|
||||
|
||||
def validate_request(self, model: str, messages: List[Dict[str, Any]]) -> None:
|
||||
"""
|
||||
Validate the incoming count tokens request.
|
||||
|
||||
Args:
|
||||
model: The model name
|
||||
messages: The messages to count tokens for
|
||||
|
||||
Raises:
|
||||
ValueError: If the request is invalid
|
||||
"""
|
||||
if not model:
|
||||
raise ValueError("model parameter is required")
|
||||
|
||||
if not messages:
|
||||
raise ValueError("messages parameter is required")
|
||||
|
||||
if not isinstance(messages, list):
|
||||
raise ValueError("messages must be a list")
|
||||
|
||||
for i, message in enumerate(messages):
|
||||
if not isinstance(message, dict):
|
||||
raise ValueError(f"Message {i} must be a dictionary")
|
||||
|
||||
if "role" not in message:
|
||||
raise ValueError(f"Message {i} must have a 'role' field")
|
||||
|
||||
if "content" not in message:
|
||||
raise ValueError(f"Message {i} must have a 'content' field")
|
||||
Reference in New Issue
Block a user