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 <chihyu.jimmy.yeh@gmail.com>

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

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}

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}

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,
         )