From 739c07332889d18b23db46a11271a495435f3044 Mon Sep 17 00:00:00 2001
From: Nil Andreu <65730003+Nil-Andreu@users.noreply.github.com>
Date: Thu, 5 Jan 2023 17:49:39 +0100
Subject: [PATCH] Session Close when called HuggingFace Client (#411)

* [FIX] Session Close

* [NEW] Async Context Managers of Response

* [NEW] Linting
---
 backend/oasst_backend/utils/hugging_face.py | 18 ++++++++++--------
 1 file changed, 10 insertions(+), 8 deletions(-)

diff --git a/backend/oasst_backend/utils/hugging_face.py b/backend/oasst_backend/utils/hugging_face.py
index ff73a1c5..0df913f5 100644
--- a/backend/oasst_backend/utils/hugging_face.py
+++ b/backend/oasst_backend/utils/hugging_face.py
@@ -35,15 +35,17 @@ class HuggingFaceAPI:
             inference: the inference we obtain from the model in HF
         """
 
-        session = aiohttp.ClientSession()
-        payload: Dict[str, str] = {"inputs": input}
-        response = await session.post(self.api_url, headers=self.headers, json=payload)
+        async with aiohttp.ClientSession() as session:
+            payload: Dict[str, str] = {"inputs": input}
 
-        # If we get a bad response
-        if response.status != 200:
-            raise OasstError("Response Error Detoxify HuggingFace", error_code=OasstErrorCode.HUGGINGFACE_API_ERROR)
+            async with session.post(self.api_url, headers=self.headers, json=payload) as response:
+                # If we get a bad response
+                if response.status != 200:
+                    raise OasstError(
+                        "Response Error Detoxify HuggingFace", error_code=OasstErrorCode.HUGGINGFACE_API_ERROR
+                    )
 
-        # Get the response from the API call
-        inference = await response.json()
+                # Get the response from the API call
+                inference = await response.json()
 
         return inference