add retries

run-llama · Feb 16, 2024 · 3a58cde · 3a58cde
1 parent bf79202
commit 3a58cde
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 1 deletion.
diff --git a/llama_hub/github_repo/README.md b/llama_hub/github_repo/README.md
@@ -76,6 +76,7 @@ if docs is None:
         verbose =                True,
         concurrent_requests =    10,
         timeout =                5,
+        retries =                0,
     )
 
     docs = loader.load_data(branch="main")

diff --git a/llama_hub/github_repo/base.py b/llama_hub/github_repo/base.py
@@ -73,6 +73,7 @@ def __init__(
         verbose: bool = False,
         concurrent_requests: int = 5,
         timeout: Optional[int] = 5,
+        retries: int = 0,
         filter_directories: Optional[Tuple[List[str], FilterType]] = None,
         filter_file_extensions: Optional[Tuple[List[str], FilterType]] = None,
     ):
@@ -89,6 +90,7 @@ def __init__(
             - concurrent_requests (int): Number of concurrent requests to
                 make to the Github API.
             - timeout (int or None): Timeout for the requests to the Github API. Default is 5.
+            - retries (int): Number of retries for the requests to the Github API. Default is 0.
             - filter_directories (Optional[Tuple[List[str], FilterType]]): Tuple
                 containing a list of directories and a FilterType. If the FilterType
                 is INCLUDE, only the files in the directories in the list will be
@@ -112,6 +114,7 @@ def __init__(
         self._verbose = verbose
         self._concurrent_requests = concurrent_requests
         self._timeout = timeout
+        self._retries = retries
         self._filter_directories = filter_directories
         self._filter_file_extensions = filter_file_extensions
 

diff --git a/llama_hub/github_repo/github_client.py b/llama_hub/github_repo/github_client.py
@@ -272,6 +272,7 @@ async def request(
         method: str,
         headers: Dict[str, Any] = {},
         timeout: Optional[int] = 5,
+        retries: int = 0,
         **kwargs: Any,
     ) -> Any:
         """
@@ -285,6 +286,7 @@ async def request(
             - `method (str)`: HTTP method to use for the request.
             - `headers (dict)`: HTTP headers to include in the request.
             - `timeout (int or None)`: Timeout for the request in seconds. Default is 5.
+            - `retries (int)`: Number of retries for the request. Default is 0.
             - `**kwargs`: Keyword arguments to pass to the endpoint URL.
 
         Returns:
@@ -297,7 +299,7 @@ async def request(
         Examples:
             >>> response = client.request("getTree", "GET",
                                 owner="owner", repo="repo",
-                                tree_sha="tree_sha", timeout=5)
+                                tree_sha="tree_sha", timeout=5, retries=2)
         """
         try:
             import httpx
@@ -314,6 +316,7 @@ async def request(
             headers=_headers,
             base_url=self._base_url,
             timeout=timeout,
+            transport=httpx.AsyncHTTPTransport(retries=retries),
         ) as _client:
             try:
                 response = await _client.request(
@@ -331,6 +334,7 @@ async def get_branch(
         branch: Optional[str] = None,
         branch_name: Optional[str] = None,
         timeout: Optional[int] = 5,
+        retries: int = 0,
     ) -> GitBranchResponseModel:
         """
         Get information about a branch. (Github API endpoint: getBranch).
@@ -339,6 +343,9 @@ async def get_branch(
             - `owner (str)`: Owner of the repository.
             - `repo (str)`: Name of the repository.
             - `branch (str)`: Name of the branch.
+            - `branch_name (str)`: Name of the branch.
+            - `timeout (int or None)`: Timeout for the request in seconds. Default is 5.
+            - `retries (int)`: Number of retries for the request. Default is 0.
 
         Returns:
             - `branch_info (GitBranchResponseModel)`: Information about the branch.
@@ -360,6 +367,7 @@ async def get_branch(
                     repo=repo,
                     branch=branch,
                     timeout=timeout,
+                    retries=retries,
                 )
             ).text
         )
@@ -370,6 +378,7 @@ async def get_tree(
         repo: str,
         tree_sha: str,
         timeout: Optional[int] = 5,
+        retries: int = 0,
     ) -> GitTreeResponseModel:
         """
         Get information about a tree. (Github API endpoint: getTree).
@@ -379,6 +388,7 @@ async def get_tree(
             - `repo (str)`: Name of the repository.
             - `tree_sha (str)`: SHA of the tree.
             - `timeout (int or None)`: Timeout for the request in seconds. Default is 5.
+            - `retries (int)`: Number of retries for the request. Default is 0.
 
         Returns:
             - `tree_info (GitTreeResponseModel)`: Information about the tree.
@@ -395,6 +405,7 @@ async def get_tree(
                     repo=repo,
                     tree_sha=tree_sha,
                     timeout=timeout,
+                    retries=retries,
                 )
             ).text
         )
@@ -405,6 +416,7 @@ async def get_blob(
         repo: str,
         file_sha: str,
         timeout: Optional[int] = 5,
+        retries: int = 0,
     ) -> GitBlobResponseModel:
         """
         Get information about a blob. (Github API endpoint: getBlob).
@@ -414,6 +426,7 @@ async def get_blob(
             - `repo (str)`: Name of the repository.
             - `file_sha (str)`: SHA of the file.
             - `timeout (int or None)`: Timeout for the request in seconds. Default is 5.
+            - `retries (int)`: Number of retries for the request. Default is 0.
 
         Returns:
             - `blob_info (GitBlobResponseModel)`: Information about the blob.
@@ -430,6 +443,7 @@ async def get_blob(
                     repo=repo,
                     file_sha=file_sha,
                     timeout=timeout,
+                    retries=retries,
                 )
             ).text
         )
@@ -440,6 +454,7 @@ async def get_commit(
         repo: str,
         commit_sha: str,
         timeout: Optional[int] = 5,
+        retries: int = 0,
     ) -> GitCommitResponseModel:
         """
         Get information about a commit. (Github API endpoint: getCommit).
@@ -449,6 +464,7 @@ async def get_commit(
             - `repo (str)`: Name of the repository.
             - `commit_sha (str)`: SHA of the commit.
             - `timeout (int or None)`: Timeout for the request in seconds. Default is 5.
+            - `retries (int)`: Number of retries for the request. Default is 0.
 
         Returns:
             - `commit_info (GitCommitResponseModel)`: Information about the commit.
@@ -465,6 +481,7 @@ async def get_commit(
                     repo=repo,
                     commit_sha=commit_sha,
                     timeout=timeout,
+                    retries=retries,
                 )
             ).text
         )