well updated

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})"