From 7decf5fbc49ded836dcf8fba28b59b0fbcfce8bf Mon Sep 17 00:00:00 2001 From: Houpus <33534878+RoacherM@users.noreply.github.com> Date: Tue, 4 Jun 2024 21:38:23 +0800 Subject: [PATCH] Support 3rd party openai api (#365) * Support for specifying third-party OPENAI proxy services * Remove unnecessary changes * Support for specifying third-party OPENAI proxy services * update get_generator method --------- Co-authored-by: Chih-Yu Yeh --- docker/.env.example | 1 + docker/docker-compose.yaml | 1 + wren-ai-service/docker/docker-compose.yml | 1 + wren-ai-service/src/providers/llm/openai.py | 28 +++++++++++++++------ 4 files changed, 24 insertions(+), 7 deletions(-) diff --git a/docker/.env.example b/docker/.env.example index 446c4f541..bdf9101a0 100644 --- a/docker/.env.example +++ b/docker/.env.example @@ -17,6 +17,7 @@ WREN_BOOTSTRAP_VERSION=0.1.4 # keys # CHANGE THIS TO YOUR OPENAI API KEY OPENAI_API_KEY=sk-1234567890 +OPENAI_API_BASE=https://api.openai.com/v1 OPENAI_GENERATION_MODEL=gpt-3.5-turbo # SQL Protocol diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml index 7943957af..5fe195f6a 100644 --- a/docker/docker-compose.yaml +++ b/docker/docker-compose.yaml @@ -45,6 +45,7 @@ services: environment: WREN_AI_SERVICE_PORT: ${WREN_AI_SERVICE_PORT} OPENAI_API_KEY: ${OPENAI_API_KEY} + OPENAI_API_BASE: ${OPENAI_API_BASE} OPENAI_GENERATION_MODEL: ${OPENAI_GENERATION_MODEL} QDRANT_HOST: qdrant WREN_ENGINE_ENDPOINT: http://wren-engine:${WREN_ENGINE_PORT} diff --git a/wren-ai-service/docker/docker-compose.yml b/wren-ai-service/docker/docker-compose.yml index a5b8a77a0..da6afdfc7 100644 --- a/wren-ai-service/docker/docker-compose.yml +++ b/wren-ai-service/docker/docker-compose.yml @@ -13,6 +13,7 @@ services: environment: WREN_AI_SERVICE_PORT: ${WREN_AI_SERVICE_PORT} OPENAI_API_KEY: ${OPENAI_API_KEY} + OPENAI_API_BASE: ${OPENAI_API_BASE} OPENAI_GENERATION_MODEL: ${OPENAI_GENERATION_MODEL} QDRANT_HOST: ${QDRANT_HOST} WREN_ENGINE_ENDPOINT: ${WREN_ENGINE_ENDPOINT} diff --git a/wren-ai-service/src/providers/llm/openai.py b/wren-ai-service/src/providers/llm/openai.py index f01905d1d..2037daaef 100644 --- a/wren-ai-service/src/providers/llm/openai.py +++ b/wren-ai-service/src/providers/llm/openai.py @@ -8,9 +8,7 @@ from haystack import component from haystack.components.embedders import OpenAIDocumentEmbedder, OpenAITextEmbedder from haystack.components.generators import OpenAIGenerator from haystack.dataclasses import ChatMessage, StreamingChunk -from haystack.utils import ( - Secret, -) +from haystack.utils import Secret from openai import AsyncOpenAI, OpenAI, Stream from openai.types.chat import ChatCompletion, ChatCompletionChunk @@ -19,6 +17,7 @@ from src.providers.loader import provider logger = logging.getLogger("wren-ai-service") +OPENAI_API_BASE = "https://api.openai.com/v1" GENERATION_MODEL_NAME = "gpt-3.5-turbo" GENERATION_MODEL_KWARGS = { "temperature": 0, @@ -121,18 +120,20 @@ class OpenAILLMProvider(LLMProvider): def __init__( self, api_key: Secret = Secret.from_env_var("OPENAI_API_KEY"), + api_base: str = os.getenv("OPENAI_API_BASE") or OPENAI_API_BASE, generation_model: str = os.getenv("OPENAI_GENERATION_MODEL") or GENERATION_MODEL_NAME, ): - def _verify_api_key(api_key: str) -> None: + def _verify_api_key(api_key: str, api_base: str) -> None: """ this is a temporary solution to verify that the required environment variables are set """ - OpenAI(api_key=api_key).models.list() + OpenAI(api_key=api_key, base_url=api_base).models.list() - _verify_api_key(api_key.resolve_value()) + _verify_api_key(api_key.resolve_value(), api_base) logger.info(f"Using OpenAI Generation Model: {generation_model}") self._api_key = api_key + self._api_base = api_base self._generation_model = generation_model def get_generator( @@ -140,11 +141,22 @@ class OpenAILLMProvider(LLMProvider): model_kwargs: Optional[Dict[str, Any]] = GENERATION_MODEL_KWARGS, system_prompt: Optional[str] = None, ): + def _get_generation_kwargs( + model_kwargs: Optional[Dict[str, Any]] = GENERATION_MODEL_KWARGS, + api_base: str = OPENAI_API_BASE, + ): + if api_base != OPENAI_API_BASE: + return model_kwargs + elif model_kwargs != GENERATION_MODEL_KWARGS: + return model_kwargs + return None + return AsyncGenerator( api_key=self._api_key, + api_base_url=self._api_base, model=self._generation_model, system_prompt=system_prompt, - generation_kwargs=model_kwargs, + generation_kwargs=_get_generation_kwargs(model_kwargs, self._api_base), ) def get_text_embedder( @@ -154,6 +166,7 @@ class OpenAILLMProvider(LLMProvider): ): return OpenAITextEmbedder( api_key=self._api_key, + api_base_url=self._api_base, model=model_name, dimensions=model_dim, ) @@ -165,6 +178,7 @@ class OpenAILLMProvider(LLMProvider): ): return OpenAIDocumentEmbedder( api_key=self._api_key, + api_base_url=self._api_base, model=model_name, dimensions=model_dim, )