tau-coding-agent 0.1.0__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.
Files changed (283) hide show
  1. tau/__init__.py +0 -0
  2. tau/agent/__init__.py +11 -0
  3. tau/agent/prompt/__init__.py +10 -0
  4. tau/agent/prompt/builder.py +302 -0
  5. tau/agent/prompt/types.py +33 -0
  6. tau/agent/service.py +369 -0
  7. tau/agent/types.py +61 -0
  8. tau/auth/manager.py +247 -0
  9. tau/auth/storage.py +82 -0
  10. tau/auth/types.py +41 -0
  11. tau/builtins/__init__.py +4 -0
  12. tau/builtins/__pycache__/__init__.cpython-313.pyc +0 -0
  13. tau/builtins/__pycache__/__init__.cpython-314.pyc +0 -0
  14. tau/builtins/commands/__init__.py +41 -0
  15. tau/builtins/commands/__pycache__/__init__.cpython-313.pyc +0 -0
  16. tau/builtins/commands/__pycache__/__init__.cpython-314.pyc +0 -0
  17. tau/builtins/commands/__pycache__/clear.cpython-313.pyc +0 -0
  18. tau/builtins/commands/__pycache__/clear.cpython-314.pyc +0 -0
  19. tau/builtins/commands/__pycache__/compact.cpython-313.pyc +0 -0
  20. tau/builtins/commands/__pycache__/compact.cpython-314.pyc +0 -0
  21. tau/builtins/commands/__pycache__/reload.cpython-313.pyc +0 -0
  22. tau/builtins/commands/__pycache__/reload.cpython-314.pyc +0 -0
  23. tau/builtins/commands/__pycache__/session.cpython-313.pyc +0 -0
  24. tau/builtins/commands/__pycache__/session.cpython-314.pyc +0 -0
  25. tau/builtins/commands/clear.py +16 -0
  26. tau/builtins/commands/compact.py +28 -0
  27. tau/builtins/commands/reload.py +27 -0
  28. tau/builtins/commands/session.py +19 -0
  29. tau/builtins/extensions/footer/__init__.py +76 -0
  30. tau/builtins/extensions/footer/__pycache__/__init__.cpython-313.pyc +0 -0
  31. tau/builtins/extensions/footer/__pycache__/git.cpython-313.pyc +0 -0
  32. tau/builtins/extensions/footer/__pycache__/model.cpython-313.pyc +0 -0
  33. tau/builtins/extensions/footer/__pycache__/utils.cpython-313.pyc +0 -0
  34. tau/builtins/extensions/footer/git.py +26 -0
  35. tau/builtins/extensions/footer/model.py +69 -0
  36. tau/builtins/extensions/footer/utils.py +44 -0
  37. tau/builtins/extensions/header/__init__.py +18 -0
  38. tau/builtins/extensions/header/__pycache__/__init__.cpython-313.pyc +0 -0
  39. tau/builtins/models/__init__.py +0 -0
  40. tau/builtins/models/__pycache__/__init__.cpython-313.pyc +0 -0
  41. tau/builtins/models/__pycache__/text.cpython-313.pyc +0 -0
  42. tau/builtins/models/audio.py +43 -0
  43. tau/builtins/models/image.py +43 -0
  44. tau/builtins/models/text.py +482 -0
  45. tau/builtins/models/video.py +40 -0
  46. tau/builtins/prompts/commit.md +7 -0
  47. tau/builtins/prompts/docs.md +7 -0
  48. tau/builtins/prompts/explain.md +7 -0
  49. tau/builtins/prompts/fix.md +7 -0
  50. tau/builtins/prompts/refactor.md +7 -0
  51. tau/builtins/prompts/review.md +7 -0
  52. tau/builtins/prompts/test.md +7 -0
  53. tau/builtins/providers/__init__.py +0 -0
  54. tau/builtins/providers/__pycache__/__init__.cpython-313.pyc +0 -0
  55. tau/builtins/providers/__pycache__/text.cpython-313.pyc +0 -0
  56. tau/builtins/providers/audio.py +10 -0
  57. tau/builtins/providers/image.py +9 -0
  58. tau/builtins/providers/text.py +33 -0
  59. tau/builtins/providers/video.py +6 -0
  60. tau/builtins/skills/code-review/SKILL.md +4 -0
  61. tau/builtins/skills/debug/SKILL.md +4 -0
  62. tau/builtins/skills/git-commit/SKILL.md +4 -0
  63. tau/builtins/themes/dark.yaml +1 -0
  64. tau/builtins/themes/light.yaml +46 -0
  65. tau/builtins/tools/__init__.py +73 -0
  66. tau/builtins/tools/__pycache__/__init__.cpython-313.pyc +0 -0
  67. tau/builtins/tools/__pycache__/__init__.cpython-314.pyc +0 -0
  68. tau/builtins/tools/__pycache__/bash.cpython-313.pyc +0 -0
  69. tau/builtins/tools/__pycache__/bash.cpython-314.pyc +0 -0
  70. tau/builtins/tools/__pycache__/edit.cpython-313.pyc +0 -0
  71. tau/builtins/tools/__pycache__/edit.cpython-314.pyc +0 -0
  72. tau/builtins/tools/__pycache__/glob.cpython-313.pyc +0 -0
  73. tau/builtins/tools/__pycache__/glob.cpython-314.pyc +0 -0
  74. tau/builtins/tools/__pycache__/grep.cpython-313.pyc +0 -0
  75. tau/builtins/tools/__pycache__/grep.cpython-314.pyc +0 -0
  76. tau/builtins/tools/__pycache__/ls.cpython-313.pyc +0 -0
  77. tau/builtins/tools/__pycache__/ls.cpython-314.pyc +0 -0
  78. tau/builtins/tools/__pycache__/read.cpython-313.pyc +0 -0
  79. tau/builtins/tools/__pycache__/read.cpython-314.pyc +0 -0
  80. tau/builtins/tools/__pycache__/terminal.cpython-313.pyc +0 -0
  81. tau/builtins/tools/__pycache__/terminal.cpython-314.pyc +0 -0
  82. tau/builtins/tools/__pycache__/write.cpython-313.pyc +0 -0
  83. tau/builtins/tools/__pycache__/write.cpython-314.pyc +0 -0
  84. tau/builtins/tools/edit.py +215 -0
  85. tau/builtins/tools/glob.py +112 -0
  86. tau/builtins/tools/grep.py +146 -0
  87. tau/builtins/tools/ls.py +135 -0
  88. tau/builtins/tools/read.py +122 -0
  89. tau/builtins/tools/terminal.py +150 -0
  90. tau/builtins/tools/write.py +105 -0
  91. tau/commands/__init__.py +10 -0
  92. tau/commands/registry.py +71 -0
  93. tau/commands/types.py +33 -0
  94. tau/console/__init__.py +0 -0
  95. tau/console/cli.py +266 -0
  96. tau/console/commands/__init__.py +0 -0
  97. tau/console/commands/auth.py +193 -0
  98. tau/console/commands/packages.py +104 -0
  99. tau/console/commands/update.py +76 -0
  100. tau/core/__init__.py +0 -0
  101. tau/core/registry.py +102 -0
  102. tau/engine/__init__.py +47 -0
  103. tau/engine/service.py +768 -0
  104. tau/engine/types.py +163 -0
  105. tau/extensions/__init__.py +28 -0
  106. tau/extensions/api.py +928 -0
  107. tau/extensions/context.py +462 -0
  108. tau/extensions/events.py +70 -0
  109. tau/extensions/loader.py +386 -0
  110. tau/extensions/runtime.py +184 -0
  111. tau/extensions/settings.py +137 -0
  112. tau/hooks/__init__.py +112 -0
  113. tau/hooks/engine.py +237 -0
  114. tau/hooks/inference.py +21 -0
  115. tau/hooks/runtime.py +126 -0
  116. tau/hooks/service.py +121 -0
  117. tau/hooks/session.py +117 -0
  118. tau/hooks/tui.py +61 -0
  119. tau/hooks/types.py +72 -0
  120. tau/inference/__init__.py +180 -0
  121. tau/inference/api/__init__.py +0 -0
  122. tau/inference/api/audio/__init__.py +0 -0
  123. tau/inference/api/audio/base.py +29 -0
  124. tau/inference/api/audio/builtins.py +15 -0
  125. tau/inference/api/audio/elevenlabs_audio.py +183 -0
  126. tau/inference/api/audio/gemini_audio.py +95 -0
  127. tau/inference/api/audio/openai_audio.py +159 -0
  128. tau/inference/api/audio/registry.py +15 -0
  129. tau/inference/api/audio/sarvam_audio.py +163 -0
  130. tau/inference/api/audio/service.py +103 -0
  131. tau/inference/api/audio/utils.py +47 -0
  132. tau/inference/api/image/__init__.py +0 -0
  133. tau/inference/api/image/base.py +17 -0
  134. tau/inference/api/image/builtins.py +8 -0
  135. tau/inference/api/image/gemini_image.py +77 -0
  136. tau/inference/api/image/openai_image.py +103 -0
  137. tau/inference/api/image/openrouter.py +144 -0
  138. tau/inference/api/image/registry.py +15 -0
  139. tau/inference/api/image/service.py +71 -0
  140. tau/inference/api/registry.py +82 -0
  141. tau/inference/api/text/__init__.py +0 -0
  142. tau/inference/api/text/anthropic_claude_code.py +222 -0
  143. tau/inference/api/text/anthropic_messages.py +196 -0
  144. tau/inference/api/text/base.py +40 -0
  145. tau/inference/api/text/builtins.py +19 -0
  146. tau/inference/api/text/gemini_generate.py +234 -0
  147. tau/inference/api/text/github_copilot_chat.py +172 -0
  148. tau/inference/api/text/google_antigravity.py +522 -0
  149. tau/inference/api/text/mistral_chat.py +284 -0
  150. tau/inference/api/text/ollama_chat.py +200 -0
  151. tau/inference/api/text/openai_codex_responses.py +497 -0
  152. tau/inference/api/text/openai_completions.py +227 -0
  153. tau/inference/api/text/openai_responses.py +235 -0
  154. tau/inference/api/text/registry.py +50 -0
  155. tau/inference/api/text/service.py +297 -0
  156. tau/inference/api/text/types.py +7 -0
  157. tau/inference/api/text/utils.py +228 -0
  158. tau/inference/api/video/__init__.py +0 -0
  159. tau/inference/api/video/base.py +26 -0
  160. tau/inference/api/video/builtins.py +7 -0
  161. tau/inference/api/video/fal_video.py +119 -0
  162. tau/inference/api/video/openrouter_video.py +142 -0
  163. tau/inference/api/video/registry.py +15 -0
  164. tau/inference/api/video/service.py +72 -0
  165. tau/inference/model/__init__.py +0 -0
  166. tau/inference/model/registry.py +102 -0
  167. tau/inference/model/types.py +65 -0
  168. tau/inference/provider/__init__.py +0 -0
  169. tau/inference/provider/oauth/__init__.py +35 -0
  170. tau/inference/provider/oauth/anthropic_claude_code.py +286 -0
  171. tau/inference/provider/oauth/github_copilot.py +333 -0
  172. tau/inference/provider/oauth/google_antigravity.py +258 -0
  173. tau/inference/provider/oauth/openai_codex.py +309 -0
  174. tau/inference/provider/oauth/pkce.py +14 -0
  175. tau/inference/provider/oauth/types.py +46 -0
  176. tau/inference/provider/oauth/utils.py +154 -0
  177. tau/inference/provider/registry.py +141 -0
  178. tau/inference/provider/types.py +114 -0
  179. tau/inference/types.py +549 -0
  180. tau/inference/utils.py +219 -0
  181. tau/message/__init__.py +0 -0
  182. tau/message/types.py +482 -0
  183. tau/message/utils.py +178 -0
  184. tau/packages/__init__.py +11 -0
  185. tau/packages/manager.py +190 -0
  186. tau/packages/types.py +20 -0
  187. tau/packages/utils.py +67 -0
  188. tau/prompts/expand.py +58 -0
  189. tau/prompts/loader.py +69 -0
  190. tau/prompts/registry.py +45 -0
  191. tau/prompts/types.py +24 -0
  192. tau/rpc/__init__.py +8 -0
  193. tau/rpc/mode.py +783 -0
  194. tau/rpc/types.py +252 -0
  195. tau/runtime/service.py +759 -0
  196. tau/runtime/types.py +303 -0
  197. tau/session/branch_summarization.py +312 -0
  198. tau/session/compaction.py +646 -0
  199. tau/session/manager.py +652 -0
  200. tau/session/types.py +188 -0
  201. tau/session/utils.py +233 -0
  202. tau/settings/manager.py +1077 -0
  203. tau/settings/paths.py +150 -0
  204. tau/settings/storage.py +63 -0
  205. tau/settings/types.py +173 -0
  206. tau/settings/utils.py +25 -0
  207. tau/skills/loader.py +91 -0
  208. tau/skills/registry.py +70 -0
  209. tau/skills/types.py +25 -0
  210. tau/themes/loader.py +238 -0
  211. tau/themes/registry.py +108 -0
  212. tau/themes/types.py +19 -0
  213. tau/tool/__init__.py +3 -0
  214. tau/tool/registry.py +117 -0
  215. tau/tool/render.py +21 -0
  216. tau/tool/types.py +244 -0
  217. tau/trust/__init__.py +13 -0
  218. tau/trust/manager.py +80 -0
  219. tau/trust/types.py +14 -0
  220. tau/trust/utils.py +72 -0
  221. tau/tui/__init__.py +54 -0
  222. tau/tui/agent_hooks.py +346 -0
  223. tau/tui/ansi.py +330 -0
  224. tau/tui/app.py +540 -0
  225. tau/tui/autocomplete.py +33 -0
  226. tau/tui/capabilities.py +119 -0
  227. tau/tui/commands/__init__.py +3 -0
  228. tau/tui/commands/appearance.py +498 -0
  229. tau/tui/commands/auth.py +232 -0
  230. tau/tui/commands/context.py +38 -0
  231. tau/tui/commands/misc.py +82 -0
  232. tau/tui/commands/model.py +118 -0
  233. tau/tui/commands/session.py +464 -0
  234. tau/tui/component.py +268 -0
  235. tau/tui/components/__init__.py +0 -0
  236. tau/tui/components/autocomplete_manager.py +267 -0
  237. tau/tui/components/autocomplete_picker.py +143 -0
  238. tau/tui/components/box.py +90 -0
  239. tau/tui/components/command_palette.py +144 -0
  240. tau/tui/components/dynamic_border.py +19 -0
  241. tau/tui/components/file_picker.py +233 -0
  242. tau/tui/components/image.py +181 -0
  243. tau/tui/components/inline_selector.py +71 -0
  244. tau/tui/components/layout.py +1194 -0
  245. tau/tui/components/message_list.py +692 -0
  246. tau/tui/components/modal.py +97 -0
  247. tau/tui/components/model_palette.py +204 -0
  248. tau/tui/components/picker_overlay.py +174 -0
  249. tau/tui/components/prompt_overlay.py +236 -0
  250. tau/tui/components/resume_modal.py +372 -0
  251. tau/tui/components/select_list.py +222 -0
  252. tau/tui/components/settings_modal.py +274 -0
  253. tau/tui/components/settings_schema.py +203 -0
  254. tau/tui/components/spinner.py +119 -0
  255. tau/tui/components/text_input.py +396 -0
  256. tau/tui/components/text_prompt.py +82 -0
  257. tau/tui/components/tree_select_list.py +580 -0
  258. tau/tui/components/trust_screen.py +97 -0
  259. tau/tui/diff.py +114 -0
  260. tau/tui/fuzzy.py +99 -0
  261. tau/tui/input.py +496 -0
  262. tau/tui/input_handler.py +716 -0
  263. tau/tui/keybindings.py +87 -0
  264. tau/tui/markdown.py +286 -0
  265. tau/tui/message_renderers.py +31 -0
  266. tau/tui/overlay.py +326 -0
  267. tau/tui/renderer.py +378 -0
  268. tau/tui/terminal.py +499 -0
  269. tau/tui/theme.py +148 -0
  270. tau/tui/tui.py +544 -0
  271. tau/tui/ui_context.py +768 -0
  272. tau/tui/utils.py +20 -0
  273. tau/utils/__init__.py +0 -0
  274. tau/utils/http_proxy.py +221 -0
  275. tau/utils/image_processing.py +172 -0
  276. tau/utils/secrets.py +59 -0
  277. tau/utils/version_check.py +60 -0
  278. tau_coding_agent-0.1.0.dist-info/METADATA +177 -0
  279. tau_coding_agent-0.1.0.dist-info/RECORD +283 -0
  280. tau_coding_agent-0.1.0.dist-info/WHEEL +5 -0
  281. tau_coding_agent-0.1.0.dist-info/entry_points.txt +2 -0
  282. tau_coding_agent-0.1.0.dist-info/licenses/LICENSE +21 -0
  283. tau_coding_agent-0.1.0.dist-info/top_level.txt +1 -0
@@ -0,0 +1,154 @@
1
+ """Shared utilities for PKCE + local-callback OAuth flows."""
2
+ from __future__ import annotations
3
+
4
+ import asyncio
5
+ import urllib.parse
6
+ from typing import Optional
7
+
8
+ from tau.inference.provider.oauth.types import OAuthLoginCallbacks
9
+
10
+ __all__ = [
11
+ "OAUTH_SUCCESS_HTML",
12
+ "OAUTH_ERROR_HTML",
13
+ "parse_authorization_input",
14
+ "start_oauth_callback_server",
15
+ "await_oauth_code",
16
+ ]
17
+
18
+ OAUTH_SUCCESS_HTML = b"""<!DOCTYPE html><html><head><title>Auth complete</title></head><body>
19
+ <h2>Authentication successful!</h2>
20
+ <p>You can close this window and return to the application.</p>
21
+ </body></html>"""
22
+
23
+ OAUTH_ERROR_HTML = b"""<!DOCTYPE html><html><head><title>Auth failed</title></head><body>
24
+ <h2>Authentication failed</h2>
25
+ <p>An error occurred. Please try again.</p>
26
+ </body></html>"""
27
+
28
+
29
+ def parse_authorization_input(value: str) -> tuple[Optional[str], Optional[str]]:
30
+ """Parse (code, state) from a redirect URL, raw query string, or bare code."""
31
+ value = value.strip()
32
+ if not value:
33
+ return None, None
34
+ try:
35
+ parsed = urllib.parse.urlparse(value)
36
+ if parsed.scheme in ("http", "https"):
37
+ params = urllib.parse.parse_qs(parsed.query)
38
+ return params.get("code", [None])[0], params.get("state", [None])[0]
39
+ except Exception:
40
+ pass
41
+ if "#" in value:
42
+ code, state = value.split("#", 1)
43
+ return code or None, state or None
44
+ if "code=" in value:
45
+ params = urllib.parse.parse_qs(value)
46
+ return params.get("code", [None])[0], params.get("state", [None])[0]
47
+ return value, None
48
+
49
+
50
+ async def start_oauth_callback_server(
51
+ callback_path: str,
52
+ expected_state: str,
53
+ host: Optional[str],
54
+ port: int,
55
+ ) -> tuple[asyncio.Server, asyncio.Future[str]]:
56
+ """Start a minimal HTTP server that captures the OAuth authorization code."""
57
+ loop = asyncio.get_running_loop()
58
+ code_future: asyncio.Future[str] = loop.create_future()
59
+
60
+ async def _handle(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
61
+ try:
62
+ raw = await reader.read(4096)
63
+ line = raw.decode(errors="replace").split("\r\n")[0]
64
+ parts = line.split(" ")
65
+ if len(parts) < 2:
66
+ writer.write(b"HTTP/1.1 400 Bad Request\r\n\r\n")
67
+ await writer.drain()
68
+ return
69
+
70
+ parsed = urllib.parse.urlparse(parts[1])
71
+ params = urllib.parse.parse_qs(parsed.query)
72
+
73
+ if parsed.path == callback_path:
74
+ recv_state = params.get("state", [None])[0]
75
+ code = params.get("code", [None])[0]
76
+ error = params.get("error", [None])[0]
77
+
78
+ if error:
79
+ writer.write(
80
+ b"HTTP/1.1 400 Bad Request\r\nContent-Type: text/html; charset=utf-8\r\n\r\n"
81
+ + OAUTH_ERROR_HTML
82
+ )
83
+ elif recv_state == expected_state and code:
84
+ writer.write(
85
+ b"HTTP/1.1 200 OK\r\nContent-Type: text/html; charset=utf-8\r\n\r\n"
86
+ + OAUTH_SUCCESS_HTML
87
+ )
88
+ if not code_future.done():
89
+ code_future.set_result(code)
90
+ else:
91
+ writer.write(
92
+ b"HTTP/1.1 400 Bad Request\r\nContent-Type: text/html; charset=utf-8\r\n\r\n"
93
+ + OAUTH_ERROR_HTML
94
+ )
95
+ else:
96
+ writer.write(b"HTTP/1.1 404 Not Found\r\n\r\n")
97
+
98
+ await writer.drain()
99
+ except Exception:
100
+ pass
101
+ finally:
102
+ writer.close()
103
+ await writer.wait_closed()
104
+
105
+ server = await asyncio.start_server(_handle, host, port)
106
+ return server, code_future
107
+
108
+
109
+ async def await_oauth_code(
110
+ code_future: asyncio.Future[str],
111
+ state: str,
112
+ server: asyncio.Server,
113
+ callbacks: OAuthLoginCallbacks,
114
+ ) -> tuple[Optional[str], Optional[str]]:
115
+ """Race browser callback vs manual paste; close the server either way.
116
+
117
+ Returns (code, recv_state). recv_state falls back to state when the browser
118
+ callback wins (the state was already validated by the server handler).
119
+ """
120
+ code: Optional[str] = None
121
+ recv_state: Optional[str] = None
122
+ try:
123
+ if callbacks.on_manual_code_input:
124
+ browser_task = asyncio.ensure_future(code_future)
125
+ manual_task = asyncio.ensure_future(callbacks.on_manual_code_input())
126
+ done, pending = await asyncio.wait(
127
+ [browser_task, manual_task],
128
+ timeout=300,
129
+ return_when=asyncio.FIRST_COMPLETED,
130
+ )
131
+ for task in pending:
132
+ task.cancel()
133
+ await asyncio.gather(*pending, return_exceptions=True)
134
+
135
+ if browser_task in done and not browser_task.cancelled() and browser_task.exception() is None:
136
+ code = browser_task.result()
137
+ recv_state = state
138
+ elif manual_task in done and not manual_task.cancelled() and manual_task.exception() is None:
139
+ raw = manual_task.result()
140
+ parsed_code, parsed_state = parse_authorization_input(raw)
141
+ if parsed_state and parsed_state != state:
142
+ raise ValueError("OAuth state mismatch")
143
+ code = parsed_code
144
+ recv_state = parsed_state or state
145
+ else:
146
+ try:
147
+ code = await asyncio.wait_for(asyncio.shield(code_future), timeout=300)
148
+ recv_state = state
149
+ except asyncio.TimeoutError:
150
+ pass
151
+ finally:
152
+ server.close()
153
+ await server.wait_closed()
154
+ return code, recv_state
@@ -0,0 +1,141 @@
1
+ from __future__ import annotations
2
+ from typing import Generic, TypeVar
3
+ from tau.inference.provider.types import APIProvider, OAuthProvider, ImageProvider, AudioProvider, VideoProvider, AuthType
4
+
5
+ TextProvider = APIProvider | OAuthProvider
6
+
7
+ _T = TypeVar("_T")
8
+
9
+
10
+ class _ProviderRegistryBase(Generic[_T]):
11
+ """Internal shared base for all provider sub-registries."""
12
+
13
+ def __init__(self) -> None:
14
+ self._providers: dict[str, _T] = {}
15
+
16
+ def _key(self, _provider: _T) -> str:
17
+ raise NotImplementedError
18
+
19
+ def register(self, provider: _T) -> None:
20
+ self._providers[self._key(provider)] = provider
21
+
22
+ def unregister(self, key: str) -> None:
23
+ self._providers.pop(key, None)
24
+
25
+ def list(self) -> list[_T]:
26
+ return list(self._providers.values())
27
+
28
+ def get(self, key: str) -> _T | None:
29
+ return self._providers.get(key)
30
+
31
+ def reset(self) -> None:
32
+ self._providers.clear()
33
+
34
+
35
+ class TextProviderRegistry(_ProviderRegistryBase[TextProvider]):
36
+ """Registry mapping provider IDs to text-LLM provider instances."""
37
+
38
+ def _key(self, provider: TextProvider) -> str:
39
+ return provider.id
40
+
41
+ def is_using_oauth(self, provider: str) -> bool:
42
+ """Return True if the named provider uses OAuth auth; raise ValueError if unknown."""
43
+ if p := self.get(provider):
44
+ return p.auth_type == AuthType.OAuth
45
+ raise ValueError(f"Provider '{provider}' not found.")
46
+
47
+ def get_oauth_providers(self) -> list[OAuthProvider]:
48
+ """Return all registered OAuth-authenticated providers."""
49
+ return [p for p in self._providers.values() if isinstance(p, OAuthProvider)]
50
+
51
+ def get_api_providers(self) -> list[APIProvider]:
52
+ """Return all registered API-key-authenticated providers."""
53
+ return [p for p in self._providers.values() if isinstance(p, APIProvider)]
54
+
55
+ def get_oauth_provider(self, provider: str) -> OAuthProvider | None:
56
+ """Return the named provider only if it is an OAuthProvider."""
57
+ p = self.get(provider)
58
+ return p if isinstance(p, OAuthProvider) else None
59
+
60
+ def get_api_provider(self, provider: str) -> APIProvider | None:
61
+ """Return the named provider only if it is an APIProvider."""
62
+ p = self.get(provider)
63
+ return p if isinstance(p, APIProvider) else None
64
+
65
+ @classmethod
66
+ def from_builtins(cls) -> TextProviderRegistry:
67
+ from tau.builtins.providers.text import providers
68
+ instance = cls()
69
+ for provider in providers:
70
+ instance.register(provider)
71
+ return instance
72
+
73
+
74
+ class ImageProviderRegistry(_ProviderRegistryBase[ImageProvider]):
75
+ """Registry mapping provider names to image-generation provider instances."""
76
+
77
+ def _key(self, provider: ImageProvider) -> str:
78
+ return provider.name
79
+
80
+ @classmethod
81
+ def from_builtins(cls) -> ImageProviderRegistry:
82
+ from tau.builtins.providers.image import providers
83
+ instance = cls()
84
+ for provider in providers:
85
+ instance.register(provider)
86
+ return instance
87
+
88
+
89
+ class AudioProviderRegistry(_ProviderRegistryBase[AudioProvider]):
90
+ """Registry mapping provider names to audio (STT/TTS) provider instances."""
91
+
92
+ def _key(self, provider: AudioProvider) -> str:
93
+ return provider.name
94
+
95
+ @classmethod
96
+ def from_builtins(cls) -> AudioProviderRegistry:
97
+ from tau.builtins.providers.audio import providers
98
+ instance = cls()
99
+ for provider in providers:
100
+ instance.register(provider)
101
+ return instance
102
+
103
+
104
+ class VideoProviderRegistry(_ProviderRegistryBase[VideoProvider]):
105
+ """Registry mapping provider names to video-generation provider instances."""
106
+
107
+ def _key(self, provider: VideoProvider) -> str:
108
+ return provider.name
109
+
110
+ @classmethod
111
+ def from_builtins(cls) -> VideoProviderRegistry:
112
+ from tau.builtins.providers.video import providers
113
+ instance = cls()
114
+ for provider in providers:
115
+ instance.register(provider)
116
+ return instance
117
+
118
+
119
+ class ProviderRegistry:
120
+ """Unified registry holding all provider types (text, image, audio, video)."""
121
+
122
+ def __init__(
123
+ self,
124
+ text: TextProviderRegistry | None = None,
125
+ image: ImageProviderRegistry | None = None,
126
+ audio: AudioProviderRegistry | None = None,
127
+ video: VideoProviderRegistry | None = None,
128
+ ) -> None:
129
+ self.text = text or TextProviderRegistry()
130
+ self.image = image or ImageProviderRegistry()
131
+ self.audio = audio or AudioProviderRegistry()
132
+ self.video = video or VideoProviderRegistry()
133
+
134
+ @classmethod
135
+ def from_builtins(cls) -> ProviderRegistry:
136
+ return cls(
137
+ text=TextProviderRegistry.from_builtins(),
138
+ image=ImageProviderRegistry.from_builtins(),
139
+ audio=AudioProviderRegistry.from_builtins(),
140
+ video=VideoProviderRegistry.from_builtins(),
141
+ )
@@ -0,0 +1,114 @@
1
+ from __future__ import annotations
2
+
3
+ import time
4
+ from abc import ABC, abstractmethod
5
+ from dataclasses import dataclass, field
6
+ from typing import Optional, Type, Union, TYPE_CHECKING
7
+
8
+ from tau.inference.types import AuthType, LLMOptions, Transport
9
+
10
+ if TYPE_CHECKING:
11
+ from tau.inference.api.text.base import BaseLLMAPI
12
+ from tau.inference.provider.oauth.types import OAuthCredential, OAuthLoginCallbacks, AbortSignal
13
+
14
+ __all__ = ["AuthType", "APIProvider", "OAuthProvider", "ImageProvider", "AudioProvider", "VideoProvider"]
15
+
16
+
17
+ @dataclass
18
+ class OAuthProvider(ABC):
19
+ """Base for providers that authenticate via OAuth rather than a static API key."""
20
+
21
+ id: str
22
+ name: str
23
+ auth_type: AuthType = AuthType.OAuth
24
+ uses_callback_server: bool = False
25
+
26
+ @property
27
+ @abstractmethod
28
+ def api(self) -> Type["BaseLLMAPI"]:
29
+ """Return the LLM API class for this provider."""
30
+ ...
31
+
32
+ @abstractmethod
33
+ async def login(self, callbacks: "OAuthLoginCallbacks") -> "OAuthCredential":
34
+ """Authenticate with the provider and return a credential."""
35
+ ...
36
+
37
+ @abstractmethod
38
+ async def refresh_token(self, credential: "OAuthCredential", signal: Optional["AbortSignal"] = None) -> "OAuthCredential":
39
+ """Refresh an expired credential."""
40
+ ...
41
+
42
+ @abstractmethod
43
+ async def logout(self, credential: "OAuthCredential") -> None:
44
+ """Revoke the credential and clean up."""
45
+ ...
46
+
47
+ def get_api_key(self, credential: "OAuthCredential") -> str:
48
+ """Return the access token as the Bearer key for API calls."""
49
+ return credential.access
50
+
51
+ @abstractmethod
52
+ async def validate(self, credential: "OAuthCredential", signal: Optional["AbortSignal"] = None) -> bool:
53
+ """Check if the credential is valid (not necessarily fresh)."""
54
+ ...
55
+
56
+ def is_expired(self, credential: "OAuthCredential") -> bool:
57
+ """Return True if the token expires within the next 30 seconds."""
58
+ # 30-second buffer prevents using a token that expires mid-request
59
+ return int(time.time() * 1000) + 30_000 >= credential.expires
60
+
61
+ async def ensure_fresh(self, credential: "OAuthCredential", signal: Optional["AbortSignal"] = None) -> "OAuthCredential":
62
+ """Return a valid credential, transparently refreshing if expired."""
63
+ if self.is_expired(credential):
64
+ return await self.refresh_token(credential=credential, signal=signal)
65
+ return credential
66
+
67
+
68
+ @dataclass
69
+ class APIProvider:
70
+ """Provider that authenticates with a static API key stored in LLMOptions."""
71
+
72
+ id: str
73
+ name: str
74
+ api: Union[str, Type["BaseLLMAPI"]]
75
+ options: LLMOptions
76
+ auth_type: AuthType = AuthType.ApiKey
77
+
78
+ def get_api_key(self) -> Optional[str]:
79
+ """Return the configured API key, or None if not set."""
80
+ return self.options.api_key
81
+
82
+ def get_base_url(self) -> Optional[str]:
83
+ """Return the configured base URL override, or None to use the default."""
84
+ return self.options.base_url
85
+
86
+
87
+ @dataclass
88
+ class ImageProvider:
89
+ """Provider descriptor for image generation APIs."""
90
+
91
+ name: str
92
+ api: str
93
+ base_url: str
94
+ auth_type: AuthType = AuthType.ApiKey
95
+
96
+
97
+ @dataclass
98
+ class AudioProvider:
99
+ """Provider descriptor for speech-to-text and text-to-speech APIs."""
100
+
101
+ name: str
102
+ api: str
103
+ base_url: Optional[str] = None
104
+ auth_type: AuthType = AuthType.ApiKey
105
+
106
+
107
+ @dataclass
108
+ class VideoProvider:
109
+ """Provider descriptor for video generation APIs."""
110
+
111
+ name: str
112
+ api: str
113
+ base_url: Optional[str] = None
114
+ auth_type: AuthType = AuthType.ApiKey