llmbb/LLMBB-Agent

1from abc import ABC, abstractmethod2from typing import Dict, Iterator, List, Optional, Union3 4from agent.log import logger5from agent.utils.utils import print_traceback6 7 8class FnCallNotImplError(NotImplementedError):9    pass10 11 12class BaseChatModel(ABC):13 14    def __init__(self):15        self._support_fn_call: Optional[bool] = None16 17    # It is okay to use the same code to handle the output18    # regardless of whether stream is True or False, as follows:19    # ```py20    # for chunk in chat_model.chat(..., stream=True/False):21    #   response += chunk22    #   yield response23    # ```24    def chat(25        self,26        prompt: Optional[str] = None,27        messages: Optional[List[Dict]] = None,28        stop: Optional[List[str]] = None,29        stream: bool = False,30    ) -> Union[str, Iterator[str]]:31        if messages is None:32            assert isinstance(prompt, str)33            messages = [{'role': 'user', 'content': prompt}]34        else:35            assert prompt is None, 'Do not pass prompt and messages at the same time.'36        logger.debug(messages)37        if stream:38            return self._chat_stream(messages, stop=stop)39        else:40            return self._chat_no_stream(messages, stop=stop)41 42    def support_function_calling(self) -> bool:43        if self._support_fn_call is None:44            functions = [{45                'name': 'get_current_weather',46                'description': 'Get the current weather in a given location.',47                'parameters': {48                    'type': 'object',49                    'properties': {50                        'location': {51                            'type':52                            'string',53                            'description':54                            'The city and state, e.g. San Francisco, CA',55                        },56                        'unit': {57                            'type': 'string',58                            'enum': ['celsius', 'fahrenheit'],59                        },60                    },61                    'required': ['location'],62                },63            }]64            messages = [{65                'role': 'user',66                'content': 'What is the weather like in Boston?'67            }]68            self._support_fn_call = False69            try:70                response = self.chat_with_functions(messages=messages,71                                                    functions=functions)72                if response.get('function_call', None):73                    logger.info('Support of function calling is detected.')74                    self._support_fn_call = True75            except FnCallNotImplError:76                pass77            except Exception:  # TODO: more specific78                print_traceback()79        return self._support_fn_call80 81    def chat_with_functions(self,82                            messages: List[Dict],83                            functions: Optional[List[Dict]] = None) -> Dict:84        raise FnCallNotImplError85 86    @abstractmethod87    def _chat_stream(88        self,89        messages: List[Dict],90        stop: Optional[List[str]] = None,91    ) -> Iterator[str]:92        raise NotImplementedError93 94    @abstractmethod95    def _chat_no_stream(96        self,97        messages: List[Dict],98        stop: Optional[List[str]] = None,99    ) -> str:100        raise NotImplementedError101