Allow parameterizing OpenAI embeddings component (api_base, key, model)

This commit is contained in:
Yevhenii Semendiak 2024-05-09 00:42:55 +03:00
parent 659e0f1a9d
commit 90dadda249
2 changed files with 20 additions and 2 deletions

View File

@ -55,8 +55,17 @@ class EmbeddingComponent:
"OpenAI dependencies not found, install with `poetry install --extras embeddings-openai`"
) from e
openai_settings = settings.openai.api_key
self.embedding_model = OpenAIEmbedding(api_key=openai_settings)
api_base = (
settings.openai.embedding_api_base or settings.openai.api_base
)
api_key = settings.openai.embedding_api_key or settings.openai.api_key
model = settings.openai.embedding_model
self.embedding_model = OpenAIEmbedding(
api_base=api_base,
api_key=api_key,
model=model,
)
case "ollama":
try:
from llama_index.embeddings.ollama import ( # type: ignore

View File

@ -206,6 +206,15 @@ class OpenAISettings(BaseModel):
"gpt-3.5-turbo",
description="OpenAI Model to use. Example: 'gpt-4'.",
)
embedding_api_base: str = Field(
None,
description="Base URL of OpenAI API. Example: 'https://api.openai.com/v1'.",
)
embedding_api_key: str
embedding_model: str = Field(
"text-embedding-ada-002",
description="OpenAI Model to use. Example: 'gpt-4'.",
)
class OllamaSettings(BaseModel):