55 lines
2 KiB
Python
55 lines
2 KiB
Python
from typing import Dict, List, Union
|
|
|
|
|
|
class BardResponse:
|
|
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', [])]
|
|
|
|
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: 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})"
|