Skip to content

Commit 77a33ed

Browse files
authored
clean: add harcoded route for whoami to remove it from publicly exposed ones (#467)
1 parent 68a4cb7 commit 77a33ed

3 files changed

Lines changed: 102 additions & 0 deletions

File tree

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
#!/usr/bin/env python
2+
3+
import asyncio
4+
import os
5+
6+
from mistralai.client import Mistral
7+
from mistralai.extra.workflows.helpers import get_scheduler_namespace
8+
9+
10+
async def main():
11+
api_key = os.environ["MISTRAL_API_KEY"]
12+
13+
client = Mistral(api_key=api_key)
14+
print(await get_scheduler_namespace(client))
15+
16+
17+
if __name__ == "__main__":
18+
asyncio.run(main())

src/mistralai/extra/workflows/__init__.py

Whitespace-only changes.
Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
from typing import Mapping, Optional
2+
3+
from pydantic import BaseModel
4+
5+
from mistralai.client import errors, models, utils
6+
from mistralai.client._hooks.types import HookContext
7+
from mistralai.client.sdk import Mistral
8+
from mistralai.client.types.basemodel import UNSET, OptionalNullable
9+
from mistralai.client.utils.security import get_security_from_env
10+
from mistralai.client.utils.unmarshal_json_response import unmarshal_json_response
11+
12+
13+
class WorkerInfo(BaseModel):
14+
scheduler_url: str
15+
namespace: str
16+
17+
18+
async def get_scheduler_namespace(
19+
client: Mistral,
20+
*,
21+
retries: OptionalNullable[utils.RetryConfig] = UNSET,
22+
server_url: Optional[str] = None,
23+
timeout_ms: Optional[int] = None,
24+
http_headers: Optional[Mapping[str, str]] = None,
25+
) -> str:
26+
base_url = None
27+
url_variables = None
28+
if timeout_ms is None:
29+
timeout_ms = client.sdk_configuration.timeout_ms
30+
31+
if server_url is not None:
32+
base_url = server_url
33+
else:
34+
base_url = client._get_url(base_url, url_variables)
35+
req = client._build_request_async(
36+
method="GET",
37+
path="/v1/workflows/workers/whoami",
38+
base_url=base_url,
39+
url_variables=url_variables,
40+
request=None,
41+
request_body_required=False,
42+
request_has_path_params=False,
43+
request_has_query_params=True,
44+
user_agent_header="user-agent",
45+
accept_header_value="application/json",
46+
http_headers=http_headers,
47+
security=client.sdk_configuration.security,
48+
allow_empty_value=None,
49+
timeout_ms=timeout_ms,
50+
)
51+
52+
if retries == UNSET:
53+
if client.sdk_configuration.retry_config is not UNSET:
54+
retries = client.sdk_configuration.retry_config
55+
56+
retry_config = None
57+
if isinstance(retries, utils.RetryConfig):
58+
retry_config = (retries, ["429", "500", "502", "503", "504"])
59+
60+
http_res = await client.do_request_async(
61+
hook_ctx=HookContext(
62+
config=client.sdk_configuration,
63+
base_url=base_url or "",
64+
operation_id="get_worker_info_v1_workflows_workers_whoami_get",
65+
oauth2_scopes=None,
66+
security_source=get_security_from_env(
67+
client.sdk_configuration.security, models.Security
68+
),
69+
),
70+
request=req,
71+
error_status_codes=["4XX", "5XX"],
72+
retry_config=retry_config,
73+
)
74+
75+
if utils.match_response(http_res, "200", "application/json"):
76+
return unmarshal_json_response(WorkerInfo, http_res).namespace
77+
if utils.match_response(http_res, "4XX", "*"):
78+
http_res_text = await utils.stream_to_text_async(http_res)
79+
raise errors.SDKError("API error occurred", http_res, http_res_text)
80+
if utils.match_response(http_res, "5XX", "*"):
81+
http_res_text = await utils.stream_to_text_async(http_res)
82+
raise errors.SDKError("API error occurred", http_res, http_res_text)
83+
84+
raise errors.SDKError("Unexpected response received", http_res)

0 commit comments

Comments
 (0)