| """Module for client.""" | |
| from __future__ import annotations | |
| from typing import Any, cast | |
| from .backends import ILLMBackend, LLMResponse | |
| class LLMClient: | |
| def __init__(self, *, default_backend: str) -> None: | |
| self._default_backend = default_backend | |
| def complete( | |
| self, | |
| *, | |
| fallback_policy: Any, | |
| prompt: str, | |
| system_prompt: str | None = None, | |
| backend: str | None = None, | |
| model: str | None = None, | |
| temperature: float = 0.7, | |
| max_tokens: int | None = None, | |
| fallback: bool = True, | |
| **kwargs: Any, | |
| ) -> LLMResponse: | |
| messages: list[dict[str, str]] = [] | |
| if system_prompt: | |
| messages.append({"role": "system", "content": system_prompt}) | |
| messages.append({"role": "user", "content": prompt}) | |
| return self.chat( | |
| fallback_policy=fallback_policy, | |
| messages=messages, | |
| backend=backend, | |
| model=model, | |
| temperature=temperature, | |
| max_tokens=max_tokens, | |
| fallback=fallback, | |
| **kwargs, | |
| ) | |
| def chat( | |
| self, | |
| *, | |
| fallback_policy: Any, | |
| messages: list[dict[str, str]], | |
| backend: str | None = None, | |
| model: str | None = None, | |
| temperature: float = 0.7, | |
| max_tokens: int | None = None, | |
| fallback: bool = True, | |
| **kwargs: Any, | |
| ) -> LLMResponse: | |
| def operation(b: ILLMBackend) -> LLMResponse: | |
| return b.chat(messages=messages, model=model, temperature=temperature, max_tokens=max_tokens, **kwargs) | |
| backend_name = backend or self._default_backend | |
| return cast(LLMResponse, fallback_policy.execute(operation=operation, backend=backend_name, fallback=fallback)) | |
| async def achat( | |
| self, | |
| *, | |
| fallback_policy: Any, | |
| messages: list[dict[str, str]], | |
| backend: str | None = None, | |
| model: str | None = None, | |
| temperature: float = 0.7, | |
| max_tokens: int | None = None, | |
| fallback: bool = True, | |
| **kwargs: Any, | |
| ) -> LLMResponse: | |
| async def operation(b: ILLMBackend) -> LLMResponse: | |
| return await b.achat( | |
| messages=messages, model=model, temperature=temperature, max_tokens=max_tokens, **kwargs | |
| ) | |
| backend_name = backend or self._default_backend | |
| return cast( | |
| LLMResponse, | |
| await fallback_policy.execute_async(operation=operation, backend=backend_name, fallback=fallback), | |
| ) | |
| def embed_texts( | |
| self, | |
| *, | |
| fallback_policy: Any, | |
| texts: list[str], | |
| backend: str | None = None, | |
| model: str | None = None, | |
| fallback: bool = True, | |
| **kwargs: Any, | |
| ) -> list[list[float]]: | |
| def operation(b: ILLMBackend) -> list[list[float]]: | |
| return b.embed_texts(texts=texts, model=model, **kwargs) | |
| backend_name = backend or self._default_backend | |
| return cast( | |
| list[list[float]], | |
| fallback_policy.execute(operation=operation, backend=backend_name, fallback=fallback), | |
| ) | |