Spaces:
Sleeping
Sleeping
| import os | |
| import re | |
| from typing import List | |
| from tenacity import retry, stop_after_attempt, wait_random_exponential | |
| from ..message import SYSTEM_NAME as SYSTEM | |
| from ..message import Message | |
| from .base import IntelligenceBackend | |
| try: | |
| import bardapi | |
| except ImportError: | |
| is_bard_available = False | |
| # logging.warning("bard package is not installed") | |
| else: | |
| bard_api_key = os.environ.get("_BARD_API_KEY") | |
| if bard_api_key is None: | |
| # logging.warning( | |
| # "Bard API key is not set. Please set the environment variable _BARD_API_KEY") | |
| is_bard_available = False | |
| else: | |
| is_bard_available = True | |
| DEFAULT_MAX_TOKENS = 4096 | |
| class Bard(IntelligenceBackend): | |
| """Interface to the Bard offered by Google.""" | |
| stateful = False | |
| type_name = "bard" | |
| def __init__(self, max_tokens: int = DEFAULT_MAX_TOKENS, **kwargs): | |
| assert ( | |
| is_bard_available | |
| ), "bard package is not installed or the API key is not set" | |
| super().__init__(max_tokens=max_tokens, **kwargs) | |
| self.max_tokens = max_tokens | |
| self.client = bardapi.core.Bard() | |
| def _get_response(self, prompt: str): | |
| response = self.client.get_answer( | |
| input_text=prompt, | |
| ) | |
| response = response["content"].strip() | |
| return response | |
| def query( | |
| self, | |
| agent_name: str, | |
| role_desc: str, | |
| history_messages: List[Message], | |
| global_prompt: str = None, | |
| request_msg: Message = None, | |
| *args, | |
| **kwargs, | |
| ) -> str: | |
| """ | |
| Format the input and call the Bard API. | |
| args: | |
| agent_name: the name of the agent | |
| role_desc: the description of the role of the agent | |
| env_desc: the description of the environment | |
| history_messages: the history of the conversation, or the observation for the agent | |
| request_msg: the request from the system to guide the agent's next response | |
| """ | |
| all_messages = ( | |
| [(SYSTEM, global_prompt), (SYSTEM, role_desc)] | |
| if global_prompt | |
| else [(SYSTEM, role_desc)] | |
| ) | |
| for message in history_messages: | |
| all_messages.append((message.agent_name, message.content)) | |
| if request_msg: | |
| all_messages.append((SYSTEM, request_msg.content)) | |
| # current bard api doesn't support role system, so just dump the raw messages as prompt | |
| response = self._get_response(str(all_messages), *args, **kwargs) | |
| # Remove the agent name if the response starts with it | |
| response = re.sub(rf"^\s*\[{agent_name}]:?", "", response).strip() | |
| return response | |