ya-oauth-provider 0.74.0__tar.gz
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.
- ya_oauth_provider-0.74.0/.gitignore +169 -0
- ya_oauth_provider-0.74.0/PKG-INFO +34 -0
- ya_oauth_provider-0.74.0/README.md +13 -0
- ya_oauth_provider-0.74.0/pyproject.toml +60 -0
- ya_oauth_provider-0.74.0/tests/test_codex_provider.py +125 -0
- ya_oauth_provider-0.74.0/ya_oauth_provider/__init__.py +6 -0
- ya_oauth_provider-0.74.0/ya_oauth_provider/codex.py +107 -0
- ya_oauth_provider-0.74.0/ya_oauth_provider/http.py +132 -0
|
@@ -0,0 +1,169 @@
|
|
|
1
|
+
docs/source
|
|
2
|
+
|
|
3
|
+
# From https://raw.githubusercontent.com/github/gitignore/main/Python.gitignore
|
|
4
|
+
|
|
5
|
+
# Byte-compiled / optimized / DLL files
|
|
6
|
+
__pycache__/
|
|
7
|
+
*.py[cod]
|
|
8
|
+
*$py.class
|
|
9
|
+
|
|
10
|
+
# C extensions
|
|
11
|
+
*.so
|
|
12
|
+
|
|
13
|
+
# Distribution / packaging
|
|
14
|
+
.Python
|
|
15
|
+
build/
|
|
16
|
+
develop-eggs/
|
|
17
|
+
dist/
|
|
18
|
+
downloads/
|
|
19
|
+
eggs/
|
|
20
|
+
.eggs/
|
|
21
|
+
lib/
|
|
22
|
+
!apps/
|
|
23
|
+
!apps/ya-claw-web/
|
|
24
|
+
!apps/ya-claw-web/src/
|
|
25
|
+
!apps/ya-claw-web/src/lib/
|
|
26
|
+
!apps/ya-claw-web/src/lib/**
|
|
27
|
+
lib64/
|
|
28
|
+
parts/
|
|
29
|
+
sdist/
|
|
30
|
+
var/
|
|
31
|
+
wheels/
|
|
32
|
+
share/python-wheels/
|
|
33
|
+
*.egg-info/
|
|
34
|
+
.installed.cfg
|
|
35
|
+
*.egg
|
|
36
|
+
MANIFEST
|
|
37
|
+
|
|
38
|
+
# PyInstaller
|
|
39
|
+
# Usually these files are written by a python script from a template
|
|
40
|
+
# before PyInstaller builds the exe, so as to inject date/other infos into it.
|
|
41
|
+
*.manifest
|
|
42
|
+
*.spec
|
|
43
|
+
|
|
44
|
+
# Installer logs
|
|
45
|
+
pip-log.txt
|
|
46
|
+
pip-delete-this-directory.txt
|
|
47
|
+
|
|
48
|
+
# Unit test / coverage reports
|
|
49
|
+
htmlcov/
|
|
50
|
+
.tox/
|
|
51
|
+
.nox/
|
|
52
|
+
.coverage
|
|
53
|
+
.coverage.*
|
|
54
|
+
.cache
|
|
55
|
+
nosetests.xml
|
|
56
|
+
coverage.xml
|
|
57
|
+
*.cover
|
|
58
|
+
*.py,cover
|
|
59
|
+
.hypothesis/
|
|
60
|
+
.pytest_cache/
|
|
61
|
+
cover/
|
|
62
|
+
|
|
63
|
+
# Translations
|
|
64
|
+
*.mo
|
|
65
|
+
*.pot
|
|
66
|
+
|
|
67
|
+
# Django stuff:
|
|
68
|
+
*.log
|
|
69
|
+
local_settings.py
|
|
70
|
+
db.sqlite3
|
|
71
|
+
db.sqlite3-journal
|
|
72
|
+
|
|
73
|
+
# Flask stuff:
|
|
74
|
+
instance/
|
|
75
|
+
.webassets-cache
|
|
76
|
+
|
|
77
|
+
# Scrapy stuff:
|
|
78
|
+
.scrapy
|
|
79
|
+
|
|
80
|
+
# Sphinx documentation
|
|
81
|
+
docs/_build/
|
|
82
|
+
|
|
83
|
+
# PyBuilder
|
|
84
|
+
.pybuilder/
|
|
85
|
+
target/
|
|
86
|
+
|
|
87
|
+
# Jupyter Notebook
|
|
88
|
+
.ipynb_checkpoints
|
|
89
|
+
|
|
90
|
+
# IPython
|
|
91
|
+
profile_default/
|
|
92
|
+
ipython_config.py
|
|
93
|
+
|
|
94
|
+
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
|
|
95
|
+
__pypackages__/
|
|
96
|
+
|
|
97
|
+
# Celery stuff
|
|
98
|
+
celerybeat-schedule
|
|
99
|
+
celerybeat.pid
|
|
100
|
+
|
|
101
|
+
# SageMath parsed files
|
|
102
|
+
*.sage.py
|
|
103
|
+
|
|
104
|
+
# Environments
|
|
105
|
+
.env
|
|
106
|
+
.yaai/
|
|
107
|
+
**/.yaai/
|
|
108
|
+
.venv
|
|
109
|
+
env/
|
|
110
|
+
venv/
|
|
111
|
+
ENV/
|
|
112
|
+
env.bak/
|
|
113
|
+
venv.bak/
|
|
114
|
+
|
|
115
|
+
# Spyder project settings
|
|
116
|
+
.spyderproject
|
|
117
|
+
.spyproject
|
|
118
|
+
|
|
119
|
+
# Rope project settings
|
|
120
|
+
.ropeproject
|
|
121
|
+
|
|
122
|
+
# mkdocs documentation
|
|
123
|
+
/site
|
|
124
|
+
|
|
125
|
+
# mypy
|
|
126
|
+
.mypy_cache/
|
|
127
|
+
.pyright/
|
|
128
|
+
.dmypy.json
|
|
129
|
+
dmypy.json
|
|
130
|
+
|
|
131
|
+
# Pyre type checker
|
|
132
|
+
.pyre/
|
|
133
|
+
|
|
134
|
+
# pytype static type analyzer
|
|
135
|
+
.pytype/
|
|
136
|
+
|
|
137
|
+
# Cython debug symbols
|
|
138
|
+
cython_debug/
|
|
139
|
+
|
|
140
|
+
# Vscode config files
|
|
141
|
+
# .vscode/
|
|
142
|
+
|
|
143
|
+
# PyCharm
|
|
144
|
+
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
|
|
145
|
+
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
|
|
146
|
+
# and can be added to the global gitignore or merged into this file. For a more nuclear
|
|
147
|
+
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
|
|
148
|
+
#.idea/
|
|
149
|
+
|
|
150
|
+
TO-DO.json
|
|
151
|
+
dev/
|
|
152
|
+
ya_agent_sdk/sandbox/shell/templates/public
|
|
153
|
+
|
|
154
|
+
# Sync automaticlly
|
|
155
|
+
yaacli/yaacli/skills/building-agents
|
|
156
|
+
|
|
157
|
+
# Frontend
|
|
158
|
+
node_modules/
|
|
159
|
+
apps/*/dist/
|
|
160
|
+
!apps/ya-desktop/
|
|
161
|
+
!apps/ya-desktop/src/
|
|
162
|
+
!apps/ya-desktop/src/**
|
|
163
|
+
!apps/ya-desktop/src-tauri/
|
|
164
|
+
!apps/ya-desktop/src-tauri/resources/
|
|
165
|
+
!apps/ya-desktop/src-tauri/resources/uv/
|
|
166
|
+
!apps/ya-desktop/src-tauri/resources/uv/.gitkeep
|
|
167
|
+
apps/ya-desktop/src-tauri/resources/uv/uv
|
|
168
|
+
apps/ya-desktop/src-tauri/resources/uv/uv.exe
|
|
169
|
+
*.tsbuildinfo
|
|
@@ -0,0 +1,34 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: ya-oauth-provider
|
|
3
|
+
Version: 0.74.0
|
|
4
|
+
Summary: Pydantic AI OAuth-backed model provider helpers for YA
|
|
5
|
+
Project-URL: Repository, https://github.com/wh1isper/ya-mono
|
|
6
|
+
Author-email: wh1isper <jizhongsheng957@gmail.com>
|
|
7
|
+
Keywords: agents,oauth,pydantic-ai
|
|
8
|
+
Classifier: Intended Audience :: Developers
|
|
9
|
+
Classifier: Programming Language :: Python
|
|
10
|
+
Classifier: Programming Language :: Python :: 3
|
|
11
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
12
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
13
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
14
|
+
Requires-Python: <3.14,>=3.11
|
|
15
|
+
Requires-Dist: httpx>=0.28.1
|
|
16
|
+
Requires-Dist: pydantic-ai>=1.94.0
|
|
17
|
+
Requires-Dist: pydantic>=2.12.0
|
|
18
|
+
Requires-Dist: tenacity>=9.0.0
|
|
19
|
+
Requires-Dist: ya-oauth==0.74.0
|
|
20
|
+
Description-Content-Type: text/markdown
|
|
21
|
+
|
|
22
|
+
# ya-oauth-provider
|
|
23
|
+
|
|
24
|
+
Pydantic AI model/provider helpers that consume OAuth token sources from `ya-oauth`.
|
|
25
|
+
|
|
26
|
+
## Codex model string
|
|
27
|
+
|
|
28
|
+
YA Agent SDK loads this package for model strings such as:
|
|
29
|
+
|
|
30
|
+
```text
|
|
31
|
+
oauth@codex:gpt-5.5
|
|
32
|
+
```
|
|
33
|
+
|
|
34
|
+
The provider attaches Codex-compatible bearer, account, originator, version, session, and thread headers.
|
|
@@ -0,0 +1,13 @@
|
|
|
1
|
+
# ya-oauth-provider
|
|
2
|
+
|
|
3
|
+
Pydantic AI model/provider helpers that consume OAuth token sources from `ya-oauth`.
|
|
4
|
+
|
|
5
|
+
## Codex model string
|
|
6
|
+
|
|
7
|
+
YA Agent SDK loads this package for model strings such as:
|
|
8
|
+
|
|
9
|
+
```text
|
|
10
|
+
oauth@codex:gpt-5.5
|
|
11
|
+
```
|
|
12
|
+
|
|
13
|
+
The provider attaches Codex-compatible bearer, account, originator, version, session, and thread headers.
|
|
@@ -0,0 +1,60 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "ya-oauth-provider"
|
|
3
|
+
dynamic = ["version", "dependencies"]
|
|
4
|
+
description = "Pydantic AI OAuth-backed model provider helpers for YA"
|
|
5
|
+
authors = [{ name = "wh1isper", email = "jizhongsheng957@gmail.com" }]
|
|
6
|
+
readme = "README.md"
|
|
7
|
+
keywords = ["oauth", "pydantic-ai", "agents"]
|
|
8
|
+
requires-python = ">=3.11,<3.14"
|
|
9
|
+
classifiers = [
|
|
10
|
+
"Intended Audience :: Developers",
|
|
11
|
+
"Programming Language :: Python",
|
|
12
|
+
"Programming Language :: Python :: 3",
|
|
13
|
+
"Programming Language :: Python :: 3.11",
|
|
14
|
+
"Programming Language :: Python :: 3.12",
|
|
15
|
+
"Programming Language :: Python :: 3.13",
|
|
16
|
+
]
|
|
17
|
+
|
|
18
|
+
[project.urls]
|
|
19
|
+
Repository = "https://github.com/wh1isper/ya-mono"
|
|
20
|
+
|
|
21
|
+
[dependency-groups]
|
|
22
|
+
dev = [
|
|
23
|
+
"pytest>=7.2.0",
|
|
24
|
+
"pytest-asyncio>=0.25.3",
|
|
25
|
+
"pytest-httpx>=0.35.0",
|
|
26
|
+
"ruff>=0.9.2",
|
|
27
|
+
"pyright>=1.1.0",
|
|
28
|
+
]
|
|
29
|
+
|
|
30
|
+
[build-system]
|
|
31
|
+
requires = ["hatchling", "uv-dynamic-versioning>=0.7.0"]
|
|
32
|
+
build-backend = "hatchling.build"
|
|
33
|
+
|
|
34
|
+
[tool.hatch.version]
|
|
35
|
+
source = "uv-dynamic-versioning"
|
|
36
|
+
|
|
37
|
+
[tool.uv-dynamic-versioning]
|
|
38
|
+
vcs = "git"
|
|
39
|
+
style = "pep440"
|
|
40
|
+
bump = true
|
|
41
|
+
|
|
42
|
+
[tool.hatch.metadata.hooks.uv-dynamic-versioning]
|
|
43
|
+
dependencies = [
|
|
44
|
+
"httpx>=0.28.1",
|
|
45
|
+
"pydantic>=2.12.0",
|
|
46
|
+
"pydantic-ai>=1.94.0",
|
|
47
|
+
"tenacity>=9.0.0",
|
|
48
|
+
"ya-oauth=={{ version }}",
|
|
49
|
+
]
|
|
50
|
+
|
|
51
|
+
[tool.hatch.build.targets.wheel]
|
|
52
|
+
packages = ["ya_oauth_provider"]
|
|
53
|
+
|
|
54
|
+
[tool.uv.sources]
|
|
55
|
+
ya-oauth = { workspace = true }
|
|
56
|
+
|
|
57
|
+
[tool.deptry]
|
|
58
|
+
ignore = ["DEP001", "DEP002"]
|
|
59
|
+
package_module_name_map = { "httpx" = "httpx", "pydantic" = "pydantic", "pydantic-ai" = "pydantic_ai", "pytest" = "pytest", "pytest-asyncio" = "pytest_asyncio", "pytest-httpx" = "pytest_httpx", "ruff" = "ruff", "pyright" = "pyright", "tenacity" = "tenacity", "ya-oauth" = "ya_oauth" }
|
|
60
|
+
per_rule_ignores = { DEP003 = ["anyio", "httpx", "pydantic-ai", "pytest", "tenacity"], DEP004 = ["pytest"] }
|
|
@@ -0,0 +1,125 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import json
|
|
4
|
+
|
|
5
|
+
import anyio
|
|
6
|
+
import httpx
|
|
7
|
+
import pytest
|
|
8
|
+
from pydantic_ai.exceptions import UserError
|
|
9
|
+
from ya_oauth.types import OAuthAccount, TokenSnapshot
|
|
10
|
+
from ya_oauth_provider.codex import CodexResponsesModel, build_codex_model, build_session_headers
|
|
11
|
+
from ya_oauth_provider.http import OAuthBearerAuth, build_codex_headers
|
|
12
|
+
|
|
13
|
+
ACCESS_TOKEN_OLD = "fixture-access-token-old" # noqa: S105
|
|
14
|
+
ACCESS_TOKEN_NEW = "fixture-access-token-new" # noqa: S105
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
class FakeTokenSource:
|
|
18
|
+
def __init__(self) -> None:
|
|
19
|
+
self.refresh_count = 0
|
|
20
|
+
|
|
21
|
+
async def get_token(self) -> TokenSnapshot:
|
|
22
|
+
return TokenSnapshot(
|
|
23
|
+
provider_name="codex",
|
|
24
|
+
access_token=ACCESS_TOKEN_OLD,
|
|
25
|
+
account=OAuthAccount(chatgpt_account_id="acct_123", chatgpt_account_is_fedramp=True),
|
|
26
|
+
)
|
|
27
|
+
|
|
28
|
+
async def refresh_token(self) -> TokenSnapshot:
|
|
29
|
+
self.refresh_count += 1
|
|
30
|
+
return TokenSnapshot(
|
|
31
|
+
provider_name="codex",
|
|
32
|
+
access_token=ACCESS_TOKEN_NEW,
|
|
33
|
+
account=OAuthAccount(chatgpt_account_id="acct_456"),
|
|
34
|
+
)
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
def test_build_codex_headers() -> None:
|
|
38
|
+
headers = build_codex_headers(
|
|
39
|
+
OAuthAccount(chatgpt_account_id="acct_123", chatgpt_account_is_fedramp=True),
|
|
40
|
+
extra_headers={"session_id": "s1", "thread-id": "t1", "x-client-request-id": "t1"},
|
|
41
|
+
version="test-version",
|
|
42
|
+
)
|
|
43
|
+
|
|
44
|
+
assert "Authorization" not in headers
|
|
45
|
+
assert headers["ChatGPT-Account-ID"] == "acct_123"
|
|
46
|
+
assert headers["X-OpenAI-Fedramp"] == "true"
|
|
47
|
+
assert headers["originator"] == "ya_agent_sdk"
|
|
48
|
+
assert headers["version"] == "test-version"
|
|
49
|
+
assert headers["session_id"] == "s1"
|
|
50
|
+
assert headers["thread-id"] == "t1"
|
|
51
|
+
assert headers["x-client-request-id"] == "t1"
|
|
52
|
+
|
|
53
|
+
|
|
54
|
+
def test_build_codex_headers_rejects_reserved_extra_headers() -> None:
|
|
55
|
+
with pytest.raises(ValueError, match="reserved OAuth/Codex header"):
|
|
56
|
+
build_codex_headers(OAuthAccount(), extra_headers={"Authorization": "Bearer other"})
|
|
57
|
+
|
|
58
|
+
|
|
59
|
+
def test_build_session_headers_uses_both_variants() -> None:
|
|
60
|
+
assert build_session_headers("session", "thread") == {
|
|
61
|
+
"session_id": "session",
|
|
62
|
+
"session-id": "session",
|
|
63
|
+
"thread_id": "thread",
|
|
64
|
+
"thread-id": "thread",
|
|
65
|
+
"x-client-request-id": "thread",
|
|
66
|
+
}
|
|
67
|
+
|
|
68
|
+
|
|
69
|
+
def test_build_codex_model_requires_streaming_for_non_stream_request() -> None:
|
|
70
|
+
model = build_codex_model("gpt-5.5", token_source=FakeTokenSource())
|
|
71
|
+
|
|
72
|
+
assert isinstance(model, CodexResponsesModel)
|
|
73
|
+
with pytest.raises(UserError, match="requires streaming"):
|
|
74
|
+
anyio.run(model.request, [], None, None) # type: ignore[arg-type]
|
|
75
|
+
|
|
76
|
+
|
|
77
|
+
@pytest.mark.asyncio
|
|
78
|
+
async def test_oauth_bearer_auth_fills_codex_responses_instructions() -> None:
|
|
79
|
+
source = FakeTokenSource()
|
|
80
|
+
seen: list[dict[str, object]] = []
|
|
81
|
+
|
|
82
|
+
def handler(request: httpx.Request) -> httpx.Response:
|
|
83
|
+
seen.append(dict(json.loads(request.content)))
|
|
84
|
+
return httpx.Response(200, json={"ok": True}, request=request)
|
|
85
|
+
|
|
86
|
+
client = httpx.AsyncClient(
|
|
87
|
+
transport=httpx.MockTransport(handler),
|
|
88
|
+
auth=OAuthBearerAuth(source, provider_name="codex"),
|
|
89
|
+
)
|
|
90
|
+
|
|
91
|
+
await client.post("https://chatgpt.com/backend-api/codex/responses", json={"model": "gpt-5.5"})
|
|
92
|
+
await client.post(
|
|
93
|
+
"https://chatgpt.com/backend-api/codex/responses",
|
|
94
|
+
json={"model": "gpt-5.5", "instructions": None},
|
|
95
|
+
)
|
|
96
|
+
await client.aclose()
|
|
97
|
+
|
|
98
|
+
assert seen == [
|
|
99
|
+
{"model": "gpt-5.5", "instructions": "", "store": False},
|
|
100
|
+
{"model": "gpt-5.5", "instructions": "", "store": False},
|
|
101
|
+
]
|
|
102
|
+
|
|
103
|
+
|
|
104
|
+
@pytest.mark.asyncio
|
|
105
|
+
async def test_oauth_bearer_auth_refreshes_once_on_401() -> None:
|
|
106
|
+
source = FakeTokenSource()
|
|
107
|
+
seen: list[str] = []
|
|
108
|
+
|
|
109
|
+
def handler(request: httpx.Request) -> httpx.Response:
|
|
110
|
+
seen.append(request.headers["Authorization"])
|
|
111
|
+
if len(seen) == 1:
|
|
112
|
+
return httpx.Response(401, request=request)
|
|
113
|
+
return httpx.Response(200, json={"ok": True}, request=request)
|
|
114
|
+
|
|
115
|
+
client = httpx.AsyncClient(
|
|
116
|
+
transport=httpx.MockTransport(handler),
|
|
117
|
+
auth=OAuthBearerAuth(source, provider_name="codex", extra_headers={"session_id": "s1"}),
|
|
118
|
+
)
|
|
119
|
+
|
|
120
|
+
response = await client.get("https://example.com/test")
|
|
121
|
+
await client.aclose()
|
|
122
|
+
|
|
123
|
+
assert response.status_code == 200
|
|
124
|
+
assert seen == [f"Bearer {ACCESS_TOKEN_OLD}", f"Bearer {ACCESS_TOKEN_NEW}"]
|
|
125
|
+
assert source.refresh_count == 1
|
|
@@ -0,0 +1,6 @@
|
|
|
1
|
+
"""OAuth-backed Pydantic AI provider helpers."""
|
|
2
|
+
|
|
3
|
+
from ya_oauth_provider.codex import build_codex_model, build_session_headers, infer_oauth_model
|
|
4
|
+
from ya_oauth_provider.http import OAuthBearerAuth, build_codex_headers
|
|
5
|
+
|
|
6
|
+
__all__ = ["OAuthBearerAuth", "build_codex_headers", "build_codex_model", "build_session_headers", "infer_oauth_model"]
|
|
@@ -0,0 +1,107 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from collections.abc import AsyncIterator
|
|
4
|
+
from contextlib import asynccontextmanager
|
|
5
|
+
from typing import Any
|
|
6
|
+
|
|
7
|
+
from pydantic_ai.exceptions import UserError
|
|
8
|
+
from pydantic_ai.messages import ModelMessage, ModelResponse
|
|
9
|
+
from pydantic_ai.models import Model, ModelRequestParameters, StreamedResponse
|
|
10
|
+
from pydantic_ai.models.openai import OpenAIResponsesModel
|
|
11
|
+
from pydantic_ai.profiles.openai import OpenAIModelProfile
|
|
12
|
+
from pydantic_ai.providers.openai import OpenAIProvider
|
|
13
|
+
from pydantic_ai.settings import ModelSettings
|
|
14
|
+
from ya_oauth.codex import CODEX_BASE_URL, create_codex_token_source
|
|
15
|
+
from ya_oauth.types import OAuthTokenSource
|
|
16
|
+
|
|
17
|
+
from ya_oauth_provider.http import OAuthBearerAuth
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
def infer_oauth_model(provider_name: str, model_name: str, *, extra_headers: dict[str, str] | None = None) -> Model:
|
|
21
|
+
"""Infer an OAuth-backed model from `oauth@provider:model` parts."""
|
|
22
|
+
if provider_name == "codex":
|
|
23
|
+
return build_codex_model(model_name, extra_headers=extra_headers)
|
|
24
|
+
raise KeyError(f"Unknown OAuth provider: {provider_name}")
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
def build_codex_model(
|
|
28
|
+
model_name: str,
|
|
29
|
+
*,
|
|
30
|
+
token_source: OAuthTokenSource | None = None,
|
|
31
|
+
extra_headers: dict[str, str] | None = None,
|
|
32
|
+
base_url: str = CODEX_BASE_URL,
|
|
33
|
+
) -> Model:
|
|
34
|
+
"""Build a Codex OAuth-backed OpenAI Responses model."""
|
|
35
|
+
import httpx
|
|
36
|
+
from pydantic_ai.models import get_user_agent
|
|
37
|
+
from pydantic_ai.retries import AsyncTenacityTransport, RetryConfig
|
|
38
|
+
from tenacity import retry_if_exception_type, stop_after_attempt, wait_exponential
|
|
39
|
+
|
|
40
|
+
source = token_source or create_codex_token_source()
|
|
41
|
+
http_client = httpx.AsyncClient(
|
|
42
|
+
auth=OAuthBearerAuth(source, provider_name="codex", extra_headers=extra_headers),
|
|
43
|
+
headers={"User-Agent": get_user_agent()},
|
|
44
|
+
timeout=httpx.Timeout(timeout=900, connect=5, read=300),
|
|
45
|
+
transport=AsyncTenacityTransport(
|
|
46
|
+
config=RetryConfig(
|
|
47
|
+
retry=retry_if_exception_type((httpx.HTTPError, httpx.StreamError)),
|
|
48
|
+
wait=wait_exponential(multiplier=1, max=10),
|
|
49
|
+
stop=stop_after_attempt(10),
|
|
50
|
+
reraise=True,
|
|
51
|
+
)
|
|
52
|
+
),
|
|
53
|
+
)
|
|
54
|
+
provider = OpenAIProvider(api_key="oauth-managed", base_url=base_url, http_client=http_client)
|
|
55
|
+
return CodexResponsesModel(model_name, provider=provider, profile=_codex_profile())
|
|
56
|
+
|
|
57
|
+
|
|
58
|
+
class CodexResponsesModel(OpenAIResponsesModel):
|
|
59
|
+
"""Codex Responses API model that requires streaming calls."""
|
|
60
|
+
|
|
61
|
+
async def request(
|
|
62
|
+
self,
|
|
63
|
+
messages: list[ModelMessage],
|
|
64
|
+
model_settings: ModelSettings | None,
|
|
65
|
+
model_request_parameters: ModelRequestParameters,
|
|
66
|
+
) -> ModelResponse:
|
|
67
|
+
raise UserError(
|
|
68
|
+
"Codex OAuth Responses API requires streaming. "
|
|
69
|
+
"Use agent.run_stream(), agent.iter(), or ya_agent_sdk.stream_agent()."
|
|
70
|
+
)
|
|
71
|
+
|
|
72
|
+
@asynccontextmanager
|
|
73
|
+
async def request_stream(
|
|
74
|
+
self,
|
|
75
|
+
messages: list[ModelMessage],
|
|
76
|
+
model_settings: ModelSettings | None,
|
|
77
|
+
model_request_parameters: ModelRequestParameters,
|
|
78
|
+
run_context: Any | None = None,
|
|
79
|
+
) -> AsyncIterator[StreamedResponse]:
|
|
80
|
+
async with super().request_stream(messages, model_settings, model_request_parameters, run_context) as response:
|
|
81
|
+
yield response
|
|
82
|
+
|
|
83
|
+
|
|
84
|
+
def _codex_profile() -> OpenAIModelProfile:
|
|
85
|
+
return OpenAIModelProfile(
|
|
86
|
+
supports_tools=True,
|
|
87
|
+
supports_json_schema_output=True,
|
|
88
|
+
supports_thinking=True,
|
|
89
|
+
thinking_always_enabled=True,
|
|
90
|
+
openai_supports_reasoning=True,
|
|
91
|
+
openai_supports_encrypted_reasoning_content=True,
|
|
92
|
+
openai_supports_strict_tool_definition=True,
|
|
93
|
+
openai_responses_requires_function_call_status_none=True,
|
|
94
|
+
)
|
|
95
|
+
|
|
96
|
+
|
|
97
|
+
def build_session_headers(session_id: str | None, thread_id: str | None) -> dict[str, str]:
|
|
98
|
+
"""Build Codex session/thread headers with underscore and hyphen variants."""
|
|
99
|
+
headers: dict[str, str] = {}
|
|
100
|
+
if session_id:
|
|
101
|
+
headers["session_id"] = session_id
|
|
102
|
+
headers["session-id"] = session_id
|
|
103
|
+
if thread_id:
|
|
104
|
+
headers["thread_id"] = thread_id
|
|
105
|
+
headers["thread-id"] = thread_id
|
|
106
|
+
headers["x-client-request-id"] = thread_id
|
|
107
|
+
return headers
|
|
@@ -0,0 +1,132 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import json
|
|
4
|
+
from collections.abc import AsyncGenerator
|
|
5
|
+
from typing import Any
|
|
6
|
+
|
|
7
|
+
import httpx
|
|
8
|
+
from ya_oauth.types import OAuthAccount, OAuthTokenSource, TokenSnapshot
|
|
9
|
+
|
|
10
|
+
_RESERVED_EXTRA_HEADERS = {
|
|
11
|
+
"authorization",
|
|
12
|
+
"proxy-authorization",
|
|
13
|
+
"chatgpt-account-id",
|
|
14
|
+
"x-openai-fedramp",
|
|
15
|
+
"originator",
|
|
16
|
+
"version",
|
|
17
|
+
}
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
class OAuthBearerAuth(httpx.Auth):
|
|
21
|
+
"""httpx auth flow that attaches OAuth bearer headers and refreshes once on 401."""
|
|
22
|
+
|
|
23
|
+
requires_response_body = True
|
|
24
|
+
|
|
25
|
+
def __init__(
|
|
26
|
+
self, token_source: OAuthTokenSource, *, provider_name: str, extra_headers: dict[str, str] | None = None
|
|
27
|
+
) -> None:
|
|
28
|
+
self.token_source = token_source
|
|
29
|
+
self.provider_name = provider_name
|
|
30
|
+
self.extra_headers = _safe_extra_headers(extra_headers)
|
|
31
|
+
|
|
32
|
+
async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.Request, httpx.Response]:
|
|
33
|
+
snapshot = await self.token_source.get_token()
|
|
34
|
+
self._prepare_request(request)
|
|
35
|
+
self._apply_headers(request, snapshot)
|
|
36
|
+
response = yield request
|
|
37
|
+
if response.status_code != 401:
|
|
38
|
+
return
|
|
39
|
+
refreshed = await self.token_source.refresh_token()
|
|
40
|
+
retry = _clone_request(request)
|
|
41
|
+
self._prepare_request(retry)
|
|
42
|
+
self._apply_headers(retry, refreshed)
|
|
43
|
+
yield retry
|
|
44
|
+
|
|
45
|
+
def _prepare_request(self, request: httpx.Request) -> None:
|
|
46
|
+
if self.provider_name == "codex":
|
|
47
|
+
_ensure_codex_responses_instructions(request)
|
|
48
|
+
|
|
49
|
+
def _apply_headers(self, request: httpx.Request, snapshot: TokenSnapshot) -> None:
|
|
50
|
+
request.headers["Authorization"] = f"Bearer {snapshot.access_token}"
|
|
51
|
+
if self.provider_name == "codex":
|
|
52
|
+
request.headers.update(build_codex_headers(snapshot.account, extra_headers=self.extra_headers))
|
|
53
|
+
else:
|
|
54
|
+
request.headers.update(self.extra_headers)
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
def build_codex_headers(
|
|
58
|
+
account: OAuthAccount,
|
|
59
|
+
*,
|
|
60
|
+
extra_headers: dict[str, str] | None = None,
|
|
61
|
+
version: str | None = None,
|
|
62
|
+
) -> dict[str, str]:
|
|
63
|
+
"""Build Codex-compatible request headers."""
|
|
64
|
+
headers: dict[str, str] = {
|
|
65
|
+
"originator": "ya_agent_sdk",
|
|
66
|
+
"version": version or _sdk_version(),
|
|
67
|
+
}
|
|
68
|
+
if account.chatgpt_account_id:
|
|
69
|
+
headers["ChatGPT-Account-ID"] = account.chatgpt_account_id
|
|
70
|
+
if account.chatgpt_account_is_fedramp:
|
|
71
|
+
headers["X-OpenAI-Fedramp"] = "true"
|
|
72
|
+
headers.update(_safe_extra_headers(extra_headers))
|
|
73
|
+
return headers
|
|
74
|
+
|
|
75
|
+
|
|
76
|
+
def _safe_extra_headers(extra_headers: dict[str, str] | None) -> dict[str, str]:
|
|
77
|
+
safe_headers: dict[str, str] = {}
|
|
78
|
+
for key, value in (extra_headers or {}).items():
|
|
79
|
+
if key.lower() in _RESERVED_EXTRA_HEADERS:
|
|
80
|
+
raise ValueError(f"extra_headers may not override reserved OAuth/Codex header: {key}")
|
|
81
|
+
safe_headers[key] = value
|
|
82
|
+
return safe_headers
|
|
83
|
+
|
|
84
|
+
|
|
85
|
+
def _ensure_codex_responses_instructions(request: httpx.Request) -> None:
|
|
86
|
+
"""Align Codex Responses API request body requirements."""
|
|
87
|
+
if request.method.upper() != "POST" or request.url.path.rstrip("/") != "/backend-api/codex/responses":
|
|
88
|
+
return
|
|
89
|
+
|
|
90
|
+
try:
|
|
91
|
+
body = json.loads(request.content or b"{}")
|
|
92
|
+
except (json.JSONDecodeError, httpx.RequestNotRead):
|
|
93
|
+
return
|
|
94
|
+
if not isinstance(body, dict):
|
|
95
|
+
return
|
|
96
|
+
|
|
97
|
+
changed = False
|
|
98
|
+
instructions = body.get("instructions")
|
|
99
|
+
if not instructions:
|
|
100
|
+
body["instructions"] = ""
|
|
101
|
+
changed = True
|
|
102
|
+
if body.get("store") is not False:
|
|
103
|
+
body["store"] = False
|
|
104
|
+
changed = True
|
|
105
|
+
if changed:
|
|
106
|
+
_replace_json_body(request, body)
|
|
107
|
+
|
|
108
|
+
|
|
109
|
+
def _replace_json_body(request: httpx.Request, body: dict[str, Any]) -> None:
|
|
110
|
+
content = json.dumps(body, separators=(",", ":")).encode()
|
|
111
|
+
request.stream = httpx.ByteStream(content)
|
|
112
|
+
request._content = content
|
|
113
|
+
request.headers["Content-Length"] = str(len(content))
|
|
114
|
+
|
|
115
|
+
|
|
116
|
+
def _clone_request(request: httpx.Request) -> httpx.Request:
|
|
117
|
+
return httpx.Request(
|
|
118
|
+
method=request.method,
|
|
119
|
+
url=request.url,
|
|
120
|
+
headers=request.headers.copy(),
|
|
121
|
+
content=request.content,
|
|
122
|
+
extensions=request.extensions.copy(),
|
|
123
|
+
)
|
|
124
|
+
|
|
125
|
+
|
|
126
|
+
def _sdk_version() -> str:
|
|
127
|
+
try:
|
|
128
|
+
from importlib.metadata import version
|
|
129
|
+
|
|
130
|
+
return version("ya-agent-sdk")
|
|
131
|
+
except Exception:
|
|
132
|
+
return "0.0.0"
|