Skip to content

Commit c77d219

Browse files
jamescalamjamescalam
authored
Merge pull request #545 from aurelio-labs/TaylorN15/azure-openai-updates
feat: azure openai updates
2 parents b16b85f + 004877a commit c77d219

File tree

7 files changed

+127
-92
lines changed

7 files changed

+127
-92
lines changed

‎docs/source/conf.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
project = "Semantic Router"
1616
copyright = "2025, Aurelio AI"
1717
author = "Aurelio AI"
18-
release = "0.1.0.dev10"
18+
release = "0.1.1"
1919

2020
# -- General configuration ---------------------------------------------------
2121
# https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration

‎pyproject.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
[project]
22
name = "semantic-router"
3-
version = "0.1.0"
3+
version = "0.1.1"
44
description = "Super fast semantic router for AI decision making"
55
authors = [{ name = "Aurelio AI", email = "[email protected]" }]
66
requires-python = ">=3.9,<3.14"

‎semantic_router/__init__.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,4 +3,4 @@
33

44
__all__ = ["SemanticRouter", "HybridRouter", "Route", "RouterConfig"]
55

6-
__version__ = "0.1.0.dev10"
6+
__version__ = "0.1.1"

‎semantic_router/encoders/__init__.py

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
11
from typing import List, Optional
22

3+
from semantic_router.encoders.base import DenseEncoder, SparseEncoder # isort: skip
34
from semantic_router.encoders.aurelio import AurelioSparseEncoder
4-
from semantic_router.encoders.base import DenseEncoder, SparseEncoder
5+
from semantic_router.encoders.azure_openai import AzureOpenAIEncoder
56
from semantic_router.encoders.bedrock import BedrockEncoder
67
from semantic_router.encoders.bm25 import BM25Encoder
78
from semantic_router.encoders.clip import CLIPEncoder
@@ -13,7 +14,6 @@
1314
from semantic_router.encoders.openai import OpenAIEncoder
1415
from semantic_router.encoders.tfidf import TfidfEncoder
1516
from semantic_router.encoders.vit import VitEncoder
16-
from semantic_router.encoders.zure import AzureOpenAIEncoder
1717
from semantic_router.schema import EncoderType, SparseEmbedding
1818

1919
__all__ = [
@@ -45,8 +45,7 @@ def __init__(self, type: str, name: Optional[str]):
4545
self.type = EncoderType(type)
4646
self.name = name
4747
if self.type == EncoderType.AZURE:
48-
# TODO should change `model` to `name` JB
49-
self.model = AzureOpenAIEncoder(model=name)
48+
self.model = AzureOpenAIEncoder(name=name)
5049
elif self.type == EncoderType.COHERE:
5150
self.model = CohereEncoder(name=name)
5251
elif self.type == EncoderType.OPENAI:

‎semantic_router/encoders/zure.py renamed to ‎semantic_router/encoders/azure_openai.py

Lines changed: 109 additions & 74 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,9 @@
11
import os
22
from asyncio import sleep as asleep
33
from time import sleep
4-
from typing import List, Optional, Union
4+
from typing import Any, Callable, Dict, List, Optional, Union
55

6+
import httpx
67
import openai
78
from openai import OpenAIError
89
from openai._types import NotGiven
@@ -24,100 +25,135 @@ class AzureOpenAIEncoder(DenseEncoder):
2425
async_client: Optional[openai.AsyncAzureOpenAI] = None
2526
dimensions: Union[int, NotGiven] = NotGiven()
2627
type: str = "azure"
27-
api_key: Optional[str] = None
28-
deployment_name: Optional[str] = None
29-
azure_endpoint: Optional[str] = None
30-
api_version: Optional[str] = None
31-
model: Optional[str] = None
28+
deployment_name: str | None = None
3229
max_retries: int = 3
3330

3431
def __init__(
3532
self,
36-
api_key: Optional[str] = None,
37-
deployment_name: Optional[str] = None,
38-
azure_endpoint: Optional[str] = None,
39-
api_version: Optional[str] = None,
40-
model: Optional[str] = None, # TODO we should change to `name` JB
33+
name: Optional[str] = None,
34+
azure_endpoint: str | None = None,
35+
api_version: str | None = None,
36+
api_key: str | None = None,
37+
azure_ad_token: str | None = None,
38+
azure_ad_token_provider: Callable[[], str] | None = None,
39+
http_client_options: Optional[Dict[str, Any]] = None,
40+
deployment_name: str = EncoderDefault.AZURE.value["deployment_name"],
4141
score_threshold: float = 0.82,
4242
dimensions: Union[int, NotGiven] = NotGiven(),
4343
max_retries: int = 3,
4444
):
4545
"""Initialize the AzureOpenAIEncoder.
4646
47-
:param api_key: The API key for the Azure OpenAI API.
48-
:type api_key: str
49-
:param deployment_name: The name of the deployment to use.
50-
:type deployment_name: str
5147
:param azure_endpoint: The endpoint for the Azure OpenAI API.
52-
:type azure_endpoint: str
48+
Example: ``https://accountname.openai.azure.com``
49+
:type azure_endpoint: str, optional
50+
5351
:param api_version: The version of the API to use.
54-
:type api_version: str
55-
:param model: The model to use.
56-
:type model: str
57-
:param score_threshold: The score threshold for the embeddings.
58-
:type score_threshold: float
59-
:param dimensions: The dimensions of the embeddings.
60-
:type dimensions: int
61-
:param max_retries: The maximum number of retries for the API call.
62-
:type max_retries: int
52+
Example: ``"2025-02-01-preview"``
53+
:type api_version: str, optional
54+
55+
:param api_key: The API key for the Azure OpenAI API.
56+
:type api_key: str, optional
57+
58+
:param azure_ad_token: The Azure AD/Entra ID token for authentication.
59+
https://www.microsoft.com/en-us/security/business/identity-access/microsoft-entra-id
60+
:type azure_ad_token: str, optional
61+
62+
:param azure_ad_token_provider: A callable function that returns an Azure AD/Entra ID token.
63+
:type azure_ad_token_provider: Callable[[], str], optional
64+
65+
:param http_client_options: Dictionary of options to configure httpx client
66+
Example:
67+
{
68+
"proxies": "http://proxy.server:8080",
69+
"timeout": 20.0,
70+
"headers": {"Authorization": "Bearer xyz"}
71+
}
72+
:type http_client_options: Dict[str, Any], optional
73+
74+
:param deployment_name: The name of the model deployment to use.
75+
:type deployment_name: str, optional
76+
77+
:param score_threshold: The score threshold for filtering embeddings.
78+
Default is ``0.82``.
79+
:type score_threshold: float, optional
80+
81+
:param dimensions: The number of dimensions for the embeddings. If not given, it defaults to the model's default setting.
82+
:type dimensions: int, optional
83+
84+
:param max_retries: The maximum number of retries for API calls in case of failures.
85+
Default is ``3``.
86+
:type max_retries: int, optional
6387
"""
64-
name = deployment_name
6588
if name is None:
66-
name = EncoderDefault.AZURE.value["embedding_model"]
89+
name = deployment_name
90+
if name is None:
91+
name = EncoderDefault.AZURE.value["embedding_model"]
6792
super().__init__(name=name, score_threshold=score_threshold)
68-
self.api_key = api_key
93+
94+
azure_endpoint = azure_endpoint or os.getenv("AZURE_OPENAI_ENDPOINT")
95+
if not azure_endpoint:
96+
raise ValueError("No Azure OpenAI endpoint provided.")
97+
98+
api_version = api_version or os.getenv("AZURE_OPENAI_API_VERSION")
99+
if not api_version:
100+
raise ValueError("No Azure OpenAI API version provided.")
101+
102+
if not (
103+
azure_ad_token
104+
or azure_ad_token_provider
105+
or api_key
106+
or os.getenv("AZURE_OPENAI_API_KEY")
107+
):
108+
raise ValueError(
109+
"No authentication method provided. Please provide either `azure_ad_token`, "
110+
"`azure_ad_token_provider`, or `api_key`."
111+
)
112+
113+
# Only check API Key if no AD token or provider is used
114+
if not azure_ad_token and not azure_ad_token_provider:
115+
api_key = api_key or os.getenv("AZURE_OPENAI_API_KEY")
116+
if not api_key:
117+
raise ValueError("No Azure OpenAI API key provided.")
118+
69119
self.deployment_name = deployment_name
70-
self.azure_endpoint = azure_endpoint
71-
self.api_version = api_version
72-
self.model = model
120+
73121
# set dimensions to support openai embed 3 dimensions param
74122
self.dimensions = dimensions
75-
if self.api_key is None:
76-
self.api_key = os.getenv("AZURE_OPENAI_API_KEY")
77-
if self.api_key is None:
78-
raise ValueError("No Azure OpenAI API key provided.")
123+
79124
if max_retries is not None:
80125
self.max_retries = max_retries
81-
if self.deployment_name is None:
82-
self.deployment_name = EncoderDefault.AZURE.value["deployment_name"]
83-
# deployment_name may still be None, but it is optional in the API
84-
if self.azure_endpoint is None:
85-
self.azure_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
86-
if self.azure_endpoint is None:
87-
raise ValueError("No Azure OpenAI endpoint provided.")
88-
if self.api_version is None:
89-
self.api_version = os.getenv("AZURE_OPENAI_API_VERSION")
90-
if self.api_version is None:
91-
raise ValueError("No Azure OpenAI API version provided.")
92-
if self.model is None:
93-
self.model = os.getenv("AZURE_OPENAI_MODEL")
94-
if self.model is None:
95-
raise ValueError("No Azure OpenAI model provided.")
96-
assert (
97-
self.api_key is not None
98-
and self.azure_endpoint is not None
99-
and self.api_version is not None
100-
and self.model is not None
126+
127+
# Only create HTTP clients if options are provided
128+
sync_http_client = (
129+
httpx.Client(**http_client_options) if http_client_options else None
130+
)
131+
async_http_client = (
132+
httpx.AsyncClient(**http_client_options) if http_client_options else None
101133
)
102134

135+
assert azure_endpoint is not None and self.deployment_name is not None
136+
103137
try:
104138
self.client = openai.AzureOpenAI(
105-
azure_deployment=(
106-
str(self.deployment_name) if self.deployment_name else None
107-
),
108-
api_key=str(self.api_key),
109-
azure_endpoint=str(self.azure_endpoint),
110-
api_version=str(self.api_version),
139+
azure_endpoint=azure_endpoint,
140+
api_version=api_version,
141+
api_key=api_key,
142+
azure_ad_token=azure_ad_token,
143+
azure_ad_token_provider=azure_ad_token_provider,
144+
http_client=sync_http_client,
111145
)
112146
self.async_client = openai.AsyncAzureOpenAI(
113-
azure_deployment=(
114-
str(self.deployment_name) if self.deployment_name else None
115-
),
116-
api_key=str(self.api_key),
117-
azure_endpoint=str(self.azure_endpoint),
118-
api_version=str(self.api_version),
147+
azure_endpoint=azure_endpoint,
148+
api_version=api_version,
149+
api_key=api_key,
150+
azure_ad_token=azure_ad_token,
151+
azure_ad_token_provider=azure_ad_token_provider,
152+
http_client=async_http_client,
119153
)
154+
120155
except Exception as e:
156+
logger.error("OpenAI API client failed to initialize. Error: %s", e)
121157
raise ValueError(
122158
f"OpenAI API client failed to initialize. Error: {e}"
123159
) from e
@@ -139,7 +175,7 @@ def __call__(self, docs: List[str]) -> List[List[float]]:
139175
try:
140176
embeds = self.client.embeddings.create(
141177
input=docs,
142-
model=str(self.model),
178+
model=str(self.deployment_name),
143179
dimensions=self.dimensions,
144180
)
145181
if embeds.data:
@@ -149,12 +185,12 @@ def __call__(self, docs: List[str]) -> List[List[float]]:
149185
if self.max_retries != 0 and j < self.max_retries:
150186
sleep(2**j)
151187
logger.warning(
152-
f"Retrying in {2**j} seconds due to OpenAIError: {e}"
188+
"Retrying in %d seconds due to OpenAIError: %s", 2**j, e
153189
)
154190
else:
155191
raise
156192
except Exception as e:
157-
logger.error(f"Azure OpenAI API call failed. Error: {e}")
193+
logger.error("Azure OpenAI API call failed. Error: %s", e)
158194
raise ValueError(f"Azure OpenAI API call failed. Error: {e}") from e
159195

160196
if (
@@ -183,23 +219,22 @@ async def acall(self, docs: List[str]) -> List[List[float]]:
183219
try:
184220
embeds = await self.async_client.embeddings.create(
185221
input=docs,
186-
model=str(self.model),
222+
model=str(self.deployment_name),
187223
dimensions=self.dimensions,
188224
)
189225
if embeds.data:
190226
break
191-
192227
except OpenAIError as e:
193228
logger.error("Exception occurred", exc_info=True)
194229
if self.max_retries != 0 and j < self.max_retries:
195230
await asleep(2**j)
196231
logger.warning(
197-
f"Retrying in {2**j} seconds due to OpenAIError: {e}"
232+
"Retrying in %d seconds due to OpenAIError: %s", 2**j, e
198233
)
199234
else:
200235
raise
201236
except Exception as e:
202-
logger.error(f"Azure OpenAI API call failed. Error: {e}")
237+
logger.error("Azure OpenAI API call failed. Error: %s", e)
203238
raise ValueError(f"Azure OpenAI API call failed. Error: {e}") from e
204239

205240
if (

‎tests/unit/encoders/test_azure.py

Lines changed: 11 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -23,11 +23,12 @@ def mock_openai_async_client():
2323
@pytest.fixture
2424
def openai_encoder(mock_openai_client, mock_openai_async_client):
2525
return AzureOpenAIEncoder(
26+
azure_endpoint="https://test-endpoint.openai.azure.com",
27+
api_version="test-version",
2628
api_key="test_api_key",
29+
http_client_options={"timeout": 10},
2730
deployment_name="test-deployment",
28-
azure_endpoint="test_endpoint",
29-
api_version="test_version",
30-
model="test_model",
31+
dimensions=1536,
3132
max_retries=2,
3233
)
3334

@@ -84,7 +85,7 @@ def test_openai_encoder_call_success(self, openai_encoder, mocker):
8485
mocker.patch.object(
8586
openai_encoder.client.embeddings, "create", side_effect=responses
8687
)
87-
with patch("semantic_router.encoders.zure.sleep", return_value=None):
88+
with patch("semantic_router.encoders.azure_openai.sleep", return_value=None):
8889
embeddings = openai_encoder(["test document"])
8990
assert embeddings == [[0.1, 0.2]]
9091

@@ -96,7 +97,7 @@ def test_openai_encoder_call_failure_non_openai_error(self, openai_encoder, mock
9697
"create",
9798
side_effect=Exception("Non-OpenAIError"),
9899
)
99-
with patch("semantic_router.encoders.zure.sleep", return_value=None):
100+
with patch("semantic_router.encoders.azure_openai.sleep", return_value=None):
100101
with pytest.raises(ValueError) as e:
101102
openai_encoder(["test document"])
102103

@@ -124,7 +125,7 @@ def test_openai_encoder_call_successful_retry(self, openai_encoder, mocker):
124125
mocker.patch.object(
125126
openai_encoder.client.embeddings, "create", side_effect=responses
126127
)
127-
with patch("semantic_router.encoders.zure.sleep", return_value=None):
128+
with patch("semantic_router.encoders.azure_openai.sleep", return_value=None):
128129
embeddings = openai_encoder(["test document"])
129130
assert embeddings == [[0.1, 0.2]]
130131

@@ -150,7 +151,7 @@ def test_retry_logic_sync(self, openai_encoder, mock_openai_client, mocker):
150151
mocker.patch("time.sleep", return_value=None) # To speed up the test
151152

152153
# Patch the sleep function in the encoder module to avoid actual sleep
153-
with patch("semantic_router.encoders.zure.sleep", return_value=None):
154+
with patch("semantic_router.encoders.azure_openai.sleep", return_value=None):
154155
result = openai_encoder(["test document"])
155156

156157
assert result == [[0.1, 0.2, 0.3]]
@@ -176,7 +177,7 @@ def test_retry_logic_sync_max_retries_exceeded(
176177
mocker.patch("time.sleep", return_value=None) # To speed up the test
177178

178179
# Patch the sleep function in the encoder module to avoid actual sleep
179-
with patch("semantic_router.encoders.zure.sleep", return_value=None):
180+
with patch("semantic_router.encoders.azure_openai.sleep", return_value=None):
180181
with pytest.raises(OpenAIError):
181182
openai_encoder(["test document"])
182183

@@ -207,7 +208,7 @@ async def test_retry_logic_async(
207208
mocker.patch("asyncio.sleep", return_value=None) # To speed up the test
208209

209210
# Patch the asleep function in the encoder module to avoid actual sleep
210-
with patch("semantic_router.encoders.zure.asleep", return_value=None):
211+
with patch("semantic_router.encoders.azure_openai.asleep", return_value=None):
211212
result = await openai_encoder.acall(["test document"])
212213

213214
assert result == [[0.1, 0.2, 0.3]]
@@ -226,7 +227,7 @@ async def raise_error(*args, **kwargs):
226227
mocker.patch("asyncio.sleep", return_value=None) # To speed up the test
227228

228229
# Patch the asleep function in the encoder module to avoid actual sleep
229-
with patch("semantic_router.encoders.zure.asleep", return_value=None):
230+
with patch("semantic_router.encoders.azure_openai.asleep", return_value=None):
230231
with pytest.raises(OpenAIError):
231232
await openai_encoder.acall(["test document"])
232233

0 commit comments

Comments
 (0)