ucbepic · staru09 · Nov 5, 2024 · Nov 5, 2024 · Nov 7, 2024 · Nov 12, 2024
diff --git a/docetl/api.py b/docetl/api.py
@@ -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,

diff --git a/docetl/operations/resolve.py b/docetl/operations/resolve.py
@@ -445,8 +445,11 @@ def merge_clusters(item1: int, item2: int) -> None:
 
         # Compute an auto-batch size based on the number of comparisons
         def auto_batch() -> int:
-            # Maximum batch size limit for 4o-mini model
-            M = 500
+            # Get model-specific rate limit from pipeline
+            model = self.config.get("comparison_model", self.default_model)
+            rate_limit = self.runner.api.get_rate_limit(model)  
+            # Use the rate limit as our maximum batch size
+            M = rate_limit["requests_per_minute"]
 
             n = len(input_data)
             m = len(blocked_pairs)
@@ -468,6 +471,7 @@ def auto_batch() -> int:
 
         # Compare pairs and update clusters in real-time
         batch_size = self.config.get("compare_batch_size", auto_batch())
+        rate_info = self.runner.pipeline.get_rate_limits(self.config.get("comparison_model", self.default_model))
         self.console.log(f"Using compare batch size: {batch_size}")
         pair_costs = 0