Skip to content

Commit 3b3e96a

Browse files
Allow parameterizing OpenAI embeddings component (api_base, key, model) (#1920)
* Allow parameterizing OpenAI embeddings component (api_base, key, model) * Update settings * Update description
1 parent 45df99f commit 3b3e96a

File tree

3 files changed

+21
-2
lines changed

3 files changed

+21
-2
lines changed

private_gpt/components/embedding/embedding_component.py

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -55,8 +55,17 @@ def __init__(self, settings: Settings) -> None:
5555
"OpenAI dependencies not found, install with `poetry install --extras embeddings-openai`"
5656
) from e
5757

58-
openai_settings = settings.openai.api_key
59-
self.embedding_model = OpenAIEmbedding(api_key=openai_settings)
58+
api_base = (
59+
settings.openai.embedding_api_base or settings.openai.api_base
60+
)
61+
api_key = settings.openai.embedding_api_key or settings.openai.api_key
62+
model = settings.openai.embedding_model
63+
64+
self.embedding_model = OpenAIEmbedding(
65+
api_base=api_base,
66+
api_key=api_key,
67+
model=model,
68+
)
6069
case "ollama":
6170
try:
6271
from llama_index.embeddings.ollama import ( # type: ignore

private_gpt/settings/settings.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -209,6 +209,15 @@ class OpenAISettings(BaseModel):
209209
120.0,
210210
description="Time elapsed until openailike server times out the request. Default is 120s. Format is float. ",
211211
)
212+
embedding_api_base: str = Field(
213+
None,
214+
description="Base URL of OpenAI API. Example: 'https://api.openai.com/v1'.",
215+
)
216+
embedding_api_key: str
217+
embedding_model: str = Field(
218+
"text-embedding-ada-002",
219+
description="OpenAI embedding Model to use. Example: 'text-embedding-3-large'.",
220+
)
212221

213222

214223
class OllamaSettings(BaseModel):

settings.yaml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,7 @@ sagemaker:
9595
openai:
9696
api_key: ${OPENAI_API_KEY:}
9797
model: gpt-3.5-turbo
98+
embedding_api_key: ${OPENAI_API_KEY:}
9899

99100
ollama:
100101
llm_model: llama2

0 commit comments

Comments
 (0)