danvi/dim1486/generalize dan perception stack

dimos/agents/mcp/mcp_adapter.py

@@ -41,8 +41,6 @@
 
 logger = setup_logger()
 
-DEFAULT_TIMEOUT = 30
-
 
 class McpError(Exception):
     """Raised when the MCP server returns a JSON-RPC error."""
@@ -55,13 +53,13 @@ def __init__(self, message: str, code: int | None = None) -> None:
 class McpAdapter:
     """Thin JSON-RPC client for a running MCP server."""
 
-    def __init__(self, url: str | None = None, timeout: int = DEFAULT_TIMEOUT) -> None:
-        if url is None:
-            from dimos.core.global_config import global_config
+    def __init__(self, url: str | None = None, timeout: int | None = None) -> None:
+        from dimos.core.global_config import global_config
 
+        if url is None:
             url = f"http://localhost:{global_config.mcp_port}/mcp"
         self.url = url
-        self.timeout = timeout
+        self.timeout = global_config.mcp_timeout if timeout is None else timeout
 
     def call(self, method: str, params: dict[str, Any] | None = None) -> dict[str, Any]:
         """Send a JSON-RPC request and return the parsed response.
@@ -137,7 +135,7 @@ def wait_for_down(self, timeout: float = 10.0, interval: float = 0.5) -> bool:
         return False
 
     @classmethod
-    def from_run_entry(cls, entry: Any | None = None, timeout: int = DEFAULT_TIMEOUT) -> McpAdapter:
+    def from_run_entry(cls, entry: Any | None = None, timeout: int | None = None) -> McpAdapter:
         """Create an adapter from a RunEntry, or discover the latest one.
 
         Falls back to the default URL if no entry is found.

dimos/cli/commands/mcp.py

@@ -27,9 +27,9 @@
 mcp_app = typer.Typer(help="Interact with the running MCP server")
 
 
-def _get_adapter() -> McpAdapter:
+def _get_adapter(timeout: int | None = None) -> McpAdapter:
     """Get an McpAdapter from the latest RunEntry or default URL."""
-    return McpAdapter.from_run_entry()
+    return McpAdapter.from_run_entry(timeout=timeout)
 
 
 @mcp_app.command("list-tools")
@@ -72,6 +72,9 @@ def mcp_call_tool(
         [], "--arg", "-a", callback=_validate_key_value_args, help="Arguments as key=value"
     ),
     json_args: str = typer.Option("", "--json-args", "-j", help="Arguments as JSON string"),
+    timeout: int = typer.Option(
+        None, "--timeout", "-t", help="Seconds to wait for the tool (default: mcp_timeout)"
+    ),
 ) -> None:
     """Call an MCP tool by name."""
     arguments: dict[str, Any] = {}
@@ -89,7 +92,7 @@ def mcp_call_tool(
             raise typer.Exit(1)
 
     try:
-        result = _get_adapter().call_tool(tool_name, arguments)
+        result = _get_adapter(timeout).call_tool(tool_name, arguments)
     except requests.ConnectionError:
         typer.echo("Error: no running MCP server (is DimOS running?)", err=True)
         raise typer.Exit(1)

dimos/core/global_config.py

@@ -104,6 +104,10 @@ class GlobalConfig(BaseSettings):
     robot_rotation_diameter: float = 0.6
     nerf_speed: float = 1.0
     mcp_port: int = 9990
+    # Seconds an MCP client waits for a tool to answer. A skill that thinks
+    # for longer than this is cut off at the client, not the server, so the
+    # caller owns the number.
+    mcp_timeout: int = 30
     # `DIMOS_TRANSPORT` (or `.env`) is the single switch read by every process
     # (dimos, humancli, agentspy, dtop). The `transport` alias keeps the bare
     # env name and the `--transport` CLI flag (which sets the field by name) working.

dimos/memory/backend.py

@@ -207,13 +207,14 @@ def _vector_search(self, query: StreamQuery) -> Iterator[Observation[T]]:
         ranked: list[Observation[T]] = []
         for obs_id, sim in hits:
             match = obs_by_id.get(obs_id)
-            if match is not None:
-                ranked.append(
-                    match.derive(data=match.data, embedding=query.search_vec, similarity=sim)
-                )
-
-        # Apply remaining query ops (skip vector search)
-        rest = replace(query, search_vec=None, search_k=None)
+            if match is None:
+                continue
+            if not all(f.matches(match) for f in query.filters):
+                continue
+            ranked.append(match.derive(data=match.data, embedding=query.search_vec, similarity=sim))
+
+        # Apply remaining query ops (filters already applied; skip vector search)
+        rest = replace(query, search_vec=None, search_k=None, filters=())
         yield from rest.apply(iter(ranked))
 
     def _iterate_live(

dimos/memory/stream.py

@@ -568,8 +568,15 @@ def drain(self) -> int:
         return n
 
     def drain_thread(self) -> DisposableBase:
-        """Drain this stream on the dimos thread pool; returns a disposable."""
-        return self.subscribe(lambda _: None)
+        """Drain this stream on the dimos thread pool; returns a disposable.
+
+        A drain has no consumer to surface an iteration error, so it logs -
+        a dead pipeline must not die silently.
+        """
+        return self.subscribe(
+            lambda _: None,
+            on_error=lambda e: logger.error("drain_thread() pipeline died: %s", e, exc_info=e),
+        )
 
     def observable(self) -> reactivex.Observable[O]:
         """Convert this stream to an RxPY Observable.

dimos/memory/utils/sqlite.py

@@ -27,6 +27,9 @@ def open_sqlite_connection(path: str | Path) -> sqlite3.Connection:
     conn = sqlite3.connect(path, check_same_thread=False)
     conn.execute("PRAGMA journal_mode=WAL")
     conn.execute("PRAGMA synchronous=NORMAL")
+    # WAL has one writer at a time; concurrent writer threads (a recorder
+    # feed plus a live embed save) can starve past the 5s default and die.
+    conn.execute("PRAGMA busy_timeout=30000")
     conn.enable_load_extension(True)
     sqlite_vec.load(conn)
     conn.enable_load_extension(False)

dimos/perception/detection/detectors/owlv2.py

@@ -16,6 +16,7 @@
 
 from __future__ import annotations
 
+from collections import OrderedDict
 from functools import cached_property
 
 import numpy as np
@@ -27,9 +28,14 @@
 from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox
 from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D
 
+# ~3.8 MB of GPU memory per cached frame
+_FEATURE_CACHE_MAX = 128
+
 
 class Owlv2Config(HuggingFaceModelConfig):
     model_name: str = "google/owlv2-base-patch16-ensemble"
+    # float16 runs the forward under autocast at roughly half the latency;
+    # scores jitter by a few thousandths, so threshold-edge boxes may flip.
     dtype: torch.dtype = torch.float32
 
 
@@ -45,6 +51,22 @@ class Owlv2Detector(HuggingFaceModel):
 
     config: Owlv2Config
 
+    def __init__(self, **kwargs: object) -> None:
+        super().__init__(**kwargs)
+        # image-tower forwards run so far; instrumentation reads deltas
+        self.forwards = 0
+        # (frame ts, label, floor) to that label's floor-filtered
+        # (boxes_kx4, scores_k); filled and read by localize
+        self.score_cache: OrderedDict[tuple[float, str, float], tuple[np.ndarray, np.ndarray]] = (
+            OrderedDict()
+        )
+        # frame ts to text-independent tower outputs; any query set scores
+        # against a cached frame without rerunning the tower
+        self._features: OrderedDict[
+            float, tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]
+        ] = OrderedDict()
+        self._text_embeds: dict[str, torch.Tensor] = {}
+
     @cached_property
     def _model(self):  # type: ignore[no-untyped-def]
         from transformers import Owlv2ForObjectDetection
@@ -62,6 +84,13 @@ def _processor(self):  # type: ignore[no-untyped-def]
 
         return Owlv2Processor.from_pretrained(self.config.model_name)
 
+    def _autocast(self) -> torch.autocast:
+        return torch.autocast(
+            device_type="cuda",
+            dtype=self.config.dtype,
+            enabled=self.config.dtype is not torch.float32 and "cuda" in str(self.config.device),
+        )
+
     def query_detections(
         self,
         image: Image,
@@ -74,38 +103,55 @@ def query_detections(
         ``confidence`` is the calibrated per-box score. ``class_id`` indexes
         into ``queries``.
         """
-        pil = PILImage.fromarray(image.to_rgb().data)
-        with torch.inference_mode():
-            inputs = self._processor(text=[queries], images=pil, return_tensors="pt").to(
-                self.config.device
-            )
+        return self.query_detections_batch([image], queries, threshold)[0]
+
+    def query_detections_batch(
+        self,
+        images: list[Image],
+        queries: list[str],
+        threshold: float = 0.1,
+    ) -> list[ImageDetections2D]:
+        """``query_detections`` over several images in one forward pass.
+
+        Per-call preprocessing, text encoding and kernel launches amortize
+        across the batch, which is what makes many-frame sweeps affordable;
+        results are per-image, in input order.
+        """
+        self.forwards += len(images)
+        pils = [PILImage.fromarray(image.to_rgb().data) for image in images]
+        with torch.inference_mode(), self._autocast():
+            inputs = self._processor(
+                text=[queries] * len(pils), images=pils, return_tensors="pt"
+            ).to(self.config.device)
             outputs = self._model(**inputs)
             results = self._processor.post_process_grounded_object_detection(
                 outputs=outputs,
-                target_sizes=torch.tensor([(pil.height, pil.width)]),
+                target_sizes=torch.tensor([(pil.height, pil.width) for pil in pils]),
                 threshold=threshold,
-            )[0]
-
-        detections: list[Detection2DBBox] = []
-        w, h = float(pil.width), float(pil.height)
-        for box, score, label in zip(
-            results["boxes"], results["scores"], results["labels"], strict=False
-        ):
-            x1, y1, x2, y2 = (float(v) for v in box)
-            bbox = (max(0.0, x1), max(0.0, y1), min(w, x2), min(h, y2))
-            det = Detection2DBBox(
-                bbox=bbox,
-                track_id=-1,
-                class_id=int(label),
-                confidence=float(score),
-                name=queries[int(label)],
-                ts=image.ts,
-                image=image,
             )
-            if det.is_valid():
-                detections.append(det)
 
-        return ImageDetections2D(image=image, detections=detections)
+        batch: list[ImageDetections2D] = []
+        for image, pil, result in zip(images, pils, results, strict=True):
+            detections: list[Detection2DBBox] = []
+            w, h = float(pil.width), float(pil.height)
+            for box, score, label in zip(
+                result["boxes"], result["scores"], result["labels"], strict=False
+            ):
+                x1, y1, x2, y2 = (float(v) for v in box)
+                bbox = (max(0.0, x1), max(0.0, y1), min(w, x2), min(h, y2))
+                det = Detection2DBBox(
+                    bbox=bbox,
+                    track_id=-1,
+                    class_id=int(label),
+                    confidence=float(score),
+                    name=queries[int(label)],
+                    ts=image.ts,
+                    image=image,
+                )
+                if det.is_valid():
+                    detections.append(det)
+            batch.append(ImageDetections2D(image=image, detections=detections))
+        return batch
 
     def query_score_rows(
         self,
@@ -122,8 +168,9 @@ def query_score_rows(
         queries and refuse. Returns pixel ``(x1, y1, x2, y2)`` boxes and
         their score rows.
         """
+        self.forwards += 1
         pil = PILImage.fromarray(image.to_rgb().data)
-        with torch.inference_mode():
+        with torch.inference_mode(), self._autocast():
             inputs = self._processor(text=[queries], images=pil, return_tensors="pt").to(
                 self.config.device
             )
@@ -143,7 +190,98 @@ def query_score_rows(
         boxes[:, 1::2] = boxes[:, 1::2].clip(0.0, float(pil.height))
         return boxes, kept
 
+    def _frame_features(
+        self, image: Image
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Text-independent tower outputs for one frame, cached by frame ts.
+
+        Returns unit-normalized per-box class embeddings, the class head's
+        logit shift and scale, and clipped pixel ``(x1, y1, x2, y2)`` boxes.
+        A miss runs the image tower once; a hit runs nothing.
+        """
+        key = image.ts
+        if key in self._features:
+            self._features.move_to_end(key)
+            return self._features[key]
+
+        from transformers.image_transforms import center_to_corners_format
+
+        self.forwards += 1
+        pil = PILImage.fromarray(image.to_rgb().data)
+        model = self._model
+        head = model.class_head
+        with torch.inference_mode(), self._autocast():
+            pixel_values = self._processor(images=pil, return_tensors="pt").pixel_values.to(
+                self.config.device
+            )
+            feature_map = model.image_embedder(pixel_values)[0]
+            batch, height, width, dim = feature_map.shape
+            image_feats = feature_map.reshape(batch, height * width, dim)
+            pred_boxes = model.box_predictor(image_feats, feature_map)
+            class_embeds = head.dense0(image_feats)
+            class_embeds = class_embeds / (
+                torch.linalg.norm(class_embeds, dim=-1, keepdim=True) + 1e-6
+            )
+            logit_shift = head.logit_shift(image_feats)
+            logit_scale = head.elu(head.logit_scale(image_feats)) + 1
+
+            # the same conversion post_process_grounded_object_detection runs:
+            # corners in model dtype, scaled to the padded square in float32
+            boxes = center_to_corners_format(pred_boxes)[0].float() * float(
+                max(pil.width, pil.height)
+            )
+            boxes[:, 0::2] = boxes[:, 0::2].clip(0.0, float(pil.width))
+            boxes[:, 1::2] = boxes[:, 1::2].clip(0.0, float(pil.height))
+
+        entry = (class_embeds[0], logit_shift[0], logit_scale[0], boxes)
+        self._features[key] = entry
+        if len(self._features) > _FEATURE_CACHE_MAX:
+            self._features.popitem(last=False)
+        return entry
+
+    def _query_embeds(self, queries: list[str]) -> torch.Tensor:
+        """Unit-normalized text embeddings, one row per query, cached per string."""
+        missing = [q for q in queries if q not in self._text_embeds]
+        if missing:
+            with torch.inference_mode(), self._autocast():
+                inputs = self._processor(text=[missing], return_tensors="pt").to(self.config.device)
+                embeds = self._model.owlv2.get_text_features(**inputs)
+                embeds = embeds / torch.linalg.norm(embeds, ord=2, dim=-1, keepdim=True)
+            for query, embed in zip(missing, embeds, strict=True):
+                self._text_embeds[query] = embed
+        return torch.stack([self._text_embeds[q] for q in queries])
+
+    def query_score_rows_batch(
+        self,
+        images: list[Image],
+        queries: list[str],
+        threshold: float = 0.1,
+    ) -> list[tuple[np.ndarray, np.ndarray]]:
+        """``query_score_rows`` over several images: one ``(boxes, rows)`` each.
+
+        The image tower runs only for frames missing from the feature cache;
+        scoring any query set against a cached frame is a matmul against its
+        stored class embeddings, so repeated windows and new labels on seen
+        frames cost no model forwards.
+        """
+        with torch.inference_mode():
+            query_embeds = self._query_embeds(queries)
+            query_embeds = query_embeds / (
+                torch.linalg.norm(query_embeds, dim=-1, keepdim=True) + 1e-6
+            )
+            out: list[tuple[np.ndarray, np.ndarray]] = []
+            for image in images:
+                class_embeds, logit_shift, logit_scale, boxes_px = self._frame_features(image)
+                logits = (class_embeds @ query_embeds.T + logit_shift) * logit_scale
+                scores = torch.sigmoid(logits.to(torch.float32))
+                keep = scores.max(dim=-1).values > threshold
+                out.append((boxes_px[keep].cpu().numpy(), scores[keep].cpu().numpy()))
+        return out
+
     def stop(self) -> None:
+        self.score_cache.clear()
+        self._features.clear()
+        self._text_embeds.clear()
         if "_processor" in self.__dict__:
             del self.__dict__["_processor"]
         super().stop()