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

Add Support for Customizing default_headers in Azure OpenAI (#1925)

Browse Source
This commit is contained in:
Jian Yu, Chen
2024-10-19 18:27:28 +08:00
committed by GitHub
parent f058cda152
commit ff7761aaf8
8 changed files with 69 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
mem0/configs/base.py
View File

@@ -63,6 +63,7 @@ class AzureConfig(BaseModel):
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.
default_headers (Dict[str, str]): Headers to include in requests to the Azure API.
"""
api_key: str = Field(
@@ -72,3 +73,4 @@ class AzureConfig(BaseModel):
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)
default_headers: Optional[Dict[str, str]] = Field(description="Headers to include in requests to the Azure API.", default=None)

2
mem0/embeddings/azure_openai.py
View File

@@ -15,6 +15,7 @@ class AzureOpenAIEmbedding(EmbeddingBase):
azure_deployment = self.config.azure_kwargs.azure_deployment or os.getenv("EMBEDDING_AZURE_DEPLOYMENT")
azure_endpoint = self.config.azure_kwargs.azure_endpoint or os.getenv("EMBEDDING_AZURE_ENDPOINT")
api_version = self.config.azure_kwargs.api_version or os.getenv("EMBEDDING_AZURE_API_VERSION")
default_headers = self.config.azure_kwargs.default_headers
self.client = AzureOpenAI(
azure_deployment=azure_deployment,
@@ -22,6 +23,7 @@ class AzureOpenAIEmbedding(EmbeddingBase):
api_version=api_version,
api_key=api_key,
http_client=self.config.http_client,
default_headers=default_headers,
)
def embed(self, text):

2
mem0/llms/azure_openai.py
View File

@@ -20,6 +20,7 @@ class AzureOpenAILLM(LLMBase):
azure_deployment = self.config.azure_kwargs.azure_deployment or os.getenv("LLM_AZURE_DEPLOYMENT")
azure_endpoint = self.config.azure_kwargs.azure_endpoint or os.getenv("LLM_AZURE_ENDPOINT")
api_version = self.config.azure_kwargs.api_version or os.getenv("LLM_AZURE_API_VERSION")
default_headers = self.config.azure_kwargs.default_headers
self.client = AzureOpenAI(
azure_deployment=azure_deployment,
@@ -27,6 +28,7 @@ class AzureOpenAILLM(LLMBase):
api_version=api_version,
api_key=api_key,
http_client=self.config.http_client,
default_headers=default_headers,
)
def _parse_response(self, response, tools):

4
mem0/llms/azure_openai_structured.py
View File

@@ -20,14 +20,16 @@ class AzureOpenAIStructuredLLM(LLMBase):
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
# Can display a warning if API version is of model and api-version
default_headers = self.config.azure_kwargs.default_headers
# Can display a warning if API version is of model and 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,
default_headers=default_headers,
)
def _parse_response(self, response, tools):