Merge pull request #951 from lucagobbi/to_async_or_not_to_async

[WIP] To async or not to async

core/cat/mad_hatter/registry.py

@@ -1,5 +1,6 @@
-import requests
+import httpx
 import random
+import aiofiles
 
 from cat.log import log
 
@@ -20,7 +21,9 @@ async def registry_search_plugins(
             # search plugins
             url = f"{registry_url}/search"
             payload = {"query": query}
-            response = requests.post(url, json=payload)
+
+            async with httpx.AsyncClient() as client:
+                response = await client.post(url, json=payload)
 
             # check the connection's status
             if response.status_code == 200:
@@ -37,7 +40,9 @@ async def registry_search_plugins(
                 "page": 1,
                 "page_size": 1000,
             }
-            response = requests.get(url, params=params)
+
+            async with httpx.AsyncClient() as client:
+                response = await client.get(url, params=params)
 
             # check the connection's status
             if response.status_code == 200:
@@ -57,16 +62,20 @@ async def registry_search_plugins(
         return []
 
 
-def registry_download_plugin(url: str) -> str:
+async def registry_download_plugin(url: str):
     log.info(f"Downloading {url}")
 
     registry_url = get_registry_url()
     payload = {"url": url}
-    response = requests.post(f"{registry_url}/download", json=payload)
-    plugin_zip_path = f"/tmp/{url.split('/')[-1]}.zip"
-    with open(plugin_zip_path, "wb") as f:
-        f.write(response.content)
 
-    log.info(f"Saved plugin as {plugin_zip_path}")
+    async with httpx.AsyncClient() as client:
+        response = await client.post(f"{registry_url}/download", json=payload)
+        response.raise_for_status()
 
+        plugin_zip_path = f"/tmp/{url.split('/')[-1]}.zip"
+
+        async with aiofiles.open(plugin_zip_path, "wb") as f:
+            await f.write(response.content)  # Write the content asynchronously
+
+    log.info(f"Saved plugin as {plugin_zip_path}")
     return plugin_zip_path

core/cat/routes/plugins.py

@@ -1,3 +1,4 @@
+import aiofiles
 import mimetypes
 from copy import deepcopy
 from typing import Dict
@@ -98,8 +99,9 @@ async def install_plugin(
 
     log.info(f"Uploading {content_type} plugin {file.filename}")
     plugin_archive_path = f"/tmp/{file.filename}"
-    with open(plugin_archive_path, "wb+") as f:
-        f.write(file.file.read())
+    async with aiofiles.open(plugin_archive_path, "wb+") as f:
+        content = await file.read()
+        await f.write(content)
     ccat.mad_hatter.install_plugin(plugin_archive_path)
 
     return {
@@ -122,7 +124,7 @@ async def install_plugin_from_registry(
 
     # download zip from registry
     try:
-        tmp_plugin_path = registry_download_plugin(payload["url"])
+        tmp_plugin_path = await registry_download_plugin(payload["url"])
         ccat.mad_hatter.install_plugin(tmp_plugin_path)
     except Exception as e:
         log.error("Could not download plugin form registry")

core/cat/routes/upload.py

@@ -1,5 +1,5 @@
 import mimetypes
-import requests
+import httpx
 import io
 import json
 from typing import Dict, List
@@ -276,9 +276,12 @@ async def upload_url(
     # check that URL is valid
     try:
         # Send a HEAD request to the specified URL
-        response = requests.head(
-            upload_config.url, headers={"User-Agent": "Magic Browser"}, allow_redirects=True
-        )
+        async with httpx.AsyncClient() as client:
+            response = await client.head(
+                upload_config.url,
+                headers={"User-Agent": "Magic Browser"},
+                follow_redirects=True,
+            )
 
         if response.status_code == 200:
             # upload file to long term memory, in the background
@@ -294,7 +297,7 @@ async def upload_url(
                 status_code=400,
                 detail={"error": "Invalid URL", "url": upload_config.url},
             )
-    except requests.exceptions.RequestException as _e:
+    except httpx.RequestError as _e:
         raise HTTPException(
             status_code=400,
             detail={"error": "Unable to reach the URL", "url": upload_config.url},

core/pyproject.toml

@@ -49,6 +49,7 @@ dependencies = [
     "rapidfuzz==3.6.1",
     "APScheduler==3.10.4",
     "ruff==0.4.7",
+    "aiofiles==24.1.0",
 ]
 
 [tool.coverage.run]

core/tests/routes/plugins/test_plugins_registry.py

@@ -35,11 +35,13 @@ def test_list_registry_plugins_by_query(client):
         plugin_text = plugin["name"] + plugin["description"]
         assert params["query"] in plugin_text  # verify searched text
 
+async def mock_registry_download_plugin(url: str):
+    return create_mock_plugin_zip(True)
 
 def test_plugin_install_from_registry(client, monkeypatch):
     # Mock the download from the registry creating a zip on-the-fly
     monkeypatch.setattr(
-        "cat.routes.plugins.registry_download_plugin", create_mock_plugin_zip
+        "cat.routes.plugins.registry_download_plugin", mock_registry_download_plugin
     )
 
     # during tests, the cat uses a different folder for plugins