diff --git a/.gitignore b/.gitignore index 82f9275..7b6caf3 100644 --- a/.gitignore +++ b/.gitignore @@ -159,4 +159,4 @@ cython_debug/ # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore # and can be added to the global gitignore or merged into this file. For a more nuclear # option (not recommended) you can uncomment the following to ignore the entire idea folder. -#.idea/ +.idea/ diff --git a/README.md b/README.md index 9c91b9b..dc42e44 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@ Inspired by [tinygrad](https://github.com/tinygrad/tinygrad) and [simpleaichat](https://github.com/minimaxir/simpleaichat/tree/main/simpleaichat), `tiny-ai-client` is the easiest way to use and switch LLMs with vision and tool usage support. It works because it is `tiny`, `simple` and most importantly `fun` to develop. -I want to change LLMs with ease, while knowing what is happening under the hood. Langchain is cool, but became bloated, complicated there is just too much chaos going on. I want to keep it simple, easy to understand and easy to use. If you want to use a LLM and have an API key, you should not need to read a 1000 lines of code and write `response.choices[0].message.content` to get the response. +I want to change LLMs with ease, while knowing what is happening under the hood. Langchain is cool, but became bloated, complicated there is just too much chaos going on. I want to keep it simple, easy to understand and easy to use. If you want to use a LLM and have an API key, you should not need to read 1000 lines of code and write `response.choices[0].message.content` to get the response. Simple and tiny, that's the goal. diff --git a/examples/anthropic_.py b/examples/anthropic_.py index 9849f64..b98f2f1 100644 --- a/examples/anthropic_.py +++ b/examples/anthropic_.py @@ -62,6 +62,6 @@ def main(): if __name__ == "__main__": - os.environ["ANTHROPIC_API_KEY"] = None + os.environ["ANTHROPIC_API_KEY"] = "" main() asyncio.run(async_ai_main()) diff --git a/examples/gemini_.py b/examples/gemini_.py index 2ccd3ee..d88cb27 100644 --- a/examples/gemini_.py +++ b/examples/gemini_.py @@ -58,6 +58,6 @@ async def async_ai_main(): if __name__ == "__main__": - os.environ["GOOGLE_API_KEY"] = None + os.environ["GOOGLE_API_KEY"] = "" main() asyncio.run(async_ai_main()) diff --git a/examples/openai_.py b/examples/openai_.py index c7fabdc..fd4efab 100644 --- a/examples/openai_.py +++ b/examples/openai_.py @@ -66,6 +66,6 @@ def main(): if __name__ == "__main__": - os.environ["OPENAI_API_KEY"] = None + os.environ["OPENAI_API_KEY"] = "" main() asyncio.run(async_ai_main()) diff --git a/requirements.txt b/requirements.txt index 40c5ca1..b68e3c8 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,5 @@ pydantic==2.7.3 openai==1.31.0 anthropic==0.28.0 +pillow~=10.3.0 +google-generativeai diff --git a/tiny_ai_client/anthropic_.py b/tiny_ai_client/anthropic_.py index 6c2cdcb..b79d098 100644 --- a/tiny_ai_client/anthropic_.py +++ b/tiny_ai_client/anthropic_.py @@ -101,7 +101,7 @@ def call_llm_provider( temperature: int | None, max_new_tokens: int | None, timeout: int, - ) -> str: + ) -> Message: kwargs = {} input_messages, system = model_input if temperature is not None: @@ -139,7 +139,7 @@ async def async_call_llm_provider( temperature: int | None, max_new_tokens: int | None, timeout: int, - ) -> str: + ) -> Message: kwargs = {} input_messages, system = model_input if temperature is not None: diff --git a/tiny_ai_client/gemini_.py b/tiny_ai_client/gemini_.py index 8bd9910..b213bfa 100644 --- a/tiny_ai_client/gemini_.py +++ b/tiny_ai_client/gemini_.py @@ -31,7 +31,6 @@ def build_model_input(self, messages: List["Message"]) -> Any: history = [] local_messages = deepcopy(messages) system = None - message = None for message in local_messages: if message.role == "system": @@ -45,6 +44,7 @@ def build_model_input(self, messages: List["Message"]) -> Any: if message.text is not None: parts.append(message.text) if message.images is not None: + # noinspection PyTypeChecker parts.extend(message.images) history.append( { @@ -53,7 +53,7 @@ def build_model_input(self, messages: List["Message"]) -> Any: } ) - return (system, history) + return system, history def call_llm_provider( self, @@ -61,7 +61,7 @@ def call_llm_provider( temperature: int | None, max_new_tokens: int | None, timeout: int, - ) -> str: + ) -> Message: system, history = model_input generation_config_kwargs = {} @@ -92,7 +92,7 @@ async def async_call_llm_provider( temperature: int | None, max_new_tokens: int | None, timeout: int, - ) -> str: + ) -> Message: system, history = model_input generation_config_kwargs = {} diff --git a/tiny_ai_client/models.py b/tiny_ai_client/models.py index 3365a69..a563fa7 100644 --- a/tiny_ai_client/models.py +++ b/tiny_ai_client/models.py @@ -5,17 +5,41 @@ from tiny_ai_client.tools import json_to_function_input +""" +Using the factory design pattern approach +provides a way to encapsulate the instantiate-and-return +The logic is code easy to code and test, much cleaner +Following the best object-oriented design practices (DRI, SRP, OCP, Factory). +""" + + +class ClientWrapperFactory: + _client_wrapper_mapping = { + "gpt": "tiny_ai_client.openai_.OpenAIClientWrapper", + "claude": "tiny_ai_client.anthropic_.AnthropicClientWrapper", + "gemini": "tiny_ai_client.gemini_.GeminiClientWrapper" + } + + @staticmethod + def create(model_name: str, tools: List[Union[Callable, Dict]]) -> "LLMClientWrapper": + for key, wrapper_path in ClientWrapperFactory._client_wrapper_mapping.items(): + if key in model_name: + module_path, wrapper_class = wrapper_path.rsplit('.', 1) + module = __import__(module_path, fromlist=[wrapper_class]) + return getattr(module, wrapper_class)(model_name, tools) + raise NotImplementedError(f"{model_name=} not supported") + class AI: def __init__( - self, - model_name: str, - system: str | None = None, - temperature: int = 1, - max_new_tokens: int | None = None, - timeout: int = 30, - tools: List[Union[Callable, Dict]] | None = None, - chat: List["Message"] | None = None, + self, + model_name: str, + system: str | None = None, + temperature: int = 1, + max_new_tokens: int | None = None, + timeout: int = 30, + tools: List[Union[Callable, Dict]] | None = None, + chat: List["Message"] | None = None, ): # llm sampling parameters self.temperature: int = temperature @@ -36,12 +60,12 @@ def __init__( ) def __call__( - self, - message: str | None = None, - max_new_tokens: int | None = None, - temperature: int | None = 1, - timeout: int | None = None, - images: List[PIL_Image.Image] | PIL_Image.Image | None = None, + self, + message: str | None = None, + max_new_tokens: int | None = None, + temperature: int | None = 1, + timeout: int | None = None, + images: List[PIL_Image.Image] | PIL_Image.Image | None = None, ) -> str: max_new_tokens = max_new_tokens or self.max_new_tokens temperature = temperature or self.temperature @@ -66,42 +90,30 @@ def __call__( response_msg.tool_call.result if response_msg.tool_call else "" ) + # noinspection PyMethodMayBeStatic def get_llm_client_wrapper( - self, model_name: str, tools: List[Union[Callable, Dict]] + self, model_name: str, tools: List[Union[Callable, Dict]] ) -> "LLMClientWrapper": - if "gpt" in model_name: - from tiny_ai_client.openai_ import OpenAIClientWrapper - - return OpenAIClientWrapper(model_name, tools) - if "claude" in model_name: - from tiny_ai_client.anthropic_ import AnthropicClientWrapper - - return AnthropicClientWrapper(model_name, tools) - - if "gemini" in model_name: - from tiny_ai_client.gemini_ import GeminiClientWrapper - - return GeminiClientWrapper(model_name, tools) - - raise NotImplementedError(f"{model_name=} not supported") + return ClientWrapperFactory.create(model_name, tools) @property def model_name(self) -> str: - return self._model_name + return self.model_name @model_name.setter def model_name(self, value: str) -> None: + self.model_name = value self.client_wrapper = self.get_llm_client_wrapper(value, self.tools) class AsyncAI(AI): async def __call__( - self, - message: str | None = None, - max_new_tokens: int | None = None, - temperature: int | None = 1, - timeout: int | None = None, - images: List[PIL_Image.Image] | PIL_Image.Image | None = None, + self, + message: str | None = None, + max_new_tokens: int | None = None, + temperature: int | None = 1, + timeout: int | None = None, + images: List[PIL_Image.Image] | PIL_Image.Image | None = None, ) -> str: max_new_tokens = max_new_tokens or self.max_new_tokens temperature = temperature or self.temperature @@ -148,25 +160,25 @@ class Config: class LLMClientWrapper: - def build_model_input(self, messages: List["Message"]) -> Any: + def build_model_input(self, messages: List[Message]) -> Any: raise NotImplementedError def call_llm_provider( - self, - model_input: Any, - temperature: int | None, - max_new_tokens: int | None, - timeout: int, - ) -> "Message": + self, + model_input: Any, + temperature: int | None, + max_new_tokens: int | None, + timeout: int, + ) -> Message: raise NotImplementedError def __call__( - self, - max_new_tokens: int | None, - temperature: int | None, - timeout: int, - chat: List["Message"], - ) -> "Message": + self, + max_new_tokens: int | None, + temperature: int | None, + timeout: int, + chat: List[Message], + ) -> Message: model_input = self.build_model_input(chat) return self.call_llm_provider( model_input, @@ -176,21 +188,21 @@ def __call__( ) async def async_call_llm_provider( - self, - model_input: Any, - temperature: int | None, - max_new_tokens: int | None, - timeout: int, - ) -> "Message": + self, + model_input: Any, + temperature: int | None, + max_new_tokens: int | None, + timeout: int, + ) -> Message: raise NotImplementedError async def acall( - self, - max_new_tokens: int | None, - temperature: int | None, - timeout: int, - chat: List["Message"], - ) -> "Message": + self, + max_new_tokens: int | None, + temperature: int | None, + timeout: int, + chat: List[Message], + ) -> Message: model_input = self.build_model_input(chat) return await self.async_call_llm_provider( model_input, diff --git a/tiny_ai_client/openai_.py b/tiny_ai_client/openai_.py index 3c72798..cab5d55 100644 --- a/tiny_ai_client/openai_.py +++ b/tiny_ai_client/openai_.py @@ -68,7 +68,7 @@ def call_llm_provider( temperature: int | None, max_new_tokens: int | None, timeout: int, - ) -> str: + ) -> Message: kwargs = {} if temperature is not None: kwargs["temperature"] = temperature @@ -103,7 +103,7 @@ async def async_call_llm_provider( temperature: int | None, max_new_tokens: int | None, timeout: int, - ) -> str: + ) -> Message: kwargs = {} if temperature is not None: kwargs["temperature"] = temperature