2 files changed, 102 insertions, 86 deletions

diff --git a/unfinished/bard/__init__.py b/unfinished/bard/__init__.py
index 2b26997f..ef8980d7 100644
--- a/unfinished/bard/__init__.py
+++ b/unfinished/bard/__init__.py
@@ -1,19 +1,19 @@
-from requests     import Session
-from re           import search
-from random       import randint
-from json         import dumps, loads
-from random       import randint
+from requests import Session
+from re import search
+from random import randint
+from json import dumps, loads
 from urllib.parse import urlencode
-from dotenv       import load_dotenv; load_dotenv()
-from os           import getenv
+from dotenv import load_dotenv
+from os import getenv
 
 from bard.typings import BardResponse
 
+load_dotenv()
 token = getenv('1psid')
 proxy = getenv('proxy')
 
 temperatures = {
-    0  : "Generate text strictly following known patterns, with no creativity.",
+    0: "Generate text strictly following known patterns, with no creativity.",
     0.1: "Produce text adhering closely to established patterns, allowing minimal creativity.",
     0.2: "Create text with modest deviations from familiar patterns, injecting a slight creative touch.",
     0.3: "Craft text with a mild level of creativity, deviating somewhat from common patterns.",
@@ -23,93 +23,70 @@ temperatures = {
     0.7: "Produce text favoring creativity over typical patterns for more original results.",
     0.8: "Create text heavily focused on creativity, with limited concern for familiar patterns.",
     0.9: "Craft text with a strong emphasis on unique and inventive ideas, largely ignoring established patterns.",
-    1  : "Generate text with maximum creativity, disregarding any constraints of known patterns or structures."
+    1: "Generate text with maximum creativity, disregarding any constraints of known patterns or structures."
 }
 
+
 class Completion:
-    # def __init__(self, _token, proxy: str or None = None) -> None:
-    #     self.client         = Session()
-    #     self.client.proxies = {
-    #         'http': f'http://{proxy}',
-    #         'https': f'http://{proxy}' } if proxy else None
-        
-    #     self.client.headers = {
-    #         'authority'     : 'bard.google.com',
-    #         'content-type'  : 'application/x-www-form-urlencoded;charset=UTF-8',
-    #         'origin'        : 'https://bard.google.com',
-    #         'referer'       : 'https://bard.google.com/',
-    #         'user-agent'    : 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/111.0.0.0 Safari/537.36',
-    #         'x-same-domain' : '1',
-    #         'cookie'        : f'__Secure-1PSID={_token}'
-    #     }
-        
-    #     self.snlm0e          = self.__init_client()
-    #     self.conversation_id = ''
-    #     self.response_id     = ''
-    #     self.choice_id       = ''
-    #     self.reqid           = randint(1111, 9999)
-    
     def create(
-        prompt          : str = 'hello world',
-        temperature     : int = None,
-        conversation_id : str = '',
-        response_id     : str = '',
-        choice_id       : str = '') -> BardResponse:
-        
+            prompt: str = 'hello world',
+            temperature: int = None,
+            conversation_id: str = '',
+            response_id: str = '',
+            choice_id: str = '') -> BardResponse:
+
         if temperature:
             prompt = f'''settings: follow these settings for your response: [temperature: {temperature} - {temperatures[temperature]}] | prompt  : {prompt}'''
-        
-        client         = Session()
+
+        client = Session()
         client.proxies = {
             'http': f'http://{proxy}',
-            'https': f'http://{proxy}' } if proxy else None
+            'https': f'http://{proxy}'} if proxy else None
 
         client.headers = {
-            'authority'     : 'bard.google.com',
-            'content-type'  : 'application/x-www-form-urlencoded;charset=UTF-8',
-            'origin'        : 'https://bard.google.com',
-            'referer'       : 'https://bard.google.com/',
-            'user-agent'    : 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/111.0.0.0 Safari/537.36',
-            'x-same-domain' : '1',
-            'cookie'        : f'__Secure-1PSID={token}'
+            'authority': 'bard.google.com',
+            'content-type': 'application/x-www-form-urlencoded;charset=UTF-8',
+            'origin': 'https://bard.google.com',
+            'referer': 'https://bard.google.com/',
+            'user-agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/111.0.0.0 Safari/537.36',
+            'x-same-domain': '1',
+            'cookie': f'__Secure-1PSID={token}'
         }
 
-        snlm0e          = search(r'SNlM0e\":\"(.*?)\"', client.get('https://bard.google.com/').text).group(1)
-    
+        snlm0e = search(r'SNlM0e\":\"(.*?)\"',
+                        client.get('https://bard.google.com/').text).group(1)
+
         params = urlencode({
-            'bl'     : 'boq_assistant-bard-web-server_20230326.21_p0',
-            '_reqid' : randint(1111, 9999),
-            'rt'     : 'c',
+            'bl': 'boq_assistant-bard-web-server_20230326.21_p0',
+            '_reqid': randint(1111, 9999),
+            'rt': 'c',
         })
 
-        response = client.post(f'https://bard.google.com/_/BardChatUi/data/assistant.lamda.BardFrontendService/StreamGenerate?{params}', 
-            data = {
-                'at': snlm0e,
-                'f.req': dumps([None, dumps([
-                        [prompt],
-                        None,
-                        [conversation_id, response_id, choice_id],
-                    ])
-                ])
-            }
-        )
-        
+        response = client.post(f'https://bard.google.com/_/BardChatUi/data/assistant.lamda.BardFrontendService/StreamGenerate?{params}',
+                               data={
+                                   'at': snlm0e,
+                                   'f.req': dumps([None, dumps([
+                                       [prompt],
+                                       None,
+                                       [conversation_id, response_id, choice_id],
+                                   ])])
+                               }
+                               )
+
         chat_data = loads(response.content.splitlines()[3])[0][2]
-        if not chat_data: print('error, retrying'); Completion.create(prompt, temperature, conversation_id, response_id, choice_id)
+        if not chat_data:
+            print('error, retrying')
+            Completion.create(prompt, temperature,
+                              conversation_id, response_id, choice_id)
 
         json_chat_data = loads(chat_data)
         results = {
-            'content'           : json_chat_data[0][0],
-            'conversation_id'   : json_chat_data[1][0],
-            'response_id'       : json_chat_data[1][1],
-            'factualityQueries' : json_chat_data[3],
-            'textQuery'         : json_chat_data[2][0] if json_chat_data[2] is not None else '',
-            'choices'           : [{'id': i[0], 'content': i[1]} for i in json_chat_data[4]],
+            'content': json_chat_data[0][0],
+            'conversation_id': json_chat_data[1][0],
+            'response_id': json_chat_data[1][1],
+            'factualityQueries': json_chat_data[3],
+            'textQuery': json_chat_data[2][0] if json_chat_data[2] is not None else '',
+            'choices': [{'id': i[0], 'content': i[1]} for i in json_chat_data[4]],
         }
-        
-        # self.conversation_id = results['conversation_id']
-        # self.response_id     = results['response_id']
-        # self.choice_id       = results['choices'][0]['id']
-        # self.reqid           += 100000
 
         return BardResponse(results)
diff --git a/unfinished/bard/typings.py b/unfinished/bard/typings.py
index 69239762..ddf803b6 100644
--- a/unfinished/bard/typings.py
+++ b/unfinished/bard/typings.py
@@ -1,15 +1,54 @@
+from typing import Dict, List, Optional, Union
+
+
 class BardResponse:
-    def __init__(self, json_dict):
-        self.json               = json_dict
-        
-        self.content            = json_dict.get('content')
-        self.conversation_id    = json_dict.get('conversation_id')
-        self.response_id        = json_dict.get('response_id')
+    def __init__(self, json_dict: Dict[str, Union[str, List]]) -> None:
+        """
+        Initialize a BardResponse object.
+
+        :param json_dict: A dictionary containing the JSON response data.
+        """
+        self.json = json_dict
+
+        self.content = json_dict.get('content')
+        self.conversation_id = json_dict.get('conversation_id')
+        self.response_id = json_dict.get('response_id')
         self.factuality_queries = json_dict.get('factualityQueries', [])
-        self.text_query         = json_dict.get('textQuery', [])
-        self.choices            = [self.BardChoice(choice) for choice in json_dict.get('choices', [])]
+        self.text_query = json_dict.get('textQuery', [])
+        self.choices = [self.BardChoice(choice)
+                        for choice in json_dict.get('choices', [])]
+
+    def __repr__(self) -> str:
+        """
+        Return a string representation of the BardResponse object.
+
+        :return: A string representation of the BardResponse object.
+        """
+        return f"BardResponse(conversation_id={self.conversation_id}, response_id={self.response_id}, content={self.content})"
+
+    def filter_choices(self, keyword: str) -> List['BardChoice']:
+        """
+        Filter the choices based on a keyword.
+
+        :param keyword: The keyword to filter choices by.
+        :return: A list of filtered BardChoice objects.
+        """
+        return [choice for choice in self.choices if keyword.lower() in choice.content.lower()]
 
     class BardChoice:
-        def __init__(self, choice_dict):
-            self.id      = choice_dict.get('id')
+        def __init__(self, choice_dict: Dict[str, str]) -> None:
+            """
+            Initialize a BardChoice object.
+
+            :param choice_dict: A dictionary containing the choice data.
+            """
+            self.id = choice_dict.get('id')
             self.content = choice_dict.get('content')[0]
+
+        def __repr__(self) -> str:
+            """
+            Return a string representation of the BardChoice object.
+
+            :return: A string representation of the BardChoice object.
+            """
+            return f"BardChoice(id={self.id}, content={self.content})"