2024-03-11 14:26:34 -04:00
|
|
|
from ..typing import Union, Iterator, Messages
|
|
|
|
from ..stubs import ChatCompletion, ChatCompletionChunk
|
2024-03-19 07:40:59 -04:00
|
|
|
from ._engine import LocalProvider
|
|
|
|
from ._models import models
|
|
|
|
from ..client import iter_response, filter_none, IterResponse
|
2024-03-11 14:26:34 -04:00
|
|
|
|
|
|
|
class LocalClient():
|
2024-03-11 14:31:13 -04:00
|
|
|
def __init__(self, **kwargs) -> None:
|
2024-03-11 14:26:34 -04:00
|
|
|
self.chat: Chat = Chat(self)
|
2024-03-11 14:31:13 -04:00
|
|
|
|
2024-03-11 14:26:34 -04:00
|
|
|
@staticmethod
|
|
|
|
def list_models():
|
|
|
|
return list(models.keys())
|
|
|
|
|
|
|
|
class Completions():
|
|
|
|
def __init__(self, client: LocalClient):
|
|
|
|
self.client: LocalClient = client
|
|
|
|
|
|
|
|
def create(
|
|
|
|
self,
|
|
|
|
messages: Messages,
|
|
|
|
model: str,
|
|
|
|
stream: bool = False,
|
|
|
|
response_format: dict = None,
|
|
|
|
max_tokens: int = None,
|
|
|
|
stop: Union[list[str], str] = None,
|
|
|
|
**kwargs
|
|
|
|
) -> Union[ChatCompletion, Iterator[ChatCompletionChunk]]:
|
|
|
|
|
|
|
|
stop = [stop] if isinstance(stop, str) else stop
|
|
|
|
response = LocalProvider.create_completion(
|
|
|
|
model, messages, stream,
|
|
|
|
**filter_none(
|
|
|
|
max_tokens=max_tokens,
|
|
|
|
stop=stop,
|
|
|
|
),
|
|
|
|
**kwargs
|
|
|
|
)
|
|
|
|
response = iter_response(response, stream, response_format, max_tokens, stop)
|
|
|
|
return response if stream else next(response)
|
2024-03-11 14:31:13 -04:00
|
|
|
|
2024-03-11 14:26:34 -04:00
|
|
|
class Chat():
|
|
|
|
completions: Completions
|
|
|
|
|
|
|
|
def __init__(self, client: LocalClient):
|
2024-03-11 14:31:13 -04:00
|
|
|
self.completions = Completions(client)
|