beamlit 0.1.0__py3-none-any.whl

Sign up to get free protection for your applications and to get access to all the features.

Potentially problematic release.


This version of beamlit might be problematic. Click here for more details.

Files changed (211) hide show
  1. beamlit/__init__.py +8 -0
  2. beamlit/api/__init__.py +1 -0
  3. beamlit/api/agents/__init__.py +0 -0
  4. beamlit/api/agents/create_agent.py +155 -0
  5. beamlit/api/agents/create_agent_release.py +146 -0
  6. beamlit/api/agents/delete_agent.py +146 -0
  7. beamlit/api/agents/delete_agent_deployment.py +163 -0
  8. beamlit/api/agents/delete_agent_deployment_history.py +172 -0
  9. beamlit/api/agents/get_agent.py +146 -0
  10. beamlit/api/agents/get_agent_deployment.py +163 -0
  11. beamlit/api/agents/get_agent_deployment_history.py +172 -0
  12. beamlit/api/agents/get_agent_deployment_logs.py +164 -0
  13. beamlit/api/agents/get_agent_deployment_metrics.py +159 -0
  14. beamlit/api/agents/get_agent_metrics.py +150 -0
  15. beamlit/api/agents/list_agent_deployment_history.py +164 -0
  16. beamlit/api/agents/list_agent_deployments.py +155 -0
  17. beamlit/api/agents/list_agents.py +127 -0
  18. beamlit/api/agents/put_agent_deployment.py +185 -0
  19. beamlit/api/agents/put_agent_deployment_history.py +198 -0
  20. beamlit/api/agents/update_agent.py +168 -0
  21. beamlit/api/authentication_providers/__init__.py +0 -0
  22. beamlit/api/authentication_providers/get_model_with_repo_for_authentication_provider.py +184 -0
  23. beamlit/api/authentication_providers/list_models_for_authentication_provider.py +163 -0
  24. beamlit/api/authentication_providers/list_organizations_for_authentication_provider.py +163 -0
  25. beamlit/api/configurations/__init__.py +0 -0
  26. beamlit/api/configurations/get_configuration.py +122 -0
  27. beamlit/api/environments/__init__.py +0 -0
  28. beamlit/api/environments/create_environment.py +167 -0
  29. beamlit/api/environments/delete_environment.py +154 -0
  30. beamlit/api/environments/get_environment.py +154 -0
  31. beamlit/api/environments/get_environment_metrics.py +158 -0
  32. beamlit/api/environments/list_environments.py +139 -0
  33. beamlit/api/environments/update_environment.py +180 -0
  34. beamlit/api/functions/__init__.py +0 -0
  35. beamlit/api/functions/create_function.py +155 -0
  36. beamlit/api/functions/create_function_release.py +150 -0
  37. beamlit/api/functions/delete_function.py +146 -0
  38. beamlit/api/functions/delete_function_deployment.py +163 -0
  39. beamlit/api/functions/get_function.py +146 -0
  40. beamlit/api/functions/get_function_deployment.py +163 -0
  41. beamlit/api/functions/get_function_deployment_logs.py +164 -0
  42. beamlit/api/functions/get_function_deployment_metrics.py +159 -0
  43. beamlit/api/functions/get_function_metrics.py +150 -0
  44. beamlit/api/functions/list_function_deployments.py +155 -0
  45. beamlit/api/functions/list_functions.py +131 -0
  46. beamlit/api/functions/put_function_deployment.py +185 -0
  47. beamlit/api/functions/update_function.py +168 -0
  48. beamlit/api/history/__init__.py +0 -0
  49. beamlit/api/history/get_agents_history.py +155 -0
  50. beamlit/api/history/list_agents_history.py +131 -0
  51. beamlit/api/invitations/__init__.py +0 -0
  52. beamlit/api/invitations/list_all_pending_invitations.py +142 -0
  53. beamlit/api/locations/__init__.py +0 -0
  54. beamlit/api/locations/list_locations.py +139 -0
  55. beamlit/api/metrics/__init__.py +0 -0
  56. beamlit/api/metrics/get_metrics.py +130 -0
  57. beamlit/api/model_providers/__init__.py +0 -0
  58. beamlit/api/model_providers/create_model_provider.py +163 -0
  59. beamlit/api/model_providers/delete_model_provider.py +154 -0
  60. beamlit/api/model_providers/get_model_provider.py +154 -0
  61. beamlit/api/model_providers/list_model_providers.py +139 -0
  62. beamlit/api/model_providers/update_model_provider.py +176 -0
  63. beamlit/api/models/__init__.py +0 -0
  64. beamlit/api/models/create_model.py +168 -0
  65. beamlit/api/models/delete_model.py +154 -0
  66. beamlit/api/models/delete_model_deployment.py +171 -0
  67. beamlit/api/models/get_model.py +154 -0
  68. beamlit/api/models/get_model_deployment.py +171 -0
  69. beamlit/api/models/get_model_deployment_logs.py +168 -0
  70. beamlit/api/models/get_model_deployment_metrics.py +163 -0
  71. beamlit/api/models/get_model_metrics.py +154 -0
  72. beamlit/api/models/list_model_deployments.py +163 -0
  73. beamlit/api/models/list_models.py +135 -0
  74. beamlit/api/models/put_model_deployment.py +193 -0
  75. beamlit/api/models/release_model.py +154 -0
  76. beamlit/api/models/update_model.py +181 -0
  77. beamlit/api/policies/__init__.py +0 -0
  78. beamlit/api/policies/create_policy.py +167 -0
  79. beamlit/api/policies/delete_policy.py +154 -0
  80. beamlit/api/policies/get_policy.py +154 -0
  81. beamlit/api/policies/list_policies.py +139 -0
  82. beamlit/api/policies/update_policy.py +180 -0
  83. beamlit/api/service_accounts/__init__.py +0 -0
  84. beamlit/api/service_accounts/create_api_key_for_service_account.py +177 -0
  85. beamlit/api/service_accounts/create_workspace_service_account.py +168 -0
  86. beamlit/api/service_accounts/delete_api_key_for_service_account.py +104 -0
  87. beamlit/api/service_accounts/delete_workspace_service_account.py +158 -0
  88. beamlit/api/service_accounts/get_workspace_service_accounts.py +139 -0
  89. beamlit/api/service_accounts/list_api_keys_for_service_account.py +163 -0
  90. beamlit/api/service_accounts/update_workspace_service_account.py +181 -0
  91. beamlit/api/store/__init__.py +0 -0
  92. beamlit/api/store/get_store_agent.py +146 -0
  93. beamlit/api/store/get_store_function.py +146 -0
  94. beamlit/api/store/list_store_agents.py +131 -0
  95. beamlit/api/store/list_store_functions.py +131 -0
  96. beamlit/api/workspaces/__init__.py +0 -0
  97. beamlit/api/workspaces/accept_workspace_invitation.py +161 -0
  98. beamlit/api/workspaces/create_worspace.py +163 -0
  99. beamlit/api/workspaces/decline_workspace_invitation.py +158 -0
  100. beamlit/api/workspaces/delete_workspace.py +154 -0
  101. beamlit/api/workspaces/get_workspace.py +154 -0
  102. beamlit/api/workspaces/invite_workspace_user.py +174 -0
  103. beamlit/api/workspaces/leave_workspace.py +161 -0
  104. beamlit/api/workspaces/list_workspace_users.py +139 -0
  105. beamlit/api/workspaces/list_workspaces.py +139 -0
  106. beamlit/api/workspaces/remove_workspace_user.py +101 -0
  107. beamlit/api/workspaces/update_workspace.py +176 -0
  108. beamlit/api/workspaces/update_workspace_user_role.py +187 -0
  109. beamlit/authentication/__init__.py +24 -0
  110. beamlit/authentication/apikey.py +14 -0
  111. beamlit/authentication/authentication.py +35 -0
  112. beamlit/authentication/credentials.py +165 -0
  113. beamlit/authentication/device_mode.py +115 -0
  114. beamlit/client.py +270 -0
  115. beamlit/errors.py +16 -0
  116. beamlit/models/__init__.py +189 -0
  117. beamlit/models/acl.py +149 -0
  118. beamlit/models/agent.py +151 -0
  119. beamlit/models/agent_chain.py +77 -0
  120. beamlit/models/agent_configuration.py +68 -0
  121. beamlit/models/agent_deployment.py +327 -0
  122. beamlit/models/agent_deployment_configuration.py +43 -0
  123. beamlit/models/agent_deployment_history.py +183 -0
  124. beamlit/models/agent_deployment_history_event.py +131 -0
  125. beamlit/models/agent_deployment_pod_template.py +43 -0
  126. beamlit/models/agent_release.py +68 -0
  127. beamlit/models/api_key.py +140 -0
  128. beamlit/models/authentication_provider_model.py +142 -0
  129. beamlit/models/authentication_provider_organization.py +86 -0
  130. beamlit/models/configuration.py +72 -0
  131. beamlit/models/continent.py +68 -0
  132. beamlit/models/country.py +68 -0
  133. beamlit/models/create_api_key_for_service_account_body.py +67 -0
  134. beamlit/models/create_workspace_service_account_body.py +69 -0
  135. beamlit/models/create_workspace_service_account_response_200.py +103 -0
  136. beamlit/models/delete_workspace_service_account_response_200.py +94 -0
  137. beamlit/models/deployment_configuration.py +68 -0
  138. beamlit/models/deployment_configurations.py +43 -0
  139. beamlit/models/deployment_serverless_config.py +129 -0
  140. beamlit/models/environment.py +162 -0
  141. beamlit/models/environment_metrics.py +75 -0
  142. beamlit/models/flavor.py +68 -0
  143. beamlit/models/function.py +151 -0
  144. beamlit/models/function_configuration.py +68 -0
  145. beamlit/models/function_deployment.py +327 -0
  146. beamlit/models/function_deployment_configuration.py +43 -0
  147. beamlit/models/function_deployment_pod_template.py +43 -0
  148. beamlit/models/function_kit.py +95 -0
  149. beamlit/models/function_provider_ref.py +68 -0
  150. beamlit/models/function_release.py +68 -0
  151. beamlit/models/get_workspace_service_accounts_response_200_item.py +94 -0
  152. beamlit/models/invite_workspace_user_body.py +58 -0
  153. beamlit/models/labels.py +43 -0
  154. beamlit/models/labels_type_0.py +43 -0
  155. beamlit/models/location.py +120 -0
  156. beamlit/models/location_response.py +111 -0
  157. beamlit/models/metric.py +68 -0
  158. beamlit/models/metrics.py +111 -0
  159. beamlit/models/model.py +151 -0
  160. beamlit/models/model_deployment.py +283 -0
  161. beamlit/models/model_deployment_log.py +68 -0
  162. beamlit/models/model_deployment_metrics.py +170 -0
  163. beamlit/models/model_deployment_metrics_inference_per_second_per_region.py +75 -0
  164. beamlit/models/model_deployment_metrics_query_per_second_per_region_per_code.py +73 -0
  165. beamlit/models/model_deployment_pod_template.py +43 -0
  166. beamlit/models/model_metrics.py +94 -0
  167. beamlit/models/model_provider.py +187 -0
  168. beamlit/models/model_provider_ref.py +68 -0
  169. beamlit/models/model_release.py +68 -0
  170. beamlit/models/model_with_deployments.py +174 -0
  171. beamlit/models/pending_invitation.py +122 -0
  172. beamlit/models/pending_invitation_accept.py +81 -0
  173. beamlit/models/pending_invitation_render.py +135 -0
  174. beamlit/models/pending_invitation_render_invited_by.py +86 -0
  175. beamlit/models/pending_invitation_render_workspace.py +68 -0
  176. beamlit/models/pending_invitation_workspace_details.py +70 -0
  177. beamlit/models/policy.py +216 -0
  178. beamlit/models/policy_location.py +68 -0
  179. beamlit/models/provider_config.py +101 -0
  180. beamlit/models/qps.py +59 -0
  181. beamlit/models/resource_deployment_log.py +68 -0
  182. beamlit/models/resource_deployment_metrics.py +172 -0
  183. beamlit/models/resource_deployment_metrics_inference_per_second_per_region.py +75 -0
  184. beamlit/models/resource_deployment_metrics_query_per_second_per_region_per_code.py +73 -0
  185. beamlit/models/resource_metrics.py +94 -0
  186. beamlit/models/runtime.py +152 -0
  187. beamlit/models/runtime_readiness_probe.py +43 -0
  188. beamlit/models/runtime_resources.py +43 -0
  189. beamlit/models/serverless_config.py +122 -0
  190. beamlit/models/standard_fields_dynamo_db.py +86 -0
  191. beamlit/models/store_agent.py +165 -0
  192. beamlit/models/store_agent_configuration.py +95 -0
  193. beamlit/models/store_agent_labels.py +43 -0
  194. beamlit/models/store_configuration.py +149 -0
  195. beamlit/models/store_configuration_option.py +77 -0
  196. beamlit/models/store_function.py +207 -0
  197. beamlit/models/store_function_configuration.py +95 -0
  198. beamlit/models/store_function_kit.py +95 -0
  199. beamlit/models/store_function_labels.py +43 -0
  200. beamlit/models/store_function_parameter.py +86 -0
  201. beamlit/models/update_workspace_service_account_body.py +67 -0
  202. beamlit/models/update_workspace_service_account_response_200.py +94 -0
  203. beamlit/models/update_workspace_user_role_body.py +58 -0
  204. beamlit/models/workspace.py +126 -0
  205. beamlit/models/workspace_labels.py +43 -0
  206. beamlit/models/workspace_user.py +113 -0
  207. beamlit/py.typed +1 -0
  208. beamlit/types.py +46 -0
  209. beamlit-0.1.0.dist-info/METADATA +59 -0
  210. beamlit-0.1.0.dist-info/RECORD +211 -0
  211. beamlit-0.1.0.dist-info/WHEEL +4 -0
@@ -0,0 +1,187 @@
1
+ from http import HTTPStatus
2
+ from typing import Any, Optional, Union, cast
3
+
4
+ import httpx
5
+
6
+ from ... import errors
7
+ from ...client import AuthenticatedClient, Client
8
+ from ...models.update_workspace_user_role_body import UpdateWorkspaceUserRoleBody
9
+ from ...models.workspace_user import WorkspaceUser
10
+ from ...types import Response
11
+
12
+
13
+ def _get_kwargs(
14
+ sub_or_email: str,
15
+ *,
16
+ body: UpdateWorkspaceUserRoleBody,
17
+ ) -> dict[str, Any]:
18
+ headers: dict[str, Any] = {}
19
+
20
+ _kwargs: dict[str, Any] = {
21
+ "method": "put",
22
+ "url": f"/users/{sub_or_email}",
23
+ }
24
+
25
+ _body = body.to_dict()
26
+
27
+ _kwargs["json"] = _body
28
+ headers["Content-Type"] = "application/json"
29
+
30
+ _kwargs["headers"] = headers
31
+ return _kwargs
32
+
33
+
34
+ def _parse_response(
35
+ *, client: Union[AuthenticatedClient, Client], response: httpx.Response
36
+ ) -> Optional[Union[Any, WorkspaceUser]]:
37
+ if response.status_code == 200:
38
+ response_200 = WorkspaceUser.from_dict(response.json())
39
+
40
+ return response_200
41
+ if response.status_code == 400:
42
+ response_400 = cast(Any, None)
43
+ return response_400
44
+ if response.status_code == 404:
45
+ response_404 = cast(Any, None)
46
+ return response_404
47
+ if client.raise_on_unexpected_status:
48
+ raise errors.UnexpectedStatus(response.status_code, response.content)
49
+ else:
50
+ return None
51
+
52
+
53
+ def _build_response(
54
+ *, client: Union[AuthenticatedClient, Client], response: httpx.Response
55
+ ) -> Response[Union[Any, WorkspaceUser]]:
56
+ return Response(
57
+ status_code=HTTPStatus(response.status_code),
58
+ content=response.content,
59
+ headers=response.headers,
60
+ parsed=_parse_response(client=client, response=response),
61
+ )
62
+
63
+
64
+ def sync_detailed(
65
+ sub_or_email: str,
66
+ *,
67
+ client: AuthenticatedClient,
68
+ body: UpdateWorkspaceUserRoleBody,
69
+ ) -> Response[Union[Any, WorkspaceUser]]:
70
+ """Update user role in workspace
71
+
72
+ Updates the role of a user in the workspace.
73
+
74
+ Args:
75
+ sub_or_email (str):
76
+ body (UpdateWorkspaceUserRoleBody):
77
+
78
+ Raises:
79
+ errors.UnexpectedStatus: If the server returns an undocumented status code and Client.raise_on_unexpected_status is True.
80
+ httpx.TimeoutException: If the request takes longer than Client.timeout.
81
+
82
+ Returns:
83
+ Response[Union[Any, WorkspaceUser]]
84
+ """
85
+
86
+ kwargs = _get_kwargs(
87
+ sub_or_email=sub_or_email,
88
+ body=body,
89
+ )
90
+
91
+ response = client.get_httpx_client().request(
92
+ **kwargs,
93
+ )
94
+
95
+ return _build_response(client=client, response=response)
96
+
97
+
98
+ def sync(
99
+ sub_or_email: str,
100
+ *,
101
+ client: AuthenticatedClient,
102
+ body: UpdateWorkspaceUserRoleBody,
103
+ ) -> Optional[Union[Any, WorkspaceUser]]:
104
+ """Update user role in workspace
105
+
106
+ Updates the role of a user in the workspace.
107
+
108
+ Args:
109
+ sub_or_email (str):
110
+ body (UpdateWorkspaceUserRoleBody):
111
+
112
+ Raises:
113
+ errors.UnexpectedStatus: If the server returns an undocumented status code and Client.raise_on_unexpected_status is True.
114
+ httpx.TimeoutException: If the request takes longer than Client.timeout.
115
+
116
+ Returns:
117
+ Union[Any, WorkspaceUser]
118
+ """
119
+
120
+ return sync_detailed(
121
+ sub_or_email=sub_or_email,
122
+ client=client,
123
+ body=body,
124
+ ).parsed
125
+
126
+
127
+ async def asyncio_detailed(
128
+ sub_or_email: str,
129
+ *,
130
+ client: AuthenticatedClient,
131
+ body: UpdateWorkspaceUserRoleBody,
132
+ ) -> Response[Union[Any, WorkspaceUser]]:
133
+ """Update user role in workspace
134
+
135
+ Updates the role of a user in the workspace.
136
+
137
+ Args:
138
+ sub_or_email (str):
139
+ body (UpdateWorkspaceUserRoleBody):
140
+
141
+ Raises:
142
+ errors.UnexpectedStatus: If the server returns an undocumented status code and Client.raise_on_unexpected_status is True.
143
+ httpx.TimeoutException: If the request takes longer than Client.timeout.
144
+
145
+ Returns:
146
+ Response[Union[Any, WorkspaceUser]]
147
+ """
148
+
149
+ kwargs = _get_kwargs(
150
+ sub_or_email=sub_or_email,
151
+ body=body,
152
+ )
153
+
154
+ response = await client.get_async_httpx_client().request(**kwargs)
155
+
156
+ return _build_response(client=client, response=response)
157
+
158
+
159
+ async def asyncio(
160
+ sub_or_email: str,
161
+ *,
162
+ client: AuthenticatedClient,
163
+ body: UpdateWorkspaceUserRoleBody,
164
+ ) -> Optional[Union[Any, WorkspaceUser]]:
165
+ """Update user role in workspace
166
+
167
+ Updates the role of a user in the workspace.
168
+
169
+ Args:
170
+ sub_or_email (str):
171
+ body (UpdateWorkspaceUserRoleBody):
172
+
173
+ Raises:
174
+ errors.UnexpectedStatus: If the server returns an undocumented status code and Client.raise_on_unexpected_status is True.
175
+ httpx.TimeoutException: If the request takes longer than Client.timeout.
176
+
177
+ Returns:
178
+ Union[Any, WorkspaceUser]
179
+ """
180
+
181
+ return (
182
+ await asyncio_detailed(
183
+ sub_or_email=sub_or_email,
184
+ client=client,
185
+ body=body,
186
+ )
187
+ ).parsed
@@ -0,0 +1,24 @@
1
+ from .apikey import ApiKeyProvider
2
+ from .authentication import (PublicProvider, RunClientWithCredentials,
3
+ new_client_with_credentials)
4
+ from .credentials import (Config, ContextConfig, Credentials, WorkspaceConfig,
5
+ load_credentials)
6
+ from .device_mode import (BearerToken, DeviceLogin, DeviceLoginFinalizeRequest,
7
+ DeviceLoginFinalizeResponse, DeviceLoginResponse)
8
+
9
+ __all__ = (
10
+ "ApiKeyProvider",
11
+ "PublicProvider",
12
+ "RunClientWithCredentials",
13
+ "new_client_with_credentials",
14
+ "Config",
15
+ "ContextConfig",
16
+ "Credentials",
17
+ "WorkspaceConfig",
18
+ "load_credentials",
19
+ "BearerToken",
20
+ "DeviceLogin",
21
+ "DeviceLoginFinalizeRequest",
22
+ "DeviceLoginFinalizeResponse",
23
+ "DeviceLoginResponse"
24
+ )
@@ -0,0 +1,14 @@
1
+ from typing import Generator
2
+
3
+ from httpx import Auth, Request, Response
4
+
5
+
6
+ class ApiKeyProvider(Auth):
7
+ def __init__(self, credentials, workspace_name: str):
8
+ self.credentials = credentials
9
+ self.workspace_name = workspace_name
10
+
11
+ def auth_flow(self, request: Request) -> Generator[Request, Response, None]:
12
+ request.headers['X-Beamlit-Api-Key'] = self.credentials.api_key
13
+ request.headers['X-Beamlit-Workspace'] = self.workspace_name
14
+ yield request
@@ -0,0 +1,35 @@
1
+ from dataclasses import dataclass
2
+ from typing import Generator
3
+
4
+ from httpx import Auth, Request, Response
5
+
6
+ from ..client import AuthenticatedClient
7
+ from .apikey import ApiKeyProvider
8
+ from .credentials import Credentials
9
+ from .device_mode import BearerToken
10
+
11
+
12
+ class PublicProvider(Auth):
13
+ def auth_flow(self, request: Request) -> Generator[Request, Response, None]:
14
+ yield request
15
+
16
+
17
+
18
+ @dataclass
19
+ class RunClientWithCredentials:
20
+ credentials: Credentials
21
+ workspace: str
22
+ api_url: str = "https://api.beamlit.dev/v0"
23
+ run_url: str = "https://run.beamlit.dev/v0"
24
+
25
+
26
+ def new_client_with_credentials(config: RunClientWithCredentials):
27
+ provider: Auth = None
28
+ if config.credentials.api_key:
29
+ provider = ApiKeyProvider(config.credentials, config.workspace)
30
+ elif config.credentials.access_token:
31
+ provider = BearerToken(config.credentials, config.workspace, config.api_url)
32
+ else:
33
+ provider = PublicProvider()
34
+
35
+ return AuthenticatedClient(base_url=config.api_url, provider=provider)
@@ -0,0 +1,165 @@
1
+ from dataclasses import dataclass
2
+ from pathlib import Path
3
+ from typing import List
4
+
5
+ import yaml
6
+
7
+
8
+ @dataclass
9
+ class Credentials:
10
+ api_key: str = ""
11
+ access_token: str = ""
12
+ refresh_token: str = ""
13
+ expires_in: int = 0
14
+ device_code: str = ""
15
+
16
+ @dataclass
17
+ class WorkspaceConfig:
18
+ name: str
19
+ credentials: Credentials
20
+
21
+
22
+ @dataclass
23
+ class ContextConfig:
24
+ workspace: str = ""
25
+ environment: str = ""
26
+
27
+
28
+ @dataclass
29
+ class Config:
30
+ workspaces: List[WorkspaceConfig] = None
31
+ context: ContextConfig = None
32
+
33
+ def __post_init__(self):
34
+ if self.workspaces is None:
35
+ self.workspaces = []
36
+ if self.context is None:
37
+ self.context = ContextConfig()
38
+
39
+
40
+ def load_config() -> Config:
41
+ config = Config()
42
+ home_dir = Path.home()
43
+ if home_dir:
44
+ config_path = home_dir / ".beamlit" / "config.yaml"
45
+ if config_path.exists():
46
+ try:
47
+ with open(config_path) as f:
48
+ data = yaml.safe_load(f)
49
+ if data:
50
+ workspaces = []
51
+ for ws in data.get("workspaces", []):
52
+ creds = Credentials(**ws.get("credentials", {}))
53
+ workspaces.append(WorkspaceConfig(name=ws["name"], credentials=creds))
54
+ config.workspaces = workspaces
55
+ if "context" in data:
56
+ config.context = ContextConfig(**data["context"])
57
+ except yaml.YAMLError:
58
+ # Invalid YAML, use empty config
59
+ pass
60
+ return config
61
+
62
+
63
+ def save_config(config: Config):
64
+ data = {
65
+ "workspaces": [
66
+ {
67
+ "name": ws.name,
68
+ "credentials": {
69
+ "access_token": ws.credentials.access_token,
70
+ "api_key": ws.credentials.api_key
71
+ }
72
+ }
73
+ for ws in config.workspaces
74
+ ],
75
+ "context": {
76
+ "workspace": config.context.workspace,
77
+ "environment": config.context.environment
78
+ }
79
+ }
80
+
81
+ home_dir = Path.home()
82
+ if not home_dir:
83
+ raise RuntimeError("Could not determine home directory")
84
+
85
+ config_dir = home_dir / ".beamlit"
86
+ config_file = config_dir / "config.yaml"
87
+
88
+ config_dir.mkdir(mode=0o700, parents=True, exist_ok=True)
89
+
90
+ with open(config_file, "w", encoding="utf-8") as f:
91
+ yaml.dump(data, f)
92
+
93
+
94
+ def list_workspaces() -> List[str]:
95
+ config = load_config()
96
+ return [workspace.name for workspace in config.workspaces]
97
+
98
+
99
+ def current_context() -> ContextConfig:
100
+ config = load_config()
101
+ return config.context
102
+
103
+
104
+ def set_current_workspace(workspace_name: str, environment: str):
105
+ config = load_config()
106
+ config.context.workspace = workspace_name
107
+ config.context.environment = environment
108
+ save_config(config)
109
+
110
+
111
+ def load_credentials(workspace_name: str) -> Credentials:
112
+ config = load_config()
113
+ for workspace in config.workspaces:
114
+ if workspace.name == workspace_name:
115
+ return workspace.credentials
116
+ return Credentials()
117
+
118
+
119
+ def create_home_dir_if_missing():
120
+ home_dir = Path.home()
121
+ if not home_dir:
122
+ print("Error getting home directory")
123
+ return
124
+
125
+ credentials_dir = home_dir / ".beamlit"
126
+ credentials_file = credentials_dir / "credentials.json"
127
+
128
+ if credentials_file.exists():
129
+ print("You are already logged in. Enter a new API key to overwrite it.")
130
+ else:
131
+ try:
132
+ credentials_dir.mkdir(mode=0o700, parents=True, exist_ok=True)
133
+ except Exception as e:
134
+ print(f"Error creating credentials directory: {e}")
135
+
136
+
137
+ def save_credentials(workspace_name: str, credentials: Credentials):
138
+ create_home_dir_if_missing()
139
+ if not credentials.access_token and not credentials.api_key:
140
+ print("No credentials to save, error")
141
+ return
142
+
143
+ config = load_config()
144
+ found = False
145
+
146
+ for i, workspace in enumerate(config.workspaces):
147
+ if workspace.name == workspace_name:
148
+ config.workspaces[i].credentials = credentials
149
+ found = True
150
+ break
151
+
152
+ if not found:
153
+ config.workspaces.append(WorkspaceConfig(name=workspace_name, credentials=credentials))
154
+
155
+ save_config(config)
156
+
157
+
158
+ def clear_credentials(workspace_name: str):
159
+ config = load_config()
160
+ config.workspaces = [ws for ws in config.workspaces if ws.name != workspace_name]
161
+
162
+ if config.context.workspace == workspace_name:
163
+ config.context.workspace = ""
164
+
165
+ save_config(config)
@@ -0,0 +1,115 @@
1
+ import base64
2
+ import json
3
+ import time
4
+ from dataclasses import dataclass
5
+ from typing import Generator, Optional
6
+
7
+ from httpx import Auth, Request, Response, post
8
+
9
+
10
+ @dataclass
11
+ class DeviceLogin:
12
+ client_id: str
13
+ scope: str
14
+
15
+
16
+ @dataclass
17
+ class DeviceLoginResponse:
18
+ client_id: str
19
+ device_code: str
20
+ user_code: str
21
+ expires_in: int
22
+ interval: int
23
+ verification_uri: str
24
+ verification_uri_complete: str
25
+
26
+
27
+ @dataclass
28
+ class DeviceLoginFinalizeRequest:
29
+ grant_type: str
30
+ client_id: str
31
+ device_code: str
32
+
33
+
34
+ @dataclass
35
+ class DeviceLoginFinalizeResponse:
36
+ access_token: str
37
+ expires_in: int
38
+ refresh_token: str
39
+ token_type: str
40
+
41
+
42
+ class BearerToken(Auth):
43
+ def __init__(self, credentials, workspace_name: str, base_url: str):
44
+ self.credentials = credentials
45
+ self.workspace_name = workspace_name
46
+ self.base_url = base_url
47
+
48
+ def refresh_if_needed(self) -> Optional[Exception]:
49
+ # Need to refresh token if expires in less than 10 minutes
50
+ parts = self.credentials.access_token.split('.')
51
+ if len(parts) != 3:
52
+ return Exception("Invalid JWT token format")
53
+
54
+ try:
55
+ claims_bytes = base64.urlsafe_b64decode(parts[1] + '=' * (-len(parts[1]) % 4))
56
+ claims = json.loads(claims_bytes)
57
+ except Exception as e:
58
+ return Exception(f"Failed to decode/parse JWT claims: {str(e)}")
59
+
60
+ exp_time = time.gmtime(claims['exp'])
61
+ # Refresh if token expires in less than 10 minutes
62
+ if time.time() + (10 * 60) > time.mktime(exp_time):
63
+ return self.do_refresh()
64
+
65
+ return None
66
+
67
+ def auth_flow(self, request: Request) -> Generator[Request, Response, None]:
68
+ err = self.refresh_if_needed()
69
+ if err:
70
+ return err
71
+
72
+ request.headers['X-Beamlit-Authorization'] = f'Bearer {self.credentials.access_token}'
73
+ request.headers['X-Beamlit-Workspace'] = self.workspace_name
74
+ yield request
75
+
76
+ def do_refresh(self) -> Optional[Exception]:
77
+ if not self.credentials.refresh_token:
78
+ return Exception("No refresh token to refresh")
79
+
80
+ url = f"{self.base_url}/oauth/token"
81
+ refresh_data = {
82
+ "grant_type": "refresh_token",
83
+ "refresh_token": self.credentials.refresh_token,
84
+ "device_code": self.credentials.device_code,
85
+ "client_id": "beamlit"
86
+ }
87
+
88
+ try:
89
+ print(refresh_data)
90
+ response = post(
91
+ url,
92
+ json=refresh_data,
93
+ headers={"Content-Type": "application/json"}
94
+ )
95
+ print(response.text)
96
+ response.raise_for_status()
97
+ finalize_response = DeviceLoginFinalizeResponse(**response.json())
98
+
99
+ if not finalize_response.refresh_token:
100
+ finalize_response.refresh_token = self.credentials.refresh_token
101
+
102
+ from .credentials import Credentials, save_credentials
103
+ creds = Credentials(
104
+ access_token=finalize_response.access_token,
105
+ refresh_token=finalize_response.refresh_token,
106
+ expires_in=finalize_response.expires_in,
107
+ device_code=self.credentials.device_code
108
+ )
109
+
110
+ self.credentials = creds
111
+ save_credentials(self.workspace_name, creds)
112
+ return None
113
+
114
+ except Exception as e:
115
+ return Exception(f"Failed to refresh token: {str(e)}")