support installing openai, and proxying requests (#44)

src/frontend/src/default_code.py

@@ -0,0 +1 @@
+print('hello world')

src/frontend/src/install_dependencies.py → src/frontend/src/prepare_env.py

@@ -18,17 +18,15 @@
 from urllib.parse import urlparse
 
 import tomllib
-from packaging.tags import parse_tag  # noqa
-from packaging.version import Version  # noqa
 
 import micropip  # noqa
 from micropip import transaction  # noqa
-from micropip.wheelinfo import WheelInfo  # noqa
+from micropip.wheelinfo import WheelInfo, Tag, Version  # noqa
 
 from pyodide.code import find_imports  # noqa
 import pyodide_js  # noqa
 
-__all__ = ('install_deps',)
+__all__ = ('prepare_env',)
 
 
 class File(TypedDict):
@@ -49,35 +47,36 @@ class Error:
     kind: Literal['error'] = 'error'
 
 
-# This is a temporary hack to install jiter from a URL until
-# https://github.com/pyodide/pyodide/pull/5388 is released.
-real_find_wheel = transaction.find_wheel
+async def prepare_env(files: list[File]) -> Success | Error:
+    # This is a temporary hack to install jiter from a URL until
+    # https://github.com/pyodide/pyodide/pull/5388 is released.
+    real_find_wheel = transaction.find_wheel
+
+    def custom_find_wheel(metadata: Any, req: Any) -> Any:
+        if metadata.name == 'jiter':
+            known_version = Version('0.8.2')
+            if known_version in metadata.releases:
+                tag = Tag('cp312', 'cp312', 'emscripten_3_1_58_wasm32')
+                filename = f'{metadata.name}-{known_version}-{tag}.whl'
+                url = f'https://files.pydantic.run/{filename}'
+                return WheelInfo(
+                    name=metadata.name,
+                    version=known_version,
+                    filename=filename,
+                    build=(),
+                    tags=frozenset({tag}),
+                    url=url,
+                    parsed_url=urlparse(url),
+                )
+        return real_find_wheel(metadata, req)
+
+    transaction.find_wheel = custom_find_wheel
+    # end `transaction.find_wheel` hack
 
+    sys.setrecursionlimit(400)
 
-def custom_find_wheel(metadata: Any, req: Any) -> Any:
-    if metadata.name == 'jiter':
-        known_version = Version('0.8.2')
-        if known_version in metadata.releases:
-            tag = 'cp312-cp312-emscripten_3_1_58_wasm32'
-            filename = f'{metadata.name}-{known_version}-{tag}.whl'
-            url = f'https://files.pydantic.run/{filename}'
-            return WheelInfo(
-                name=metadata.name,
-                version=known_version,
-                filename=filename,
-                build=(),
-                tags=frozenset({parse_tag(tag)}),
-                url=url,
-                parsed_url=urlparse(url),
-            )
-    return real_find_wheel(metadata, req)
-
-
-transaction.find_wheel = custom_find_wheel
-
+    os.environ.update(OPENAI_BASE_URL='https://proxy.pydantic.run/proxy/openai', OPENAI_API_KEY='proxy-key')
 
-async def install_deps(files: list[File]) -> Success | Error:
-    sys.setrecursionlimit(400)
     cwd = Path.cwd()
     for file in files:
         (cwd / file['name']).write_text(file['content'])
@@ -102,22 +101,7 @@ async def install_deps(files: list[File]) -> Success | Error:
             dependencies = await _find_import_dependencies(python_code)
 
     if dependencies:
-        # pygments seems to be required to get rich to work properly, ssl is required for FastAPI and HTTPX
-        install_pygments = False
-        install_ssl = False
-        for d in dependencies:
-            if d.startswith(('logfire', 'rich')):
-                install_pygments = True
-            elif d.startswith(('fastapi', 'httpx')):
-                install_ssl = True
-            if install_pygments and install_ssl:
-                break
-
-        install_dependencies = dependencies.copy()
-        if install_pygments:
-            install_dependencies.append('pygments')
-        if install_ssl:
-            install_dependencies.append('ssl')
+        install_dependencies = _add_extra_dependencies(dependencies)
 
         with _micropip_logging() as logs_filename:
             try:
@@ -128,9 +112,57 @@ async def install_deps(files: list[File]) -> Success | Error:
                     logs = f.read()
                 return Error(message=f'{logs}\n{traceback.format_exc()}')
 
+    # temporary hack until the debug prints in https://github.com/encode/httpx/pull/3330 are used/merged
+    try:
+        from httpx import AsyncClient
+    except ImportError:
+        pass
+    else:
+        original_send = AsyncClient.send
+
+        def print_monkeypatch(*args, **kwargs):
+            pass
+
+        async def send_monkeypatch_print(self, *args, **kwargs):
+            import builtins
+            original_print = builtins.print
+            builtins.print = print_monkeypatch
+            try:
+                return await original_send(self, *args, **kwargs)
+            finally:
+                builtins.print = original_print
+
+        AsyncClient.send = send_monkeypatch_print
+    # end temporary hack for httpx debug prints
+
     return Success(message=', '.join(dependencies))
 
 
+def _add_extra_dependencies(dependencies: list[str]) -> list[str]:
+    """Add extra dependencies we know some packages need.
+
+    Workaround for micropip not installing some required transitive dependencies.
+    See https://github.com/pyodide/micropip/issues/204
+
+    pygments seems to be required to get rich to work properly, ssl is required for FastAPI and HTTPX,
+    pydantic_ai requires newest typing_extensions.
+    """
+    extras = []
+    for d in dependencies:
+        if d.startswith(('logfire', 'rich')):
+            extras.append('pygments')
+        elif d.startswith(('fastapi', 'httpx', 'pydantic_ai')):
+            extras.append('ssl')
+
+        if d.startswith('pydantic_ai'):
+            extras.append('typing_extensions>=4.12')
+
+        if len(extras) == 3:
+            break
+
+    return dependencies + extras
+
+
 @contextmanager
 def _micropip_logging() -> Iterable[str]:
     from micropip import logging as micropip_logging  # noqa

src/frontend/src/store.ts

@@ -1,4 +1,5 @@
 import type { CodeFile } from './types'
+import defaultPythonCode from './default_code.py?raw'
 
 interface StoreHttpResponse {
   readKey: string
@@ -80,7 +81,7 @@ export async function retrieve(): Promise<InitialState> {
   const files = [
     {
       name: 'main.py',
-      content: `print('hello world')`,
+      content: defaultPythonCode,
       activeIndex: 0,
     },
   ]

src/frontend/src/worker.ts

@@ -1,13 +1,13 @@
 /* eslint @typescript-eslint/no-explicit-any: off */
 import { loadPyodide, PyodideInterface, version as pyodideVersion } from 'pyodide'
-import installPythonCode from './install_dependencies.py?raw'
+import preparePythonEnvCode from './prepare_env.py?raw'
 import type { CodeFile, RunCode, WorkerResponse } from './types'
 
-interface InstallSuccess {
+interface PrepareSuccess {
   kind: 'success'
   message: string
 }
-interface InstallError {
+interface PrepareError {
   kind: 'error'
   message: string
 }
@@ -16,24 +16,22 @@ self.onmessage = async ({ data }: { data: RunCode }) => {
   const { files } = data
   let msg = ''
   try {
-    const [setupTime, { pyodide, installDeps }] = await time(getPyodideEnv())
+    const [setupTime, { pyodide, preparePyEnv }] = await time(getPyodideEnv())
     if (setupTime > 50) {
       msg += `Started Python in ${asMs(setupTime)}, `
     }
     post({ kind: 'status', message: `${msg}Installing dependencies…` })
     const sys = pyodide.pyimport('sys')
 
-    const [installTime, installStatus]: [number, InstallSuccess | InstallError] = await time(
-      installDeps.install_deps(pyodide.toPy(files)),
-    )
+    const [installTime, prepareStatus] = await time(preparePyEnv.prepare_env(pyodide.toPy(files)))
     sys.stdout.flush()
     sys.stderr.flush()
-    if (installStatus.kind == 'error') {
+    if (prepareStatus.kind == 'error') {
       post({ kind: 'status', message: `${msg}Error occurred` })
-      post({ kind: 'error', message: installStatus.message })
+      post({ kind: 'error', message: prepareStatus.message })
       return
     }
-    post({ kind: 'installed', message: installStatus.message })
+    post({ kind: 'installed', message: prepareStatus.message })
     if (installTime > 50) {
       msg += `Installed dependencies in ${asMs(installTime)}, `
     }
@@ -87,14 +85,19 @@ async function time<T>(promise: Promise<T>): Promise<[number, T]> {
 
 interface PyodideEnv {
   pyodide: PyodideInterface
-  installDeps: any
+  // matches the signature of the `prepare_env` function in prepare_env.py
+  preparePyEnv: { prepare_env: (files: any) => Promise<PrepareSuccess | PrepareError> }
 }
 
+// see https://github.com/pyodide/micropip/issues/201
+const micropipV09 =
+  'https://files.pythonhosted.org/packages/27/6d/195810e3e73e5f351dc6082cada41bb4d5b0746a6804155ba6bae4304612/micropip-0.9.0-py3-none-any.whl'
+
 // we rerun this on every invocation to avoid issues with conflicting packages
 async function getPyodideEnv(): Promise<PyodideEnv> {
   const pyodide = await loadPyodide({
     indexURL: `https://cdn.jsdelivr.net/pyodide/v${pyodideVersion}/full/`,
-    packages: ['micropip'],
+    packages: [micropipV09],
   })
   const sys = pyodide.pyimport('sys')
   const pv = sys.version_info
@@ -109,12 +112,12 @@ async function getPyodideEnv(): Promise<PyodideEnv> {
   sys.path.append(dirPath)
   const pathlib = pyodide.pyimport('pathlib')
   pathlib.Path(dirPath).mkdir()
-  const moduleName = '_install_dependencies'
-  pathlib.Path(`${dirPath}/${moduleName}.py`).write_text(installPythonCode)
+  const moduleName = '_prepare_env'
+  pathlib.Path(`${dirPath}/${moduleName}.py`).write_text(preparePythonEnvCode)
 
   return {
     pyodide,
-    installDeps: pyodide.pyimport(moduleName),
+    preparePyEnv: pyodide.pyimport(moduleName),
   }
 }