refactor: clean code

bardapi/core.py

@@ -1,16 +1,20 @@
-import json
+import os
 import random
-import re
 import string
-import os
 import requests
+import json
+import re
 
 
 class Bard:
-    def __init__(self, timeout=6, proxies=None, session=None):
+    def __init__(
+        self, token=os.getenv("_BARD_API_KEY"), timeout=6, proxies=None, session=None
+    ):
         """
         Initialize Bard
 
+        :param token: (`str`, *optional*)
+            Enter the __Secure-1PSID value. defaults to os.getenv("_BARD_API_KEY")
         :param timeout: (`int`, *optional*)
             Timeout in seconds when connecting bard server. The timeout is used on each request.
         :param proxies: (`Dict[str, str]`, *optional*)
@@ -19,8 +23,6 @@ def __init__(self, timeout=6, proxies=None, session=None):
         :param session: (`requests.Session`, *optional*)
             An existing requests.Session object to be used for making HTTP requests.
         """
-        self.proxies = proxies
-        self.timeout = timeout
         headers = {
             "Host": "bard.google.com",
             "X-Same-Domain": "1",
@@ -30,27 +32,36 @@ def __init__(self, timeout=6, proxies=None, session=None):
             "Origin": "https://bard.google.com",
             "Referer": "https://bard.google.com/",
         }
+        self.token, self.proxies, self.timeout = token, proxies, timeout
         self._reqid = int("".join(random.choices(string.digits, k=4)))
-        self.conversation_id = ""
-        self.response_id = ""
-        self.choice_id = ""
-
-        if session is None:
-            self.session = requests.Session()
-            self.session.headers = headers
-            self.session.cookies.set("__Secure-1PSID", os.environ["_BARD_API_KEY"])
-        else:
-            self.session = session
-
+        self.conversation_id = self.response_id = self.choice_id = ""
+        self.session = requests.Session() if session is None else session
+        self.session.headers = headers
+        self.session.cookies.set("__Secure-1PSID", self.token)
         self.SNlM0e = self._get_snim0e()
 
     def _get_snim0e(self):
+        if self.token and self.token[-1] != ".":
+            raise Exception(
+                "__Secure-1PSID value must end with a single dot. Enter correct __Secure-1PSID value."
+            )
+        elif not self.token:
+            raise Exception("No token value entered.")
         resp = self.session.get(
-            url="https://bard.google.com/", timeout=self.timeout, proxies=self.proxies
+            "https://bard.google.com/", timeout=self.timeout, proxies=self.proxies
         )
         if resp.status_code != 200:
-            raise Exception(f"Response Status: {resp.status_code}")
-        return re.search(r"SNlM0e\":\"(.*?)\"", resp.text).group(1)
+            raise Exception(
+                f"Response code not 200. \n - Try Bard on https://bard.google.com/ \n "
+                f"- Check service area on https://support.google.com/bard/answer/13575153?hl=en. \n "
+                f"- Response Status is {resp.status_code}"
+            )
+        snim0e = re.search(r"SNlM0e\":\"(.*?)\"", resp.text)
+        if not snim0e:
+            raise Exception(
+                "SNlM0e value not found in response. Check __Secure-1PSID value."
+            )
+        return snim0e.group(1)
 
     def get_answer(self, input_text: str) -> dict:
         params = {
@@ -75,20 +86,24 @@ def get_answer(self, input_text: str) -> dict:
             proxies=self.proxies,
         )
         resp_dict = json.loads(resp.content.splitlines()[3])[0][2]
-        if resp_dict is None:
+
+        if not resp_dict:
             return {"content": f"Response Error: {resp.content}."}
+
         parsed_answer = json.loads(resp_dict)
         bard_answer = {
             "content": parsed_answer[0][0],
             "conversation_id": parsed_answer[1][0],
             "response_id": parsed_answer[1][1],
             "factualityQueries": parsed_answer[3],
-            "textQuery": parsed_answer[2][0] if parsed_answer[2] is not None else "",
+            "textQuery": parsed_answer[2][0] if parsed_answer[2] else "",
             "choices": [{"id": i[0], "content": i[1]} for i in parsed_answer[4]],
         }
-        self.conversation_id = bard_answer["conversation_id"]
-        self.response_id = bard_answer["response_id"]
-        self.choice_id = bard_answer["choices"][0]["id"]
+        self.conversation_id, self.response_id, self.choice_id = (
+            bard_answer["conversation_id"],
+            bard_answer["response_id"],
+            bard_answer["choices"][0]["id"],
+        )
         self._reqid += 100000
 
         return bard_answer