From f049adbf6938d119917fc96180f97b4679fbf449 Mon Sep 17 00:00:00 2001
From: EarningsCall <dev@earningscall.biz>
Date: Mon, 27 Jan 2025 09:47:30 -0600
Subject: [PATCH] Add customized retry strategies for different plan types

* Implement plan-specific retry strategies with varying base delays and max attempts
* Update default retry strategy to use 0.5s base delay and 3 max attempts
* Add function to dynamically select retry strategy based on API key
* Update README and version to reflect new retry configuration
---
 CHANGELOG.md        |  5 ++++
 README.md           |  4 +--
 earningscall/api.py | 62 ++++++++++++++++++++++++++++++++++++++++++---
 pyproject.toml      |  2 +-
 4 files changed, 67 insertions(+), 6 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7920850..c8cf14e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,10 @@
 # Changelog
 
+## Release `1.2.0` - 2025-01-27
+
+* Add customized retry strategy for different plan types.
+* Update default retry strategy to use 0.5s base delay and 5 max attempts.
+
 ## Release `1.1.1` - 2025-01-26
 
 * Modify default retry strategy to use 1s base delay and 10 max attempts (necessary for starter plan).
diff --git a/README.md b/README.md
index b876311..08696a5 100644
--- a/README.md
+++ b/README.md
@@ -284,8 +284,8 @@ import earningscall
 
 earningscall.retry_strategy = {
     "strategy": "exponential",
-    "base_delay": 1,
-    "max_attempts": 10,
+    "base_delay": 0.5,
+    "max_attempts": 3,
 }
 ```
 
diff --git a/earningscall/api.py b/earningscall/api.py
index 746c907..f466219 100644
--- a/earningscall/api.py
+++ b/earningscall/api.py
@@ -19,8 +19,50 @@
 API_BASE = f"https://v2.api.{DOMAIN}"
 DEFAULT_RETRY_STRATEGY: Dict[str, Union[str, int, float]] = {
     "strategy": "exponential",
-    "base_delay": 1,
-    "max_attempts": 10,
+    "base_delay": 0.5,
+    "max_attempts": 3,
+}
+RETRY_STRATEGIES_BY_PLAN = {
+    "starter": {
+        "strategy": "exponential",
+        "base_delay": 1,
+        "max_attempts": 10,
+    },
+    "premium": {
+        "strategy": "exponential",
+        "base_delay": 1,
+        "max_attempts": 5,
+    },
+    "ultimate": {
+        "strategy": "exponential",
+        "base_delay": 0.001,
+        "max_attempts": 5,
+    },
+    "ultimate_plus": {
+        "strategy": "exponential",
+        "base_delay": 0.0001,
+        "max_attempts": 5,
+    },
+    "startup": {
+        "strategy": "exponential",
+        "base_delay": 0.00001,
+        "max_attempts": 5,
+    },
+    "seed": {
+        "strategy": "exponential",
+        "base_delay": 0.000001,
+        "max_attempts": 5,
+    },
+    "enterprise": {
+        "strategy": "exponential",
+        "base_delay": 0.1,
+        "max_attempts": 5,
+    },
+    "enterprise_plus": {
+        "strategy": "exponential",
+        "base_delay": 0.1,
+        "max_attempts": 5,
+    },
 }
 
 
@@ -102,6 +144,20 @@ def is_success(response: requests.Response) -> bool:
     return response.status_code == 200
 
 
+def get_retry_strategy_from_api_key(api_key: str) -> Dict[str, Union[str, int, float]]:
+    for key, value in RETRY_STRATEGIES_BY_PLAN.items():
+        if api_key.startswith(key):
+            return value
+    return DEFAULT_RETRY_STRATEGY
+
+
+def get_retry_strategy(api_key: str) -> Dict[str, Union[str, int, float]]:
+    if earningscall.retry_strategy:
+        return earningscall.retry_strategy
+    # Customize retry strategy for different plan types
+    return get_retry_strategy_from_api_key(get_api_key())
+
+
 def do_get(
     path: str,
     use_cache: bool = False,
@@ -127,7 +183,7 @@ def do_get(
         full_url = f"{url}?{urllib.parse.urlencode(params)}"
         log.debug(f"GET: {full_url}")
 
-    retry_strategy = earningscall.retry_strategy or DEFAULT_RETRY_STRATEGY
+    retry_strategy = get_retry_strategy(get_api_key())
     delay = retry_strategy["base_delay"]
     max_attempts = int(retry_strategy["max_attempts"])
 
diff --git a/pyproject.toml b/pyproject.toml
index 4ed2ef1..324cf39 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "earningscall"
-version = "1.1.1"
+version = "1.2.0"
 description = "The EarningsCall Python library provides convenient access to the EarningsCall API.  It includes a pre-defined set of classes for API resources that initialize themselves dynamically from API responses."
 readme = "README.md"
 authors = [{ name = "EarningsCall", email = "dev@earningscall.biz" }]