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.
- tau/__init__.py +0 -0
- tau/agent/__init__.py +11 -0
- tau/agent/prompt/__init__.py +10 -0
- tau/agent/prompt/builder.py +302 -0
- tau/agent/prompt/types.py +33 -0
- tau/agent/service.py +369 -0
- tau/agent/types.py +61 -0
- tau/auth/manager.py +247 -0
- tau/auth/storage.py +82 -0
- tau/auth/types.py +41 -0
- tau/builtins/__init__.py +4 -0
- tau/builtins/__pycache__/__init__.cpython-313.pyc +0 -0
- tau/builtins/__pycache__/__init__.cpython-314.pyc +0 -0
- tau/builtins/commands/__init__.py +41 -0
- tau/builtins/commands/__pycache__/__init__.cpython-313.pyc +0 -0
- tau/builtins/commands/__pycache__/__init__.cpython-314.pyc +0 -0
- tau/builtins/commands/__pycache__/clear.cpython-313.pyc +0 -0
- tau/builtins/commands/__pycache__/clear.cpython-314.pyc +0 -0
- tau/builtins/commands/__pycache__/compact.cpython-313.pyc +0 -0
- tau/builtins/commands/__pycache__/compact.cpython-314.pyc +0 -0
- tau/builtins/commands/__pycache__/reload.cpython-313.pyc +0 -0
- tau/builtins/commands/__pycache__/reload.cpython-314.pyc +0 -0
- tau/builtins/commands/__pycache__/session.cpython-313.pyc +0 -0
- tau/builtins/commands/__pycache__/session.cpython-314.pyc +0 -0
- tau/builtins/commands/clear.py +16 -0
- tau/builtins/commands/compact.py +28 -0
- tau/builtins/commands/reload.py +27 -0
- tau/builtins/commands/session.py +19 -0
- tau/builtins/extensions/footer/__init__.py +76 -0
- tau/builtins/extensions/footer/__pycache__/__init__.cpython-313.pyc +0 -0
- tau/builtins/extensions/footer/__pycache__/git.cpython-313.pyc +0 -0
- tau/builtins/extensions/footer/__pycache__/model.cpython-313.pyc +0 -0
- tau/builtins/extensions/footer/__pycache__/utils.cpython-313.pyc +0 -0
- tau/builtins/extensions/footer/git.py +26 -0
- tau/builtins/extensions/footer/model.py +69 -0
- tau/builtins/extensions/footer/utils.py +44 -0
- tau/builtins/extensions/header/__init__.py +18 -0
- tau/builtins/extensions/header/__pycache__/__init__.cpython-313.pyc +0 -0
- tau/builtins/models/__init__.py +0 -0
- tau/builtins/models/__pycache__/__init__.cpython-313.pyc +0 -0
- tau/builtins/models/__pycache__/text.cpython-313.pyc +0 -0
- tau/builtins/models/audio.py +43 -0
- tau/builtins/models/image.py +43 -0
- tau/builtins/models/text.py +482 -0
- tau/builtins/models/video.py +40 -0
- tau/builtins/prompts/commit.md +7 -0
- tau/builtins/prompts/docs.md +7 -0
- tau/builtins/prompts/explain.md +7 -0
- tau/builtins/prompts/fix.md +7 -0
- tau/builtins/prompts/refactor.md +7 -0
- tau/builtins/prompts/review.md +7 -0
- tau/builtins/prompts/test.md +7 -0
- tau/builtins/providers/__init__.py +0 -0
- tau/builtins/providers/__pycache__/__init__.cpython-313.pyc +0 -0
- tau/builtins/providers/__pycache__/text.cpython-313.pyc +0 -0
- tau/builtins/providers/audio.py +10 -0
- tau/builtins/providers/image.py +9 -0
- tau/builtins/providers/text.py +33 -0
- tau/builtins/providers/video.py +6 -0
- tau/builtins/skills/code-review/SKILL.md +4 -0
- tau/builtins/skills/debug/SKILL.md +4 -0
- tau/builtins/skills/git-commit/SKILL.md +4 -0
- tau/builtins/themes/dark.yaml +1 -0
- tau/builtins/themes/light.yaml +46 -0
- tau/builtins/tools/__init__.py +73 -0
- tau/builtins/tools/__pycache__/__init__.cpython-313.pyc +0 -0
- tau/builtins/tools/__pycache__/__init__.cpython-314.pyc +0 -0
- tau/builtins/tools/__pycache__/bash.cpython-313.pyc +0 -0
- tau/builtins/tools/__pycache__/bash.cpython-314.pyc +0 -0
- tau/builtins/tools/__pycache__/edit.cpython-313.pyc +0 -0
- tau/builtins/tools/__pycache__/edit.cpython-314.pyc +0 -0
- tau/builtins/tools/__pycache__/glob.cpython-313.pyc +0 -0
- tau/builtins/tools/__pycache__/glob.cpython-314.pyc +0 -0
- tau/builtins/tools/__pycache__/grep.cpython-313.pyc +0 -0
- tau/builtins/tools/__pycache__/grep.cpython-314.pyc +0 -0
- tau/builtins/tools/__pycache__/ls.cpython-313.pyc +0 -0
- tau/builtins/tools/__pycache__/ls.cpython-314.pyc +0 -0
- tau/builtins/tools/__pycache__/read.cpython-313.pyc +0 -0
- tau/builtins/tools/__pycache__/read.cpython-314.pyc +0 -0
- tau/builtins/tools/__pycache__/terminal.cpython-313.pyc +0 -0
- tau/builtins/tools/__pycache__/terminal.cpython-314.pyc +0 -0
- tau/builtins/tools/__pycache__/write.cpython-313.pyc +0 -0
- tau/builtins/tools/__pycache__/write.cpython-314.pyc +0 -0
- tau/builtins/tools/edit.py +215 -0
- tau/builtins/tools/glob.py +112 -0
- tau/builtins/tools/grep.py +146 -0
- tau/builtins/tools/ls.py +135 -0
- tau/builtins/tools/read.py +122 -0
- tau/builtins/tools/terminal.py +150 -0
- tau/builtins/tools/write.py +105 -0
- tau/commands/__init__.py +10 -0
- tau/commands/registry.py +71 -0
- tau/commands/types.py +33 -0
- tau/console/__init__.py +0 -0
- tau/console/cli.py +266 -0
- tau/console/commands/__init__.py +0 -0
- tau/console/commands/auth.py +193 -0
- tau/console/commands/packages.py +104 -0
- tau/console/commands/update.py +76 -0
- tau/core/__init__.py +0 -0
- tau/core/registry.py +102 -0
- tau/engine/__init__.py +47 -0
- tau/engine/service.py +768 -0
- tau/engine/types.py +163 -0
- tau/extensions/__init__.py +28 -0
- tau/extensions/api.py +928 -0
- tau/extensions/context.py +462 -0
- tau/extensions/events.py +70 -0
- tau/extensions/loader.py +386 -0
- tau/extensions/runtime.py +184 -0
- tau/extensions/settings.py +137 -0
- tau/hooks/__init__.py +112 -0
- tau/hooks/engine.py +237 -0
- tau/hooks/inference.py +21 -0
- tau/hooks/runtime.py +126 -0
- tau/hooks/service.py +121 -0
- tau/hooks/session.py +117 -0
- tau/hooks/tui.py +61 -0
- tau/hooks/types.py +72 -0
- tau/inference/__init__.py +180 -0
- tau/inference/api/__init__.py +0 -0
- tau/inference/api/audio/__init__.py +0 -0
- tau/inference/api/audio/base.py +29 -0
- tau/inference/api/audio/builtins.py +15 -0
- tau/inference/api/audio/elevenlabs_audio.py +183 -0
- tau/inference/api/audio/gemini_audio.py +95 -0
- tau/inference/api/audio/openai_audio.py +159 -0
- tau/inference/api/audio/registry.py +15 -0
- tau/inference/api/audio/sarvam_audio.py +163 -0
- tau/inference/api/audio/service.py +103 -0
- tau/inference/api/audio/utils.py +47 -0
- tau/inference/api/image/__init__.py +0 -0
- tau/inference/api/image/base.py +17 -0
- tau/inference/api/image/builtins.py +8 -0
- tau/inference/api/image/gemini_image.py +77 -0
- tau/inference/api/image/openai_image.py +103 -0
- tau/inference/api/image/openrouter.py +144 -0
- tau/inference/api/image/registry.py +15 -0
- tau/inference/api/image/service.py +71 -0
- tau/inference/api/registry.py +82 -0
- tau/inference/api/text/__init__.py +0 -0
- tau/inference/api/text/anthropic_claude_code.py +222 -0
- tau/inference/api/text/anthropic_messages.py +196 -0
- tau/inference/api/text/base.py +40 -0
- tau/inference/api/text/builtins.py +19 -0
- tau/inference/api/text/gemini_generate.py +234 -0
- tau/inference/api/text/github_copilot_chat.py +172 -0
- tau/inference/api/text/google_antigravity.py +522 -0
- tau/inference/api/text/mistral_chat.py +284 -0
- tau/inference/api/text/ollama_chat.py +200 -0
- tau/inference/api/text/openai_codex_responses.py +497 -0
- tau/inference/api/text/openai_completions.py +227 -0
- tau/inference/api/text/openai_responses.py +235 -0
- tau/inference/api/text/registry.py +50 -0
- tau/inference/api/text/service.py +297 -0
- tau/inference/api/text/types.py +7 -0
- tau/inference/api/text/utils.py +228 -0
- tau/inference/api/video/__init__.py +0 -0
- tau/inference/api/video/base.py +26 -0
- tau/inference/api/video/builtins.py +7 -0
- tau/inference/api/video/fal_video.py +119 -0
- tau/inference/api/video/openrouter_video.py +142 -0
- tau/inference/api/video/registry.py +15 -0
- tau/inference/api/video/service.py +72 -0
- tau/inference/model/__init__.py +0 -0
- tau/inference/model/registry.py +102 -0
- tau/inference/model/types.py +65 -0
- tau/inference/provider/__init__.py +0 -0
- tau/inference/provider/oauth/__init__.py +35 -0
- tau/inference/provider/oauth/anthropic_claude_code.py +286 -0
- tau/inference/provider/oauth/github_copilot.py +333 -0
- tau/inference/provider/oauth/google_antigravity.py +258 -0
- tau/inference/provider/oauth/openai_codex.py +309 -0
- tau/inference/provider/oauth/pkce.py +14 -0
- tau/inference/provider/oauth/types.py +46 -0
- tau/inference/provider/oauth/utils.py +154 -0
- tau/inference/provider/registry.py +141 -0
- tau/inference/provider/types.py +114 -0
- tau/inference/types.py +549 -0
- tau/inference/utils.py +219 -0
- tau/message/__init__.py +0 -0
- tau/message/types.py +482 -0
- tau/message/utils.py +178 -0
- tau/packages/__init__.py +11 -0
- tau/packages/manager.py +190 -0
- tau/packages/types.py +20 -0
- tau/packages/utils.py +67 -0
- tau/prompts/expand.py +58 -0
- tau/prompts/loader.py +69 -0
- tau/prompts/registry.py +45 -0
- tau/prompts/types.py +24 -0
- tau/rpc/__init__.py +8 -0
- tau/rpc/mode.py +783 -0
- tau/rpc/types.py +252 -0
- tau/runtime/service.py +759 -0
- tau/runtime/types.py +303 -0
- tau/session/branch_summarization.py +312 -0
- tau/session/compaction.py +646 -0
- tau/session/manager.py +652 -0
- tau/session/types.py +188 -0
- tau/session/utils.py +233 -0
- tau/settings/manager.py +1077 -0
- tau/settings/paths.py +150 -0
- tau/settings/storage.py +63 -0
- tau/settings/types.py +173 -0
- tau/settings/utils.py +25 -0
- tau/skills/loader.py +91 -0
- tau/skills/registry.py +70 -0
- tau/skills/types.py +25 -0
- tau/themes/loader.py +238 -0
- tau/themes/registry.py +108 -0
- tau/themes/types.py +19 -0
- tau/tool/__init__.py +3 -0
- tau/tool/registry.py +117 -0
- tau/tool/render.py +21 -0
- tau/tool/types.py +244 -0
- tau/trust/__init__.py +13 -0
- tau/trust/manager.py +80 -0
- tau/trust/types.py +14 -0
- tau/trust/utils.py +72 -0
- tau/tui/__init__.py +54 -0
- tau/tui/agent_hooks.py +346 -0
- tau/tui/ansi.py +330 -0
- tau/tui/app.py +540 -0
- tau/tui/autocomplete.py +33 -0
- tau/tui/capabilities.py +119 -0
- tau/tui/commands/__init__.py +3 -0
- tau/tui/commands/appearance.py +498 -0
- tau/tui/commands/auth.py +232 -0
- tau/tui/commands/context.py +38 -0
- tau/tui/commands/misc.py +82 -0
- tau/tui/commands/model.py +118 -0
- tau/tui/commands/session.py +464 -0
- tau/tui/component.py +268 -0
- tau/tui/components/__init__.py +0 -0
- tau/tui/components/autocomplete_manager.py +267 -0
- tau/tui/components/autocomplete_picker.py +143 -0
- tau/tui/components/box.py +90 -0
- tau/tui/components/command_palette.py +144 -0
- tau/tui/components/dynamic_border.py +19 -0
- tau/tui/components/file_picker.py +233 -0
- tau/tui/components/image.py +181 -0
- tau/tui/components/inline_selector.py +71 -0
- tau/tui/components/layout.py +1194 -0
- tau/tui/components/message_list.py +692 -0
- tau/tui/components/modal.py +97 -0
- tau/tui/components/model_palette.py +204 -0
- tau/tui/components/picker_overlay.py +174 -0
- tau/tui/components/prompt_overlay.py +236 -0
- tau/tui/components/resume_modal.py +372 -0
- tau/tui/components/select_list.py +222 -0
- tau/tui/components/settings_modal.py +274 -0
- tau/tui/components/settings_schema.py +203 -0
- tau/tui/components/spinner.py +119 -0
- tau/tui/components/text_input.py +396 -0
- tau/tui/components/text_prompt.py +82 -0
- tau/tui/components/tree_select_list.py +580 -0
- tau/tui/components/trust_screen.py +97 -0
- tau/tui/diff.py +114 -0
- tau/tui/fuzzy.py +99 -0
- tau/tui/input.py +496 -0
- tau/tui/input_handler.py +716 -0
- tau/tui/keybindings.py +87 -0
- tau/tui/markdown.py +286 -0
- tau/tui/message_renderers.py +31 -0
- tau/tui/overlay.py +326 -0
- tau/tui/renderer.py +378 -0
- tau/tui/terminal.py +499 -0
- tau/tui/theme.py +148 -0
- tau/tui/tui.py +544 -0
- tau/tui/ui_context.py +768 -0
- tau/tui/utils.py +20 -0
- tau/utils/__init__.py +0 -0
- tau/utils/http_proxy.py +221 -0
- tau/utils/image_processing.py +172 -0
- tau/utils/secrets.py +59 -0
- tau/utils/version_check.py +60 -0
- tau_coding_agent-0.1.0.dist-info/METADATA +177 -0
- tau_coding_agent-0.1.0.dist-info/RECORD +283 -0
- tau_coding_agent-0.1.0.dist-info/WHEEL +5 -0
- tau_coding_agent-0.1.0.dist-info/entry_points.txt +2 -0
- tau_coding_agent-0.1.0.dist-info/licenses/LICENSE +21 -0
- 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
|