openaivec 0.13.0__py3-none-any.whl → 0.13.2__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- openaivec/embeddings.py +8 -8
- openaivec/model.py +37 -1
- openaivec/pandas_ext.py +30 -17
- openaivec/provider.py +67 -15
- openaivec/responses.py +79 -31
- openaivec/spark.py +15 -15
- openaivec/util.py +18 -12
- {openaivec-0.13.0.dist-info → openaivec-0.13.2.dist-info}/METADATA +39 -7
- {openaivec-0.13.0.dist-info → openaivec-0.13.2.dist-info}/RECORD +11 -11
- {openaivec-0.13.0.dist-info → openaivec-0.13.2.dist-info}/WHEEL +0 -0
- {openaivec-0.13.0.dist-info → openaivec-0.13.2.dist-info}/licenses/LICENSE +0 -0
openaivec/embeddings.py
CHANGED
|
@@ -4,7 +4,7 @@ from typing import List
|
|
|
4
4
|
|
|
5
5
|
import numpy as np
|
|
6
6
|
from numpy.typing import NDArray
|
|
7
|
-
from openai import AsyncOpenAI, OpenAI, RateLimitError
|
|
7
|
+
from openai import AsyncOpenAI, InternalServerError, OpenAI, RateLimitError
|
|
8
8
|
|
|
9
9
|
from .log import observe
|
|
10
10
|
from .proxy import AsyncBatchingMapProxy, BatchingMapProxy
|
|
@@ -24,7 +24,7 @@ class BatchEmbeddings:
|
|
|
24
24
|
|
|
25
25
|
Attributes:
|
|
26
26
|
client (OpenAI): Configured OpenAI client.
|
|
27
|
-
model_name (str):
|
|
27
|
+
model_name (str): For Azure OpenAI, use your deployment name. For OpenAI, use the model name (e.g., ``"text-embedding-3-small"``).
|
|
28
28
|
cache (BatchingMapProxy[str, NDArray[np.float32]]): Batching proxy for ordered, cached mapping.
|
|
29
29
|
"""
|
|
30
30
|
|
|
@@ -38,7 +38,7 @@ class BatchEmbeddings:
|
|
|
38
38
|
|
|
39
39
|
Args:
|
|
40
40
|
client (OpenAI): OpenAI client.
|
|
41
|
-
model_name (str):
|
|
41
|
+
model_name (str): For Azure OpenAI, use your deployment name. For OpenAI, use the model name.
|
|
42
42
|
batch_size (int, optional): Max unique inputs per API call. Defaults to 128.
|
|
43
43
|
|
|
44
44
|
Returns:
|
|
@@ -47,7 +47,7 @@ class BatchEmbeddings:
|
|
|
47
47
|
return cls(client=client, model_name=model_name, cache=BatchingMapProxy(batch_size=batch_size))
|
|
48
48
|
|
|
49
49
|
@observe(_LOGGER)
|
|
50
|
-
@backoff(
|
|
50
|
+
@backoff(exceptions=[RateLimitError, InternalServerError], scale=1, max_retries=12)
|
|
51
51
|
def _embed_chunk(self, inputs: List[str]) -> List[NDArray[np.float32]]:
|
|
52
52
|
"""Embed one minibatch of strings.
|
|
53
53
|
|
|
@@ -90,7 +90,7 @@ class AsyncBatchEmbeddings:
|
|
|
90
90
|
import asyncio
|
|
91
91
|
import numpy as np
|
|
92
92
|
from openai import AsyncOpenAI
|
|
93
|
-
|
|
93
|
+
from openaivec import AsyncBatchEmbeddings
|
|
94
94
|
|
|
95
95
|
# Assuming openai_async_client is an initialized AsyncOpenAI client
|
|
96
96
|
openai_async_client = AsyncOpenAI() # Replace with your actual client initialization
|
|
@@ -119,7 +119,7 @@ class AsyncBatchEmbeddings:
|
|
|
119
119
|
|
|
120
120
|
Attributes:
|
|
121
121
|
client (AsyncOpenAI): Configured OpenAI async client.
|
|
122
|
-
model_name (str):
|
|
122
|
+
model_name (str): For Azure OpenAI, use your deployment name. For OpenAI, use the model name.
|
|
123
123
|
cache (AsyncBatchingMapProxy[str, NDArray[np.float32]]): Async batching proxy.
|
|
124
124
|
"""
|
|
125
125
|
|
|
@@ -141,7 +141,7 @@ class AsyncBatchEmbeddings:
|
|
|
141
141
|
|
|
142
142
|
Args:
|
|
143
143
|
client (AsyncOpenAI): OpenAI async client.
|
|
144
|
-
model_name (str):
|
|
144
|
+
model_name (str): For Azure OpenAI, use your deployment name. For OpenAI, use the model name.
|
|
145
145
|
batch_size (int, optional): Max unique inputs per API call. Defaults to 128.
|
|
146
146
|
max_concurrency (int, optional): Max concurrent API calls. Defaults to 8.
|
|
147
147
|
|
|
@@ -155,7 +155,7 @@ class AsyncBatchEmbeddings:
|
|
|
155
155
|
)
|
|
156
156
|
|
|
157
157
|
@observe(_LOGGER)
|
|
158
|
-
@backoff_async(
|
|
158
|
+
@backoff_async(exceptions=[RateLimitError, InternalServerError], scale=1, max_retries=12)
|
|
159
159
|
async def _embed_chunk(self, inputs: List[str]) -> List[NDArray[np.float32]]:
|
|
160
160
|
"""Embed one minibatch of strings asynchronously.
|
|
161
161
|
|
openaivec/model.py
CHANGED
|
@@ -59,29 +59,65 @@ class PreparedTask:
|
|
|
59
59
|
|
|
60
60
|
@dataclass(frozen=True)
|
|
61
61
|
class ResponsesModelName:
|
|
62
|
+
"""Container for responses model name configuration.
|
|
63
|
+
|
|
64
|
+
Attributes:
|
|
65
|
+
value (str): The model name for OpenAI responses API.
|
|
66
|
+
"""
|
|
67
|
+
|
|
62
68
|
value: str
|
|
63
69
|
|
|
64
70
|
|
|
65
71
|
@dataclass(frozen=True)
|
|
66
72
|
class EmbeddingsModelName:
|
|
73
|
+
"""Container for embeddings model name configuration.
|
|
74
|
+
|
|
75
|
+
Attributes:
|
|
76
|
+
value (str): The model name for OpenAI embeddings API.
|
|
77
|
+
"""
|
|
78
|
+
|
|
67
79
|
value: str
|
|
68
80
|
|
|
69
81
|
|
|
70
82
|
@dataclass(frozen=True)
|
|
71
83
|
class OpenAIAPIKey:
|
|
84
|
+
"""Container for OpenAI API key configuration.
|
|
85
|
+
|
|
86
|
+
Attributes:
|
|
87
|
+
value (str): The API key for OpenAI services.
|
|
88
|
+
"""
|
|
89
|
+
|
|
72
90
|
value: str
|
|
73
91
|
|
|
74
92
|
|
|
75
93
|
@dataclass(frozen=True)
|
|
76
94
|
class AzureOpenAIAPIKey:
|
|
95
|
+
"""Container for Azure OpenAI API key configuration.
|
|
96
|
+
|
|
97
|
+
Attributes:
|
|
98
|
+
value (str): The API key for Azure OpenAI services.
|
|
99
|
+
"""
|
|
100
|
+
|
|
77
101
|
value: str
|
|
78
102
|
|
|
79
103
|
|
|
80
104
|
@dataclass(frozen=True)
|
|
81
|
-
class
|
|
105
|
+
class AzureOpenAIBaseURL:
|
|
106
|
+
"""Container for Azure OpenAI base URL configuration.
|
|
107
|
+
|
|
108
|
+
Attributes:
|
|
109
|
+
value (str): The base URL for Azure OpenAI services.
|
|
110
|
+
"""
|
|
111
|
+
|
|
82
112
|
value: str
|
|
83
113
|
|
|
84
114
|
|
|
85
115
|
@dataclass(frozen=True)
|
|
86
116
|
class AzureOpenAIAPIVersion:
|
|
117
|
+
"""Container for Azure OpenAI API version configuration.
|
|
118
|
+
|
|
119
|
+
Attributes:
|
|
120
|
+
value (str): The API version for Azure OpenAI services.
|
|
121
|
+
"""
|
|
122
|
+
|
|
87
123
|
value: str
|
openaivec/pandas_ext.py
CHANGED
|
@@ -7,7 +7,7 @@ from openaivec import pandas_ext
|
|
|
7
7
|
|
|
8
8
|
# Option 1: Use environment variables (automatic detection)
|
|
9
9
|
# Set OPENAI_API_KEY or Azure OpenAI environment variables
|
|
10
|
-
# (AZURE_OPENAI_API_KEY,
|
|
10
|
+
# (AZURE_OPENAI_API_KEY, AZURE_OPENAI_BASE_URL, AZURE_OPENAI_API_VERSION)
|
|
11
11
|
# No explicit setup needed - clients are automatically created
|
|
12
12
|
|
|
13
13
|
# Option 2: Use an existing OpenAI client instance
|
|
@@ -17,14 +17,18 @@ pandas_ext.use(client)
|
|
|
17
17
|
# Option 3: Use an existing Azure OpenAI client instance
|
|
18
18
|
azure_client = AzureOpenAI(
|
|
19
19
|
api_key="your-azure-key",
|
|
20
|
-
|
|
21
|
-
api_version="
|
|
20
|
+
base_url="https://YOUR-RESOURCE-NAME.services.ai.azure.com/openai/v1/",
|
|
21
|
+
api_version="preview"
|
|
22
22
|
)
|
|
23
23
|
pandas_ext.use(azure_client)
|
|
24
24
|
|
|
25
|
-
# Option 4: Use async
|
|
26
|
-
|
|
27
|
-
|
|
25
|
+
# Option 4: Use async Azure OpenAI client instance
|
|
26
|
+
async_azure_client = AsyncAzureOpenAI(
|
|
27
|
+
api_key="your-azure-key",
|
|
28
|
+
base_url="https://YOUR-RESOURCE-NAME.services.ai.azure.com/openai/v1/",
|
|
29
|
+
api_version="preview"
|
|
30
|
+
)
|
|
31
|
+
pandas_ext.use_async(async_azure_client)
|
|
28
32
|
|
|
29
33
|
# Set up model names (optional, defaults shown)
|
|
30
34
|
pandas_ext.responses_model("gpt-4.1-mini")
|
|
@@ -48,7 +52,7 @@ from pydantic import BaseModel
|
|
|
48
52
|
|
|
49
53
|
from .embeddings import AsyncBatchEmbeddings, BatchEmbeddings
|
|
50
54
|
from .model import EmbeddingsModelName, PreparedTask, ResponseFormat, ResponsesModelName
|
|
51
|
-
from .provider import CONTAINER
|
|
55
|
+
from .provider import CONTAINER, _check_azure_v1_api_url
|
|
52
56
|
from .proxy import AsyncBatchingMapProxy, BatchingMapProxy
|
|
53
57
|
from .responses import AsyncBatchResponses, BatchResponses
|
|
54
58
|
from .task.table import FillNaResponse, fillna
|
|
@@ -74,6 +78,10 @@ def use(client: OpenAI) -> None:
|
|
|
74
78
|
`openai.AzureOpenAI` instance.
|
|
75
79
|
The same instance is reused by every helper in this module.
|
|
76
80
|
"""
|
|
81
|
+
# Check Azure v1 API URL if using AzureOpenAI client
|
|
82
|
+
if client.__class__.__name__ == "AzureOpenAI" and hasattr(client, "base_url"):
|
|
83
|
+
_check_azure_v1_api_url(str(client.base_url))
|
|
84
|
+
|
|
77
85
|
CONTAINER.register(OpenAI, lambda: client)
|
|
78
86
|
|
|
79
87
|
|
|
@@ -85,6 +93,10 @@ def use_async(client: AsyncOpenAI) -> None:
|
|
|
85
93
|
`openai.AsyncAzureOpenAI` instance.
|
|
86
94
|
The same instance is reused by every helper in this module.
|
|
87
95
|
"""
|
|
96
|
+
# Check Azure v1 API URL if using AsyncAzureOpenAI client
|
|
97
|
+
if client.__class__.__name__ == "AsyncAzureOpenAI" and hasattr(client, "base_url"):
|
|
98
|
+
_check_azure_v1_api_url(str(client.base_url))
|
|
99
|
+
|
|
88
100
|
CONTAINER.register(AsyncOpenAI, lambda: client)
|
|
89
101
|
|
|
90
102
|
|
|
@@ -92,7 +104,7 @@ def responses_model(name: str) -> None:
|
|
|
92
104
|
"""Override the model used for text responses.
|
|
93
105
|
|
|
94
106
|
Args:
|
|
95
|
-
name (str):
|
|
107
|
+
name (str): For Azure OpenAI, use your deployment name. For OpenAI, use the model name
|
|
96
108
|
(for example, ``gpt-4.1-mini``).
|
|
97
109
|
"""
|
|
98
110
|
CONTAINER.register(ResponsesModelName, lambda: ResponsesModelName(name))
|
|
@@ -102,7 +114,8 @@ def embeddings_model(name: str) -> None:
|
|
|
102
114
|
"""Override the model used for text embeddings.
|
|
103
115
|
|
|
104
116
|
Args:
|
|
105
|
-
name (str):
|
|
117
|
+
name (str): For Azure OpenAI, use your deployment name. For OpenAI, use the model name,
|
|
118
|
+
e.g. ``text-embedding-3-small``.
|
|
106
119
|
"""
|
|
107
120
|
CONTAINER.register(EmbeddingsModelName, lambda: EmbeddingsModelName(name))
|
|
108
121
|
|
|
@@ -143,7 +156,7 @@ class OpenAIVecSeriesAccessor:
|
|
|
143
156
|
instructions: str,
|
|
144
157
|
cache: BatchingMapProxy[str, ResponseFormat],
|
|
145
158
|
response_format: Type[ResponseFormat] = str,
|
|
146
|
-
temperature: float = 0.0,
|
|
159
|
+
temperature: float | None = 0.0,
|
|
147
160
|
top_p: float = 1.0,
|
|
148
161
|
) -> pd.Series:
|
|
149
162
|
client: BatchResponses = BatchResponses(
|
|
@@ -205,7 +218,7 @@ class OpenAIVecSeriesAccessor:
|
|
|
205
218
|
instructions: str,
|
|
206
219
|
response_format: Type[ResponseFormat] = str,
|
|
207
220
|
batch_size: int = 128,
|
|
208
|
-
temperature: float = 0.0,
|
|
221
|
+
temperature: float | None = 0.0,
|
|
209
222
|
top_p: float = 1.0,
|
|
210
223
|
) -> pd.Series:
|
|
211
224
|
"""Call an LLM once for every Series element.
|
|
@@ -438,7 +451,7 @@ class OpenAIVecDataFrameAccessor:
|
|
|
438
451
|
instructions: str,
|
|
439
452
|
cache: BatchingMapProxy[str, ResponseFormat],
|
|
440
453
|
response_format: Type[ResponseFormat] = str,
|
|
441
|
-
temperature: float = 0.0,
|
|
454
|
+
temperature: float | None = 0.0,
|
|
442
455
|
top_p: float = 1.0,
|
|
443
456
|
) -> pd.Series:
|
|
444
457
|
"""Generate a response for each row after serialising it to JSON using a provided cache.
|
|
@@ -496,7 +509,7 @@ class OpenAIVecDataFrameAccessor:
|
|
|
496
509
|
instructions: str,
|
|
497
510
|
response_format: Type[ResponseFormat] = str,
|
|
498
511
|
batch_size: int = 128,
|
|
499
|
-
temperature: float = 0.0,
|
|
512
|
+
temperature: float | None = 0.0,
|
|
500
513
|
top_p: float = 1.0,
|
|
501
514
|
) -> pd.Series:
|
|
502
515
|
"""Generate a response for each row after serialising it to JSON.
|
|
@@ -681,7 +694,7 @@ class AsyncOpenAIVecSeriesAccessor:
|
|
|
681
694
|
instructions: str,
|
|
682
695
|
cache: AsyncBatchingMapProxy[str, ResponseFormat],
|
|
683
696
|
response_format: Type[ResponseFormat] = str,
|
|
684
|
-
temperature: float = 0.0,
|
|
697
|
+
temperature: float | None = 0.0,
|
|
685
698
|
top_p: float = 1.0,
|
|
686
699
|
) -> pd.Series:
|
|
687
700
|
"""Call an LLM once for every Series element using a provided cache (asynchronously).
|
|
@@ -848,7 +861,7 @@ class AsyncOpenAIVecSeriesAccessor:
|
|
|
848
861
|
instructions: str,
|
|
849
862
|
response_format: Type[ResponseFormat] = str,
|
|
850
863
|
batch_size: int = 128,
|
|
851
|
-
temperature: float = 0.0,
|
|
864
|
+
temperature: float | None = 0.0,
|
|
852
865
|
top_p: float = 1.0,
|
|
853
866
|
max_concurrency: int = 8,
|
|
854
867
|
) -> pd.Series:
|
|
@@ -975,7 +988,7 @@ class AsyncOpenAIVecDataFrameAccessor:
|
|
|
975
988
|
instructions: str,
|
|
976
989
|
cache: AsyncBatchingMapProxy[str, ResponseFormat],
|
|
977
990
|
response_format: Type[ResponseFormat] = str,
|
|
978
|
-
temperature: float = 0.0,
|
|
991
|
+
temperature: float | None = 0.0,
|
|
979
992
|
top_p: float = 1.0,
|
|
980
993
|
) -> pd.Series:
|
|
981
994
|
"""Generate a response for each row after serialising it to JSON using a provided cache (asynchronously).
|
|
@@ -1040,7 +1053,7 @@ class AsyncOpenAIVecDataFrameAccessor:
|
|
|
1040
1053
|
instructions: str,
|
|
1041
1054
|
response_format: Type[ResponseFormat] = str,
|
|
1042
1055
|
batch_size: int = 128,
|
|
1043
|
-
temperature: float = 0.0,
|
|
1056
|
+
temperature: float | None = 0.0,
|
|
1044
1057
|
top_p: float = 1.0,
|
|
1045
1058
|
max_concurrency: int = 8,
|
|
1046
1059
|
) -> pd.Series:
|
openaivec/provider.py
CHANGED
|
@@ -1,4 +1,5 @@
|
|
|
1
1
|
import os
|
|
2
|
+
import warnings
|
|
2
3
|
|
|
3
4
|
import tiktoken
|
|
4
5
|
from openai import AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI
|
|
@@ -7,7 +8,7 @@ from . import di
|
|
|
7
8
|
from .model import (
|
|
8
9
|
AzureOpenAIAPIKey,
|
|
9
10
|
AzureOpenAIAPIVersion,
|
|
10
|
-
|
|
11
|
+
AzureOpenAIBaseURL,
|
|
11
12
|
EmbeddingsModelName,
|
|
12
13
|
OpenAIAPIKey,
|
|
13
14
|
ResponsesModelName,
|
|
@@ -17,51 +18,102 @@ from .util import TextChunker
|
|
|
17
18
|
CONTAINER = di.Container()
|
|
18
19
|
|
|
19
20
|
|
|
21
|
+
def _check_azure_v1_api_url(base_url: str) -> None:
|
|
22
|
+
"""Check if Azure OpenAI base URL uses the recommended v1 API format.
|
|
23
|
+
|
|
24
|
+
Issues a warning if the URL doesn't end with '/openai/v1/' to encourage
|
|
25
|
+
migration to the v1 API format as recommended by Microsoft.
|
|
26
|
+
|
|
27
|
+
Reference: https://learn.microsoft.com/en-us/azure/ai-foundry/openai/api-version-lifecycle
|
|
28
|
+
|
|
29
|
+
Args:
|
|
30
|
+
base_url (str): The Azure OpenAI base URL to check.
|
|
31
|
+
"""
|
|
32
|
+
if base_url and not base_url.rstrip("/").endswith("/openai/v1"):
|
|
33
|
+
warnings.warn(
|
|
34
|
+
"⚠️ Azure OpenAI v1 API is recommended. Your base URL should end with '/openai/v1/'. "
|
|
35
|
+
f"Current URL: '{base_url}'. "
|
|
36
|
+
"Consider updating to: 'https://YOUR-RESOURCE-NAME.services.ai.azure.com/openai/v1/' "
|
|
37
|
+
"for better performance and future compatibility. "
|
|
38
|
+
"See: https://learn.microsoft.com/en-us/azure/ai-foundry/openai/api-version-lifecycle",
|
|
39
|
+
UserWarning,
|
|
40
|
+
stacklevel=3,
|
|
41
|
+
)
|
|
42
|
+
|
|
43
|
+
|
|
20
44
|
def provide_openai_client() -> OpenAI:
|
|
21
|
-
"""Provide OpenAI client based on environment variables.
|
|
45
|
+
"""Provide OpenAI client based on environment variables.
|
|
46
|
+
|
|
47
|
+
Automatically detects and prioritizes OpenAI over Azure OpenAI configuration.
|
|
48
|
+
Checks the following environment variables in order:
|
|
49
|
+
1. OPENAI_API_KEY - if set, creates standard OpenAI client
|
|
50
|
+
2. Azure OpenAI variables (AZURE_OPENAI_API_KEY, AZURE_OPENAI_BASE_URL,
|
|
51
|
+
AZURE_OPENAI_API_VERSION) - if all set, creates Azure OpenAI client
|
|
52
|
+
|
|
53
|
+
Returns:
|
|
54
|
+
OpenAI: Configured OpenAI or AzureOpenAI client instance.
|
|
55
|
+
|
|
56
|
+
Raises:
|
|
57
|
+
ValueError: If no valid environment variables are found for either service.
|
|
58
|
+
"""
|
|
22
59
|
openai_api_key = CONTAINER.resolve(OpenAIAPIKey)
|
|
23
60
|
if openai_api_key.value:
|
|
24
61
|
return OpenAI()
|
|
25
62
|
|
|
26
63
|
azure_api_key = CONTAINER.resolve(AzureOpenAIAPIKey)
|
|
27
|
-
|
|
64
|
+
azure_base_url = CONTAINER.resolve(AzureOpenAIBaseURL)
|
|
28
65
|
azure_api_version = CONTAINER.resolve(AzureOpenAIAPIVersion)
|
|
29
66
|
|
|
30
|
-
if all(param.value for param in [azure_api_key,
|
|
67
|
+
if all(param.value for param in [azure_api_key, azure_base_url, azure_api_version]):
|
|
68
|
+
_check_azure_v1_api_url(azure_base_url.value)
|
|
31
69
|
return AzureOpenAI(
|
|
32
70
|
api_key=azure_api_key.value,
|
|
33
|
-
|
|
71
|
+
base_url=azure_base_url.value,
|
|
34
72
|
api_version=azure_api_version.value,
|
|
35
73
|
)
|
|
36
74
|
|
|
37
75
|
raise ValueError(
|
|
38
76
|
"No valid OpenAI or Azure OpenAI environment variables found. "
|
|
39
77
|
"Please set either OPENAI_API_KEY or AZURE_OPENAI_API_KEY, "
|
|
40
|
-
"
|
|
78
|
+
"AZURE_OPENAI_BASE_URL, and AZURE_OPENAI_API_VERSION."
|
|
41
79
|
)
|
|
42
80
|
|
|
43
81
|
|
|
44
82
|
def provide_async_openai_client() -> AsyncOpenAI:
|
|
45
|
-
"""Provide
|
|
83
|
+
"""Provide asynchronous OpenAI client based on environment variables.
|
|
84
|
+
|
|
85
|
+
Automatically detects and prioritizes OpenAI over Azure OpenAI configuration.
|
|
86
|
+
Checks the following environment variables in order:
|
|
87
|
+
1. OPENAI_API_KEY - if set, creates standard AsyncOpenAI client
|
|
88
|
+
2. Azure OpenAI variables (AZURE_OPENAI_API_KEY, AZURE_OPENAI_BASE_URL,
|
|
89
|
+
AZURE_OPENAI_API_VERSION) - if all set, creates AsyncAzureOpenAI client
|
|
90
|
+
|
|
91
|
+
Returns:
|
|
92
|
+
AsyncOpenAI: Configured AsyncOpenAI or AsyncAzureOpenAI client instance.
|
|
93
|
+
|
|
94
|
+
Raises:
|
|
95
|
+
ValueError: If no valid environment variables are found for either service.
|
|
96
|
+
"""
|
|
46
97
|
openai_api_key = CONTAINER.resolve(OpenAIAPIKey)
|
|
47
98
|
if openai_api_key.value:
|
|
48
99
|
return AsyncOpenAI()
|
|
49
100
|
|
|
50
101
|
azure_api_key = CONTAINER.resolve(AzureOpenAIAPIKey)
|
|
51
|
-
|
|
102
|
+
azure_base_url = CONTAINER.resolve(AzureOpenAIBaseURL)
|
|
52
103
|
azure_api_version = CONTAINER.resolve(AzureOpenAIAPIVersion)
|
|
53
104
|
|
|
54
|
-
if all(param.value for param in [azure_api_key,
|
|
105
|
+
if all(param.value for param in [azure_api_key, azure_base_url, azure_api_version]):
|
|
106
|
+
_check_azure_v1_api_url(azure_base_url.value)
|
|
55
107
|
return AsyncAzureOpenAI(
|
|
56
108
|
api_key=azure_api_key.value,
|
|
57
|
-
|
|
109
|
+
base_url=azure_base_url.value,
|
|
58
110
|
api_version=azure_api_version.value,
|
|
59
111
|
)
|
|
60
112
|
|
|
61
113
|
raise ValueError(
|
|
62
114
|
"No valid OpenAI or Azure OpenAI environment variables found. "
|
|
63
115
|
"Please set either OPENAI_API_KEY or AZURE_OPENAI_API_KEY, "
|
|
64
|
-
"
|
|
116
|
+
"AZURE_OPENAI_BASE_URL, and AZURE_OPENAI_API_VERSION."
|
|
65
117
|
)
|
|
66
118
|
|
|
67
119
|
|
|
@@ -69,10 +121,10 @@ CONTAINER.register(ResponsesModelName, lambda: ResponsesModelName("gpt-4.1-mini"
|
|
|
69
121
|
CONTAINER.register(EmbeddingsModelName, lambda: EmbeddingsModelName("text-embedding-3-small"))
|
|
70
122
|
CONTAINER.register(OpenAIAPIKey, lambda: OpenAIAPIKey(os.getenv("OPENAI_API_KEY")))
|
|
71
123
|
CONTAINER.register(AzureOpenAIAPIKey, lambda: AzureOpenAIAPIKey(os.getenv("AZURE_OPENAI_API_KEY")))
|
|
72
|
-
CONTAINER.register(
|
|
124
|
+
CONTAINER.register(AzureOpenAIBaseURL, lambda: AzureOpenAIBaseURL(os.getenv("AZURE_OPENAI_BASE_URL")))
|
|
73
125
|
CONTAINER.register(
|
|
74
126
|
cls=AzureOpenAIAPIVersion,
|
|
75
|
-
provider=lambda: AzureOpenAIAPIVersion(os.getenv("AZURE_OPENAI_API_VERSION", "
|
|
127
|
+
provider=lambda: AzureOpenAIAPIVersion(os.getenv("AZURE_OPENAI_API_VERSION", "preview")),
|
|
76
128
|
)
|
|
77
129
|
CONTAINER.register(OpenAI, provide_openai_client)
|
|
78
130
|
CONTAINER.register(AsyncOpenAI, provide_async_openai_client)
|
|
@@ -89,10 +141,10 @@ def reset_environment_registrations():
|
|
|
89
141
|
"""
|
|
90
142
|
CONTAINER.register(OpenAIAPIKey, lambda: OpenAIAPIKey(os.getenv("OPENAI_API_KEY")))
|
|
91
143
|
CONTAINER.register(AzureOpenAIAPIKey, lambda: AzureOpenAIAPIKey(os.getenv("AZURE_OPENAI_API_KEY")))
|
|
92
|
-
CONTAINER.register(
|
|
144
|
+
CONTAINER.register(AzureOpenAIBaseURL, lambda: AzureOpenAIBaseURL(os.getenv("AZURE_OPENAI_BASE_URL")))
|
|
93
145
|
CONTAINER.register(
|
|
94
146
|
cls=AzureOpenAIAPIVersion,
|
|
95
|
-
provider=lambda: AzureOpenAIAPIVersion(os.getenv("AZURE_OPENAI_API_VERSION", "
|
|
147
|
+
provider=lambda: AzureOpenAIAPIVersion(os.getenv("AZURE_OPENAI_API_VERSION", "preview")),
|
|
96
148
|
)
|
|
97
149
|
CONTAINER.register(OpenAI, provide_openai_client)
|
|
98
150
|
CONTAINER.register(AsyncOpenAI, provide_async_openai_client)
|
openaivec/responses.py
CHANGED
|
@@ -1,8 +1,9 @@
|
|
|
1
|
+
import warnings
|
|
1
2
|
from dataclasses import dataclass, field
|
|
2
3
|
from logging import Logger, getLogger
|
|
3
4
|
from typing import Generic, List, Type, cast
|
|
4
5
|
|
|
5
|
-
from openai import AsyncOpenAI, OpenAI, RateLimitError
|
|
6
|
+
from openai import AsyncOpenAI, BadRequestError, InternalServerError, OpenAI, RateLimitError
|
|
6
7
|
from openai.types.responses import ParsedResponse
|
|
7
8
|
from pydantic import BaseModel
|
|
8
9
|
|
|
@@ -19,6 +20,35 @@ __all__ = [
|
|
|
19
20
|
_LOGGER: Logger = getLogger(__name__)
|
|
20
21
|
|
|
21
22
|
|
|
23
|
+
def _handle_temperature_error(error: BadRequestError, model_name: str, temperature: float) -> None:
|
|
24
|
+
"""Handle temperature-related errors for reasoning models.
|
|
25
|
+
|
|
26
|
+
Detects when a model doesn't support temperature parameter and provides guidance.
|
|
27
|
+
|
|
28
|
+
Args:
|
|
29
|
+
error (BadRequestError): The OpenAI API error.
|
|
30
|
+
model_name (str): The model that caused the error.
|
|
31
|
+
temperature (float): The temperature value that was rejected.
|
|
32
|
+
"""
|
|
33
|
+
error_message = str(error)
|
|
34
|
+
if "temperature" in error_message.lower() and "not supported" in error_message.lower():
|
|
35
|
+
guidance_message = (
|
|
36
|
+
f"🔧 Model '{model_name}' rejected temperature parameter (value: {temperature}). "
|
|
37
|
+
f"This typically happens with reasoning models (o1-preview, o1-mini, o3, etc.). "
|
|
38
|
+
f"To fix this, you MUST explicitly set temperature=None:\n"
|
|
39
|
+
f"• For pandas: df.col.ai.responses('prompt', temperature=None)\n"
|
|
40
|
+
f"• For Spark UDFs: responses_udf('prompt', temperature=None)\n"
|
|
41
|
+
f"• For direct API: BatchResponses.of(client, model, temperature=None)\n"
|
|
42
|
+
f"• Original error: {error_message}\n"
|
|
43
|
+
f"See: https://learn.microsoft.com/en-us/azure/ai-foundry/openai/how-to/reasoning"
|
|
44
|
+
)
|
|
45
|
+
warnings.warn(guidance_message, UserWarning, stacklevel=5)
|
|
46
|
+
|
|
47
|
+
# Re-raise with enhanced message
|
|
48
|
+
enhanced_message = f"{error_message}\n\nSUGGESTION: Set temperature=None to resolve this error."
|
|
49
|
+
raise BadRequestError(message=enhanced_message, response=error.response, body=error.body)
|
|
50
|
+
|
|
51
|
+
|
|
22
52
|
def _vectorize_system_message(system_message: str) -> str:
|
|
23
53
|
"""Build a system prompt that instructs the model to work on batched inputs.
|
|
24
54
|
|
|
@@ -116,7 +146,7 @@ class BatchResponses(Generic[ResponseFormat]):
|
|
|
116
146
|
|
|
117
147
|
Attributes:
|
|
118
148
|
client (OpenAI): Initialised OpenAI client.
|
|
119
|
-
model_name (str):
|
|
149
|
+
model_name (str): For Azure OpenAI, use your deployment name. For OpenAI, use the model name.
|
|
120
150
|
system_message (str): System prompt prepended to every request.
|
|
121
151
|
temperature (float): Sampling temperature.
|
|
122
152
|
top_p (float): Nucleus‑sampling parameter.
|
|
@@ -131,9 +161,9 @@ class BatchResponses(Generic[ResponseFormat]):
|
|
|
131
161
|
"""
|
|
132
162
|
|
|
133
163
|
client: OpenAI
|
|
134
|
-
model_name: str #
|
|
164
|
+
model_name: str # For Azure: deployment name, for OpenAI: model name
|
|
135
165
|
system_message: str
|
|
136
|
-
temperature: float = 0.0
|
|
166
|
+
temperature: float | None = 0.0
|
|
137
167
|
top_p: float = 1.0
|
|
138
168
|
response_format: Type[ResponseFormat] = str
|
|
139
169
|
cache: BatchingMapProxy[str, ResponseFormat] = field(default_factory=lambda: BatchingMapProxy(batch_size=128))
|
|
@@ -146,7 +176,7 @@ class BatchResponses(Generic[ResponseFormat]):
|
|
|
146
176
|
client: OpenAI,
|
|
147
177
|
model_name: str,
|
|
148
178
|
system_message: str,
|
|
149
|
-
temperature: float = 0.0,
|
|
179
|
+
temperature: float | None = 0.0,
|
|
150
180
|
top_p: float = 1.0,
|
|
151
181
|
response_format: Type[ResponseFormat] = str,
|
|
152
182
|
batch_size: int = 128,
|
|
@@ -155,7 +185,7 @@ class BatchResponses(Generic[ResponseFormat]):
|
|
|
155
185
|
|
|
156
186
|
Args:
|
|
157
187
|
client (OpenAI): OpenAI client.
|
|
158
|
-
model_name (str):
|
|
188
|
+
model_name (str): For Azure OpenAI, use your deployment name. For OpenAI, use the model name.
|
|
159
189
|
system_message (str): System prompt for the model.
|
|
160
190
|
temperature (float, optional): Sampling temperature. Defaults to 0.0.
|
|
161
191
|
top_p (float, optional): Nucleus sampling parameter. Defaults to 1.0.
|
|
@@ -181,7 +211,7 @@ class BatchResponses(Generic[ResponseFormat]):
|
|
|
181
211
|
|
|
182
212
|
Args:
|
|
183
213
|
client (OpenAI): OpenAI client.
|
|
184
|
-
model_name (str):
|
|
214
|
+
model_name (str): For Azure OpenAI, use your deployment name. For OpenAI, use the model name.
|
|
185
215
|
task (PreparedTask): Prepared task with instructions and response format.
|
|
186
216
|
batch_size (int, optional): Max unique prompts per API call. Defaults to 128.
|
|
187
217
|
|
|
@@ -206,7 +236,7 @@ class BatchResponses(Generic[ResponseFormat]):
|
|
|
206
236
|
)
|
|
207
237
|
|
|
208
238
|
@observe(_LOGGER)
|
|
209
|
-
@backoff(
|
|
239
|
+
@backoff(exceptions=[RateLimitError, InternalServerError], scale=1, max_retries=12)
|
|
210
240
|
def _request_llm(self, user_messages: List[Message[str]]) -> ParsedResponse[Response[ResponseFormat]]:
|
|
211
241
|
"""Make a single call to the OpenAI JSON‑mode endpoint.
|
|
212
242
|
|
|
@@ -231,14 +261,23 @@ class BatchResponses(Generic[ResponseFormat]):
|
|
|
231
261
|
class ResponseT(BaseModel):
|
|
232
262
|
assistant_messages: List[MessageT]
|
|
233
263
|
|
|
234
|
-
|
|
235
|
-
|
|
236
|
-
|
|
237
|
-
|
|
238
|
-
|
|
239
|
-
top_p
|
|
240
|
-
text_format
|
|
241
|
-
|
|
264
|
+
# Prepare API parameters, excluding temperature if None (for reasoning models)
|
|
265
|
+
api_params = {
|
|
266
|
+
"model": self.model_name,
|
|
267
|
+
"instructions": self._vectorized_system_message,
|
|
268
|
+
"input": Request(user_messages=user_messages).model_dump_json(),
|
|
269
|
+
"top_p": self.top_p,
|
|
270
|
+
"text_format": ResponseT,
|
|
271
|
+
}
|
|
272
|
+
if self.temperature is not None:
|
|
273
|
+
api_params["temperature"] = self.temperature
|
|
274
|
+
|
|
275
|
+
try:
|
|
276
|
+
completion: ParsedResponse[ResponseT] = self.client.responses.parse(**api_params)
|
|
277
|
+
except BadRequestError as e:
|
|
278
|
+
_handle_temperature_error(e, self.model_name, self.temperature or 0.0)
|
|
279
|
+
raise # Re-raise if it wasn't a temperature error
|
|
280
|
+
|
|
242
281
|
return cast(ParsedResponse[Response[ResponseFormat]], completion)
|
|
243
282
|
|
|
244
283
|
@observe(_LOGGER)
|
|
@@ -309,7 +348,7 @@ class AsyncBatchResponses(Generic[ResponseFormat]):
|
|
|
309
348
|
|
|
310
349
|
Attributes:
|
|
311
350
|
client (AsyncOpenAI): Initialised OpenAI async client.
|
|
312
|
-
model_name (str):
|
|
351
|
+
model_name (str): For Azure OpenAI, use your deployment name. For OpenAI, use the model name.
|
|
313
352
|
system_message (str): System prompt prepended to every request.
|
|
314
353
|
temperature (float): Sampling temperature.
|
|
315
354
|
top_p (float): Nucleus‑sampling parameter.
|
|
@@ -318,9 +357,9 @@ class AsyncBatchResponses(Generic[ResponseFormat]):
|
|
|
318
357
|
"""
|
|
319
358
|
|
|
320
359
|
client: AsyncOpenAI
|
|
321
|
-
model_name: str #
|
|
360
|
+
model_name: str # For Azure: deployment name, for OpenAI: model name
|
|
322
361
|
system_message: str
|
|
323
|
-
temperature: float = 0.0
|
|
362
|
+
temperature: float | None = 0.0
|
|
324
363
|
top_p: float = 1.0
|
|
325
364
|
response_format: Type[ResponseFormat] = str
|
|
326
365
|
cache: AsyncBatchingMapProxy[str, ResponseFormat] = field(
|
|
@@ -335,7 +374,7 @@ class AsyncBatchResponses(Generic[ResponseFormat]):
|
|
|
335
374
|
client: AsyncOpenAI,
|
|
336
375
|
model_name: str,
|
|
337
376
|
system_message: str,
|
|
338
|
-
temperature: float = 0.0,
|
|
377
|
+
temperature: float | None = 0.0,
|
|
339
378
|
top_p: float = 1.0,
|
|
340
379
|
response_format: Type[ResponseFormat] = str,
|
|
341
380
|
batch_size: int = 128,
|
|
@@ -345,7 +384,7 @@ class AsyncBatchResponses(Generic[ResponseFormat]):
|
|
|
345
384
|
|
|
346
385
|
Args:
|
|
347
386
|
client (AsyncOpenAI): OpenAI async client.
|
|
348
|
-
model_name (str):
|
|
387
|
+
model_name (str): For Azure OpenAI, use your deployment name. For OpenAI, use the model name.
|
|
349
388
|
system_message (str): System prompt.
|
|
350
389
|
temperature (float, optional): Sampling temperature. Defaults to 0.0.
|
|
351
390
|
top_p (float, optional): Nucleus sampling parameter. Defaults to 1.0.
|
|
@@ -374,7 +413,7 @@ class AsyncBatchResponses(Generic[ResponseFormat]):
|
|
|
374
413
|
|
|
375
414
|
Args:
|
|
376
415
|
client (AsyncOpenAI): OpenAI async client.
|
|
377
|
-
model_name (str):
|
|
416
|
+
model_name (str): For Azure OpenAI, use your deployment name. For OpenAI, use the model name.
|
|
378
417
|
task (PreparedTask): Prepared task with instructions and response format.
|
|
379
418
|
batch_size (int, optional): Max unique prompts per API call. Defaults to 128.
|
|
380
419
|
max_concurrency (int, optional): Max concurrent API calls. Defaults to 8.
|
|
@@ -400,7 +439,7 @@ class AsyncBatchResponses(Generic[ResponseFormat]):
|
|
|
400
439
|
)
|
|
401
440
|
|
|
402
441
|
@observe(_LOGGER)
|
|
403
|
-
@backoff_async(
|
|
442
|
+
@backoff_async(exceptions=[RateLimitError, InternalServerError], scale=1, max_retries=12)
|
|
404
443
|
async def _request_llm(self, user_messages: List[Message[str]]) -> ParsedResponse[Response[ResponseFormat]]:
|
|
405
444
|
"""Make a single async call to the OpenAI JSON‑mode endpoint.
|
|
406
445
|
|
|
@@ -422,14 +461,23 @@ class AsyncBatchResponses(Generic[ResponseFormat]):
|
|
|
422
461
|
class ResponseT(BaseModel):
|
|
423
462
|
assistant_messages: List[MessageT]
|
|
424
463
|
|
|
425
|
-
|
|
426
|
-
|
|
427
|
-
|
|
428
|
-
|
|
429
|
-
|
|
430
|
-
top_p
|
|
431
|
-
text_format
|
|
432
|
-
|
|
464
|
+
# Prepare API parameters, excluding temperature if None (for reasoning models)
|
|
465
|
+
api_params = {
|
|
466
|
+
"model": self.model_name,
|
|
467
|
+
"instructions": self._vectorized_system_message,
|
|
468
|
+
"input": Request(user_messages=user_messages).model_dump_json(),
|
|
469
|
+
"top_p": self.top_p,
|
|
470
|
+
"text_format": ResponseT,
|
|
471
|
+
}
|
|
472
|
+
if self.temperature is not None:
|
|
473
|
+
api_params["temperature"] = self.temperature
|
|
474
|
+
|
|
475
|
+
try:
|
|
476
|
+
completion: ParsedResponse[ResponseT] = await self.client.responses.parse(**api_params)
|
|
477
|
+
except BadRequestError as e:
|
|
478
|
+
_handle_temperature_error(e, self.model_name, self.temperature or 0.0)
|
|
479
|
+
raise # Re-raise if it wasn't a temperature error
|
|
480
|
+
|
|
433
481
|
return cast(ParsedResponse[Response[ResponseFormat]], completion)
|
|
434
482
|
|
|
435
483
|
@observe(_LOGGER)
|
openaivec/spark.py
CHANGED
|
@@ -28,8 +28,8 @@ sc.environment["OPENAI_API_KEY"] = "your-openai-api-key"
|
|
|
28
28
|
|
|
29
29
|
# Option 2: Using Azure OpenAI
|
|
30
30
|
# sc.environment["AZURE_OPENAI_API_KEY"] = "your-azure-openai-api-key"
|
|
31
|
-
# sc.environment["
|
|
32
|
-
# sc.environment["AZURE_OPENAI_API_VERSION"] = "
|
|
31
|
+
# sc.environment["AZURE_OPENAI_BASE_URL"] = "https://YOUR-RESOURCE-NAME.services.ai.azure.com/openai/v1/"
|
|
32
|
+
# sc.environment["AZURE_OPENAI_API_VERSION"] = "preview"
|
|
33
33
|
```
|
|
34
34
|
|
|
35
35
|
Next, create UDFs and register them:
|
|
@@ -50,7 +50,7 @@ spark.udf.register(
|
|
|
50
50
|
responses_udf(
|
|
51
51
|
instructions="Translate the text to multiple languages.",
|
|
52
52
|
response_format=Translation,
|
|
53
|
-
model_name="gpt-4.1-mini", #
|
|
53
|
+
model_name="gpt-4.1-mini", # For Azure: deployment name, for OpenAI: model name
|
|
54
54
|
batch_size=64, # Rows per API request within partition
|
|
55
55
|
max_concurrency=8 # Concurrent requests PER EXECUTOR
|
|
56
56
|
),
|
|
@@ -67,7 +67,7 @@ spark.udf.register(
|
|
|
67
67
|
spark.udf.register(
|
|
68
68
|
"embed_async",
|
|
69
69
|
embeddings_udf(
|
|
70
|
-
model_name="text-embedding-3-small", #
|
|
70
|
+
model_name="text-embedding-3-small", # For Azure: deployment name, for OpenAI: model name
|
|
71
71
|
batch_size=128, # Larger batches for embeddings
|
|
72
72
|
max_concurrency=8 # Concurrent requests PER EXECUTOR
|
|
73
73
|
),
|
|
@@ -224,7 +224,7 @@ def responses_udf(
|
|
|
224
224
|
response_format: Type[ResponseFormat] = str,
|
|
225
225
|
model_name: str = "gpt-4.1-mini",
|
|
226
226
|
batch_size: int = 128,
|
|
227
|
-
temperature: float = 0.0,
|
|
227
|
+
temperature: float | None = 0.0,
|
|
228
228
|
top_p: float = 1.0,
|
|
229
229
|
max_concurrency: int = 8,
|
|
230
230
|
) -> UserDefinedFunction:
|
|
@@ -245,15 +245,15 @@ def responses_udf(
|
|
|
245
245
|
|
|
246
246
|
For Azure OpenAI:
|
|
247
247
|
sc.environment["AZURE_OPENAI_API_KEY"] = "your-azure-openai-api-key"
|
|
248
|
-
sc.environment["
|
|
249
|
-
sc.environment["AZURE_OPENAI_API_VERSION"] = "
|
|
248
|
+
sc.environment["AZURE_OPENAI_BASE_URL"] = "https://YOUR-RESOURCE-NAME.services.ai.azure.com/openai/v1/"
|
|
249
|
+
sc.environment["AZURE_OPENAI_API_VERSION"] = "preview"
|
|
250
250
|
|
|
251
251
|
Args:
|
|
252
252
|
instructions (str): The system prompt or instructions for the model.
|
|
253
253
|
response_format (Type[ResponseFormat]): The desired output format. Either `str` for plain text
|
|
254
254
|
or a Pydantic `BaseModel` for structured JSON output. Defaults to `str`.
|
|
255
|
-
model_name (str):
|
|
256
|
-
Defaults to "gpt-4.1-mini".
|
|
255
|
+
model_name (str): For Azure OpenAI, use your deployment name (e.g., "my-gpt4-deployment").
|
|
256
|
+
For OpenAI, use the model name (e.g., "gpt-4.1-mini"). Defaults to "gpt-4.1-mini".
|
|
257
257
|
batch_size (int): Number of rows per async batch request within each partition.
|
|
258
258
|
Larger values reduce API call overhead but increase memory usage.
|
|
259
259
|
Recommended: 32-128 depending on data complexity. Defaults to 128.
|
|
@@ -357,8 +357,8 @@ def task_udf(
|
|
|
357
357
|
Args:
|
|
358
358
|
task (PreparedTask): A predefined task configuration containing instructions,
|
|
359
359
|
response format, temperature, and top_p settings.
|
|
360
|
-
model_name (str):
|
|
361
|
-
Defaults to "gpt-4.1-mini".
|
|
360
|
+
model_name (str): For Azure OpenAI, use your deployment name (e.g., "my-gpt4-deployment").
|
|
361
|
+
For OpenAI, use the model name (e.g., "gpt-4.1-mini"). Defaults to "gpt-4.1-mini".
|
|
362
362
|
batch_size (int): Number of rows per async batch request within each partition.
|
|
363
363
|
Larger values reduce API call overhead but increase memory usage.
|
|
364
364
|
Recommended: 32-128 depending on task complexity. Defaults to 128.
|
|
@@ -474,12 +474,12 @@ def embeddings_udf(
|
|
|
474
474
|
|
|
475
475
|
For Azure OpenAI:
|
|
476
476
|
sc.environment["AZURE_OPENAI_API_KEY"] = "your-azure-openai-api-key"
|
|
477
|
-
sc.environment["
|
|
478
|
-
sc.environment["AZURE_OPENAI_API_VERSION"] = "
|
|
477
|
+
sc.environment["AZURE_OPENAI_BASE_URL"] = "https://YOUR-RESOURCE-NAME.services.ai.azure.com/openai/v1/"
|
|
478
|
+
sc.environment["AZURE_OPENAI_API_VERSION"] = "preview"
|
|
479
479
|
|
|
480
480
|
Args:
|
|
481
|
-
model_name (str):
|
|
482
|
-
Defaults to "text-embedding-3-small".
|
|
481
|
+
model_name (str): For Azure OpenAI, use your deployment name (e.g., "my-embedding-deployment").
|
|
482
|
+
For OpenAI, use the model name (e.g., "text-embedding-3-small"). Defaults to "text-embedding-3-small".
|
|
483
483
|
batch_size (int): Number of rows per async batch request within each partition.
|
|
484
484
|
Larger values reduce API call overhead but increase memory usage.
|
|
485
485
|
Embeddings typically handle larger batches efficiently.
|
openaivec/util.py
CHANGED
|
@@ -3,7 +3,7 @@ import functools
|
|
|
3
3
|
import re
|
|
4
4
|
import time
|
|
5
5
|
from dataclasses import dataclass
|
|
6
|
-
from typing import Awaitable, Callable, List, TypeVar
|
|
6
|
+
from typing import Awaitable, Callable, List, Type, TypeVar
|
|
7
7
|
|
|
8
8
|
import numpy as np
|
|
9
9
|
import tiktoken
|
|
@@ -34,24 +34,28 @@ def get_exponential_with_cutoff(scale: float) -> float:
|
|
|
34
34
|
return v
|
|
35
35
|
|
|
36
36
|
|
|
37
|
-
def backoff(
|
|
37
|
+
def backoff(
|
|
38
|
+
exceptions: List[Type[Exception]],
|
|
39
|
+
scale: int | None = None,
|
|
40
|
+
max_retries: int | None = None,
|
|
41
|
+
) -> Callable[..., V]:
|
|
38
42
|
"""Decorator implementing exponential back‑off retry logic.
|
|
39
43
|
|
|
40
44
|
Args:
|
|
41
|
-
|
|
45
|
+
exceptions (List[Type[Exception]]): List of exception types that trigger a retry.
|
|
42
46
|
scale (int | None): Initial scale parameter for the exponential jitter.
|
|
43
47
|
This scale is used as the mean for the first delay's exponential
|
|
44
48
|
distribution and doubles with each subsequent retry. If ``None``,
|
|
45
49
|
an initial scale of 1.0 is used.
|
|
46
|
-
max_retries (
|
|
50
|
+
max_retries (int | None): Maximum number of retries. ``None`` means
|
|
47
51
|
retry indefinitely.
|
|
48
52
|
|
|
49
53
|
Returns:
|
|
50
54
|
Callable[..., V]: A decorated function that retries on the specified
|
|
51
|
-
|
|
55
|
+
exceptions with exponential back‑off.
|
|
52
56
|
|
|
53
57
|
Raises:
|
|
54
|
-
|
|
58
|
+
Exception: Re‑raised when the maximum number of retries is exceeded.
|
|
55
59
|
"""
|
|
56
60
|
|
|
57
61
|
def decorator(func: Callable[..., V]) -> Callable[..., V]:
|
|
@@ -65,7 +69,7 @@ def backoff(exception: type[Exception], scale: int | None = None, max_retries: i
|
|
|
65
69
|
while True:
|
|
66
70
|
try:
|
|
67
71
|
return func(*args, **kwargs)
|
|
68
|
-
except
|
|
72
|
+
except tuple(exceptions):
|
|
69
73
|
attempt += 1
|
|
70
74
|
if max_retries is not None and attempt >= max_retries:
|
|
71
75
|
raise
|
|
@@ -83,12 +87,14 @@ def backoff(exception: type[Exception], scale: int | None = None, max_retries: i
|
|
|
83
87
|
|
|
84
88
|
|
|
85
89
|
def backoff_async(
|
|
86
|
-
|
|
90
|
+
exceptions: List[Type[Exception]],
|
|
91
|
+
scale: int | None = None,
|
|
92
|
+
max_retries: int | None = None,
|
|
87
93
|
) -> Callable[..., Awaitable[V]]:
|
|
88
94
|
"""Asynchronous version of the backoff decorator.
|
|
89
95
|
|
|
90
96
|
Args:
|
|
91
|
-
|
|
97
|
+
exceptions (List[Type[Exception]]): List of exception types that trigger a retry.
|
|
92
98
|
scale (int | None): Initial scale parameter for the exponential jitter.
|
|
93
99
|
This scale is used as the mean for the first delay's exponential
|
|
94
100
|
distribution and doubles with each subsequent retry. If ``None``,
|
|
@@ -98,10 +104,10 @@ def backoff_async(
|
|
|
98
104
|
|
|
99
105
|
Returns:
|
|
100
106
|
Callable[..., Awaitable[V]]: A decorated asynchronous function that
|
|
101
|
-
retries on the specified
|
|
107
|
+
retries on the specified exceptions with exponential back‑off.
|
|
102
108
|
|
|
103
109
|
Raises:
|
|
104
|
-
|
|
110
|
+
Exception: Re‑raised when the maximum number of retries is exceeded.
|
|
105
111
|
"""
|
|
106
112
|
|
|
107
113
|
def decorator(func: Callable[..., Awaitable[V]]) -> Callable[..., Awaitable[V]]:
|
|
@@ -115,7 +121,7 @@ def backoff_async(
|
|
|
115
121
|
while True:
|
|
116
122
|
try:
|
|
117
123
|
return await func(*args, **kwargs)
|
|
118
|
-
except
|
|
124
|
+
except tuple(exceptions):
|
|
119
125
|
attempt += 1
|
|
120
126
|
if max_retries is not None and attempt >= max_retries:
|
|
121
127
|
raise
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: openaivec
|
|
3
|
-
Version: 0.13.
|
|
3
|
+
Version: 0.13.2
|
|
4
4
|
Summary: Generative mutation for tabular calculation
|
|
5
5
|
Project-URL: Homepage, https://microsoft.github.io/openaivec/
|
|
6
6
|
Project-URL: Repository, https://github.com/microsoft/openaivec
|
|
@@ -180,8 +180,8 @@ from openaivec import pandas_ext
|
|
|
180
180
|
os.environ["OPENAI_API_KEY"] = "your-api-key-here"
|
|
181
181
|
# Or for Azure OpenAI:
|
|
182
182
|
# os.environ["AZURE_OPENAI_API_KEY"] = "your-azure-key"
|
|
183
|
-
# os.environ["
|
|
184
|
-
# os.environ["AZURE_OPENAI_API_VERSION"] = "
|
|
183
|
+
# os.environ["AZURE_OPENAI_BASE_URL"] = "https://YOUR-RESOURCE-NAME.services.ai.azure.com/openai/v1/"
|
|
184
|
+
# os.environ["AZURE_OPENAI_API_VERSION"] = "preview"
|
|
185
185
|
|
|
186
186
|
# Authentication Option 2: Custom client (optional)
|
|
187
187
|
# from openai import OpenAI, AsyncOpenAI
|
|
@@ -190,6 +190,7 @@ os.environ["OPENAI_API_KEY"] = "your-api-key-here"
|
|
|
190
190
|
# pandas_ext.use_async(AsyncOpenAI())
|
|
191
191
|
|
|
192
192
|
# Configure model (optional - defaults to gpt-4.1-mini)
|
|
193
|
+
# For Azure OpenAI: use your deployment name, for OpenAI: use model name
|
|
193
194
|
pandas_ext.responses_model("gpt-4.1-mini")
|
|
194
195
|
|
|
195
196
|
# Create your data
|
|
@@ -211,6 +212,27 @@ result = df.assign(
|
|
|
211
212
|
|
|
212
213
|
📓 **[Interactive pandas examples →](https://microsoft.github.io/openaivec/examples/pandas/)**
|
|
213
214
|
|
|
215
|
+
### Using with Reasoning Models
|
|
216
|
+
|
|
217
|
+
When using reasoning models (o1-preview, o1-mini, o3-mini, etc.), you must set `temperature=None` to avoid API errors:
|
|
218
|
+
|
|
219
|
+
```python
|
|
220
|
+
# For reasoning models like o1-preview, o1-mini, o3-mini
|
|
221
|
+
pandas_ext.responses_model("o1-mini") # Set your reasoning model
|
|
222
|
+
|
|
223
|
+
# MUST use temperature=None with reasoning models
|
|
224
|
+
result = df.assign(
|
|
225
|
+
analysis=lambda df: df.text.ai.responses(
|
|
226
|
+
"Analyze this text step by step",
|
|
227
|
+
temperature=None # Required for reasoning models
|
|
228
|
+
)
|
|
229
|
+
)
|
|
230
|
+
```
|
|
231
|
+
|
|
232
|
+
**Why this is needed**: Reasoning models don't support temperature parameters and will return an error if temperature is specified. The library automatically detects these errors and provides guidance on how to fix them.
|
|
233
|
+
|
|
234
|
+
**Reference**: [Azure OpenAI Reasoning Models](https://learn.microsoft.com/en-us/azure/ai-foundry/openai/how-to/reasoning)
|
|
235
|
+
|
|
214
236
|
### Using Pre-configured Tasks
|
|
215
237
|
|
|
216
238
|
For common text processing operations, openaivec provides ready-to-use tasks that eliminate the need to write custom prompts:
|
|
@@ -322,7 +344,7 @@ sc.environment["OPENAI_API_KEY"] = os.environ.get("OPENAI_API_KEY")
|
|
|
322
344
|
|
|
323
345
|
# Option 2: Using Azure OpenAI
|
|
324
346
|
# sc.environment["AZURE_OPENAI_API_KEY"] = os.environ.get("AZURE_OPENAI_API_KEY")
|
|
325
|
-
# sc.environment["
|
|
347
|
+
# sc.environment["AZURE_OPENAI_BASE_URL"] = os.environ.get("AZURE_OPENAI_BASE_URL")
|
|
326
348
|
# sc.environment["AZURE_OPENAI_API_VERSION"] = os.environ.get("AZURE_OPENAI_API_VERSION")
|
|
327
349
|
```
|
|
328
350
|
|
|
@@ -380,6 +402,16 @@ spark.udf.register(
|
|
|
380
402
|
)
|
|
381
403
|
)
|
|
382
404
|
|
|
405
|
+
# --- Register UDF for Reasoning Models ---
|
|
406
|
+
# For reasoning models (o1-preview, o1-mini, o3, etc.), set temperature=None
|
|
407
|
+
spark.udf.register(
|
|
408
|
+
"reasoning_analysis",
|
|
409
|
+
responses_udf(
|
|
410
|
+
instructions="Analyze this step by step with detailed reasoning",
|
|
411
|
+
temperature=None # Required for reasoning models
|
|
412
|
+
)
|
|
413
|
+
)
|
|
414
|
+
|
|
383
415
|
```
|
|
384
416
|
|
|
385
417
|
You can now use these UDFs in Spark SQL:
|
|
@@ -666,15 +698,15 @@ steps:
|
|
|
666
698
|
|
|
667
699
|
# Configure Azure OpenAI authentication
|
|
668
700
|
sc.environment["AZURE_OPENAI_API_KEY"] = "<your-api-key>"
|
|
669
|
-
sc.environment["
|
|
670
|
-
sc.environment["AZURE_OPENAI_API_VERSION"] = "
|
|
701
|
+
sc.environment["AZURE_OPENAI_BASE_URL"] = "https://YOUR-RESOURCE-NAME.services.ai.azure.com/openai/v1/"
|
|
702
|
+
sc.environment["AZURE_OPENAI_API_VERSION"] = "preview"
|
|
671
703
|
|
|
672
704
|
# Register UDFs
|
|
673
705
|
spark.udf.register(
|
|
674
706
|
"analyze_text",
|
|
675
707
|
responses_udf(
|
|
676
708
|
instructions="Analyze the sentiment of the text",
|
|
677
|
-
model_name="
|
|
709
|
+
model_name="gpt-4.1-mini" # Use your Azure deployment name here
|
|
678
710
|
)
|
|
679
711
|
)
|
|
680
712
|
```
|
|
@@ -1,16 +1,16 @@
|
|
|
1
1
|
openaivec/__init__.py,sha256=CuUAtLtX5RhFUbgF94UmXjurgL2VaerHTFPjMl0rRlE,236
|
|
2
2
|
openaivec/di.py,sha256=RSAiS9PqtoeobX-MZvAN8VIgaQa5EUMysj9aJwGuM9Y,10646
|
|
3
|
-
openaivec/embeddings.py,sha256=
|
|
3
|
+
openaivec/embeddings.py,sha256=psWNHnkmiw9Zb8bNAwrf8yWzKq2d6DZXdRXivT9APJ4,7171
|
|
4
4
|
openaivec/log.py,sha256=GofgzUpv_xDVuGC-gYmit5Oyu06it1SBXRck6COR5go,1439
|
|
5
|
-
openaivec/model.py,sha256=
|
|
6
|
-
openaivec/pandas_ext.py,sha256=
|
|
5
|
+
openaivec/model.py,sha256=wu1UGetqLbUGvGqmOiQna4SJnO5VvyMoCHdAQhSG6MY,3295
|
|
6
|
+
openaivec/pandas_ext.py,sha256=VvHqGRYjob2VHCFpopa_kuGhmMbzcB2o0FX13JWk0wE,51749
|
|
7
7
|
openaivec/prompt.py,sha256=p_AyCLIhlrxp7JXLoywPuDuBiOVmhuAp0XWlKcRlAGY,17715
|
|
8
|
-
openaivec/provider.py,sha256=
|
|
8
|
+
openaivec/provider.py,sha256=HdGaxSBmJtuFiYYHtPH-BjSyS2FiAmwWZC3LC9L8x7U,6192
|
|
9
9
|
openaivec/proxy.py,sha256=36d5Rb69HOhENBpS0xzx0eeJZuIfFBzKH6vEyurK-CA,24337
|
|
10
|
-
openaivec/responses.py,sha256=
|
|
10
|
+
openaivec/responses.py,sha256=HVS2Ih9igR_5HF_p2GjH2lAPl2-hMBxch_9jUomMPsc,20409
|
|
11
11
|
openaivec/serialize.py,sha256=YbYOAxhw67ELcq48wCCelqp7J4vX42liyIWdl9pjJwU,7325
|
|
12
|
-
openaivec/spark.py,sha256=
|
|
13
|
-
openaivec/util.py,sha256=
|
|
12
|
+
openaivec/spark.py,sha256=IfCehUAj99a7F0AbOJAm-iRMF7ZwYsYdzSdruGAA0NE,23117
|
|
13
|
+
openaivec/util.py,sha256=cyiATh_IGnkMU0IBPcLpFMUHjjFB6TjLAug5TdC3vS8,6350
|
|
14
14
|
openaivec/task/__init__.py,sha256=26Kd2eRRfUqRJnYH83sRrbuEuRwshypRRSkTWtT50Tw,6177
|
|
15
15
|
openaivec/task/customer_support/__init__.py,sha256=p_OPWdONia_62F5PS-w83-Rngql_4r8Xhn0p2ArzSUk,1067
|
|
16
16
|
openaivec/task/customer_support/customer_sentiment.py,sha256=2DdJHmG5OcLGC-u12F4b2EaFhy5jEjYunCHU2OUrSus,7587
|
|
@@ -28,7 +28,7 @@ openaivec/task/nlp/sentiment_analysis.py,sha256=-P56GL_krrthu2g6oQiQ7MeA9QR8-m76
|
|
|
28
28
|
openaivec/task/nlp/translation.py,sha256=bp_yr4_0_CwrX9Lw3UQMXdmlXBkc6hsAapAEZudCK00,6618
|
|
29
29
|
openaivec/task/table/__init__.py,sha256=qJ-oPcCUZzn2la_NFO1eGRUTNHgTO2jCR4L5xHp6qpM,83
|
|
30
30
|
openaivec/task/table/fillna.py,sha256=ZKYfcweiA6ODGna986pxG9RXYJZthbfp5XXOeT-WUjg,6557
|
|
31
|
-
openaivec-0.13.
|
|
32
|
-
openaivec-0.13.
|
|
33
|
-
openaivec-0.13.
|
|
34
|
-
openaivec-0.13.
|
|
31
|
+
openaivec-0.13.2.dist-info/METADATA,sha256=EJ0P0RBxo6_1ClL20UX4yL0beY7s_BrxzsHTEvakO3I,27286
|
|
32
|
+
openaivec-0.13.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
|
|
33
|
+
openaivec-0.13.2.dist-info/licenses/LICENSE,sha256=ws_MuBL-SCEBqPBFl9_FqZkaaydIJmxHrJG2parhU4M,1141
|
|
34
|
+
openaivec-0.13.2.dist-info/RECORD,,
|
|
File without changes
|
|
File without changes
|