base.py•4.99 kB
import os
import httpx
from selfmemory.configs.base import AzureConfig
class BaseEmbedderConfig:
"""
Config for Embeddings.
"""
def __init__(
self,
model: str | None = None,
api_key: str | None = None,
embedding_dims: int | None = None,
# Ollama specific
ollama_base_url: str | None = None,
# Openai specific
openai_base_url: str | None = None,
# Huggingface specific
model_kwargs: dict | None = None,
huggingface_base_url: str | None = None,
# AzureOpenAI specific
azure_kwargs: AzureConfig | None = None,
http_client_proxies: dict | str | None = None,
# VertexAI specific
vertex_credentials_json: str | None = None,
memory_add_embedding_type: str | None = None,
memory_update_embedding_type: str | None = None,
memory_search_embedding_type: str | None = None,
# Gemini specific
output_dimensionality: str | None = None,
# LM Studio specific
lmstudio_base_url: str | None = "http://localhost:1234/v1",
# AWS Bedrock specific
aws_access_key_id: str | None = None,
aws_secret_access_key: str | None = None,
aws_region: str | None = None,
):
"""
Initializes a configuration class instance for the Embeddings.
:param model: Embedding model to use, defaults to None
:type model: Optional[str], optional
:param api_key: API key to be use, defaults to None
:type api_key: Optional[str], optional
:param embedding_dims: The number of dimensions in the embedding, defaults to None
:type embedding_dims: Optional[int], optional
:param ollama_base_url: Base URL for the Ollama API, defaults to None
:type ollama_base_url: Optional[str], optional
:param model_kwargs: key-value arguments for the huggingface embedding model, defaults a dict inside init
:type model_kwargs: Optional[Dict[str, Any]], defaults a dict inside init
:param huggingface_base_url: Huggingface base URL to be use, defaults to None
:type huggingface_base_url: Optional[str], optional
:param openai_base_url: Openai base URL to be use, defaults to "https://api.openai.com/v1"
:type openai_base_url: Optional[str], optional
:param azure_kwargs: key-value arguments for the AzureOpenAI embedding model, defaults a dict inside init
:type azure_kwargs: Optional[Dict[str, Any]], defaults a dict inside init
:param http_client_proxies: The proxy server settings used to create self.http_client, defaults to None
:type http_client_proxies: Optional[Dict | str], optional
:param vertex_credentials_json: The path to the Vertex AI credentials JSON file, defaults to None
:type vertex_credentials_json: Optional[str], optional
:param memory_add_embedding_type: The type of embedding to use for the add memory action, defaults to None
:type memory_add_embedding_type: Optional[str], optional
:param memory_update_embedding_type: The type of embedding to use for the update memory action, defaults to None
:type memory_update_embedding_type: Optional[str], optional
:param memory_search_embedding_type: The type of embedding to use for the search memory action, defaults to None
:type memory_search_embedding_type: Optional[str], optional
:param lmstudio_base_url: LM Studio base URL to be use, defaults to "http://localhost:1234/v1"
:type lmstudio_base_url: Optional[str], optional
"""
self.model = model
self.api_key = api_key
self.openai_base_url = openai_base_url
self.embedding_dims = embedding_dims
# AzureOpenAI specific
self.http_client = (
httpx.Client(proxies=http_client_proxies) if http_client_proxies else None
)
# Ollama specific
self.ollama_base_url = ollama_base_url
# Huggingface specific
self.model_kwargs = model_kwargs or {}
self.huggingface_base_url = huggingface_base_url
# AzureOpenAI specific
self.azure_kwargs = (
AzureConfig(**azure_kwargs) if azure_kwargs else AzureConfig()
)
# VertexAI specific
self.vertex_credentials_json = vertex_credentials_json
self.memory_add_embedding_type = memory_add_embedding_type
self.memory_update_embedding_type = memory_update_embedding_type
self.memory_search_embedding_type = memory_search_embedding_type
# Gemini specific
self.output_dimensionality = output_dimensionality
# LM Studio specific
self.lmstudio_base_url = lmstudio_base_url
# AWS Bedrock specific
self.aws_access_key_id = aws_access_key_id
self.aws_secret_access_key = aws_secret_access_key
self.aws_region = aws_region or os.environ.get("AWS_REGION") or "us-west-2"