Skip to content

Commit

Permalink
custom rate limits added
Browse files Browse the repository at this point in the history
  • Loading branch information
staru09 committed Nov 5, 2024
1 parent 169b601 commit bed0b86
Showing 1 changed file with 39 additions and 0 deletions.
39 changes: 39 additions & 0 deletions docetl/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -134,6 +134,45 @@ def custom_parser(text: str) -> List[str]:
This example shows a complete pipeline configuration with datasets, operations,
steps, and output settings.
"""
DEFAULT_RATE_LIMITS = {
# OpenAI models
"gpt-4o": 1000,
"gpt-4o-mini": 200,
"gpt-3.5-turbo": 500,

# Anthropic models
"claude 3.5-sonnet": 1000,
"claude-3-opus": 500,
"claude-3-sonnet": 400,
"claude-3-haiku": 200,
}

def get_rate_limits(self, model: str) -> dict:
"""Get rate limits for a specific model.
Args:
model: The model identifier (e.g., 'gpt-4o', 'claude-3-sonnet')
Returns:
dict: Rate limit information including requests_per_minute
"""
if self.rate_limits and model in self.rate_limits:
return {
"requests_per_minute": self.rate_limits[model],
"source": "custom"
}

if model in self.DEFAULT_RATE_LIMITS:
return {
"requests_per_minute": self.DEFAULT_RATE_LIMITS[model],
"source": "default"
}

return {
"requests_per_minute": 200,
"source": "fallback"
}


def __init__(
self,
Expand Down

0 comments on commit bed0b86

Please sign in to comment.