From b2037302e984a214014e32b00b2453a69adb8ef7 Mon Sep 17 00:00:00 2001 From: Andrew Tsegaye Date: Thu, 27 Apr 2023 09:17:07 -0700 Subject: well updated --- unfinished/bard/typings.py | 59 ++++++++++++++++++++++++++++++++++++++-------- 1 file changed, 49 insertions(+), 10 deletions(-) (limited to 'unfinished/bard/typings.py') 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})" -- cgit v1.2.3 From f1594cfddecf2bd47a25d54852834b392eaa4efd Mon Sep 17 00:00:00 2001 From: Aymane Hrouch Date: Thu, 27 Apr 2023 15:43:59 +0100 Subject: Reformat code using PyCharm --- unfinished/bard/typings.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) (limited to 'unfinished/bard/typings.py') diff --git a/unfinished/bard/typings.py b/unfinished/bard/typings.py index ddf803b6..75b73bf9 100644 --- a/unfinished/bard/typings.py +++ b/unfinished/bard/typings.py @@ -1,4 +1,4 @@ -from typing import Dict, List, Optional, Union +from typing import Dict, List, Union class BardResponse: -- cgit v1.2.3