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