klas/t6_mem0
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

AzureOpenAI Embedding Model and LLM Model Initialisation from Config. (#1773)

Browse Source
This commit is contained in:
k10
2024-09-01 02:09:00 +05:30
committed by GitHub
parent ad233034ef
commit 077d0c47f9
10 changed files with 88 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
mem0/configs/base.py
View File

@@ -55,4 +55,20 @@ class MemoryConfig(BaseModel):
description="The version of the API",
default="v1.0",
)
class AzureConfig(BaseModel):
"""
Configuration settings for Azure.
Args:
api_key (str): The API key used for authenticating with the Azure service.
azure_deployment (str): The name of the Azure deployment.
azure_endpoint (str): The endpoint URL for the Azure service.
api_version (str): The version of the Azure API being used.
"""
api_key: str = Field(description="The API key used for authenticating with the Azure service.", default=None)
azure_deployment : str = Field(description="The name of the Azure deployment.", default=None)
azure_endpoint : str = Field(description="The endpoint URL for the Azure service.", default=None)
api_version : str = Field(description="The version of the Azure API being used.", default=None)

7
mem0/configs/embeddings/base.py
View File

@@ -1,4 +1,5 @@
from abc import ABC
from mem0.configs.base import AzureConfig
from typing import Optional, Union, Dict
import httpx
@@ -21,6 +22,7 @@ class BaseEmbedderConfig(ABC):
# Huggingface specific
model_kwargs: Optional[dict] = None,
# AzureOpenAI specific
azure_kwargs: Optional[AzureConfig] = {},
http_client_proxies: Optional[Union[Dict, str]] = None,
):
"""
@@ -38,6 +40,8 @@ class BaseEmbedderConfig(ABC):
:type model_kwargs: Optional[Dict[str, Any]], defaults a dict inside init
: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
"""
@@ -55,3 +59,6 @@ class BaseEmbedderConfig(ABC):
# Huggingface specific
self.model_kwargs = model_kwargs or {}
# AzureOpenAI specific
self.azure_kwargs = AzureConfig(**azure_kwargs) or {}

9
mem0/configs/llms/base.py
View File

@@ -1,4 +1,5 @@
from abc import ABC
from mem0.configs.base import AzureConfig
from typing import Optional, Union, Dict
import httpx
@@ -27,7 +28,8 @@ class BaseLlmConfig(ABC):
app_name: Optional[str] = None,
# Ollama specific
ollama_base_url: Optional[str] = None,
# AzureOpenAI specific
azure_kwargs: Optional[AzureConfig] = {},
# AzureOpenAI specific
http_client_proxies: Optional[Union[Dict, str]] = None,
):
@@ -62,6 +64,8 @@ class BaseLlmConfig(ABC):
:type ollama_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 LLM 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(s) settings used to create self.http_client, defaults to None
:type http_client_proxies: Optional[Dict | str], optional
"""
@@ -86,3 +90,6 @@ class BaseLlmConfig(ABC):
# Ollama specific
self.ollama_base_url = ollama_base_url
# AzureOpenAI specific
self.azure_kwargs = AzureConfig(**azure_kwargs) or {}

19
mem0/embeddings/azure_openai.py
View File

@@ -11,13 +11,18 @@ class AzureOpenAIEmbedding(EmbeddingBase):
def __init__(self, config: Optional[BaseEmbedderConfig] = None):
super().__init__(config)
if self.config.model is None:
self.config.model = "text-embedding-3-small"
if self.config.embedding_dims is None:
self.config.embedding_dims = 1536
api_key = os.getenv("AZURE_OPENAI_API_KEY") or self.config.api_key
self.client = AzureOpenAI(api_key=api_key, http_client=self.config.http_client)
api_key = os.getenv("EMBEDDING_AZURE_OPENAI_API_KEY") or self.config.azure_kwargs.api_key
azure_deployment = os.getenv("EMBEDDING_AZURE_DEPLOYMENT") or self.config.azure_kwargs.azure_deployment
azure_endpoint = os.getenv("EMBEDDING_AZURE_ENDPOINT") or self.config.azure_kwargs.azure_endpoint
api_version = os.getenv("EMBEDDING_AZURE_API_VERSION") or self.config.azure_kwargs.api_version
self.client = AzureOpenAI(
azure_deployment=azure_deployment,
azure_endpoint=azure_endpoint,
api_version=api_version,
api_key=api_key,
http_client=self.config.http_client
)
def embed(self, text):
"""

16
mem0/llms/azure_openai.py
View File

@@ -15,10 +15,20 @@ class AzureOpenAILLM(LLMBase):
# Model name should match the custom deployment name chosen for it.
if not self.config.model:
self.config.model = "gpt-4o"
api_key = os.getenv("AZURE_OPENAI_API_KEY") or self.config.api_key
self.client = AzureOpenAI(api_key=api_key, http_client=self.config.http_client)
api_key = os.getenv("LLM_AZURE_OPENAI_API_KEY") or self.config.azure_kwargs.api_key
azure_deployment = os.getenv("LLM_AZURE_DEPLOYMENT") or self.config.azure_kwargs.azure_deployment
azure_endpoint = os.getenv("LLM_AZURE_ENDPOINT") or self.config.azure_kwargs.azure_endpoint
api_version = os.getenv("LLM_AZURE_API_VERSION") or self.config.azure_kwargs.api_version
self.client = AzureOpenAI(
azure_deployment=azure_deployment,
azure_endpoint=azure_endpoint,
api_version=api_version,
api_key=api_key,
http_client=self.config.http_client
)
def _parse_response(self, response, tools):
"""
Process the response based on whether tools are used or not.