@pentatonic-ai/ai-agent-sdk 0.7.1 → 0.7.3

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@pentatonic-ai/ai-agent-sdk",
-  "version": "0.7.1",
+  "version": "0.7.3",
   "description": "TES SDK — LLM observability and lifecycle tracking via Pentatonic Thing Event System. Track token usage, tool calls, and conversations. Manage things through event-sourced lifecycle stages with AI enrichment and vector search.",
   "type": "module",
   "main": "./dist/index.cjs",

package/packages/memory-engine/compat/server.py

@@ -85,6 +85,19 @@ class SearchRequest(BaseModel):
     query: str
     limit: Optional[int] = 10
     min_score: Optional[float] = 0.001
+    # Tenant scope. Required for multi-tenant deployments. Forwarded to
+    # layers that support arena filtering natively (L6); applied as a
+    # post-filter on the shim for layers that don't yet (L2, L4, L5).
+    # When unset, search is global — same behaviour as v0.7.x; safe for
+    # single-tenant deployments. Multi-tenant callers MUST set this.
+    arena: Optional[str] = None
+    # Arbitrary metadata equality filters, applied as a post-filter on
+    # the shim. Useful for `kind`, `layer_type`, `source_repo`, etc.
+    # Keys not present on a result's metadata are treated as no-match.
+    # Each pair is exact string equality. Engine doesn't currently
+    # forward these to underlying stores, so over-fetch happens; the
+    # shim trims to the requested limit after filtering.
+    metadata_filter: Optional[dict[str, Any]] = None
 
 
 class ForgetRequest(BaseModel):
@@ -424,6 +437,51 @@ async def store_batch(req: StoreBatchRequest):
     }
 
 
+def _apply_metadata_filters(results: list[dict[str, Any]], req: SearchRequest) -> list[dict[str, Any]]:
+    """Post-filter results by arena + arbitrary metadata equality.
+
+    Many layer searches don't yet honour arena/metadata at the storage
+    level, so the shim enforces tenant isolation here as defence in
+    depth. Even if the underlying layer leaks across arenas, the shim
+    drops cross-tenant rows before returning.
+    """
+    arena = req.arena
+    extra = req.metadata_filter or {}
+    if not arena and not extra:
+        return results
+    out: list[dict[str, Any]] = []
+    for item in results:
+        meta = item.get("metadata") or {}
+        if arena:
+            row_arena = meta.get("arena") or item.get("arena")
+            if row_arena and row_arena != arena:
+                continue
+            # If row has no arena tag at all, drop on multi-tenant
+            # safety: a row without arena predates the multi-tenant
+            # plumbing and could belong to anyone.
+            if arena and not row_arena:
+                continue
+        ok = True
+        for k, v in extra.items():
+            if str(meta.get(k, "")) != str(v):
+                ok = False
+                break
+        if ok:
+            out.append(item)
+    return out
+
+
+def _search_overfetch(req: SearchRequest) -> int:
+    """Decide how many results to over-fetch from layers.
+
+    Post-filtering can drop many rows; we ask layers for more than the
+    user's limit so we have headroom after filtering. 5x is a balance
+    between accuracy and latency.
+    """
+    base = req.limit or 10
+    return base * 5 if (req.arena or req.metadata_filter) else base * 3
+
+
 @app.post("/search")
 async def search(req: SearchRequest):
     """
@@ -431,6 +489,12 @@ async def search(req: SearchRequest):
     queries L0 BM25, L4 vec, L5 Milvus, L6 doc-store in parallel and fuses
     the results with Reciprocal Rank Fusion. L3 KG adds entity-aware
     boosting for graph queries.
+
+    Multi-tenancy: pass `arena` to scope results to a single tenant.
+    Underlying layers may or may not honour arena natively (L6 does;
+    L2/L4/L5 don't yet — engine TODO); the shim applies arena as a
+    post-filter regardless, so cross-tenant leakage is prevented even
+    when a layer is non-compliant.
     """
     if not req.query:
         return {"results": []}
@@ -452,10 +516,19 @@ async def search(req: SearchRequest):
         import asyncio
         async def _q_l6(query: str):
             try:
+                params: dict[str, Any] = {
+                    "q": query,
+                    "limit": _search_overfetch(req),
+                    "method": "hybrid",
+                }
+                if req.arena:
+                    # L6 supports arena natively (l6-document-store.py:837).
+                    # Forward it so the underlying Milvus query and FTS
+                    # query both filter to this tenant before returning.
+                    params["arena"] = req.arena
                 r = await _client().get(
                     f"{L6_DOC_URL}/search",
-                    params={"q": query, "limit": (req.limit or 10) * 3,
-                            "method": "hybrid"},
+                    params=params,
                     timeout=30.0,
                 )
                 r.raise_for_status()
@@ -534,21 +607,35 @@ async def search(req: SearchRequest):
                         if m:
                             attached_meta = m
                             break
+            # Some layers surface arena/kind/layer_type as top-level
+            # fields on the raw item rather than nested under metadata.
+            # Pull them through into the formatted metadata so the
+            # post-filter can see them even on cache miss (e.g., after
+            # a compat container restart wipes _META_CACHE).
+            raw_top_level = {
+                k: item[k]
+                for k in ("arena", "kind", "layer_type", "doc_type", "source_file")
+                if item.get(k)
+            }
+            merged_meta = {**raw_top_level, **(attached_meta or item.get("metadata") or {})}
             out_results.append({
                 "id": key,
                 "content": item.get("text") or item.get("content") or item.get("snippet") or "",
-                "metadata": attached_meta or item.get("metadata") or {},
+                "metadata": merged_meta,
                 "similarity": float(rrf_scores[key]),
                 "layer_id": f"ml_{CLIENT_ID}_episodic",
                 "client_id": CLIENT_ID,
                 "source": item.get("source_file") or item.get("path") or "",
                 "engine_layer": "+".join(sorted(set(layer_provenance.get(key, [])))),
             })
-        return {"results": out_results}
+        # Defense-in-depth post-filter (arena + arbitrary metadata),
+        # then trim to the requested limit.
+        out_results = _apply_metadata_filters(out_results, req)
+        return {"results": out_results[: req.limit or 10]}
     try:
         r = await _client().get(
             f"{L2_PROXY_URL}/search",
-            params={"q": req.query, "limit": req.limit or 10},
+            params={"q": req.query, "limit": _search_overfetch(req)},
             timeout=30.0,
         )
         r.raise_for_status()
@@ -558,7 +645,7 @@ async def search(req: SearchRequest):
         try:
             r = await _client().post(
                 f"{L2_PROXY_URL}/v1/search",
-                json={"query": req.query, "limit": req.limit or 10,
+                json={"query": req.query, "limit": _search_overfetch(req),
                       "min_score": req.min_score or 0.001},
                 timeout=30.0,
             )
@@ -567,9 +654,14 @@ async def search(req: SearchRequest):
         except Exception as exc2:
             last_err = exc2
             try:
+                params: dict[str, Any] = {"q": req.query, "limit": _search_overfetch(req)}
+                # L6 supports arena natively; forward it on the
+                # last-resort fallback path too.
+                if req.arena:
+                    params["arena"] = req.arena
                 r = await _client().get(
                     f"{L6_DOC_URL}/search",
-                    params={"q": req.query, "limit": req.limit or 10},
+                    params=params,
                     timeout=10.0,
                 )
                 r.raise_for_status()
@@ -611,17 +703,28 @@ async def search(req: SearchRequest):
         if not chosen_id:
             chosen_id = (item.get("id") or item.get("doc_id")
                          or item.get("path") or item.get("source_file") or "")
+        # Pull arena/kind/layer_type from raw item top-level if present
+        # (L6 returns these as top-level columns; cache may be empty).
+        raw_top_level = {
+            k: item[k]
+            for k in ("arena", "kind", "layer_type", "doc_type", "source_file")
+            if item.get(k)
+        }
+        merged_meta = {**raw_top_level, **(attached_meta or item.get("metadata") or {})}
         out_results.append({
             "id": chosen_id,
             "content": item.get("text") or item.get("content") or item.get("snippet") or "",
-            "metadata": attached_meta or item.get("metadata") or {},
+            "metadata": merged_meta,
             "similarity": float(item.get("score") or item.get("similarity") or 0.0),
             "layer_id": f"ml_{CLIENT_ID}_episodic",
             "client_id": CLIENT_ID,
             "source": item.get("source", item.get("source_file", "")),
             "engine_layer": item.get("layer", item.get("source_layer", "")),
         })
-    return {"results": out_results}
+    # Defense-in-depth post-filter (arena + arbitrary metadata) on L2/L6
+    # fallback paths. Same logic as the BYPASS branch above.
+    out_results = _apply_metadata_filters(out_results, req)
+    return {"results": out_results[: req.limit or 10]}
 
 
 @app.post("/forget")