Add retry mechanism to LLMFactory and BaseLLM - #827
Conversation
Related to #651 Add retry mechanism to `LLMFactory` and `LiteLLM` classes to increase robustness. * **LLMFactory**: - Add `max_retry` parameter to `LLMFactory.build` method in `promptulate/llms/factory.py`. - Pass `max_retry` parameter to `LiteLLM` constructor. * **LiteLLM**: - Add `max_retry` attribute to `LiteLLM` class in `promptulate/llms/_litellm.py`. - Modify `_predict` method to include retry logic. * **AIChat**: - Add `max_retry` parameter to `AIChat` constructor in `promptulate/chat.py`. - Pass `max_retry` parameter to `LLMFactory.build` method. * **Tests**: - Add tests in `tests/llms/test_factory.py` to verify retry mechanism. - Mock `LiteLLM` to simulate failure and retry behavior. --- For more details, open the [Copilot Workspace session](https://copilot-workspace.githubnext.com/Undertone0809/promptulate/issues/651?shareId=XXXX-XXXX-XXXX-XXXX).
WalkthroughThe recent changes introduce a Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant AIChat
participant LiteLLM
User->>AIChat: Initialize chat with max_retry
AIChat->>LiteLLM: Create instance with max_retry
LiteLLM->>LiteLLM: Call _predict with retries
LiteLLM-->>LiteLLM: If success, return result
LiteLLM-->>AIChat: Return prediction result
AIChat-->>User: Provide response
Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media? TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (invoked as PR comments)
Additionally, you can add CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 8
Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Files selected for processing (4)
- promptulate/chat.py (7 hunks)
- promptulate/llms/_litellm.py (1 hunks)
- promptulate/llms/factory.py (2 hunks)
- tests/llms/test_factory.py (2 hunks)
Additional context used
Ruff
promptulate/llms/factory.py
10-10: Line too long (108 > 88)
(E501)
20-20: Line too long (102 > 88)
(E501)
24-24: Line too long (104 > 88)
(E501)
28-28: Line too long (98 > 88)
(E501)
tests/llms/test_factory.py
65-65: Line too long (101 > 88)
(E501)
73-73: Line too long (101 > 88)
(E501)
promptulate/llms/_litellm.py
34-34: Line too long (102 > 88)
(E501)
69-69: Line too long (94 > 88)
(E501)
Additional comments not posted (9)
promptulate/chat.py (9)
47-47: Parameter Addition: max_retryThe
max_retryparameter has been added to the_get_llmfunction, which enhances the function's flexibility by allowing users to specify the maximum number of retries.
55-55: Docstring Update: max_retryThe docstring has been updated to include the
max_retryparameter, ensuring that the function documentation is accurate and complete.
86-86: Parameter Addition: max_retryThe
max_retryparameter has been added to theAIChat.__init__method, which enhances the class's flexibility by allowing users to specify the maximum number of retries.
98-98: Docstring Update: max_retryThe docstring has been updated to include the
max_retryparameter, ensuring that the class documentation is accurate and complete.
214-214: Parameter Addition: max_retryThe
max_retryparameter has been added to thechatfunction, which enhances the function's flexibility by allowing users to specify the maximum number of retries.
234-234: Docstring Update: max_retryThe docstring has been updated to include the
max_retryparameter, ensuring that the function documentation is accurate and complete.
74-74: Verify max_retry usage in LLMFactory.buildEnsure that the
max_retryparameter is correctly utilized in theLLMFactory.buildmethod.Verification successful
The
max_retryparameter is correctly utilized in theLLMFactory.buildmethod.
- The parameter is passed to the constructors of
ZhiPu,QianFan, andLiteLLMclasses within thebuildmethod.Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash # Description: Verify that the `max_retry` parameter is correctly utilized in the `LLMFactory.build` method. # Test: Search for the `LLMFactory.build` method implementation. Expect: Usage of `max_retry` parameter. rg --type python 'def build' -A 10Length of output: 67
Script:
#!/bin/bash # Correcting the previous command to search for the `build` method implementation in Python files and check the usage of `max_retry`. # Search for the `build` method in Python files within the `LLMFactory` class. rg 'def build' -t py -A 20 # Use `ast-grep` to locate the exact implementation of the `build` method in the `LLMFactory` class. ast-grep --lang python --pattern $'class LLMFactory:\n $$$\n def build($_, $_, model_config, max_retry):\n $$$\n $$$'Length of output: 3294
100-100: Verify max_retry usage in _get_llmEnsure that the
max_retryparameter is correctly utilized in the_get_llmfunction.Verification successful
max_retryis correctly utilized in the_get_llmfunction.
- The
max_retryparameter is passed to theLLMFactory.buildmethod within the_get_llmfunction.Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash # Description: Verify that the `max_retry` parameter is correctly utilized in the `_get_llm` function. # Test: Search for the `_get_llm` function implementation. Expect: Usage of `max_retry` parameter. rg --type python 'def _get_llm' -A 10Length of output: 70
Script:
#!/bin/bash # Description: Verify that the `max_retry` parameter is correctly utilized in the `_get_llm` function. # Test: Search for the `_get_llm` function implementation. Expect: Usage of `max_retry` parameter. rg --type py 'def _get_llm' -A 20Length of output: 1022
Script:
#!/bin/bash # Description: Verify that the `max_retry` parameter is correctly utilized in the `_get_llm` function. # Test: Search for the `_get_llm` function implementation. Expect: Usage of `max_retry` parameter. rg --type py 'def _get_llm' -A 50Length of output: 2579
249-249: Verify max_retry usage in AIChat constructorEnsure that the
max_retryparameter is correctly utilized in theAIChatconstructor.Verification successful
The
max_retryparameter is correctly utilized in theAIChatconstructor.
- The
max_retryparameter is passed from theAIChatconstructor to the_get_llmfunction.- The
_get_llmfunction then passesmax_retryto theLLMFactory.buildmethod, which uses it to initialize various LLM instances.Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash # Description: Verify that the `max_retry` parameter is correctly utilized in the `AIChat` constructor. # Test: Search for the `AIChat` class implementation. Expect: Usage of `max_retry` parameter. rg --type python 'class AIChat' -A 10Length of output: 70
Script:
#!/bin/bash # Description: Verify that the `max_retry` parameter is correctly utilized in the `AIChat` constructor. # Step 1: Locate the `AIChat` class definition. rg 'class AIChat' -A 20 # Step 2: Locate the constructor of the `AIChat` class and check for `max_retry` usage. rg 'def __init__' -A 20Length of output: 72708
Script:
#!/bin/bash # Description: Verify that the `max_retry` parameter is correctly utilized in the `AIChat` class. # Step 1: Locate the `AIChat` class definition and its methods. rg 'class AIChat' -A 50 # Step 2: Search for `max_retry` usage within the `promptulate/chat.py` file. rg 'max_retry' promptulate/chat.pyLength of output: 3571
Script:
#!/bin/bash # Description: Verify that the `max_retry` parameter is correctly utilized in the `_get_llm` function and `LLMFactory.build` method. # Step 1: Locate the `_get_llm` function definition. rg 'def _get_llm' -A 20 # Step 2: Locate the `LLMFactory.build` method definition. rg 'def build' -A 20Length of output: 4180
| from promptulate.llms import ZhiPu | ||
|
|
||
| return ZhiPu(model=model_id, model_config=model_config, **kwargs) | ||
| return ZhiPu(model=model_id, model_config=model_config, max_retry=max_retry, **kwargs) |
There was a problem hiding this comment.
Refactor long line.
The line is too long (102 > 88 characters). Consider breaking it into multiple lines for better readability.
- return ZhiPu(model=model_id, model_config=model_config, max_retry=max_retry, **kwargs)
+ return ZhiPu(
+ model=model_id,
+ model_config=model_config,
+ max_retry=max_retry,
+ **kwargs
+ )Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| return ZhiPu(model=model_id, model_config=model_config, max_retry=max_retry, **kwargs) | |
| return ZhiPu( | |
| model=model_id, | |
| model_config=model_config, | |
| max_retry=max_retry, | |
| **kwargs | |
| ) |
Tools
Ruff
20-20: Line too long (102 > 88)
(E501)
| pass | ||
|
|
||
| return LiteLLM(model=model_name, model_config=model_config, **kwargs) | ||
| return LiteLLM(model=model_name, model_config=model_config, max_retry=max_retry, **kwargs) |
There was a problem hiding this comment.
Refactor long line.
The line is too long (98 > 88 characters). Consider breaking it into multiple lines for better readability.
- return LiteLLM(model=model_name, model_config=model_config, max_retry=max_retry, **kwargs)
+ return LiteLLM(
+ model=model_name,
+ model_config=model_config,
+ max_retry=max_retry,
+ **kwargs
+ )Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| return LiteLLM(model=model_name, model_config=model_config, max_retry=max_retry, **kwargs) | |
| return LiteLLM( | |
| model=model_name, | |
| model_config=model_config, | |
| max_retry=max_retry, | |
| **kwargs | |
| ) |
Tools
Ruff
28-28: Line too long (98 > 88)
(E501)
| @classmethod | ||
| def build( | ||
| cls, model_name: str, *, model_config: Optional[Dict[str, Any]] = None, **kwargs | ||
| cls, model_name: str, *, model_config: Optional[Dict[str, Any]] = None, max_retry: int = 5, **kwargs |
There was a problem hiding this comment.
Refactor long line.
The line is too long (108 > 88 characters). Consider breaking it into multiple lines for better readability.
- cls, model_name: str, *, model_config: Optional[Dict[str, Any]] = None, max_retry: int = 5, **kwargs
+ cls,
+ model_name: str,
+ *,
+ model_config: Optional[Dict[str, Any]] = None,
+ max_retry: int = 5,
+ **kwargsCommittable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| cls, model_name: str, *, model_config: Optional[Dict[str, Any]] = None, max_retry: int = 5, **kwargs | |
| cls, | |
| model_name: str, | |
| *, | |
| model_config: Optional[Dict[str, Any]] = None, | |
| max_retry: int = 5, | |
| **kwargs |
Tools
Ruff
10-10: Line too long (108 > 88)
(E501)
| from promptulate.llms import QianFan | ||
|
|
||
| return QianFan(model=model_id, model_config=model_config, **kwargs) | ||
| return QianFan(model=model_id, model_config=model_config, max_retry=max_retry, **kwargs) |
There was a problem hiding this comment.
Refactor long line.
The line is too long (104 > 88 characters). Consider breaking it into multiple lines for better readability.
- return QianFan(model=model_id, model_config=model_config, max_retry=max_retry, **kwargs)
+ return QianFan(
+ model=model_id,
+ model_config=model_config,
+ max_retry=max_retry,
+ **kwargs
+ )Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| return QianFan(model=model_id, model_config=model_config, max_retry=max_retry, **kwargs) | |
| return QianFan( | |
| model=model_id, | |
| model_config=model_config, | |
| max_retry=max_retry, | |
| **kwargs | |
| ) |
Tools
Ruff
24-24: Line too long (104 > 88)
(E501)
|
|
||
|
|
||
| def test_litellm_retry(): | ||
| with mock.patch.object(LiteLLM, "_predict", side_effect=Exception("Test error")) as mock_predict: |
There was a problem hiding this comment.
Refactor long line.
The line is too long (101 > 88 characters). Consider breaking it into multiple lines for better readability.
- with mock.patch.object(LiteLLM, "_predict", side_effect=Exception("Test error")) as mock_predict:
+ with mock.patch.object(
+ LiteLLM,
+ "_predict",
+ side_effect=Exception("Test error")
+ ) as mock_predict:Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| with mock.patch.object(LiteLLM, "_predict", side_effect=Exception("Test error")) as mock_predict: | |
| with mock.patch.object( | |
| LiteLLM, | |
| "_predict", | |
| side_effect=Exception("Test error") | |
| ) as mock_predict: |
Tools
Ruff
73-73: Line too long (101 > 88)
(E501)
|
|
||
|
|
||
| def test_retry_mechanism(): | ||
| with mock.patch.object(LiteLLM, "_predict", side_effect=Exception("Test error")) as mock_predict: |
There was a problem hiding this comment.
Refactor long line.
The line is too long (101 > 88 characters). Consider breaking it into multiple lines for better readability.
- with mock.patch.object(LiteLLM, "_predict", side_effect=Exception("Test error")) as mock_predict:
+ with mock.patch.object(
+ LiteLLM,
+ "_predict",
+ side_effect=Exception("Test error")
+ ) as mock_predict:Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| with mock.patch.object(LiteLLM, "_predict", side_effect=Exception("Test error")) as mock_predict: | |
| with mock.patch.object( | |
| LiteLLM, | |
| "_predict", | |
| side_effect=Exception("Test error") | |
| ) as mock_predict: |
Tools
Ruff
65-65: Line too long (101 > 88)
(E501)
| else json.loads(temp_response.json()), | ||
| ) | ||
| logger.debug( | ||
| f"[pne chat] response: {json.dumps(response.additional_kwargs, indent=2)}" |
There was a problem hiding this comment.
Refactor long line.
The line is too long (94 > 88 characters). Consider breaking it into multiple lines for better readability.
- f"[pne chat] response: {json.dumps(response.additional_kwargs, indent=2)}"
+ f"[pne chat] response: {json.dumps(
+ response.additional_kwargs,
+ indent=2
+ )}"Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| f"[pne chat] response: {json.dumps(response.additional_kwargs, indent=2)}" | |
| f"[pne chat] response: {json.dumps( | |
| response.additional_kwargs, | |
| indent=2 | |
| )}" |
Tools
Ruff
69-69: Line too long (94 > 88)
(E501)
|
|
||
| class LiteLLM(BaseLLM): | ||
| def __init__(self, model: str, model_config: Optional[dict] = None, **kwargs): | ||
| def __init__(self, model: str, model_config: Optional[dict] = None, max_retry: int = 5, **kwargs): |
There was a problem hiding this comment.
Refactor long line.
The line is too long (102 > 88 characters). Consider breaking it into multiple lines for better readability.
- def __init__(self, model: str, model_config: Optional[dict] = None, max_retry: int = 5, **kwargs):
+ def __init__(
+ self,
+ model: str,
+ model_config: Optional[dict] = None,
+ max_retry: int = 5,
+ **kwargs
+ ):Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def __init__(self, model: str, model_config: Optional[dict] = None, max_retry: int = 5, **kwargs): | |
| def __init__( | |
| self, | |
| model: str, | |
| model_config: Optional[dict] = None, | |
| max_retry: int = 5, | |
| **kwargs | |
| ): |
Tools
Ruff
34-34: Line too long (102 > 88)
(E501)
1 participant
Related to #651
Add retry mechanism to
LLMFactoryandLiteLLMclasses to increase robustness.LLMFactory:
max_retryparameter toLLMFactory.buildmethod inpromptulate/llms/factory.py.max_retryparameter toLiteLLMconstructor.LiteLLM:
max_retryattribute toLiteLLMclass inpromptulate/llms/_litellm.py._predictmethod to include retry logic.AIChat:
max_retryparameter toAIChatconstructor inpromptulate/chat.py.max_retryparameter toLLMFactory.buildmethod.Tests:
tests/llms/test_factory.pyto verify retry mechanism.LiteLLMto simulate failure and retry behavior.For more details, open the Copilot Workspace session.
Summary by CodeRabbit
New Features
max_retryparameter across various functionalities, enhancing user control over retry attempts during LLM initialization and chat operations.LiteLLMprediction process with a retry mechanism.Bug Fixes
Tests
LiteLLMmodel, ensuring robust error handling.