openai/openai-python
Publicmirrored from https://github.com/openai/openai-pythonAvailable
scripts/utils/validate-bedrock-wheel.py
123lines · modecode
| 1 | from __future__ import annotations |
| 2 | |
| 3 | import os |
| 4 | import sys |
| 5 | import email |
| 6 | import zipfile |
| 7 | import tempfile |
| 8 | import subprocess |
| 9 | from pathlib import Path |
| 10 | |
| 11 | _SMOKE_TEST = r""" |
| 12 | import os |
| 13 | import sys |
| 14 | import importlib.abc |
| 15 | from pathlib import Path |
| 16 | |
| 17 | import httpx |
| 18 | |
| 19 | |
| 20 | class BlockBotocore(importlib.abc.MetaPathFinder): |
| 21 | def find_spec(self, fullname, path=None, target=None): |
| 22 | if fullname == "botocore" or fullname.startswith("botocore."): |
| 23 | raise ImportError(f"unexpected AWS import: {fullname}") |
| 24 | return None |
| 25 | |
| 26 | |
| 27 | wheel_root = Path(os.environ["OPENAI_WHEEL_ROOT"]).resolve() |
| 28 | blocker = BlockBotocore() |
| 29 | sys.meta_path.insert(0, blocker) |
| 30 | |
| 31 | import openai |
| 32 | from openai import OpenAI |
| 33 | from openai.providers import bedrock |
| 34 | |
| 35 | assert Path(openai.__file__).resolve().is_relative_to(wheel_root), openai.__file__ |
| 36 | |
| 37 | requests = [] |
| 38 | |
| 39 | |
| 40 | def handler(request): |
| 41 | requests.append(request) |
| 42 | return httpx.Response(200, request=request, json={}) |
| 43 | |
| 44 | |
| 45 | http_client = httpx.Client(transport=httpx.MockTransport(handler), trust_env=False) |
| 46 | with OpenAI( |
| 47 | provider=bedrock(region="us-east-1", api_key="bearer-token"), |
| 48 | http_client=http_client, |
| 49 | ) as client: |
| 50 | client.get("/models", cast_to=httpx.Response) |
| 51 | |
| 52 | assert requests[0].headers["Authorization"] == "Bearer bearer-token" |
| 53 | assert not any(name == "botocore" or name.startswith("botocore.") for name in sys.modules) |
| 54 | |
| 55 | sys.meta_path.remove(blocker) |
| 56 | requests.clear() |
| 57 | http_client = httpx.Client(transport=httpx.MockTransport(handler), trust_env=False) |
| 58 | with OpenAI( |
| 59 | provider=bedrock( |
| 60 | region="us-east-1", |
| 61 | access_key_id="fixture-access-key", |
| 62 | secret_access_key="fixture-secret-key", |
| 63 | session_token="fixture-session-token", |
| 64 | ), |
| 65 | http_client=http_client, |
| 66 | ) as client: |
| 67 | client.get("/models", cast_to=httpx.Response) |
| 68 | |
| 69 | assert "Credential=fixture-access-key/" in requests[0].headers["Authorization"] |
| 70 | assert requests[0].headers["X-Amz-Security-Token"] == "fixture-session-token" |
| 71 | """ |
| 72 | |
| 73 | |
| 74 | def main() -> None: |
| 75 | wheels = list(Path("dist").glob("*.whl")) |
| 76 | if len(wheels) != 1: |
| 77 | raise RuntimeError(f"Expected exactly one wheel in dist/, found: {wheels}") |
| 78 | |
| 79 | wheel = wheels[0] |
| 80 | with tempfile.TemporaryDirectory() as directory: |
| 81 | wheel_root = Path(directory) |
| 82 | with zipfile.ZipFile(wheel) as archive: |
| 83 | metadata_names = [name for name in archive.namelist() if name.endswith(".dist-info/METADATA")] |
| 84 | if len(metadata_names) != 1: |
| 85 | raise RuntimeError(f"Expected exactly one METADATA file in {wheel}, found: {metadata_names}") |
| 86 | |
| 87 | metadata = email.message_from_bytes(archive.read(metadata_names[0])) |
| 88 | archive.extractall(wheel_root) |
| 89 | |
| 90 | requirements = metadata.get_all("Requires-Dist", []) |
| 91 | botocore_requirements = [requirement for requirement in requirements if requirement.startswith("botocore")] |
| 92 | if len(botocore_requirements) != 2: |
| 93 | raise RuntimeError(f"Expected two Python-version-specific botocore requirements: {botocore_requirements}") |
| 94 | if any("[crt]" in requirement for requirement in botocore_requirements): |
| 95 | raise RuntimeError( |
| 96 | f"The Bedrock extra must not install the unused botocore CRT extra: {botocore_requirements}" |
| 97 | ) |
| 98 | if not all("extra == 'bedrock'" in requirement for requirement in botocore_requirements): |
| 99 | raise RuntimeError(f"Botocore requirements must belong to the Bedrock extra: {botocore_requirements}") |
| 100 | |
| 101 | environment = os.environ.copy() |
| 102 | for name in ( |
| 103 | "AWS_ACCESS_KEY_ID", |
| 104 | "AWS_SECRET_ACCESS_KEY", |
| 105 | "AWS_SESSION_TOKEN", |
| 106 | "AWS_PROFILE", |
| 107 | "AWS_REGION", |
| 108 | "AWS_DEFAULT_REGION", |
| 109 | "AWS_BEARER_TOKEN_BEDROCK", |
| 110 | "AWS_BEDROCK_BASE_URL", |
| 111 | "OPENAI_API_KEY", |
| 112 | "OPENAI_ORG_ID", |
| 113 | "OPENAI_PROJECT_ID", |
| 114 | "OPENAI_CUSTOM_HEADERS", |
| 115 | ): |
| 116 | environment.pop(name, None) |
| 117 | environment["OPENAI_WHEEL_ROOT"] = str(wheel_root) |
| 118 | environment["PYTHONPATH"] = str(wheel_root) |
| 119 | subprocess.run([sys.executable, "-c", _SMOKE_TEST], cwd=wheel_root, env=environment, check=True) |
| 120 | |
| 121 | |
| 122 | if __name__ == "__main__": |
| 123 | main() |
| 124 | |