|
from abc import ABC
|
|
from abc import abstractmethod
|
|
|
|
|
|
class Provider(ABC):
|
|
"""Base class for providers"""
|
|
|
|
@abstractmethod
|
|
def ask(
|
|
self,
|
|
prompt: str,
|
|
stream: bool = False,
|
|
raw: bool = False,
|
|
optimizer: str = None,
|
|
conversationally: bool = False,
|
|
) -> dict:
|
|
"""Chat with AI
|
|
|
|
Args:
|
|
prompt (str): Prompt to be sent
|
|
stream (bool, optional): Flag for streaming response. Defaults to False.
|
|
raw (bool, optional): Stream back raw response as received
|
|
optimizer (str, optional): Prompt optimizer name - `[code, shell_command]`
|
|
conversationally (bool, optional): Chat conversationally when using optimizer. Defaults to False.
|
|
Returns:
|
|
dict : {}
|
|
```json
|
|
{
|
|
"completion": "\nNext: domestic cat breeds with short hair >>",
|
|
"stop_reason": null,
|
|
"truncated": false,
|
|
"stop": null,
|
|
"model": "llama-2-13b-chat",
|
|
"log_id": "cmpl-3kYiYxSNDvgMShSzFooz6t",
|
|
"exception": null
|
|
}
|
|
```
|
|
"""
|
|
raise NotImplementedError("Method needs to be implemented in subclass")
|
|
|
|
@abstractmethod
|
|
def chat(
|
|
self,
|
|
prompt: str,
|
|
stream: bool = False,
|
|
optimizer: str = None,
|
|
conversationally: bool = False,
|
|
) -> str:
|
|
"""Generate response `str`
|
|
Args:
|
|
prompt (str): Prompt to be sent
|
|
stream (bool, optional): Flag for streaming response. Defaults to False.
|
|
optimizer (str, optional): Prompt optimizer name - `[code, shell_command]`
|
|
conversationally (bool, optional): Chat conversationally when using optimizer. Defaults to False.
|
|
Returns:
|
|
str: Response generated
|
|
"""
|
|
raise NotImplementedError("Method needs to be implemented in subclass")
|
|
|
|
@abstractmethod
|
|
def get_message(self, response: dict) -> str:
|
|
"""Retrieves message only from response
|
|
|
|
Args:
|
|
response (dict): Response generated by `self.ask`
|
|
|
|
Returns:
|
|
str: Message extracted
|
|
"""
|
|
raise NotImplementedError("Method needs to be implemented in subclass")
|
|
|
|
|
|
class AsyncProvider(ABC):
|
|
"""Asynchronous base class for providers"""
|
|
|
|
@abstractmethod
|
|
async def ask(
|
|
self,
|
|
prompt: str,
|
|
stream: bool = False,
|
|
raw: bool = False,
|
|
optimizer: str = None,
|
|
conversationally: bool = False,
|
|
) -> dict:
|
|
"""Asynchronously chat with AI
|
|
|
|
Args:
|
|
prompt (str): Prompt to be sent
|
|
stream (bool, optional): Flag for streaming response. Defaults to False.
|
|
raw (bool, optional): Stream back raw response as received
|
|
optimizer (str, optional): Prompt optimizer name - `[code, shell_command]`
|
|
conversationally (bool, optional): Chat conversationally when using optimizer. Defaults to False.
|
|
Returns:
|
|
dict : {}
|
|
```json
|
|
{
|
|
"completion": "\nNext: domestic cat breeds with short hair >>",
|
|
"stop_reason": null,
|
|
"truncated": false,
|
|
"stop": null,
|
|
"model": "llama-2-13b-chat",
|
|
"log_id": "cmpl-3kYiYxSNDvgMShSzFooz6t",
|
|
"exception": null
|
|
}
|
|
```
|
|
"""
|
|
raise NotImplementedError("Method needs to be implemented in subclass")
|
|
|
|
@abstractmethod
|
|
async def chat(
|
|
self,
|
|
prompt: str,
|
|
stream: bool = False,
|
|
optimizer: str = None,
|
|
conversationally: bool = False,
|
|
) -> str:
|
|
"""Asynchronously generate response `str`
|
|
Args:
|
|
prompt (str): Prompt to be sent
|
|
stream (bool, optional): Flag for streaming response. Defaults to False.
|
|
optimizer (str, optional): Prompt optimizer name - `[code, shell_command]`
|
|
conversationally (bool, optional): Chat conversationally when using optimizer. Defaults to False.
|
|
Returns:
|
|
str: Response generated
|
|
"""
|
|
raise NotImplementedError("Method needs to be implemented in subclass")
|
|
|
|
@abstractmethod
|
|
async def get_message(self, response: dict) -> str:
|
|
"""Asynchronously retrieves message only from response
|
|
|
|
Args:
|
|
response (dict): Response generated by `self.ask`
|
|
|
|
Returns:
|
|
str: Message extracted
|
|
"""
|
|
raise NotImplementedError("Method needs to be implemented in subclass") |
