sfn_llm_client 0.3.3__tar.gz

sfn_llm_client-0.3.3/MANIFEST.in

@@ -0,0 +1 @@
+

sfn_llm_client-0.3.3/PKG-INFO

@@ -0,0 +1,91 @@
+Metadata-Version: 2.4
+Name: sfn_llm_client
+Version: 0.3.3
+Summary: SDK for using LLM
+Author-email: Rajesh Darak <rajesh@stepfuction.ai>
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/iamrajeshdaraksfn/llm-client-sdk.git
+Classifier: Operating System :: OS Independent
+Classifier: Intended Audience :: Developers
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+Requires-Dist: aiohttp
+Requires-Dist: openai
+Requires-Dist: tiktoken
+Requires-Dist: anthropic
+Requires-Dist: snowflake-connector-python>=3.17.0
+Requires-Dist: snowflake-snowpark-python>=1.26.0
+Requires-Dist: snowflake-ml-python>=1.9.0
+Requires-Dist: transformers>=4.46.2
+Requires-Dist: langchain-openai
+Requires-Dist: langchain-community
+Requires-Dist: pydantic
+Requires-Dist: langchain-core
+Requires-Dist: StrEnum
+Provides-Extra: test
+Requires-Dist: pytest; extra == "test"
+Requires-Dist: pytest-aiohttp; extra == "test"
+Requires-Dist: pytest-asyncio; extra == "test"
+Requires-Dist: pytest-mock; extra == "test"
+Requires-Dist: aioresponses; extra == "test"
+Provides-Extra: openai
+Requires-Dist: openai>=1.54.3; extra == "openai"
+Requires-Dist: tiktoken>=0.3.3; extra == "openai"
+Provides-Extra: anthropic
+Requires-Dist: anthropic>=0.39.0; extra == "anthropic"
+Provides-Extra: cortex-langchain
+Requires-Dist: snowflake-snowpark-python>=1.0.0; extra == "cortex-langchain"
+Requires-Dist: langchain-community>=0.0.30; extra == "cortex-langchain"
+Requires-Dist: pydantic>=2.6.0; extra == "cortex-langchain"
+Requires-Dist: langchain-core>=0.3; extra == "cortex-langchain"
+Provides-Extra: google
+Requires-Dist: google-generativeai>=0.1.0; extra == "google"
+Provides-Extra: api
+Requires-Dist: sfn_llm_client[anthropic,google,openai]; extra == "api"
+Provides-Extra: local
+Requires-Dist: transformers>=4.0.0; extra == "local"
+Provides-Extra: sync
+Requires-Dist: async_to_sync>=0.2.0; extra == "sync"
+Provides-Extra: all
+Requires-Dist: sfn_llm_client[api,local,sync]; extra == "all"
+
+# SFN_LLM_Client
+
+This is an enhanced and improved version with latest llm provider chat completion feature The `sfn_llm_client` now includes:
+
+- **Updated to the latest version of OpenAI**.
+- **Integrated Cortex LLM provider support**.
+- **Latest improvements and updates to the codebase** for better performance and compatibility.
+
+## Features
+
+- Supports multiple LLM providers, including OpenAI and Cortex.
+- Easily extensible to include new LLM providers by implementing base client classes.
+- Well-documented and tested.
+
+### Adding a New LLM Client
+
+To add a new LLM client, follow these steps:
+
+1. **Implement `BaseLLMClient` or `BaseLLMAPIClient`:**  
+   If you're adding a new LLM provider, you'll need to implement either the `BaseLLMClient` or `BaseLLMAPIClient` interfaces.
+   
+2. **Register in `LLMAPIClientFactory`:**  
+   If you're adding a client based on `BaseLLMAPIClient`, don't forget to register it in the `LLMAPIClientFactory` so that it's available for use.
+
+### Adding Dependencies
+
+If your LLM client requires additional dependencies, you can add them to the `pyproject.toml` file under the appropriate section.
+
+## Contributing
+Contributions are welcome! If you'd like to help improve this SDK, please check out the todos or open an issue or pull request.
+
+### Credits
+the core forked functionality taken from `llm-client-sdk` created by uripeled2.
+
+## Contact:
+For any queries or issues, please contact the maintainer at: `rajesh@stepfunction.ai`

sfn_llm_client-0.3.3/README.md

@@ -0,0 +1,36 @@
+# SFN_LLM_Client
+
+This is an enhanced and improved version with latest llm provider chat completion feature The `sfn_llm_client` now includes:
+
+- **Updated to the latest version of OpenAI**.
+- **Integrated Cortex LLM provider support**.
+- **Latest improvements and updates to the codebase** for better performance and compatibility.
+
+## Features
+
+- Supports multiple LLM providers, including OpenAI and Cortex.
+- Easily extensible to include new LLM providers by implementing base client classes.
+- Well-documented and tested.
+
+### Adding a New LLM Client
+
+To add a new LLM client, follow these steps:
+
+1. **Implement `BaseLLMClient` or `BaseLLMAPIClient`:**  
+   If you're adding a new LLM provider, you'll need to implement either the `BaseLLMClient` or `BaseLLMAPIClient` interfaces.
+   
+2. **Register in `LLMAPIClientFactory`:**  
+   If you're adding a client based on `BaseLLMAPIClient`, don't forget to register it in the `LLMAPIClientFactory` so that it's available for use.
+
+### Adding Dependencies
+
+If your LLM client requires additional dependencies, you can add them to the `pyproject.toml` file under the appropriate section.
+
+## Contributing
+Contributions are welcome! If you'd like to help improve this SDK, please check out the todos or open an issue or pull request.
+
+### Credits
+the core forked functionality taken from `llm-client-sdk` created by uripeled2.
+
+## Contact:
+For any queries or issues, please contact the maintainer at: `rajesh@stepfunction.ai`

sfn_llm_client-0.3.3/pyproject.toml

@@ -0,0 +1,96 @@
+[build-system]
+requires = ["setuptools", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "sfn_llm_client"
+version = "0.3.3"
+description = "SDK for using LLM"
+readme = "README.md"
+requires-python = ">=3.9"
+authors = [
+    { name = "Rajesh Darak", email = "rajesh@stepfuction.ai" },
+]
+license = "MIT" 
+classifiers = [
+    "Operating System :: OS Independent",
+    "Intended Audience :: Developers",
+    # "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3 :: Only",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+]
+dependencies = [
+    "aiohttp",
+    "openai",
+    "tiktoken",
+    "anthropic",
+    "snowflake-connector-python >=3.17.0",
+    "snowflake-snowpark-python >=1.26.0",
+    "snowflake-ml-python>=1.9.0",
+    "transformers >= 4.46.2",
+    "langchain-openai",
+    "langchain-community ",
+    "pydantic",
+    "langchain-core",
+    "StrEnum"
+]
+# dynamic = ["version"]
+[tool.setuptools.package-data]
+sfn_llm_client = ["README.md"]
+
+[project.urls]
+Homepage = "https://github.com/iamrajeshdaraksfn/llm-client-sdk.git"
+
+[project.optional-dependencies]
+test = [
+    "pytest",
+    "pytest-aiohttp",
+    "pytest-asyncio",
+    "pytest-mock",
+    "aioresponses"
+]
+openai = [
+    "openai >=1.54.3",
+    "tiktoken >=0.3.3",
+]
+anthropic = [
+    "anthropic >= 0.39.0"
+]
+
+cortex_langchain = [
+    "snowflake-snowpark-python >=1.0.0",
+    "langchain-community >= 0.0.30",
+    "pydantic >=2.6.0",
+    "langchain-core>=0.3",
+]
+google = [
+    "google-generativeai >= 0.1.0"
+]
+api = [
+    "sfn_llm_client[openai,anthropic,google]"
+]
+local = [
+    "transformers >= 4.0.0"
+]
+sync =  [
+    "async_to_sync >= 0.2.0"
+]
+all = [
+    "sfn_llm_client[api,local,sync]"
+]
+
+# [tool.setuptools.dynamic]
+# version = {attr = "sfn_llm_client.__version__"}
+[tool.setuptools.packages.find]
+where = ["."]
+include = ["sfn_llm_client*"]
+
+
+[tool.coverage.run]
+parallel = true
+source = [
+    "sfn_llm_client"
+]
+context = '${CONTEXT}'

sfn_llm_client-0.3.3/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

sfn_llm_client-0.3.3/sfn_llm_client/__init__.py

@@ -0,0 +1,75 @@
+
+# __version__ = "0.2.0a1"  
+
+
+# load utils
+try:
+    from .utils.base_llm_client import BaseLLMClient
+    from .utils.consts import BaseLLMClient
+    from .utils.base_llm_client import BaseLLMClient
+    from .utils.base_llm_client import BaseLLMClient
+
+except ImportError:
+    pass
+
+# load api clients
+try:
+    from .llm_api_client.base_llm_api_client import BaseLLMAPIClient, LLMAPIClientConfig, ChatMessage, Role
+    from .llm_api_client.llm_api_client_factory import LLMAPIClientFactory, LLMAPIClientType
+    # load base-api clients
+    try:
+        from .llm_api_client.ai21_client import AI21Client
+        from .llm_api_client.aleph_alpha_client import AlephAlphaClient
+        from .llm_api_client.google_client import GoogleClient, MessagePrompt
+    except ImportError:
+        pass
+    # load apis with different dependencies
+    try:
+        from .llm_api_client.openai_client import OpenAIClient
+    except ImportError:
+        pass
+    try:
+        from .llm_api_client.huggingface_client import HuggingFaceClient
+    except ImportError:
+        pass
+    try:
+        from .llm_api_client.anthropic_client import AnthropicClient
+    except ImportError:
+        pass
+    try:
+        from .llm_api_client.cortex_client import CortexClient
+    except ImportError:
+        pass
+    try:
+        from .llm_cost_calculation.openai_cost_calculation import openai_cost_calculation
+    except ImportError:
+        pass
+    try:
+        from .llm_cost_calculation.snowflake_cortex_cost_calculation import snowflake_cortex_cost_calculation
+    except ImportError:
+        pass
+    try:
+        from .llm_cost_calculation.anthropic_cost_calculation import anthropic_cost_calculation
+    except ImportError:
+        pass
+except ImportError:
+    pass
+# load local clients
+try:
+    from .llm_client.local_client import LocalClient, LocalClientConfig
+except ImportError:
+    pass
+# load sync support
+try:
+    from .sync.sync_llm_api_client_factory import init_sync_llm_api_client
+except ImportError:
+    pass
+
+
+
+try:
+    from .llm_api_client.core.llm import load_model
+    from .llm_api_client.core.model_schema import  LLMConfig
+    from .llm_cost_calculation.cost_tracker import CostCallbackHandler
+except ImportError:
+    pass

sfn_llm_client-0.3.3/sfn_llm_client/agent/agent_executor.py

@@ -0,0 +1,39 @@
+# from langchain.agents import AgentExecutor, create_openai_tools_agent
+# from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
+# from langchain_community.callbacks import get_openai_callback
+
+
+
+# def call_agent(llm_client, tools, configuration):
+#     raw_prompt = configuration["system_prompt"]
+#     user_query = configuration["user_query"]
+#     chat_history = configuration.get("chat_history", None)
+
+#     # clean prompt
+#     _prompt = raw_prompt.rstrip("\n").splitlines()
+#     _prompt= " ".join(_prompt).replace("'", " ").replace('"', " ")
+#     print(_prompt)
+#     # Create prompt template with explicit input variables
+#     prompt = ChatPromptTemplate.from_messages([
+#         ("system", _prompt),
+#         MessagesPlaceholder(variable_name="chat_history"),
+#         ("human", "{{query}}"),  # Use query instead of input
+#         MessagesPlaceholder(variable_name="agent_scratchpad"),
+#     ])
+
+#     agent = create_openai_tools_agent(llm=llm_client, tools=tools, prompt=prompt)
+#     agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=False, handle_parsing_errors=True, return_intermediate_steps=False, max_iterations=5)
+#     # Run the agent
+#     # TODO: for cortex?
+#     with get_openai_callback() as cb:
+#         response = agent_executor.invoke({
+#             "input": user_query,
+#             "chat_history": chat_history
+#         })
+#     token_cost_summary = {
+#             "prompt_tokens": cb.prompt_tokens,
+#             "completion_tokens": cb.completion_tokens,
+#             "total_tokens": cb.total_tokens,
+#             "total_cost_usd": cb.total_cost,
+#         }
+#     return response, token_cost_summary

sfn_llm_client-0.3.3/sfn_llm_client/llm_api_client/ai21_client.py

@@ -0,0 +1,45 @@
+from typing import Optional
+
+from sfn_llm_client.llm_api_client.base_llm_api_client import BaseLLMAPIClient, LLMAPIClientConfig
+from sfn_llm_client.utils.consts import PROMPT_KEY
+
+COMPLETE_PATH = "complete"
+TOKENIZE_PATH = "tokenize"
+BASE_URL = "https://api.ai21.com/studio/v1/"
+COMPLETIONS_KEY = "completions"
+DATA_KEY = "data"
+TEXT_KEY = "text"
+TOKENS_KEY = "tokens"
+AUTH_HEADER = "Authorization"
+BEARER_TOKEN = "Bearer "
+
+
+class AI21Client(BaseLLMAPIClient):
+    def __init__(self, config: LLMAPIClientConfig):
+        super().__init__(config)
+        if self._base_url is None:
+            self._base_url = BASE_URL
+        self._headers[AUTH_HEADER] = BEARER_TOKEN + self._api_key
+
+    async def text_completion(self, prompt: str, model: Optional[str] = None, max_tokens: int = 16,
+                              temperature: float = 0.7, top_p: float = 1,**kwargs) -> list[str]:
+        model = model or self._default_model
+        kwargs[PROMPT_KEY] = prompt
+        kwargs["topP"] = kwargs.pop("topP", top_p)
+        kwargs["maxTokens"] = kwargs.pop("maxTokens", max_tokens)
+        kwargs["temperature"] = temperature
+        response = await self._session.post(self._base_url + model + "/" + COMPLETE_PATH,
+                                            json=kwargs,
+                                            headers=self._headers,
+                                            raise_for_status=True)
+        response_json = await response.json()
+        completions = response_json[COMPLETIONS_KEY]
+        return [completion[DATA_KEY][TEXT_KEY] for completion in completions]
+
+    async def get_tokens_count(self, text: str, **kwargs) -> int:
+        response = await self._session.post(self._base_url + TOKENIZE_PATH,
+                                            json={TEXT_KEY: text},
+                                            headers=self._headers,
+                                            raise_for_status=True)
+        response_json = await response.json()
+        return len(response_json[TOKENS_KEY])

sfn_llm_client-0.3.3/sfn_llm_client/llm_api_client/aleph_alpha_client.py

@@ -0,0 +1,70 @@
+from typing import Optional
+
+from sfn_llm_client.llm_api_client.base_llm_api_client import BaseLLMAPIClient, LLMAPIClientConfig
+from sfn_llm_client.utils.consts import PROMPT_KEY
+
+COMPLETE_PATH = "complete"
+TOKENIZE_PATH = "tokenize"
+EMBEDDING_PATH = "semantic_embed"
+BASE_URL = "https://api.aleph-alpha.com/"
+COMPLETIONS_KEY = "completions"
+TEXT_KEY = "completion"
+TOKENS_IDS_KEY = "token_ids"
+TOKENS_KEY = "tokens"
+REPRESENTATION_KEY = "representation"
+REPRESENTATION_DEFAULT_VALUE = "symmetric"
+EMBEDDING_KEY = "embedding"
+AUTH_HEADER = "Authorization"
+BEARER_TOKEN = "Bearer "
+MAX_TOKENS_KEY = "maximum_tokens"
+
+
+class AlephAlphaClient(BaseLLMAPIClient):
+    def __init__(self, config: LLMAPIClientConfig):
+        super().__init__(config)
+        if self._base_url is None:
+            self._base_url = BASE_URL
+        self._headers[AUTH_HEADER] = BEARER_TOKEN + self._api_key
+
+    async def text_completion(self, prompt: str, model: Optional[str] = None, max_tokens: Optional[int] = None,
+                              temperature: float = 0,top_p: float = 0, **kwargs) -> \
+            list[str]:
+        self._set_model_in_kwargs(kwargs, model)
+        if max_tokens is None:
+            raise ValueError("max_tokens must be specified")
+        kwargs[PROMPT_KEY] = prompt
+        kwargs["top_p"] = top_p
+        kwargs["maximum_tokens"] = kwargs.pop("maximum_tokens", max_tokens)
+        kwargs["temperature"] = temperature
+        response = await self._session.post(self._base_url + COMPLETE_PATH,
+                                            json=kwargs,
+                                            headers=self._headers,
+                                            raise_for_status=True)
+        response_json = await response.json()
+        completions = response_json[COMPLETIONS_KEY]
+        return [completion[TEXT_KEY] for completion in completions]
+
+    async def embedding(self, text: str, model: Optional[str] = None,
+                        representation: str = REPRESENTATION_DEFAULT_VALUE,
+                        **kwargs) -> list[float]:
+        self._set_model_in_kwargs(kwargs, model)
+        kwargs[REPRESENTATION_KEY] = representation
+        kwargs[PROMPT_KEY] = text
+        response = await self._session.post(self._base_url + EMBEDDING_PATH,
+                                            json=kwargs,
+                                            headers=self._headers,
+                                            raise_for_status=True)
+        response_json = await response.json()
+        return response_json[EMBEDDING_KEY]
+
+    async def get_tokens_count(self, text: str, model: Optional[str] = None, **kwargs) -> int:
+        self._set_model_in_kwargs(kwargs, model)
+        kwargs[TOKENS_KEY] = False
+        kwargs[TOKENS_IDS_KEY] = True
+        kwargs[PROMPT_KEY] = text
+        response = await self._session.post(self._base_url + TOKENIZE_PATH,
+                                            json=kwargs,
+                                            headers=self._headers,
+                                            raise_for_status=True)
+        response_json = await response.json()
+        return len(response_json[TOKENS_IDS_KEY])

sfn_llm_client-0.3.3/sfn_llm_client/llm_api_client/anthropic_client.py

@@ -0,0 +1,110 @@
+from typing import Optional
+import time
+from anthropic import Anthropic
+
+from sfn_llm_client.llm_api_client.base_llm_api_client import BaseLLMAPIClient, LLMAPIClientConfig, ChatMessage, Role
+from sfn_llm_client.utils.consts import PROMPT_KEY
+from sfn_llm_client.utils.logging import setup_logger
+from sfn_llm_client.utils.retry_with import retry_with
+from sfn_llm_client.llm_cost_calculation.anthropic_cost_calculation import anthropic_cost_calculation
+
+COMPLETE_PATH = "complete"
+BASE_URL = "https://api.anthropic.com/v1/"
+COMPLETIONS_KEY = "completion"
+AUTH_HEADER = "x-api-key"
+ACCEPT_HEADER = "Accept"
+VERSION_HEADER = "anthropic-version"
+ACCEPT_VALUE = "application/json"
+MAX_TOKENS_KEY = "max_tokens_to_sample"
+USER_PREFIX = "Human:"
+ASSISTANT_PREFIX = "Assistant:"
+START_PREFIX = "\n\n"
+SYSTEM_START_PREFIX = "<admin>"
+SYSTEM_END_PREFIX = "</admin>"
+
+
+class AnthropicClient(BaseLLMAPIClient):
+    def __init__(self, config: LLMAPIClientConfig):
+        super().__init__(config)
+        self._base_url = config.base_url or BASE_URL
+        self._anthropic = Anthropic(api_key=config.api_key)
+        self.logger, _ = setup_logger(logger_name="OpenAIClient")
+        self._headers = {
+            VERSION_HEADER: self._anthropic.default_headers[VERSION_HEADER],
+            ACCEPT_HEADER: ACCEPT_VALUE,
+            AUTH_HEADER: config.api_key,
+        }
+
+    @retry_with(retries=3, retry_delay=3.0, backoff=True)
+    def chat_completion(self, messages: list[ChatMessage], model: Optional[str] = None,
+                    max_tokens: Optional[int] = None, temperature: float = 1.0, retries: int = 3,
+                    retry_delay: float = 3.0, **kwargs) -> list[str]:
+        """
+        This method performs chat completion with retry logic for handling exceptions or empty responses.
+        It also calculates token consumption based on the API response.
+
+        :param messages: List of ChatMessage objects representing the conversation history.
+        :param model: Optional model name to be used.
+        :param max_tokens: Maximum number of tokens to generate in the response.
+        :param temperature: Controls the creativity of the response.
+        :param retries: Number of retries in case of failure.
+        :param retry_delay: Delay in seconds between retries.
+
+        :return: ChatCompletion object containing the generated text and other metadata,
+                or None if all retries fail.
+        """
+
+        self.logger.info("Started running llm client sdk chat completion...")
+        del kwargs['session'] #deleting session key not required for anthropic
+        response = self._anthropic.messages.create(model=model, max_tokens=max_tokens, messages=messages, **kwargs)
+
+        # Calculate token consumption cost
+        token_cost_summary = anthropic_cost_calculation(
+            response.usage.input_tokens,
+            response.usage.output_tokens,
+            model=model,
+        )
+
+        return response, token_cost_summary
+
+
+    def text_completion(self, prompt: str, model: Optional[str] = None, max_tokens: Optional[int] = None,
+                         temperature: float = 1.0, top_p: Optional[float] = None, **kwargs) -> list[str]:
+        if max_tokens is None and kwargs.get(MAX_TOKENS_KEY) is None:
+            raise ValueError(f"max_tokens or {MAX_TOKENS_KEY} must be specified")
+
+        kwargs[PROMPT_KEY] = prompt
+        kwargs[MAX_TOKENS_KEY] = kwargs.pop(MAX_TOKENS_KEY, max_tokens)
+        kwargs["temperature"] = temperature
+        if top_p:
+            kwargs["top_p"] = top_p
+
+        self._set_model_in_kwargs(kwargs, model)
+        response = self._anthropic.messages.create(model=model, messages=[{"content": prompt}], **kwargs)
+        return [response[COMPLETIONS_KEY]]
+
+    def get_chat_tokens_count(self, messages: list[ChatMessage], **kwargs) -> int:
+        return self.get_tokens_count(self.messages_to_text(messages), **kwargs)
+
+    def get_tokens_count(self, text: str, **kwargs) -> int:
+        return self._anthropic.count_tokens(text)
+        # return sum(len(word.split()) for word in text.split("\n"))  # Approximate token count based on words
+
+    def messages_to_text(self, messages: list[ChatMessage]) -> str:
+        prompt = START_PREFIX
+        prompt += START_PREFIX.join(map(self._message_to_prompt, messages))
+        if messages[-1].role != Role.ASSISTANT:
+            prompt += START_PREFIX
+            prompt += self._message_to_prompt(ChatMessage(role=Role.ASSISTANT, content=""))
+        return prompt.rstrip()
+
+    @staticmethod
+    def _message_to_prompt(message: ChatMessage) -> str:
+        if message.role == Role.USER:
+            return f"{USER_PREFIX} {message.content}"
+        elif message.role == Role.ASSISTANT:
+            return f"{ASSISTANT_PREFIX} {message.content}"
+        elif message.role == Role.SYSTEM:
+            return f"{USER_PREFIX} {SYSTEM_START_PREFIX}{message.content}{SYSTEM_END_PREFIX}"
+        else:
+            raise ValueError(f"Unknown role: {message.role}")

sfn_llm_client-0.3.3/sfn_llm_client/llm_api_client/base_llm_api_client.py

@@ -0,0 +1,67 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Optional
+
+from dataclasses_json import dataclass_json, config
+
+# try:
+#     from aiohttp import ClientSession
+# except ImportError:
+#     ClientSession = Any
+
+from sfn_llm_client import BaseLLMClient
+from sfn_llm_client.utils.consts import MODEL_KEY
+
+
+class Role(Enum):
+    SYSTEM = "system"
+    USER = "user"
+    ASSISTANT = "assistant"
+
+
+@dataclass_json
+@dataclass
+class ChatMessage:
+    role: Role = field(metadata=config(encoder=lambda role: role.value, decoder=Role))
+    content: str
+    name: Optional[str] = field(default=None, metadata=config(exclude=lambda name: name is None))
+    example: bool = field(default=False, metadata=config(exclude=lambda _: True))
+
+
+@dataclass
+class LLMAPIClientConfig:
+    api_key: str
+    # session: ClientSession
+    base_url: Optional[str] = None
+    default_model: Optional[str] = None
+    headers: dict[str, Any] = field(default_factory=dict)
+
+
+class BaseLLMAPIClient(BaseLLMClient, ABC):
+    def __init__(self, config: LLMAPIClientConfig):
+        self._api_key: str = config.api_key
+        # self._session: ClientSession = config.session
+        self._base_url: str = config.base_url
+        self._default_model: str = config.default_model
+        self._headers: dict[str, str] = config.headers
+
+    @abstractmethod
+    async def text_completion(self, prompt: str, model: Optional[str] = None, max_tokens: Optional[int] = None,
+                              temperature: Optional[float] = None, top_p: Optional[float] = None, **kwargs) -> list[str]:
+        raise NotImplementedError()
+
+    async def chat_completion(self, messages: list[ChatMessage], temperature: float = 0,
+                              max_tokens: int = 16, model: Optional[str] = None, **kwargs) -> list[str]:
+        raise NotImplementedError()
+
+    async def embedding(self, text: str, model: Optional[str] = None, **kwargs) -> list[float]:
+        raise NotImplementedError()
+
+    async def get_chat_tokens_count(self, messages: list[ChatMessage], **kwargs) -> int:
+        raise NotImplementedError()
+
+    def _set_model_in_kwargs(self, kwargs, model: Optional[str]) -> None:
+        if model is not None:
+            kwargs[MODEL_KEY] = model
+        kwargs.setdefault(MODEL_KEY, self._default_model)