diff --git a/docs/docs/workers/api.md b/docs/docs/workers/api.md index a0681ba..982f701 100644 --- a/docs/docs/workers/api.md +++ b/docs/docs/workers/api.md @@ -391,6 +391,7 @@ above. Each returns Python values. | Class | Config section | Methods | | ------------------------ | --------------------------- | ----------------------------------------------------------------------------------------- | +| `AI` | `ai` | `async run(model, inputs, *, gateway, return_raw_response, extra_headers)`, `async models(**params)`, and `gateway(gateway_id)`, which returns the AI Gateway as a `Binding`. With `"stream": True` in `inputs`, `run` returns a `js.ReadableStream`. With `return_raw_response=True`, it returns a `Response`. | | `Vectorize` | `vectorize` | `async describe()`, `async query(vector, *, top_k, namespace, return_values, return_metadata, filter)`, `async query_by_id(vector_id, ...)` with the same keywords, `async insert(vectors)`, `async upsert(vectors)`, `async delete_by_ids(ids)`, `async get_by_ids(ids)`. `ids` is one id or a list. | | `AnalyticsEngineDataset` | `analytics_engine_datasets` | `write_data_point(event=None, *, indexes, doubles, blobs)`, not async. `indexes` and `blobs` must be lists, and a bare `str` or `bytes` raises `TypeError`. | | `ImagesBinding` | `images` | `async info(stream)`, and `input(stream)`, which returns an `ImageTransformer`. Chain `transform(**options)` and `draw(image, ...)`, then `await output(format=..., quality=...)`. The result has `response()`, `content_type()` and `image()`. | diff --git a/docs/docs/workers/bindings.md b/docs/docs/workers/bindings.md index 3936124..c3859a1 100644 --- a/docs/docs/workers/bindings.md +++ b/docs/docs/workers/bindings.md @@ -90,9 +90,9 @@ runtime: | `durable_objects.bindings` | [`DurableObjectNamespace`](api.md#durableobjectnamespace) | | `services`, `assets` | [`Fetcher`](api.md#fetcher) | | `workflows` | [`Workflow`](api.md#workflow) | -| `vectorize`, `analytics_engine_datasets`, `images`, `ratelimits` | `Vectorize`, `AnalyticsEngineDataset`, `ImagesBinding`, `RateLimit`. See [Other binding classes](api.md#other-binding-classes). | +| `ai`, `vectorize`, `analytics_engine_datasets`, `images`, `ratelimits` | `AI`, `Vectorize`, `AnalyticsEngineDataset`, `ImagesBinding`, `RateLimit`. See [Other binding classes](api.md#other-binding-classes). | | `vars` and secrets | `str`, or a Python value for a JSON var | -| Anything else, such as `ai` | [`Binding`](api.md#binding) | +| Anything else, such as `browser` | [`Binding`](api.md#binding) | Run `uvx wynd types` after you change the config, so `Env` in `src/env.py` lists the new bindings. `wynd dev` and `wynd deploy` also regenerate it. @@ -258,6 +258,27 @@ class Default(WorkerEntrypoint[Env]): `TypeError: Default.queue() takes 2 positional arguments but 4 were given` in the `wynd dev` log, and the runtime delivers the batch again. +## Workers AI + +Add `"ai": { "binding": "AI" }` to the config. `AI.run` takes the model name +and a `dict` of inputs, and returns the model's output as a `dict`: + +```python +out = await self.env.AI.run( + "@cf/meta/llama-4-scout-17b-16e-instruct", + {"messages": [{"role": "user", "content": "Hello"}]}, +) +out["response"] +``` + +With `"stream": True` in the inputs, `run` returns a `js.ReadableStream` of +server-sent events. Return it in a `Response` to stream it to the client: + +```python +stream = await self.env.AI.run(model, {"messages": messages, "stream": True}) +return Response(stream, headers={"content-type": "text/event-stream"}) +``` + ## Service bindings A service binding is a [`Fetcher`](api.md#fetcher). Its `fetch` sends an HTTP diff --git a/docs/docs/wynd/types.md b/docs/docs/wynd/types.md index c5f600a..414c35f 100644 --- a/docs/docs/wynd/types.md +++ b/docs/docs/wynd/types.md @@ -152,8 +152,9 @@ runtime. | `analytics_engine_datasets` | `AnalyticsEngineDataset` | | `images` | `ImagesBinding` | | `ratelimits` | `RateLimit` | +| `ai` | `AI` | | `version_metadata` | `JsDict` | -| Any other binding, such as `ai` or `browser` | `Binding` | +| Any other binding, such as `browser` | `Binding` | | A text var, and each name in `.dev.vars` | `str` | | A JSON var | The type of its value, such as `int`, `list[str]` or `dict[str, int]` | diff --git a/workers/src/workers/bindings/__init__.py b/workers/src/workers/bindings/__init__.py index ec4c6ad..034c940 100644 --- a/workers/src/workers/bindings/__init__.py +++ b/workers/src/workers/bindings/__init__.py @@ -7,6 +7,7 @@ keep working. """ +from .ai import AI from .analytics import AnalyticsEngineDataset from .base import Binding from .base import _convert as _convert @@ -24,6 +25,7 @@ from .workflow import Workflow, WorkflowInstance __all__ = ( + "AI", "AnalyticsEngineDataset", "Binding", "D1Database", diff --git a/workers/src/workers/bindings/ai.py b/workers/src/workers/bindings/ai.py new file mode 100644 index 0000000..386d731 --- /dev/null +++ b/workers/src/workers/bindings/ai.py @@ -0,0 +1,37 @@ +from typing import Any + +from ..rpc import python_from_rpc, python_to_rpc +from .base import Binding, _js_options + +__all__ = ("AI",) + + +class AI(Binding): + async def run( + self, + model: str, + inputs: dict[str, Any], + options: Any = None, + /, + *, + gateway: dict[str, Any] | None = None, + return_raw_response: bool | None = None, + extra_headers: dict[str, str] | None = None, + **extra: Any, + ) -> Any: + js_options = _js_options( + options, + extra, + gateway=gateway, + return_raw_response=return_raw_response, + extra_headers=extra_headers, + ) + result = await self.js_object.run(model, python_to_rpc(inputs), js_options) + # A stream stays a js.ReadableStream and a raw response becomes a Response. + return python_from_rpc(result) + + async def models(self, params: Any = None, /, **extra: Any) -> list[dict[str, Any]]: + return python_from_rpc(await self.js_object.models(_js_options(params, extra))) + + def gateway(self, gateway_id: str) -> Binding: + return Binding(self.js_object.gateway(gateway_id)) diff --git a/workers/src/workers/bindings/env.py b/workers/src/workers/bindings/env.py index c7f60b4..38af23a 100644 --- a/workers/src/workers/bindings/env.py +++ b/workers/src/workers/bindings/env.py @@ -5,6 +5,7 @@ from .._js import js_name from ..rpc import python_from_rpc +from .ai import AI from .analytics import AnalyticsEngineDataset from .base import Binding from .d1 import D1Database @@ -33,6 +34,7 @@ "LocalAnalyticsEngineDataset": AnalyticsEngineDataset, "ImagesBindingImpl": ImagesBinding, "Ratelimit": RateLimit, + "Ai": AI, } diff --git a/workers/tests/test_bindings.py b/workers/tests/test_bindings.py index 366ac1e..6e41f72 100644 --- a/workers/tests/test_bindings.py +++ b/workers/tests/test_bindings.py @@ -5,6 +5,7 @@ from workers import DurableObject, Response, WorkerEntrypoint, python_to_rpc from workers.bindings import ( + AI, AnalyticsEngineDataset, Binding, D1Database, @@ -59,7 +60,7 @@ def test_env_wraps_each_binding_by_its_js_class() -> None: assert isinstance(env.LOCAL_EVENTS, AnalyticsEngineDataset) assert isinstance(env.IMAGES, ImagesBinding) assert isinstance(env.LIMITER, RateLimit) - assert type(env.AI) is Binding + assert isinstance(env.AI, AI) assert type(env.HYPERDRIVE) is Binding assert env.TEXT == "hello" assert env.FLAGS == {"beta": True} @@ -176,6 +177,33 @@ def test_vectorize_sends_camel_case_query_options_and_lists() -> None: ] +def test_ai_run_converts_inputs_and_sends_camel_case_options() -> None: + js_ai = JsBinding("Ai", returns=JsValue("Object", {"response": "hi"})) + ai = AI(js_ai) + messages = [{"role": "user", "content": "hello"}] + + result = run( + ai.run( + "@cf/model", + {"messages": messages}, + gateway={"id": "g"}, + return_raw_response=False, + ) + ) + + assert result == {"response": "hi"} + assert js_ai.calls == [ + ( + "run", + ( + "@cf/model", + {"messages": messages}, + {"gateway": {"id": "g"}, "returnRawResponse": False}, + ), + ) + ] + + def test_analytics_engine_writes_a_data_point_synchronously() -> None: js_dataset = JsBinding("AnalyticsEngineDataset", sync={"writeDataPoint": None}) diff --git a/wynd/src/wynd/typegen.py b/wynd/src/wynd/typegen.py index 64a7dfe..50c309d 100644 --- a/wynd/src/wynd/typegen.py +++ b/wynd/src/wynd/typegen.py @@ -28,6 +28,7 @@ "analytics_engine_datasets": "AnalyticsEngineDataset", "images": "ImagesBinding", "ratelimits": "RateLimit", + "ai": "AI", } DEV_VARS_KEY = re.compile(r"^\s*(?:export\s+)?([A-Za-z_]\w*)\s*=", re.MULTILINE) diff --git a/wynd/tests/test_types.py b/wynd/tests/test_types.py index 430fd82..b87455d 100644 --- a/wynd/tests/test_types.py +++ b/wynd/tests/test_types.py @@ -43,8 +43,8 @@ from typing import TYPE_CHECKING, Protocol from workers.bindings import ( + AI, AnalyticsEngineDataset, - Binding, D1Database, DurableObjectNamespace, Fetcher, @@ -72,7 +72,7 @@ class Env(Protocol): AUTH: Fetcher ASSETS: Fetcher FLOW: Workflow - AI: Binding + AI: AI LIMITER: RateLimit INDEX: Vectorize EVENTS: AnalyticsEngineDataset