v1.3.4.dev0 Full Collab Integration and bug fixed

agent_config.yaml

@@ -23,5 +23,6 @@ agent:
   # openrouter_api_key: <your_openrouter_api_key>  # Open router API key (can also be set as env variable OPENROUTER_API_KEY)
   # openrouter_base_url: https://openrouter.ai/api/v1  # Open router base URL (can also be set as env variable OPENROUTER_BASE_URL)
   # openrouter_request_timeout: 15.0  # Timeout for OpenRouter API requests in seconds (can also be set as env variable OPENROUTER_REQUEST_TIMEOUT)
+  # openrouter_max_tokens: 2048  # Max completion length. OpenRouter reserves max_tokens*price against the key budget BEFORE generating, so an uncapped value can 402 ("more credits, or fewer max_tokens") on a credit/weekly-limited key. Raise only if responses get truncated (env: OPENROUTER_MAX_TOKENS)
   # openrouter_provider_sort: throughput  # Bias OpenRouter's upstream routing to avoid slow providers: 'throughput' | 'latency' | 'price'. Empty string = let OpenRouter choose (env: OPENROUTER_PROVIDER_SORT)
   # openrouter_structured_output: false  # Ask the model for a schema-validated plan directly (skips free-form JSON + regex repair). More reliable, but only works on models whose OpenRouter route supports structured/JSON-schema output (e.g. Gemini, GPT-4o). env: OPENROUTER_STRUCTURED_OUTPUT=1

docs/examples/pytorch/generation.rst

@@ -101,12 +101,3 @@ filter by pair distance to find the hardest negatives.
 
       weightslab ui launch           # 1. deploy the studio
       weightslab start example --gen # 2. start the generation / anomaly demo
-
-
-.. raw:: html
-
-   <div style="text-align:right; margin-top:2rem;">
-     <a href="https://colab.research.google.com/github/GrayboxTech/weightslab/blob/main/weightslab/examples/Notebooks/PyTorch/ws-generation.ipynb" target="_blank" rel="noopener noreferrer">
-       <img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab">
-     </a>
-   </div>

docs/usage/parameters.rst

@@ -511,6 +511,16 @@ LLM / agent integration (optional)
    * - ``OPENROUTER_REQUEST_TIMEOUT``
      - *(unset)*
      - Per-request timeout in seconds for OpenRouter calls.
+   * - ``OPENROUTER_MAX_TOKENS``
+     - ``2048``
+     - Maximum completion length requested from OpenRouter. OpenRouter
+       pre-authorizes ``max_tokens × completion_price`` against the key's
+       remaining budget *before* generating, so leaving this uncapped makes
+       the model request its full output window and can fail with a ``402``
+       ("requires more credits, or fewer max_tokens") on a credit- or
+       weekly-limited key — even though the model is otherwise usable. The
+       default is ample for intent planning; raise it only if you see
+       truncated responses.
 
 Telemetry
 ~~~~~~~~~~

pyproject.toml

@@ -33,7 +33,11 @@ dependencies = [
   # versions pip resolves. Range left open (<3) so Colab's numpy 2.x installs
   # without a downgrade. Verified: import + torch/pandas interop + gRPC on np 2.4.
   "numpy>=1.24,<3",
-  "pandas>=2.2.3,<3",
+  # Floor at 2.2.2 (not 2.2.3) so Colab's pinned pandas==2.2.2 is accepted
+  # as-is: google-colab 1.0.0 requires pandas==2.2.2, and a >=2.2.3 floor forced
+  # an upgrade to 2.3.x that breaks that pin. Clean installs still pick the
+  # newest numpy-2-safe pandas (2.3.x).
+  "pandas>=2.2.2,<3",
   "duckdb>=1.1,<2",  # signal/sample/instance history store
   "PyYAML>=6.0.3,<7",
   "dill>=0.3.8,<0.5",
@@ -53,10 +57,14 @@ dependencies = [
 
   # Serving / schema
   "grpcio>=1.80,<2",
-  # protobuf <8 so Colab's protobuf 7.x is accepted (its <7 cap forced the
-  # downgrade cascade). Verified: all 70 generated gRPC messages serialize/parse
-  # and the descriptor pool builds under protobuf 7.35.
-  "protobuf>=4.25,<8",
+  # Floor tracks the checked-in gencode: the *_pb2.py are regenerated with
+  # grpcio-tools ~=1.68 (protoc gencode 5.28.1), so any runtime >=5.28.1 satisfies
+  # protobuf's runtime>=gencode rule — including Colab's stock protobuf 5.29.6, so
+  # NO upgrade is forced on Colab. IMPORTANT: regenerate the stubs with
+  # grpcio-tools ~=1.68 to keep the gencode Colab-compatible; a newer grpcio-tools
+  # bumps the gencode and reintroduces the VersionError on Colab. Upper bound <8
+  # keeps modern runtimes (7.x) working.
+  "protobuf>=5.28.1,<8",
   "pydantic>=2.7,<3",
 
   # Imaging

tests/backend/test_data_loader_interface.py

@@ -353,6 +353,35 @@ def test_multiple_sequential_epochs_with_auto_reset(self):
             f"Each epoch should yield the same batch count, got {epoch_counts}",
         )
 
+    def test_reset_is_callable_through_ledger_proxy(self):
+        """`reset()` must exist and be callable via the ledger Proxy handle.
+
+        Ultralytics' trainer calls ``train_loader.reset()`` after closing mosaic
+        augmentation. The loader handed to it is a ledger Proxy wrapping this
+        interface, so a real ``reset()`` method (mapped to ``_reset_iterator``)
+        must be reachable through the proxy — otherwise the call resolves to a
+        non-callable None. Regression for the Ultralytics YOLO notebook crash
+        ``TypeError: 'NoneType' object is not callable``.
+        """
+        DataLoaderInterface(
+            self.train_ds, batch_size=self.batch_size,
+            loader_name="reset_proxy_loader", register=True, compute_hash=True,
+        )
+        loader = ledgers.get_dataloader("reset_proxy_loader")
+
+        self.assertTrue(hasattr(loader, "reset"))
+        self.assertTrue(callable(loader.reset))
+
+        # A full epoch before reset.
+        batches_before = sum(1 for _ in loader)
+        self.assertGreater(batches_before, 0)
+
+        # reset() must not raise, and the loader stays fully usable afterwards
+        # (a fresh full epoch with the same batch count can be iterated).
+        loader.reset()
+        batches_after = sum(1 for _ in loader)
+        self.assertEqual(batches_after, batches_before)
+
 
 class TestDataLoaderReproducibility(unittest.TestCase):
     """Test RNG and iteration state reproducibility for dataloaders."""

tests/backend/test_ledgers.py

@@ -672,6 +672,49 @@ def __init__(self):
         delattr(proxy, "flag")
         self.assertFalse(hasattr(wrapped, "flag"))
 
+    def test_proxy_missing_attr_raises_attributeerror(self):
+        """Missing attributes on the wrapped object must raise AttributeError.
+
+        Regression: __getattr__ used to `return None`, which made
+        `hasattr(proxy, missing)` report True and silently resolved
+        `proxy.missing` to a non-callable None. That broke capability probes
+        such as Ultralytics' `if not hasattr(loader, "reset")` shim, leaving
+        `loader.reset` as None -> `'NoneType' object is not callable`.
+        """
+        proxy = Proxy(Dummy("x"))
+
+        # Existing attribute still forwards.
+        self.assertEqual(proxy.name, "x")
+
+        # Missing attribute raises (does not return None).
+        with self.assertRaises(AttributeError):
+            _ = proxy.definitely_not_here
+
+        # hasattr therefore reports False, and getattr honors the default.
+        self.assertFalse(hasattr(proxy, "definitely_not_here"))
+        self.assertIsNone(getattr(proxy, "definitely_not_here", None))
+        self.assertEqual(getattr(proxy, "definitely_not_here", 42), 42)
+
+    def test_proxy_forwards_callable_capability(self):
+        """A method present on the wrapped object is discoverable and callable.
+
+        Mirrors how a ledger Proxy wraps a DataLoaderInterface whose real
+        `reset()` must be reachable through the proxy.
+        """
+        class _HasReset:
+            def __init__(self):
+                self.calls = 0
+
+            def reset(self):
+                self.calls += 1
+
+        wrapped = _HasReset()
+        proxy = Proxy(wrapped)
+
+        self.assertTrue(hasattr(proxy, "reset"))
+        self.assertTrue(callable(proxy.reset))
+        proxy.reset()
+        self.assertEqual(wrapped.calls, 1)
 
 
 if __name__ == "__main__":

tests/backend/test_write_dataframe.py

@@ -48,6 +48,19 @@ def _call(path, manager, **kwargs):
         return write_dataframe(path, **kwargs)
 
 
+def _records(path):
+    """Read a JSON export back into a list of row records.
+
+    write_dataframe defaults to ``orient="columns"`` (``{column: {row: value}}``)
+    — compact (~6x smaller than records for wide sparse tables) and round-trips
+    with ``pd.read_json``'s default orient. These tests assert row-level
+    behavior, so normalize the columnar layout back to records here. Empty
+    exports (``{}`` / ``{"index": {}}``) collapse to ``[]``.
+    """
+    raw = json.loads(open(path, encoding="utf-8").read())
+    return pd.DataFrame(raw).to_dict(orient="records")
+
+
 # ---------------------------------------------------------------------------
 # Fixtures
 # ---------------------------------------------------------------------------
@@ -89,20 +102,22 @@ def test_flush_failure_does_not_abort(self, tmp_json):
 # ---------------------------------------------------------------------------
 
 class TestWriteDataframeJsonStructure:
-    def test_json_is_list_of_records(self, mgr, tmp_json):
+    def test_json_is_columnar_dict(self, mgr, tmp_json):
         _call(tmp_json, mgr)
-        data = json.loads(open(tmp_json).read())
-        assert isinstance(data, list)
+        raw = json.loads(open(tmp_json).read())
+        # Default orient="columns": {column: {row_index: value}}
+        assert isinstance(raw, dict)
+        assert isinstance(raw["sample_id"], dict)
 
     def test_json_includes_index_as_columns(self, mgr, tmp_json):
         _call(tmp_json, mgr)
-        data = json.loads(open(tmp_json).read())
+        data = _records(tmp_json)
         assert "sample_id" in data[0]
         assert "annotation_id" in data[0]
 
     def test_json_all_rows_present_by_default(self, mgr, tmp_json):
         _call(tmp_json, mgr)
-        data = json.loads(open(tmp_json).read())
+        data = _records(tmp_json)
         assert len(data) == 4
 
     def test_returns_written_path(self, mgr, tmp_json):
@@ -207,51 +222,51 @@ def test_returns_path_when_no_manager(self, tmp_json):
 class TestWriteDataframeColumnFilters:
     def test_columns_all_keeps_everything(self, mgr, tmp_json):
         _call(tmp_json, mgr, columns="all")
-        data = json.loads(open(tmp_json).read())
+        data = _records(tmp_json)
         assert "signals_loss" in data[0] or any("signals_loss" in r for r in data)
 
     def test_columns_tags_only(self, mgr, tmp_json):
         _call(tmp_json, mgr, columns="tags")
-        data = json.loads(open(tmp_json).read())
+        data = _records(tmp_json)
         non_index = [k for k in data[0] if k not in ("sample_id", "annotation_id")]
         assert all(k.startswith("tag:") or k.startswith("TAG:") for k in non_index)
 
     def test_columns_signals_only(self, mgr, tmp_json):
         _call(tmp_json, mgr, columns="signals")
-        data = json.loads(open(tmp_json).read())
+        data = _records(tmp_json)
         non_index = [k for k in data[0] if k not in ("sample_id", "annotation_id")]
         assert all(str(k).lower().startswith("signals") for k in non_index)
 
     def test_columns_discarded_only(self, mgr, tmp_json):
         _call(tmp_json, mgr, columns="discarded")
-        data = json.loads(open(tmp_json).read())
+        data = _records(tmp_json)
         non_index = [k for k in data[0] if k not in ("sample_id", "annotation_id")]
         assert non_index == ["discarded"]
 
     def test_columns_list_of_groups(self, mgr, tmp_json):
         _call(tmp_json, mgr, columns=["tags", "discarded"])
-        data = json.loads(open(tmp_json).read())
+        data = _records(tmp_json)
         non_index = set(k for r in data for k in r if k not in ("sample_id", "annotation_id"))
         assert "discarded" in non_index
         assert any(k.startswith("tag:") for k in non_index)
         assert not any(str(k).lower().startswith("signals") for k in non_index)
 
     def test_columns_exact_name(self, mgr, tmp_json):
         _call(tmp_json, mgr, columns=["signals_loss"])
-        data = json.loads(open(tmp_json).read())
+        data = _records(tmp_json)
         non_index = [k for k in data[0] if k not in ("sample_id", "annotation_id")]
         assert non_index == ["signals_loss"]
 
     def test_columns_nonexistent_name_yields_empty_cols(self, mgr, tmp_json):
         _call(tmp_json, mgr, columns=["nonexistent_col"])
-        data = json.loads(open(tmp_json).read())
+        data = _records(tmp_json)
         # Index columns always present; no extra columns
         for row in data:
             assert set(row.keys()) <= {"sample_id", "annotation_id"}
 
     def test_columns_none_keeps_all(self, mgr, tmp_json):
         _call(tmp_json, mgr, columns=None)
-        data = json.loads(open(tmp_json).read())
+        data = _records(tmp_json)
         assert "signals_loss" in data[0] or any("signals_loss" in r for r in data)
 
 
@@ -262,38 +277,38 @@ def test_columns_none_keeps_all(self, mgr, tmp_json):
 class TestWriteDataframeIndexFilters:
     def test_sample_id_single(self, mgr, tmp_json):
         _call(tmp_json, mgr, sample_id="s1")
-        data = json.loads(open(tmp_json).read())
+        data = _records(tmp_json)
         assert all(r["sample_id"] == "s1" for r in data)
         assert len(data) == 2 # s1 has annotation_ids 0 and 1
 
     def test_sample_id_list(self, mgr, tmp_json):
         _call(tmp_json, mgr, sample_id=["s1", "s2"])
-        data = json.loads(open(tmp_json).read())
+        data = _records(tmp_json)
         sids = {r["sample_id"] for r in data}
         assert sids == {"s1", "s2"}
         assert len(data) == 4
 
     def test_instance_id_zero_keeps_sample_rows(self, mgr, tmp_json):
         _call(tmp_json, mgr, instance_id=0)
-        data = json.loads(open(tmp_json).read())
+        data = _records(tmp_json)
         assert all(r["annotation_id"] == 0 for r in data)
         assert len(data) == 2 # s1 and s2 both have annotation_id=0
 
     def test_instance_id_list(self, mgr, tmp_json):
         _call(tmp_json, mgr, instance_id=[1, 2])
-        data = json.loads(open(tmp_json).read())
+        data = _records(tmp_json)
         assert all(r["annotation_id"] in (1, 2) for r in data)
 
     def test_sample_and_instance_combined(self, mgr, tmp_json):
         _call(tmp_json, mgr, sample_id="s1", instance_id=1)
-        data = json.loads(open(tmp_json).read())
+        data = _records(tmp_json)
         assert len(data) == 1
         assert data[0]["sample_id"] == "s1"
         assert data[0]["annotation_id"] == 1
 
     def test_empty_result_when_no_match(self, mgr, tmp_json):
         _call(tmp_json, mgr, sample_id="nonexistent")
-        data = json.loads(open(tmp_json).read())
+        data = _records(tmp_json)
         assert data == []
 
 
@@ -305,7 +320,7 @@ class TestWriteDataframeEmpty:
     def test_empty_df_writes_empty_json(self, tmp_json):
         mgr = _make_manager(df=pd.DataFrame())
         _call(tmp_json, mgr)
-        data = json.loads(open(tmp_json).read())
+        data = _records(tmp_json)
         assert data == []
 
     def test_empty_df_writes_empty_csv(self, tmp_csv):

tests/gRPC/test_grpc_user_actions.py

@@ -61,7 +61,7 @@ def get_combined_df(self):
     def get_df_view(self):
         return self.df.reset_index().copy()\
 
-    def get_collapse_annotations_to_samples_df(self):
+    def get_collapse_annotations_to_samples_df(self, df=None):
         # This fixture is already one row per sample (no instance rows), so the
         # per-sample collapsed view is just the frame with its index restored to
         # (origin, sample_id) columns — matching what the real manager returns.

weightslab/backend/dataloader_interface.py

@@ -1045,6 +1045,26 @@ def reset_iterator(self) -> None:
         """
         self._reset_iterator()
 
+    def reset(self) -> None:
+        """Recreate the internal iterator (Ultralytics / ``InfiniteDataLoader``-compatible).
+
+        Some training frameworks call ``loader.reset()`` to force a fresh
+        iterator after mutating the dataset mid-run — e.g. Ultralytics closes
+        mosaic augmentation near the end of training and then calls
+        ``train_loader.reset()`` so the change takes effect on the next epoch.
+        We map it to ``_reset_iterator()``.
+
+        Exposing this as a real method (rather than relying on a caller-side
+        shim like ``if not hasattr(loader, "reset"): loader.reset = ...``) is
+        important: this loader is handed to callers wrapped in a ledger
+        ``Proxy`` whose ``__getattr__`` returns ``None`` for genuinely-missing
+        attributes instead of raising ``AttributeError``. That makes
+        ``hasattr(proxy, "reset")`` return True and silently resolves
+        ``proxy.reset`` to ``None``, so any such shim is skipped and
+        ``reset()`` blows up with ``'NoneType' object is not callable``.
+        """
+        self._reset_iterator()
+
     # -------------------------------------------------------------------------
     # Iteration state capture/restore for deterministic resume
     # -------------------------------------------------------------------------

weightslab/backend/ledgers.py

@@ -500,7 +500,16 @@ def __getattr__(self, item):
             except (KeyError, TypeError):
                 pass
 
-        return None
+        # Attribute genuinely absent on the wrapped object. Raise AttributeError
+        # (as the class docstring promises) rather than returning None: returning
+        # None makes `hasattr(proxy, x)` report True for missing attributes, which
+        # breaks capability probes (e.g. `if not hasattr(loader, "reset")`) and
+        # silently resolves `proxy.missing` to a non-callable None. Callers that
+        # want a soft default should use `getattr(proxy, x, default)`, which now
+        # works correctly.
+        raise AttributeError(
+            f"'{type(obj).__name__}' object has no attribute {item!r}"
+        )
 
     def __delattr__(self, name: str) -> None:
         """Delete an attribute from the proxy or wrapped object safely.