Skip to content

Commit 2388dc9

Browse files
committed
chore: Update auth_token(s) as auth_token_getter(s) and add_auth_token(s) as add_auth_token_getter(s)
This is to align with the toolbox-core APIs and also because the suffix of _getter(s) make these APIs more descriptive and accurate. It actively encourages best practices regarding token security and lifecycle management (like refresh logic), which is crucial for production systems.
1 parent 6693407 commit 2388dc9

File tree

11 files changed

+255
-145
lines changed

11 files changed

+255
-145
lines changed

README.md

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -225,21 +225,21 @@ async def get_auth_token():
225225
toolbox = ToolboxClient("http://127.0.0.1:5000")
226226
tools = toolbox.load_toolset()
227227

228-
auth_tool = tools[0].add_auth_token("my_auth", get_auth_token) # Single token
228+
auth_tool = tools[0].add_auth_token_getter("my_auth", get_auth_token) # Single token
229229

230-
multi_auth_tool = tools[0].add_auth_tokens({"my_auth", get_auth_token}) # Multiple tokens
230+
multi_auth_tool = tools[0].add_auth_token_getters({"my_auth", get_auth_token}) # Multiple tokens
231231

232232
# OR
233233

234-
auth_tools = [tool.add_auth_token("my_auth", get_auth_token) for tool in tools]
234+
auth_tools = [tool.add_auth_token_getter("my_auth", get_auth_token) for tool in tools]
235235
```
236236

237237
#### Add Authentication While Loading
238238

239239
```py
240-
auth_tool = toolbox.load_tool(auth_tokens={"my_auth": get_auth_token})
240+
auth_tool = toolbox.load_tool(auth_token_getters={"my_auth": get_auth_token})
241241

242-
auth_tools = toolbox.load_toolset(auth_tokens={"my_auth": get_auth_token})
242+
auth_tools = toolbox.load_toolset(auth_token_getters={"my_auth": get_auth_token})
243243
```
244244

245245
> [!NOTE]
@@ -260,7 +260,7 @@ async def get_auth_token():
260260
toolbox = ToolboxClient("http://127.0.0.1:5000")
261261
tool = toolbox.load_tool("my-tool")
262262

263-
auth_tool = tool.add_auth_token("my_auth", get_auth_token)
263+
auth_tool = tool.add_auth_token_getter("my_auth", get_auth_token)
264264
result = auth_tool.invoke({"input": "some input"})
265265
print(result)
266266
```

packages/toolbox-langchain/README.md

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -225,21 +225,21 @@ async def get_auth_token():
225225
toolbox = ToolboxClient("http://127.0.0.1:5000")
226226
tools = toolbox.load_toolset()
227227

228-
auth_tool = tools[0].add_auth_token("my_auth", get_auth_token) # Single token
228+
auth_tool = tools[0].add_auth_token_getter("my_auth", get_auth_token) # Single token
229229

230-
multi_auth_tool = tools[0].add_auth_tokens({"my_auth", get_auth_token}) # Multiple tokens
230+
multi_auth_tool = tools[0].add_auth_token_getters({"my_auth", get_auth_token}) # Multiple tokens
231231

232232
# OR
233233

234-
auth_tools = [tool.add_auth_token("my_auth", get_auth_token) for tool in tools]
234+
auth_tools = [tool.add_auth_token_getter("my_auth", get_auth_token) for tool in tools]
235235
```
236236

237237
#### Add Authentication While Loading
238238

239239
```py
240-
auth_tool = toolbox.load_tool(auth_tokens={"my_auth": get_auth_token})
240+
auth_tool = toolbox.load_tool(auth_token_getters={"my_auth": get_auth_token})
241241

242-
auth_tools = toolbox.load_toolset(auth_tokens={"my_auth": get_auth_token})
242+
auth_tools = toolbox.load_toolset(auth_token_getters={"my_auth": get_auth_token})
243243
```
244244

245245
> [!NOTE]
@@ -260,7 +260,7 @@ async def get_auth_token():
260260
toolbox = ToolboxClient("http://127.0.0.1:5000")
261261
tool = toolbox.load_tool("my-tool")
262262

263-
auth_tool = tool.add_auth_token("my_auth", get_auth_token)
263+
auth_tool = tool.add_auth_token_getter("my_auth", get_auth_token)
264264
result = auth_tool.invoke({"input": "some input"})
265265
print(result)
266266
```

packages/toolbox-langchain/src/toolbox_langchain/async_client.py

Lines changed: 52 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,8 @@ def __init__(
4444
async def aload_tool(
4545
self,
4646
tool_name: str,
47-
auth_tokens: dict[str, Callable[[], str]] = {},
47+
auth_token_getters: dict[str, Callable[[], str]] = {},
48+
auth_tokens: Optional[dict[str, Callable[[], str]]] = None,
4849
auth_headers: Optional[dict[str, Callable[[], str]]] = None,
4950
bound_params: dict[str, Union[Any, Callable[[], Any]]] = {},
5051
strict: bool = True,
@@ -54,9 +55,10 @@ async def aload_tool(
5455
5556
Args:
5657
tool_name: The name of the tool to load.
57-
auth_tokens: An optional mapping of authentication source names to
58-
functions that retrieve ID tokens.
59-
auth_headers: Deprecated. Use `auth_tokens` instead.
58+
auth_token_getters: An optional mapping of authentication source
59+
names to functions that retrieve ID tokens.
60+
auth_tokens: Deprecated. Use `auth_token_getters` instead.
61+
auth_headers: Deprecated. Use `auth_token_getters` instead.
6062
bound_params: An optional mapping of parameter names to their
6163
bound values.
6264
strict: If True, raises a ValueError if any of the given bound
@@ -67,17 +69,30 @@ async def aload_tool(
6769
A tool loaded from the Toolbox.
6870
"""
6971
if auth_headers:
70-
if auth_tokens:
72+
if auth_token_getters:
7173
warn(
72-
"Both `auth_tokens` and `auth_headers` are provided. `auth_headers` is deprecated, and `auth_tokens` will be used.",
74+
"Both `auth_token_getters` and `auth_headers` are provided. `auth_headers` is deprecated, and `auth_token_getters` will be used.",
7375
DeprecationWarning,
7476
)
7577
else:
7678
warn(
77-
"Argument `auth_headers` is deprecated. Use `auth_tokens` instead.",
79+
"Argument `auth_headers` is deprecated. Use `auth_token_getters` instead.",
7880
DeprecationWarning,
7981
)
80-
auth_tokens = auth_headers
82+
auth_token_getters = auth_headers
83+
84+
if auth_tokens:
85+
if auth_token_getters:
86+
warn(
87+
"Both `auth_token_getters` and `auth_tokens` are provided. `auth_tokens` is deprecated, and `auth_token_getters` will be used.",
88+
DeprecationWarning,
89+
)
90+
else:
91+
warn(
92+
"Argument `auth_tokens` is deprecated. Use `auth_token_getters` instead.",
93+
DeprecationWarning,
94+
)
95+
auth_token_getters = auth_tokens
8196

8297
url = f"{self.__url}/api/tool/{tool_name}"
8398
manifest: ManifestSchema = await _load_manifest(url, self.__session)
@@ -87,15 +102,16 @@ async def aload_tool(
87102
manifest.tools[tool_name],
88103
self.__url,
89104
self.__session,
90-
auth_tokens,
105+
auth_token_getters,
91106
bound_params,
92107
strict,
93108
)
94109

95110
async def aload_toolset(
96111
self,
97112
toolset_name: Optional[str] = None,
98-
auth_tokens: dict[str, Callable[[], str]] = {},
113+
auth_token_getters: dict[str, Callable[[], str]] = {},
114+
auth_tokens: Optional[dict[str, Callable[[], str]]] = None,
99115
auth_headers: Optional[dict[str, Callable[[], str]]] = None,
100116
bound_params: dict[str, Union[Any, Callable[[], Any]]] = {},
101117
strict: bool = True,
@@ -107,9 +123,10 @@ async def aload_toolset(
107123
Args:
108124
toolset_name: The name of the toolset to load. If not provided,
109125
all tools are loaded.
110-
auth_tokens: An optional mapping of authentication source names to
111-
functions that retrieve ID tokens.
112-
auth_headers: Deprecated. Use `auth_tokens` instead.
126+
auth_token_getters: An optional mapping of authentication source
127+
names to functions that retrieve ID tokens.
128+
auth_tokens: Deprecated. Use `auth_token_getters` instead.
129+
auth_headers: Deprecated. Use `auth_token_getters` instead.
113130
bound_params: An optional mapping of parameter names to their
114131
bound values.
115132
strict: If True, raises a ValueError if any of the given bound
@@ -120,17 +137,30 @@ async def aload_toolset(
120137
A list of all tools loaded from the Toolbox.
121138
"""
122139
if auth_headers:
123-
if auth_tokens:
140+
if auth_token_getters:
141+
warn(
142+
"Both `auth_token_getters` and `auth_headers` are provided. `auth_headers` is deprecated, and `auth_token_getters` will be used.",
143+
DeprecationWarning,
144+
)
145+
else:
146+
warn(
147+
"Argument `auth_headers` is deprecated. Use `auth_token_getters` instead.",
148+
DeprecationWarning,
149+
)
150+
auth_token_getters = auth_headers
151+
152+
if auth_tokens:
153+
if auth_token_getters:
124154
warn(
125-
"Both `auth_tokens` and `auth_headers` are provided. `auth_headers` is deprecated, and `auth_tokens` will be used.",
155+
"Both `auth_token_getters` and `auth_tokens` are provided. `auth_tokens` is deprecated, and `auth_token_getters` will be used.",
126156
DeprecationWarning,
127157
)
128158
else:
129159
warn(
130-
"Argument `auth_headers` is deprecated. Use `auth_tokens` instead.",
160+
"Argument `auth_tokens` is deprecated. Use `auth_token_getters` instead.",
131161
DeprecationWarning,
132162
)
133-
auth_tokens = auth_headers
163+
auth_token_getters = auth_tokens
134164

135165
url = f"{self.__url}/api/toolset/{toolset_name or ''}"
136166
manifest: ManifestSchema = await _load_manifest(url, self.__session)
@@ -143,7 +173,7 @@ async def aload_toolset(
143173
tool_schema,
144174
self.__url,
145175
self.__session,
146-
auth_tokens,
176+
auth_token_getters,
147177
bound_params,
148178
strict,
149179
)
@@ -153,7 +183,8 @@ async def aload_toolset(
153183
def load_tool(
154184
self,
155185
tool_name: str,
156-
auth_tokens: dict[str, Callable[[], str]] = {},
186+
auth_token_getters: dict[str, Callable[[], str]] = {},
187+
auth_tokens: Optional[dict[str, Callable[[], str]]] = None,
157188
auth_headers: Optional[dict[str, Callable[[], str]]] = None,
158189
bound_params: dict[str, Union[Any, Callable[[], Any]]] = {},
159190
strict: bool = True,
@@ -163,7 +194,8 @@ def load_tool(
163194
def load_toolset(
164195
self,
165196
toolset_name: Optional[str] = None,
166-
auth_tokens: dict[str, Callable[[], str]] = {},
197+
auth_token_getters: dict[str, Callable[[], str]] = {},
198+
auth_tokens: Optional[dict[str, Callable[[], str]]] = None,
167199
auth_headers: Optional[dict[str, Callable[[], str]]] = None,
168200
bound_params: dict[str, Union[Any, Callable[[], Any]]] = {},
169201
strict: bool = True,

packages/toolbox-langchain/src/toolbox_langchain/async_tools.py

Lines changed: 20 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ def __init__(
4545
schema: ToolSchema,
4646
url: str,
4747
session: ClientSession,
48-
auth_tokens: dict[str, Callable[[], str]] = {},
48+
auth_token_getters: dict[str, Callable[[], str]] = {},
4949
bound_params: dict[str, Union[Any, Callable[[], Any]]] = {},
5050
strict: bool = True,
5151
) -> None:
@@ -57,8 +57,8 @@ def __init__(
5757
schema: The tool schema.
5858
url: The base URL of the Toolbox service.
5959
session: The HTTP client session.
60-
auth_tokens: A mapping of authentication source names to functions
61-
that retrieve ID tokens.
60+
auth_token_getters: A mapping of authentication source names to
61+
functions that retrieve ID tokens.
6262
bound_params: A mapping of parameter names to their bound
6363
values.
6464
strict: If True, raises a ValueError if any of the given bound
@@ -132,7 +132,7 @@ def __init__(
132132
self.__schema = schema
133133
self.__url = url
134134
self.__session = session
135-
self.__auth_tokens = auth_tokens
135+
self.__auth_token_getters = auth_token_getters
136136
self.__auth_params = auth_params
137137
self.__bound_params = bound_params
138138

@@ -172,7 +172,7 @@ async def _arun(self, **kwargs: Any) -> dict[str, Any]:
172172
kwargs.update(evaluated_params)
173173

174174
return await _invoke_tool(
175-
self.__url, self.__session, self.__name, kwargs, self.__auth_tokens
175+
self.__url, self.__session, self.__name, kwargs, self.__auth_token_getters
176176
)
177177

178178
def __validate_auth(self, strict: bool = True) -> None:
@@ -199,7 +199,7 @@ def __validate_auth(self, strict: bool = True) -> None:
199199

200200
# Check tool for at least 1 required auth source
201201
for src in self.__schema.authRequired:
202-
if src in self.__auth_tokens:
202+
if src in self.__auth_token_getters:
203203
is_authenticated = True
204204
break
205205

@@ -211,7 +211,7 @@ def __validate_auth(self, strict: bool = True) -> None:
211211
for src in param.authSources:
212212

213213
# Find first auth source that is specified
214-
if src in self.__auth_tokens:
214+
if src in self.__auth_token_getters:
215215
has_auth = True
216216
break
217217
if not has_auth:
@@ -238,7 +238,7 @@ def __validate_auth(self, strict: bool = True) -> None:
238238
def __create_copy(
239239
self,
240240
*,
241-
auth_tokens: dict[str, Callable[[], str]] = {},
241+
auth_token_getters: dict[str, Callable[[], str]] = {},
242242
bound_params: dict[str, Union[Any, Callable[[], Any]]] = {},
243243
strict: bool,
244244
) -> "AsyncToolboxTool":
@@ -253,8 +253,8 @@ def __create_copy(
253253
original instance, ensuring immutability.
254254
255255
Args:
256-
auth_tokens: A dictionary of auth source names to functions that
257-
retrieve ID tokens. These tokens will be merged with the
256+
auth_token_getters: A dictionary of auth source names to functions
257+
that retrieve ID tokens. These tokens will be merged with the
258258
existing auth tokens.
259259
bound_params: A dictionary of parameter names to their
260260
bound values or functions to retrieve the values. These params
@@ -281,21 +281,21 @@ def __create_copy(
281281
schema=new_schema,
282282
url=self.__url,
283283
session=self.__session,
284-
auth_tokens={**self.__auth_tokens, **auth_tokens},
284+
auth_token_getters={**self.__auth_token_getters, **auth_token_getters},
285285
bound_params={**self.__bound_params, **bound_params},
286286
strict=strict,
287287
)
288288

289-
def add_auth_tokens(
290-
self, auth_tokens: dict[str, Callable[[], str]], strict: bool = True
289+
def add_auth_token_getters(
290+
self, auth_token_getters: dict[str, Callable[[], str]], strict: bool = True
291291
) -> "AsyncToolboxTool":
292292
"""
293293
Registers functions to retrieve ID tokens for the corresponding
294294
authentication sources.
295295
296296
Args:
297-
auth_tokens: A dictionary of authentication source names to the
298-
functions that return corresponding ID token.
297+
auth_token_getters: A dictionary of authentication source names to
298+
the functions that return corresponding ID token getters.
299299
strict: If True, a ValueError is raised if any of the provided auth
300300
parameters is already bound. If False, only a warning is issued.
301301
@@ -313,18 +313,18 @@ def add_auth_tokens(
313313

314314
# Check if the authentication source is already registered.
315315
dupe_tokens: list[str] = []
316-
for auth_token, _ in auth_tokens.items():
317-
if auth_token in self.__auth_tokens:
316+
for auth_token, _ in auth_token_getters.items():
317+
if auth_token in self.__auth_token_getters:
318318
dupe_tokens.append(auth_token)
319319

320320
if dupe_tokens:
321321
raise ValueError(
322322
f"Authentication source(s) `{', '.join(dupe_tokens)}` already registered in tool `{self.__name}`."
323323
)
324324

325-
return self.__create_copy(auth_tokens=auth_tokens, strict=strict)
325+
return self.__create_copy(auth_token_getters=auth_token_getters, strict=strict)
326326

327-
def add_auth_token(
327+
def add_auth_token_getter(
328328
self, auth_source: str, get_id_token: Callable[[], str], strict: bool = True
329329
) -> "AsyncToolboxTool":
330330
"""
@@ -346,7 +346,7 @@ def add_auth_token(
346346
ValueError: If the provided auth parameter is already bound and
347347
strict is True.
348348
"""
349-
return self.add_auth_tokens({auth_source: get_id_token}, strict=strict)
349+
return self.add_auth_token_getters({auth_source: get_id_token}, strict=strict)
350350

351351
def bind_params(
352352
self,

0 commit comments

Comments
 (0)