From 903f7eb78a4dd23291ac599271c5eb219d85e9b5 Mon Sep 17 00:00:00 2001 From: PranavPuranik Date: Mon, 12 Aug 2024 22:00:06 -0500 Subject: [PATCH 1/4] final --- embedchain/docs/components/llms.mdx | 54 +++++++++++++++++++++++++++++ embedchain/embedchain/factory.py | 1 + embedchain/embedchain/llm/octoai.py | 34 ++++++++++++++++++ embedchain/embedchain/utils/misc.py | 1 + embedchain/tests/llm/test_octoai.py | 36 +++++++++++++++++++ 5 files changed, 126 insertions(+) create mode 100644 embedchain/embedchain/llm/octoai.py create mode 100644 embedchain/tests/llm/test_octoai.py diff --git a/embedchain/docs/components/llms.mdx b/embedchain/docs/components/llms.mdx index c0034dc126..11b4cc7089 100644 --- a/embedchain/docs/components/llms.mdx +++ b/embedchain/docs/components/llms.mdx @@ -25,6 +25,7 @@ Embedchain comes with built-in support for various popular large language models + ## OpenAI @@ -840,6 +841,59 @@ answer = app.query("What is the net worth of Elon Musk today?") ``` + +## Octo AI + +[Octo AI](https://octo.ai/) supports a lot of models from meta, microsoft, mistral, etc... These models are available in the [Octo AI TextGen](https://octoai.cloud/text) and are ready to use in production. + + +### Usage + +In order to use LLMs from Octo AI, sign up on [Octo AI](https://octo.ai/). + +Generate token from account settings. Set the token `OCTOAI_API_TOKEN` environment variable. + +Below is an example of how to use LLM model and embedding model from NVIDIA AI: + + + +```python main.py +import os +from embedchain import App + +os.environ["OCTOAI_API_TOKEN"] = "enter_token" + +config = { + "llm": { + "provider": "octoai", + "config": { + "model": "llama-2-13b-chat-fp16", + "max_tokens": 200, + "temperature": 0.1, + "top_p": 0.9, + } + }, + "embedder": { + "provider": "huggingface", + "config": { + "model": 'nomic-ai/nomic-embed-text-v1', + "model_kwargs": { + "trust_remote_code": True, + } + } + } +} + +app = App.from_config(config=config) + +app.add("https://www.forbes.com/profile/elon-musk") +answer = app.query("What is the net worth of Elon Musk today?") +# Answer: Elon Musk's net worth is $222.6 billion as of August 12, 2024. This reflects a change since the previous trading day. He remains the richest person in the world today. +``` + + + + ## Token Usage You can get the cost of the query by setting `token_usage` to `True` in the config file. This will return the token details: `prompt_tokens`, `completion_tokens`, `total_tokens`, `total_cost`, `cost_currency`. diff --git a/embedchain/embedchain/factory.py b/embedchain/embedchain/factory.py index 69636286cf..b61672a3af 100644 --- a/embedchain/embedchain/factory.py +++ b/embedchain/embedchain/factory.py @@ -27,6 +27,7 @@ class LlmFactory: "groq": "embedchain.llm.groq.GroqLlm", "nvidia": "embedchain.llm.nvidia.NvidiaLlm", "vllm": "embedchain.llm.vllm.VLLM", + "octoai": "embedchain.llm.octoai.OctoAILlm", } provider_to_config_class = { "embedchain": "embedchain.config.llm.base.BaseLlmConfig", diff --git a/embedchain/embedchain/llm/octoai.py b/embedchain/embedchain/llm/octoai.py new file mode 100644 index 0000000000..343539c6ca --- /dev/null +++ b/embedchain/embedchain/llm/octoai.py @@ -0,0 +1,34 @@ +import os +from typing import Optional + +from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler +from langchain_community.llms.octoai_endpoint import OctoAIEndpoint + +from embedchain.config import BaseLlmConfig +from embedchain.helpers.json_serializable import register_deserializable +from embedchain.llm.base import BaseLlm + + +@register_deserializable +class OctoAILlm(BaseLlm): + def __init__(self, config: Optional[BaseLlmConfig] = None): + assert "OCTOAI_API_TOKEN" in os.environ, "Please set OCTOAI_API_TOKEN as environment variable." + super().__init__(config=config) + + def get_llm_model_answer(self, prompt): + return self._get_answer(prompt, self.config) + + @staticmethod + def _get_answer(prompt: str, config: BaseLlmConfig) -> str: + chat = OctoAIEndpoint( + model_name=config.model, + max_tokens=config.max_tokens, + temperature=config.temperature, + top_p=config.top_p, + streaming=config.stream, + callbacks=config.callbacks + if (not config.stream) or (config.stream and config.callbacks) + else [StreamingStdOutCallbackHandler()], + ) + + return chat.invoke(prompt) diff --git a/embedchain/embedchain/utils/misc.py b/embedchain/embedchain/utils/misc.py index 7c5468ec93..d787d2685b 100644 --- a/embedchain/embedchain/utils/misc.py +++ b/embedchain/embedchain/utils/misc.py @@ -418,6 +418,7 @@ def validate_config(config_data): "vllm", "groq", "nvidia", + "octoai", ), Optional("config"): { Optional("model"): str, diff --git a/embedchain/tests/llm/test_octoai.py b/embedchain/tests/llm/test_octoai.py new file mode 100644 index 0000000000..18fe3b3915 --- /dev/null +++ b/embedchain/tests/llm/test_octoai.py @@ -0,0 +1,36 @@ +import os +import pytest +from embedchain.config import BaseLlmConfig + +from embedchain.llm.octoai import OctoAILlm + +@pytest.fixture +def octoai_env(): + os.environ["OCTOAI_API_TOKEN"] = "test_api_token" + +@pytest.fixture +def octoai_llm_config(): + config = BaseLlmConfig( + temperature=0.7, + model="llama-2-13b-chat-fp16", + max_tokens=50, + top_p=0.9, + ) + return config + + +def test_get_answer(octoai_llm_config, octoai_env, mocker): + mocked_get_answer = mocker.patch("embedchain.llm.octoai.OctoAILlm._get_answer", return_value="Test answer") + + octoai_llm = OctoAILlm(octoai_llm_config) + answer = octoai_llm.get_llm_model_answer("Test query") + + assert answer == "Test answer" + mocked_get_answer.assert_called_once() + +def test_octo_env_variable(octoai_llm_config): + + with pytest.raises(AssertionError): + octoai_llm = OctoAILlm(octoai_llm_config) + + From 39e6b72056ace7d7129ce116947a0049591242b5 Mon Sep 17 00:00:00 2001 From: PranavPuranik Date: Mon, 12 Aug 2024 22:49:19 -0500 Subject: [PATCH 2/4] ci tests --- embedchain/tests/llm/test_octoai.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/embedchain/tests/llm/test_octoai.py b/embedchain/tests/llm/test_octoai.py index 18fe3b3915..7af85c6d89 100644 --- a/embedchain/tests/llm/test_octoai.py +++ b/embedchain/tests/llm/test_octoai.py @@ -7,6 +7,8 @@ @pytest.fixture def octoai_env(): os.environ["OCTOAI_API_TOKEN"] = "test_api_token" + yield + del os.environ["OCTOAI_API_TOKEN"] @pytest.fixture def octoai_llm_config(): @@ -31,6 +33,4 @@ def test_get_answer(octoai_llm_config, octoai_env, mocker): def test_octo_env_variable(octoai_llm_config): with pytest.raises(AssertionError): - octoai_llm = OctoAILlm(octoai_llm_config) - - + _ = OctoAILlm(octoai_llm_config) From 0c47009e636e66070f8cfde46bc36cbb654612ae Mon Sep 17 00:00:00 2001 From: PranavPuranik Date: Fri, 30 Aug 2024 19:20:36 -0500 Subject: [PATCH 3/4] checking config for api_key --- embedchain/embedchain/llm/octoai.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/embedchain/embedchain/llm/octoai.py b/embedchain/embedchain/llm/octoai.py index 343539c6ca..5bb6c40886 100644 --- a/embedchain/embedchain/llm/octoai.py +++ b/embedchain/embedchain/llm/octoai.py @@ -12,7 +12,8 @@ @register_deserializable class OctoAILlm(BaseLlm): def __init__(self, config: Optional[BaseLlmConfig] = None): - assert "OCTOAI_API_TOKEN" in os.environ, "Please set OCTOAI_API_TOKEN as environment variable." + assert "OCTOAI_API_TOKEN" in os.environ or config.api_key, \ + "Please set OCTOAI_API_TOKEN as environment variable." super().__init__(config=config) def get_llm_model_answer(self, prompt): From 23c40013995fb43e2c029685ea163cda6ac32c57 Mon Sep 17 00:00:00 2001 From: PranavPuranik Date: Sun, 13 Oct 2024 12:10:18 -0500 Subject: [PATCH 4/4] adding octoai api token param --- embedchain/embedchain/llm/octoai.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/embedchain/embedchain/llm/octoai.py b/embedchain/embedchain/llm/octoai.py index 5bb6c40886..08ea90757f 100644 --- a/embedchain/embedchain/llm/octoai.py +++ b/embedchain/embedchain/llm/octoai.py @@ -21,15 +21,20 @@ def get_llm_model_answer(self, prompt): @staticmethod def _get_answer(prompt: str, config: BaseLlmConfig) -> str: + + + octoai_api_key = os.getenv("OCTOAI_API_TOKEN") or config.api_key + callbacks = config.callbacks if (not config.stream) or (config.stream and config.callbacks) \ + else [StreamingStdOutCallbackHandler()] + chat = OctoAIEndpoint( + octoai_api_token=octoai_api_key, model_name=config.model, max_tokens=config.max_tokens, temperature=config.temperature, top_p=config.top_p, streaming=config.stream, - callbacks=config.callbacks - if (not config.stream) or (config.stream and config.callbacks) - else [StreamingStdOutCallbackHandler()], + callbacks=callbacks, ) return chat.invoke(prompt)