cortexgrid-infer 0.3.8__tar.gz → 0.3.9__tar.gz

{cortexgrid_infer-0.3.8 → cortexgrid_infer-0.3.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: cortexgrid-infer
-Version: 0.3.8
+Version: 0.3.9
 Summary: Serve apps per task (text-to-text, text-to-image, image-to-mesh) and importers per source (HuggingFace) for running models on cortexgrid
 Project-URL: Homepage, https://github.com/robodatalab/cortexgrid-infer
 Project-URL: Repository, https://github.com/robodatalab/cortexgrid-infer
@@ -9,7 +9,7 @@ License-File: LICENSE
 Requires-Python: <3.12,>=3.11
 Requires-Dist: accelerate>=0.33
 Requires-Dist: anthropic>=0.40.0
-Requires-Dist: cortexgrid>=0.3.16
+Requires-Dist: cortexgrid>=0.3.18
 Requires-Dist: diffusers>=0.38
 Requires-Dist: fastapi>=0.133
 Requires-Dist: google-genai>=1.0
@@ -23,6 +23,7 @@ Requires-Dist: sentencepiece>=0.2
 Requires-Dist: torch>=2.0.0
 Requires-Dist: transformers>=4.47.0
 Provides-Extra: dev
+Requires-Dist: parameterized>=0.9.0; extra == 'dev'
 Requires-Dist: pytest>=8.0; extra == 'dev'
 Description-Content-Type: text/markdown
 
@@ -50,23 +51,24 @@ Upgrading from 0.2.x? See [MIGRATION.md](https://github.com/robodatalab/cortexgr
 A model is two things, and the library keeps them apart:
 
 - **What runs it** — a serve app in `cortexgrid_infer.serve_apps`, named for the task
-  it runs, not for where the model came from. It loads the weights, answers the task's
-  routes (defined in `cortexgrid_infer.protocols`), and names the client that speaks
-  them. One serve app runs every model of its task, whichever source staged it.
+  it runs, not for where the model came from. It loads the weights and serves the
+  task's endpoints, declared by the task's type (`CompletingModel`, say); cortexgrid
+  generates the client that calls them. One serve app runs every model of its task,
+  whichever source staged it.
 - **Where its weights come from** — an importer in `cortexgrid_infer.importers`, one
   per source. It names the model in the registry, downloads the weights, and sizes
   them from the source's metadata. It is handed the serve app, so one importer takes
   any model its source holds.
 
-| Serve app | Task | Runs | Client |
+| Serve app | Task | Runs | Task type |
 |---|---|---|---|
-| `Text2Text` | text to text | any causal LM in the transformers layout, via `AutoModelForCausalLM` | `ServedCompletingModel` (`complete`) |
-| `Text2Image` | text (and image) to image | any diffusers pipeline | `ServedGeneratingModel` (`generate`) |
-| `Image2Mesh` — a base: subclass it with the model's `load` and `make_mesh` | image to mesh | the model the subclass loads, behind the task's `POST /mesh` | `ServedMeshingModel` (`mesh`: one picture in, a `GeneratedMesh` out) |
-| `TextRewriter` | text to text, rewritten | any encoder-decoder LM (T5, BART, Marian, ...) in the transformers layout, via `AutoModelForSeq2SeqLM` | `ServedRewritingModel` (`rewrite`: one text in, one text out) |
-| `AnthropicText2Text` | text to text | the Anthropic API, forwarded from the cluster | `ServedCompletingModel` (`complete`) |
-| `GeminiText2Text` | text to text | the Gemini API, forwarded from the cluster | `ServedCompletingModel` (`complete`) |
-| `GeminiText2Image` | text to image | the Gemini API, forwarded from the cluster | `ServedGeneratingModel` (`generate`) |
+| `Text2Text` | text to text | any causal LM in the transformers layout, via `AutoModelForCausalLM` | `CompletingModel` (`complete`), plus its own `loglikelihoods` and `last_hidden_states` |
+| `Text2Image` | text (and image) to image | any diffusers pipeline | `GeneratingModel` (`generate`) |
+| `Image2Mesh` — a base: subclass it with the model's `load` and `make_mesh` | image to mesh | the model the subclass loads, behind the task's `mesh` endpoint | `MeshingModel` (`mesh`: one picture in, a `GeneratedMesh` out) |
+| `TextRewriter` | text to text, rewritten | any encoder-decoder LM (T5, BART, Marian, ...) in the transformers layout, via `AutoModelForSeq2SeqLM` | `RewritingModel` (`rewrite`: one text in, one text out) |
+| `AnthropicText2Text` | text to text | the Anthropic API, forwarded from the cluster | `CompletingModel` (`complete`) |
+| `GeminiText2Text` | text to text | the Gemini API, forwarded from the cluster | `CompletingModel` (`complete`) |
+| `GeminiText2Image` | text to image | the Gemini API, forwarded from the cluster | `GeneratingModel` (`generate`) |
 
 | Importer | Source |
 |---|---|
@@ -89,8 +91,9 @@ need:
 
 One with weights goes through `cortexgrid.import_model`, one without through
 `cortexgrid.register_model`. Once deployed, the `cortexgrid.Deployment` hands back
-the client that speaks the app's routes: `deployment.client()` once the app serves,
-or `deployment.client_async()` at once, with `await model.is_ready()` to ask.
+the client that calls the app's endpoints, typed as the serve app and so as its task
+type: `deployment.client()` once the app serves, or `deployment.client_async()` at
+once, with `await model.is_ready()` to ask.
 
 ## Quick start
 
@@ -122,7 +125,7 @@ cortexgrid.remote(
 deployment = cortexgrid.deploy_model(imp.family, imp.suffix, cortexgrid.IMPORTED, timeout=1800)
 
 # 3. Inference — the client blocks until the Ray Serve app is running.
-model = deployment.client()
+model: mg.GeneratingModel = deployment.client()
 
 async def run():
     result = await mg.generate(model, "a red bicycle on a beach at sunrise")
@@ -277,10 +280,13 @@ it is the floor; only quantizing the weights or running fewer forwards moves it.
 
 ## Inference
 
-`complete` streams `CompletionChunk`s (`.content`, `.thinking`, `.tool_calls`, `.finish_reason`):
+`complete` streams `CompletionChunk`s (`.content`, `.thinking`, `.tool_calls`, `.finish_reason`).
+`tools` are Python functions or tool specs; a function goes to the model as its spec,
+and each `ToolCall` comes back as data (`.id`, `.name`, `.arguments`) for the caller
+to run:
 
 ```python
-model = deployment.client()
+model: mg.CompletingModel = deployment.client()
 messages = [{"role": "user", "content": "Explain RAG in one sentence."}]
 
 async for chunk in mg.complete(model, messages, max_new_tokens=512, temperature=0.7):
@@ -321,9 +327,9 @@ a re-deploy rather than a re-registration.
 
 A replica asks for no hardware at all, so it is placed on any node, CPU-only
 included. From there it deploys and streams exactly like a cluster-served model:
-`deployment.client()` returns the same `ServedCompletingModel`. Anthropic
+`deployment.client()` is a `CompletingModel` like `Text2Text`'s. Anthropic
 reports tool calls as structured blocks rather than as generated text, so the
-serve app re-encodes them into the text form the client parses.
+serve app turns them into `ToolCall`s as they come.
 
 ### Gemini
 
@@ -345,16 +351,14 @@ cortexgrid.register_model(
 Everything said above for Anthropic holds: `config` carries the model name and
 the name of the secret (`api_key_secret="..."` to use another), both editable on
 the model card; a replica needs no hardware; and `deployment.client()`
-is a `ServedCompletingModel`, with Gemini's function calls re-encoded into the
-text form the client parses.
+is a `CompletingModel`, with Gemini's function calls turned into `ToolCall`s.
 
-`GeminiText2Image` answers the same `/generate` route as `Text2Image`, so its
-client is a `ServedGeneratingModel`:
+`GeminiText2Image` is a `GeneratingModel` like `Text2Image`:
 
 ```python
 entry = mg.Hosted("gemini-2.5-flash-image", mg.GeminiText2Image)
 ...
-model = deployment.client()
+model: mg.GeneratingModel = deployment.client()
 picture = await mg.generate(model, "a fox in the snow", size=2048)
 ```
 
@@ -373,14 +377,11 @@ diffusion settings with no Gemini counterpart, and come back as `None`.
 | `Text2Text`, `Text2Image`, `TextRewriter`, `Image2Mesh` | Serve apps that run weights (`LocalModel`s). |
 | `AnthropicText2Text` | Serve app that forwards to the Anthropic API (a `HostedModel`). |
 | `GeminiText2Text`, `GeminiText2Image` | Serve apps that forward to the Gemini API (`HostedModel`s sharing the `GeminiModel` base). |
-| `LocalModel` | Base of the serve apps that run weights: `ignore_patterns()`, `bytes_per_param`, `min_vram_gb`, `requirements(weights)`, `config()` (the model card's defaults, `compile: "false"` among them), `client(deployment)`. |
-| `HostedModel` | Base of the serve apps that forward: `config(model_id, **settings)`, `requirements()`, `client(deployment)`. |
+| `LocalModel` | Base of the serve apps that run weights: `ignore_patterns()`, `bytes_per_param`, `min_vram_gb`, `requirements(weights)`, `config()` (the model card's defaults, `compile: "false"` among them). |
+| `HostedModel` | Base of the serve apps that forward: `config(model_id, **settings)`, `requirements()`. |
 | `Weights(params=None, file_bytes=0)` | What an importer reads about the weights, handed to the serve app to size a replica. |
-| `complete(model, messages, tools=None, max_new_tokens=2048, temperature=0.7, **kw)` | Async stream of `CompletionChunk` for a `CompletingModel`. |
-| `generate(model, prompt, *, image=None, **kw) -> GeneratedImage` | One image from a `GeneratingModel`. |
-| `ServedCompletingModel(key, url)` | Client for any deployed completion app. |
-| `ServedGeneratingModel(key, url)` | Client for any deployed text-to-image app. |
-| `ServedMeshingModel(key, url)` | Client for any deployed image-to-mesh app. |
+| `complete(model, messages, tools=None, max_new_tokens=None, temperature=0.7)` | Async stream of `CompletionChunk` for a `CompletingModel`. |
+| `generate(model, prompt, *, image=None, steps=None, guidance=None, size=1024, seed=None) -> GeneratedImage` | One image from a `GeneratingModel`. |
 | `split_model_id(model_id) -> (family, suffix)` | The registry identity an importer derives. |
 | `detect_device()` | The torch device a serve app should load onto. |
 | `compiling.Mode.GRAPHED` / `compiling.Mode.FUSED` | The two modes a serve app chooses between. See [Compilation](#compilation). |
@@ -388,19 +389,19 @@ diffusion settings with no Gemini counterpart, and come back as `None`.
 | `ModelDeployFailed` | Re-exported from cortexgrid; subclasses `RuntimeError`. |
 
 Types: `CompletingModel`, `GeneratingModel`, `MeshingModel`, `RewritingModel` (each
-a `cortexgrid.DeploymentClient`), `CompletionChunk`, `GeneratedImage`, `GeneratedMesh`, `ToolCall`, `Message`, `Tool`,
-`ToolSpec`.
+the abstract base of its task's serve apps, declaring the task's endpoints),
+`CompletionChunk`, `GeneratedImage`, `GeneratedMesh`, `ToolCall`, `Message`, `Tool`,
+`ToolSpec`, `Tensor`.
 
 Extending it:
 
 - **A new source** is one importer: subclass `Importer` with `download(local_dir)`
   and `weights()`, and every serve app whose weight format it delivers runs its
   models unchanged.
-- **A new hosted API** is one serve app: subclass `HostedModel`, and register it
-  with `Hosted`. A text-to-text one only has to stream text from `POST /complete`,
-  inlining tool calls as `<tool_call>{"name": ..., "arguments": {...}}</tool_call>`
-  (there is an `encode_tool_call` for upstreams that report them structurally) —
-  then `ServedCompletingModel` is its client, unchanged.
+- **A new hosted API** is one serve app: subclass `HostedModel` and the task's type,
+  and register it with `Hosted`. A text-to-text one subclasses `CompletingModel` and
+  streams `CompletionChunk`s from its `complete` endpoint, with
+  `StreamedCompletionWriter` turning the upstream's text and tool calls into them.
 - **A model no task's serve app can load** gets its own serve app — an `Image2Mesh`
   subclass, say, with `load` and `make_mesh`, and `compile` if compiling it pays —
   and, if its source is unusual too, its own importer.

{cortexgrid_infer-0.3.8 → cortexgrid_infer-0.3.9}/README.md

@@ -22,23 +22,24 @@ Upgrading from 0.2.x? See [MIGRATION.md](https://github.com/robodatalab/cortexgr
 A model is two things, and the library keeps them apart:
 
 - **What runs it** — a serve app in `cortexgrid_infer.serve_apps`, named for the task
-  it runs, not for where the model came from. It loads the weights, answers the task's
-  routes (defined in `cortexgrid_infer.protocols`), and names the client that speaks
-  them. One serve app runs every model of its task, whichever source staged it.
+  it runs, not for where the model came from. It loads the weights and serves the
+  task's endpoints, declared by the task's type (`CompletingModel`, say); cortexgrid
+  generates the client that calls them. One serve app runs every model of its task,
+  whichever source staged it.
 - **Where its weights come from** — an importer in `cortexgrid_infer.importers`, one
   per source. It names the model in the registry, downloads the weights, and sizes
   them from the source's metadata. It is handed the serve app, so one importer takes
   any model its source holds.
 
-| Serve app | Task | Runs | Client |
+| Serve app | Task | Runs | Task type |
 |---|---|---|---|
-| `Text2Text` | text to text | any causal LM in the transformers layout, via `AutoModelForCausalLM` | `ServedCompletingModel` (`complete`) |
-| `Text2Image` | text (and image) to image | any diffusers pipeline | `ServedGeneratingModel` (`generate`) |
-| `Image2Mesh` — a base: subclass it with the model's `load` and `make_mesh` | image to mesh | the model the subclass loads, behind the task's `POST /mesh` | `ServedMeshingModel` (`mesh`: one picture in, a `GeneratedMesh` out) |
-| `TextRewriter` | text to text, rewritten | any encoder-decoder LM (T5, BART, Marian, ...) in the transformers layout, via `AutoModelForSeq2SeqLM` | `ServedRewritingModel` (`rewrite`: one text in, one text out) |
-| `AnthropicText2Text` | text to text | the Anthropic API, forwarded from the cluster | `ServedCompletingModel` (`complete`) |
-| `GeminiText2Text` | text to text | the Gemini API, forwarded from the cluster | `ServedCompletingModel` (`complete`) |
-| `GeminiText2Image` | text to image | the Gemini API, forwarded from the cluster | `ServedGeneratingModel` (`generate`) |
+| `Text2Text` | text to text | any causal LM in the transformers layout, via `AutoModelForCausalLM` | `CompletingModel` (`complete`), plus its own `loglikelihoods` and `last_hidden_states` |
+| `Text2Image` | text (and image) to image | any diffusers pipeline | `GeneratingModel` (`generate`) |
+| `Image2Mesh` — a base: subclass it with the model's `load` and `make_mesh` | image to mesh | the model the subclass loads, behind the task's `mesh` endpoint | `MeshingModel` (`mesh`: one picture in, a `GeneratedMesh` out) |
+| `TextRewriter` | text to text, rewritten | any encoder-decoder LM (T5, BART, Marian, ...) in the transformers layout, via `AutoModelForSeq2SeqLM` | `RewritingModel` (`rewrite`: one text in, one text out) |
+| `AnthropicText2Text` | text to text | the Anthropic API, forwarded from the cluster | `CompletingModel` (`complete`) |
+| `GeminiText2Text` | text to text | the Gemini API, forwarded from the cluster | `CompletingModel` (`complete`) |
+| `GeminiText2Image` | text to image | the Gemini API, forwarded from the cluster | `GeneratingModel` (`generate`) |
 
 | Importer | Source |
 |---|---|
@@ -61,8 +62,9 @@ need:
 
 One with weights goes through `cortexgrid.import_model`, one without through
 `cortexgrid.register_model`. Once deployed, the `cortexgrid.Deployment` hands back
-the client that speaks the app's routes: `deployment.client()` once the app serves,
-or `deployment.client_async()` at once, with `await model.is_ready()` to ask.
+the client that calls the app's endpoints, typed as the serve app and so as its task
+type: `deployment.client()` once the app serves, or `deployment.client_async()` at
+once, with `await model.is_ready()` to ask.
 
 ## Quick start
 
@@ -94,7 +96,7 @@ cortexgrid.remote(
 deployment = cortexgrid.deploy_model(imp.family, imp.suffix, cortexgrid.IMPORTED, timeout=1800)
 
 # 3. Inference — the client blocks until the Ray Serve app is running.
-model = deployment.client()
+model: mg.GeneratingModel = deployment.client()
 
 async def run():
     result = await mg.generate(model, "a red bicycle on a beach at sunrise")
@@ -249,10 +251,13 @@ it is the floor; only quantizing the weights or running fewer forwards moves it.
 
 ## Inference
 
-`complete` streams `CompletionChunk`s (`.content`, `.thinking`, `.tool_calls`, `.finish_reason`):
+`complete` streams `CompletionChunk`s (`.content`, `.thinking`, `.tool_calls`, `.finish_reason`).
+`tools` are Python functions or tool specs; a function goes to the model as its spec,
+and each `ToolCall` comes back as data (`.id`, `.name`, `.arguments`) for the caller
+to run:
 
 ```python
-model = deployment.client()
+model: mg.CompletingModel = deployment.client()
 messages = [{"role": "user", "content": "Explain RAG in one sentence."}]
 
 async for chunk in mg.complete(model, messages, max_new_tokens=512, temperature=0.7):
@@ -293,9 +298,9 @@ a re-deploy rather than a re-registration.
 
 A replica asks for no hardware at all, so it is placed on any node, CPU-only
 included. From there it deploys and streams exactly like a cluster-served model:
-`deployment.client()` returns the same `ServedCompletingModel`. Anthropic
+`deployment.client()` is a `CompletingModel` like `Text2Text`'s. Anthropic
 reports tool calls as structured blocks rather than as generated text, so the
-serve app re-encodes them into the text form the client parses.
+serve app turns them into `ToolCall`s as they come.
 
 ### Gemini
 
@@ -317,16 +322,14 @@ cortexgrid.register_model(
 Everything said above for Anthropic holds: `config` carries the model name and
 the name of the secret (`api_key_secret="..."` to use another), both editable on
 the model card; a replica needs no hardware; and `deployment.client()`
-is a `ServedCompletingModel`, with Gemini's function calls re-encoded into the
-text form the client parses.
+is a `CompletingModel`, with Gemini's function calls turned into `ToolCall`s.
 
-`GeminiText2Image` answers the same `/generate` route as `Text2Image`, so its
-client is a `ServedGeneratingModel`:
+`GeminiText2Image` is a `GeneratingModel` like `Text2Image`:
 
 ```python
 entry = mg.Hosted("gemini-2.5-flash-image", mg.GeminiText2Image)
 ...
-model = deployment.client()
+model: mg.GeneratingModel = deployment.client()
 picture = await mg.generate(model, "a fox in the snow", size=2048)
 ```
 
@@ -345,14 +348,11 @@ diffusion settings with no Gemini counterpart, and come back as `None`.
 | `Text2Text`, `Text2Image`, `TextRewriter`, `Image2Mesh` | Serve apps that run weights (`LocalModel`s). |
 | `AnthropicText2Text` | Serve app that forwards to the Anthropic API (a `HostedModel`). |
 | `GeminiText2Text`, `GeminiText2Image` | Serve apps that forward to the Gemini API (`HostedModel`s sharing the `GeminiModel` base). |
-| `LocalModel` | Base of the serve apps that run weights: `ignore_patterns()`, `bytes_per_param`, `min_vram_gb`, `requirements(weights)`, `config()` (the model card's defaults, `compile: "false"` among them), `client(deployment)`. |
-| `HostedModel` | Base of the serve apps that forward: `config(model_id, **settings)`, `requirements()`, `client(deployment)`. |
+| `LocalModel` | Base of the serve apps that run weights: `ignore_patterns()`, `bytes_per_param`, `min_vram_gb`, `requirements(weights)`, `config()` (the model card's defaults, `compile: "false"` among them). |
+| `HostedModel` | Base of the serve apps that forward: `config(model_id, **settings)`, `requirements()`. |
 | `Weights(params=None, file_bytes=0)` | What an importer reads about the weights, handed to the serve app to size a replica. |
-| `complete(model, messages, tools=None, max_new_tokens=2048, temperature=0.7, **kw)` | Async stream of `CompletionChunk` for a `CompletingModel`. |
-| `generate(model, prompt, *, image=None, **kw) -> GeneratedImage` | One image from a `GeneratingModel`. |
-| `ServedCompletingModel(key, url)` | Client for any deployed completion app. |
-| `ServedGeneratingModel(key, url)` | Client for any deployed text-to-image app. |
-| `ServedMeshingModel(key, url)` | Client for any deployed image-to-mesh app. |
+| `complete(model, messages, tools=None, max_new_tokens=None, temperature=0.7)` | Async stream of `CompletionChunk` for a `CompletingModel`. |
+| `generate(model, prompt, *, image=None, steps=None, guidance=None, size=1024, seed=None) -> GeneratedImage` | One image from a `GeneratingModel`. |
 | `split_model_id(model_id) -> (family, suffix)` | The registry identity an importer derives. |
 | `detect_device()` | The torch device a serve app should load onto. |
 | `compiling.Mode.GRAPHED` / `compiling.Mode.FUSED` | The two modes a serve app chooses between. See [Compilation](#compilation). |
@@ -360,19 +360,19 @@ diffusion settings with no Gemini counterpart, and come back as `None`.
 | `ModelDeployFailed` | Re-exported from cortexgrid; subclasses `RuntimeError`. |
 
 Types: `CompletingModel`, `GeneratingModel`, `MeshingModel`, `RewritingModel` (each
-a `cortexgrid.DeploymentClient`), `CompletionChunk`, `GeneratedImage`, `GeneratedMesh`, `ToolCall`, `Message`, `Tool`,
-`ToolSpec`.
+the abstract base of its task's serve apps, declaring the task's endpoints),
+`CompletionChunk`, `GeneratedImage`, `GeneratedMesh`, `ToolCall`, `Message`, `Tool`,
+`ToolSpec`, `Tensor`.
 
 Extending it:
 
 - **A new source** is one importer: subclass `Importer` with `download(local_dir)`
   and `weights()`, and every serve app whose weight format it delivers runs its
   models unchanged.
-- **A new hosted API** is one serve app: subclass `HostedModel`, and register it
-  with `Hosted`. A text-to-text one only has to stream text from `POST /complete`,
-  inlining tool calls as `<tool_call>{"name": ..., "arguments": {...}}</tool_call>`
-  (there is an `encode_tool_call` for upstreams that report them structurally) —
-  then `ServedCompletingModel` is its client, unchanged.
+- **A new hosted API** is one serve app: subclass `HostedModel` and the task's type,
+  and register it with `Hosted`. A text-to-text one subclasses `CompletingModel` and
+  streams `CompletionChunk`s from its `complete` endpoint, with
+  `StreamedCompletionWriter` turning the upstream's text and tool calls into them.
 - **A model no task's serve app can load** gets its own serve app — an `Image2Mesh`
   subclass, say, with `load` and `make_mesh`, and `compile` if compiling it pays —
   and, if its source is unusual too, its own importer.

{cortexgrid_infer-0.3.8 → cortexgrid_infer-0.3.9}/cortexgrid_infer/__init__.py

@@ -7,9 +7,9 @@ nothing else, split along the two things a model is:
 
 - where its weights come from - an importer per source (`importers`), which
   names, fetches and sizes them;
-- what runs it - a serve app per task (`serve_apps`), which loads the weights,
-  answers the task's routes (`protocols`), and names the client that speaks
-  them.
+- what runs it - a serve app per task (`serve_apps`), which loads the weights
+  and serves the task's endpoints, declared by the task's type in `core`
+  (`CompletingModel`, say); cortexgrid generates the client that calls them.
 
     imp = cortexgrid_infer.HuggingFaceImporter("Qwen/Qwen2.5-0.5B-Instruct", cortexgrid_infer.Text2Text)
     with imp:
@@ -19,7 +19,7 @@ nothing else, split along the two things a model is:
             requirements=imp.requirements(),
         )
     deployment = cortexgrid.deploy_model(imp.family, imp.suffix, cortexgrid.IMPORTED)
-    model = deployment.client()
+    model: cortexgrid_infer.CompletingModel = deployment.client()
 
 A model hosted elsewhere has no weights and so no importer: its serve app
 forwards to it, and a `Hosted` entry registers it.
@@ -38,6 +38,7 @@ from cortexgrid_infer.core import (
     MeshingModel,
     Message,
     RewritingModel,
+    Tensor,
     Tool,
     ToolCall,
     ToolSpec,
@@ -46,12 +47,6 @@ from cortexgrid_infer.core import (
     mesh,
     rewrite,
 )
-from cortexgrid_infer.protocols import (
-    ServedCompletingModel,
-    ServedGeneratingModel,
-    ServedMeshingModel,
-    ServedRewritingModel,
-)
 from cortexgrid_infer.device import detect_device
 from cortexgrid_infer.serve_apps import (
     AnthropicText2Text,
@@ -77,6 +72,7 @@ __all__ = [
     "MeshingModel",
     "RewritingModel",
     "Message",
+    "Tensor",
     "Tool",
     "ToolCall",
     "ToolSpec",
@@ -86,10 +82,6 @@ __all__ = [
     "rewrite",
     "ModelDeployFailed",
     "detect_device",
-    "ServedCompletingModel",
-    "ServedGeneratingModel",
-    "ServedMeshingModel",
-    "ServedRewritingModel",
     "LocalModel",
     "HostedModel",
     "Weights",

{cortexgrid_infer-0.3.8 → cortexgrid_infer-0.3.9}/cortexgrid_infer/core.py

@@ -3,13 +3,24 @@
 from __future__ import annotations
 
 import abc
+import base64
 from dataclasses import dataclass, field
-from functools import partial
 from collections.abc import AsyncIterator
-from typing import Any, Callable, Sequence
+from functools import partial
+from typing import Annotated, Any
 
-import cortexgrid
 import numpy as np
+import torch
+from pydantic import Base64Bytes, PlainSerializer, PlainValidator
+
+from cortexgrid_infer.utils import Tool, ToolSpec
+
+
+Tensor = Annotated[
+    torch.Tensor,
+    PlainValidator(torch.tensor),
+    PlainSerializer(torch.Tensor.tolist),
+]
 
 
 @dataclass
@@ -17,10 +28,6 @@ class ToolCall:
     id: str
     name: str
     arguments: dict[str, Any]
-    _func: partial[Any] = field(repr=False)
-
-    def __call__(self) -> Any:
-        return self._func(**self.arguments)
 
 
 @dataclass
@@ -38,42 +45,18 @@ class CompletionChunk:
     def has_tool_calls(self) -> bool:
         return bool(self.tool_calls)
 
-    @classmethod
-    def from_wire(
-        cls, answered: dict[str, Any], tools_by_name: dict[str, Callable[..., Any]]
-    ) -> CompletionChunk:
-        tool_calls = []
-        for called in answered["tool_calls"]:
-            tool = tools_by_name.get(called["name"])
-            if tool is None:
-                continue
-            bound = partial(tool, **called["arguments"])
-            tool_call = ToolCall(
-                id=called["id"], name=called["name"], arguments=called["arguments"], _func=bound
-            )
-            tool_calls.append(tool_call)
-        return cls(
-            content=answered["content"],
-            thinking=answered["thinking"],
-            tool_calls=tool_calls,
-            finish_reason=answered["finish_reason"],
-        )
-
 
 Message = dict[str, Any]
-ToolSpec = dict[str, Any]
-Tool = Callable[..., Any] | ToolSpec
 
 
-class CompletingModel(cortexgrid.DeploymentClient, abc.ABC):
+class CompletingModel(abc.ABC):
     @abc.abstractmethod
     def complete(
         self,
         messages: list[Message],
-        tools: Sequence[Tool] | None = None,
-        max_new_tokens: int = 2048,
+        tools: list[Tool] | None = None,
+        max_new_tokens: int | None = None,
         temperature: float = 0.7,
-        **kwargs: Any,
     ) -> AsyncIterator[CompletionChunk]: ...
 
 
@@ -82,13 +65,13 @@ class GeneratedImage:
     """One generated image: the encoded file bytes (PNG) plus the resolved
     parameters the deployment used to produce it."""
 
-    image: bytes
+    image: Base64Bytes
     width: int
     height: int
     params: dict[str, Any] = field(default_factory=dict)
 
 
-class GeneratingModel(cortexgrid.DeploymentClient, abc.ABC):
+class GeneratingModel(abc.ABC):
     """A deployed image model. Unlike completion, generation is a single
     request/response (no token streaming), so `generate` returns one result."""
 
@@ -97,15 +80,43 @@ class GeneratingModel(cortexgrid.DeploymentClient, abc.ABC):
         self,
         prompt: str,
         *,
-        image: bytes | None = None,
+        image: Base64Bytes | None = None,
         steps: int | None = None,
         guidance: float | None = None,
         size: int = 1024,
         seed: int | None = None,
-        **kwargs: Any,
     ) -> GeneratedImage: ...
 
 
+def _encoded_array(array: np.ndarray, dtype: str) -> str:
+    contiguous = np.ascontiguousarray(array, dtype=dtype)
+    encoded = base64.b64encode(contiguous.tobytes())
+    return encoded.decode("ascii")
+
+
+def _decoded_array(data: str, dtype: str, columns: int) -> np.ndarray:
+    raw = base64.b64decode(data)
+    flat = np.frombuffer(raw, dtype=dtype)
+    return flat.reshape(-1, columns)
+
+
+Vertices = Annotated[
+    np.ndarray,
+    PlainValidator(partial(_decoded_array, dtype="<f4", columns=3)),
+    PlainSerializer(partial(_encoded_array, dtype="<f4")),
+]
+Faces = Annotated[
+    np.ndarray,
+    PlainValidator(partial(_decoded_array, dtype="<i4", columns=3)),
+    PlainSerializer(partial(_encoded_array, dtype="<i4")),
+]
+Colours = Annotated[
+    np.ndarray,
+    PlainValidator(partial(_decoded_array, dtype="u1", columns=3)),
+    PlainSerializer(partial(_encoded_array, dtype="u1")),
+]
+
+
 @dataclass
 class GeneratedMesh:
     """One mesh: `vertices` (N, 3) float32, x to the right of the picture it was
@@ -113,35 +124,36 @@ class GeneratedMesh:
     into them, counter-clockwise seen from outside; `colours` (N, 3) uint8, each
     vertex's colour. Plus the resolved parameters the deployment used."""
 
-    vertices: np.ndarray
-    faces: np.ndarray
-    colours: np.ndarray
+    vertices: Vertices
+    faces: Faces
+    colours: Colours
     params: dict[str, Any] = field(default_factory=dict)
 
 
-class MeshingModel(cortexgrid.DeploymentClient, abc.ABC):
+class MeshingModel(abc.ABC):
     """A deployed model that makes a mesh of the object in one picture."""
 
     @abc.abstractmethod
-    async def mesh(self, image: bytes, **kwargs: Any) -> GeneratedMesh: ...
+    async def mesh(
+        self, image: Base64Bytes, options: dict[str, Any] | None = None
+    ) -> GeneratedMesh: ...
 
 
-class RewritingModel(cortexgrid.DeploymentClient, abc.ABC):
+class RewritingModel(abc.ABC):
     """A deployed model that rewrites one text into another - corrects,
     paraphrases, summarises or translates it. Like generation, a single
     request/response."""
 
     @abc.abstractmethod
-    async def rewrite(self, text: str, max_new_tokens: int = 512, **kwargs: Any) -> str: ...
+    async def rewrite(self, text: str, max_new_tokens: int = 512) -> str: ...
 
 
 async def complete(
     deployed_model: CompletingModel,
     messages: list[Message],
-    tools: Sequence[Tool] | None = None,
-    max_new_tokens: int = 2048,
+    tools: list[Tool] | None = None,
+    max_new_tokens: int | None = None,
     temperature: float = 0.7,
-    **kwargs: Any,
 ) -> AsyncIterator[CompletionChunk]:
     """Async completion request to the deployed model."""
     async for chunk in deployed_model.complete(
@@ -149,7 +161,6 @@ async def complete(
         tools,
         max_new_tokens=max_new_tokens,
         temperature=temperature,
-        **kwargs,
     ):
         yield chunk
 
@@ -159,22 +170,28 @@ async def generate(
     prompt: str,
     *,
     image: bytes | None = None,
-    **kwargs: Any,
+    steps: int | None = None,
+    guidance: float | None = None,
+    size: int = 1024,
+    seed: int | None = None,
 ) -> GeneratedImage:
     """Single image-generation request to the deployed model."""
-    return await deployed_model.generate(prompt, image=image, **kwargs)
+    return await deployed_model.generate(
+        prompt, image=image, steps=steps, guidance=guidance, size=size, seed=seed
+    )
 
 
-async def mesh(deployed_model: MeshingModel, image: bytes, **kwargs: Any) -> GeneratedMesh:
+async def mesh(
+    deployed_model: MeshingModel, image: bytes, options: dict[str, Any] | None = None
+) -> GeneratedMesh:
     """Single image-to-mesh request to the deployed model."""
-    return await deployed_model.mesh(image, **kwargs)
+    return await deployed_model.mesh(image, options)
 
 
 async def rewrite(
     deployed_model: RewritingModel,
     text: str,
     max_new_tokens: int = 512,
-    **kwargs: Any,
 ) -> str:
     """Single rewriting request to the deployed model."""
-    return await deployed_model.rewrite(text, max_new_tokens=max_new_tokens, **kwargs)
+    return await deployed_model.rewrite(text, max_new_tokens=max_new_tokens)