From 2d9e74bffe52145af53bb79dcc122c9b39a5b9e4 Mon Sep 17 00:00:00 2001
From: Kerem Yilmaz <kerem@skyvern.com>
Date: Tue, 9 Jul 2024 17:45:51 -0700
Subject: [PATCH] migrate aiohttp helper from cloud to open source (#572)

---
 skyvern/exceptions.py                    |  5 +++
 skyvern/forge/sdk/core/aiohttp_helper.py | 45 ++++++++++++++++++++++++
 2 files changed, 50 insertions(+)
 create mode 100644 skyvern/forge/sdk/core/aiohttp_helper.py

diff --git a/skyvern/exceptions.py b/skyvern/exceptions.py
index fb69672f..4ab7ad52 100644
--- a/skyvern/exceptions.py
+++ b/skyvern/exceptions.py
@@ -416,3 +416,8 @@ class FailedToGetCurrentValueOfDropdown(SkyvernException):
         super().__init__(
             f"Failed to get current value of {dropdowm_type} dropdown. element_id={element_id}, failure_reason={fail_reason}"
         )
+
+
+class HttpException(SkyvernException):
+    def __init__(self, status_code: int, url: str, msg: str | None = None) -> None:
+        super().__init__(f"HTTP Exception, status_code={status_code}, url={url}" + (f", msg={msg}" if msg else ""))
diff --git a/skyvern/forge/sdk/core/aiohttp_helper.py b/skyvern/forge/sdk/core/aiohttp_helper.py
new file mode 100644
index 00000000..fa7b1f96
--- /dev/null
+++ b/skyvern/forge/sdk/core/aiohttp_helper.py
@@ -0,0 +1,45 @@
+import asyncio
+from typing import Any
+
+import aiohttp
+import structlog
+
+from skyvern.exceptions import HttpException
+
+LOG = structlog.get_logger()
+DEFAULT_REQUEST_TIMEOUT = 30
+
+
+async def aiohttp_get_json(
+    url: str,
+    params: dict[str, Any] | None = None,
+    headers: dict[str, str] | None = None,
+    cookies: dict[str, str] | None = None,
+    retry: int = 0,
+    proxy: str | None = None,
+    timeout: int = DEFAULT_REQUEST_TIMEOUT,
+    raise_exception: bool = True,
+    retry_timeout: float = 0,
+) -> dict[str, Any]:
+    async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=timeout)) as session:
+        count = 0
+        while count <= retry:
+            try:
+                async with session.get(
+                    url,
+                    params=params,
+                    headers=headers,
+                    cookies=cookies,
+                    proxy=proxy,
+                ) as response:
+                    if response.status == 200:
+                        return await response.json()
+                    if raise_exception:
+                        raise HttpException(response.status, url)
+                    LOG.error(f"Failed to fetch data from {url}", status_code=response.status)
+                    return {}
+            except Exception:
+                if retry_timeout > 0:
+                    await asyncio.sleep(retry_timeout)
+                count += 1
+        raise Exception(f"Failed to fetch data from {url}")