Feature/add learning agent guard

apps/agent-guard/python-service/embedder-container/pytest.ini

@@ -0,0 +1,2 @@
+[pytest]
+testpaths = tests

apps/agent-guard/python-service/embedder-container/src/embedder_service.py

@@ -7,6 +7,11 @@
 not have to host an ONNX/torch runtime and a model. The worker owns the cache
 logic (Redis vector store, rule hashing, comparison/alerting) and calls
 POST /embed here whenever it needs a vector.
+
+Stateless and shared across every customer — this is the only kind of
+computation safe to centralize into one multi-tenant fleet (see
+intent-classifier-container for the per-agent classifier, which is stateful
+and therefore deployed per customer instead).
 """
 
 import os

apps/agent-guard/python-service/embedder-container/tests/conftest.py

@@ -0,0 +1,8 @@
+"""Make the embedder container's `src/` importable as top-level modules in tests."""
+
+import pathlib
+import sys
+
+SRC = pathlib.Path(__file__).resolve().parent.parent / "src"
+if str(SRC) not in sys.path:
+    sys.path.insert(0, str(SRC))

apps/agent-guard/python-service/intent-classifier-container/Dockerfile

@@ -0,0 +1,17 @@
+FROM python:3.12-slim
+
+WORKDIR /app
+
+COPY src/requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY src/ .
+
+ENV PYTHONUNBUFFERED=1
+ENV PORT=8095
+EXPOSE 8095
+
+# One instance per customer — traffic volume per tenant is far below the
+# shared embedder-container's, so a single worker is the right default;
+# override with CLASSIFIER_WORKERS if a tenant's load warrants more.
+CMD exec uvicorn intent_classifier_service:app --host 0.0.0.0 --port 8095 --workers ${CLASSIFIER_WORKERS:-1}

apps/agent-guard/python-service/intent-classifier-container/pytest.ini

@@ -0,0 +1,2 @@
+[pytest]
+testpaths = tests

apps/agent-guard/python-service/intent-classifier-container/src/intent_classifier.py

@@ -0,0 +1,197 @@
+
+import logging
+import os
+import pickle
+import threading
+from typing import Any, Dict, List, Optional, Tuple
+
+import numpy as np
+from sklearn.calibration import CalibratedClassifierCV
+from sklearn.linear_model import LogisticRegression
+
+logger = logging.getLogger(__name__)
+
+# Minimum samples in the smallest class before we train at all. Below this a
+# discriminative model is meaningless → stay None and the caller ESCALATEs.
+_MIN_PER_CLASS = 2
+# Calibration needs ≥ cv samples per class; cap cv at 3 (the user-provided shape).
+_MAX_CV = 3
+
+
+def _debug_enabled() -> bool:
+    return os.getenv("INTENT_CLASSIFY_DEBUG", "").strip().lower() in ("1", "true", "yes", "on")
+
+_models: Dict[str, "_AgentModel"] = {}
+_lock = threading.Lock()
+
+
+class _AgentModel:
+    """A fitted multi-class classifier plus per-class centroids and risk map."""
+
+    __slots__ = ("clf", "centroids", "class_risk", "n_samples")
+
+    def __init__(self, clf, centroids: Dict[str, np.ndarray],
+                class_risk: Dict[str, str], n_samples: int):
+        self.clf = clf
+        self.centroids = centroids
+        self.class_risk = class_risk
+        self.n_samples = n_samples
+
+    def _cosine_to_centroid(self, x: np.ndarray, cls: str) -> float:
+        c = self.centroids.get(cls)
+        if c is None:
+            return 0.0
+        xn = float(np.linalg.norm(x))
+        if xn == 0.0:
+            return 0.0
+        return float(np.dot(x, c) / xn)  # c is already unit-norm
+
+    def predict_intent(self, X: np.ndarray, agent_host: str = "") -> List[Dict[str, Any]]:
+        proba = self.clf.predict_proba(X)
+        classes = self.clf.classes_
+        debug = _debug_enabled()
+        out: List[Dict[str, Any]] = []
+        for i in range(X.shape[0]):
+            row = proba[i]
+            order = np.argsort(row)[::-1]
+            top1_idx = int(order[0])
+            top2_idx = int(order[1]) if len(order) > 1 else top1_idx
+            top1_class = str(classes[top1_idx])
+            result = {
+                "intent": top1_class,
+                "confidence": float(row[top1_idx]),
+                "margin": float(row[top1_idx] - row[top2_idx]),
+                "centroid_similarity": self._cosine_to_centroid(X[i], top1_class),
+                "risk_category": self.class_risk.get(top1_class, "unknown"),
+            }
+            if debug:
+                top_k = [(str(classes[j]), round(float(row[j]), 4)) for j in order[:5]]
+                logger.info(f"[intent_clf] agent={agent_host!r} unit={i} n_classes={len(classes)} "
+                           f"top5={top_k} centroid_sim={result['centroid_similarity']:.4f}")
+            out.append(result)
+        return out
+
+
+def _filter_min_class(X: np.ndarray, y: np.ndarray) -> Tuple[np.ndarray, np.ndarray, List[str]]:
+    """Drop rows belonging to any class with fewer than _MIN_PER_CLASS samples,
+    so one singleton/rare intent doesn't sink training for every other intent
+    the agent already has enough examples for. Returns (X, y, dropped_classes)
+    — dropped classes just won't be predicted, so those requests ESCALATE
+    (safe) instead of the agent staying cold entirely."""
+    classes, counts = np.unique(y, return_counts=True)
+    keep = set(classes[counts >= _MIN_PER_CLASS].tolist())
+    dropped = sorted(str(c) for c in classes.tolist() if c not in keep)
+    if not dropped:
+        return X, y, []
+    mask = np.isin(y, list(keep))
+    return X[mask], y[mask], dropped
+
+
+def _fit(X: np.ndarray, y: np.ndarray):
+    """Fit a calibrated multi-class LogReg when there are enough per-class
+    samples, else return None (untrainable — caller stays cold)."""
+    classes, counts = np.unique(y, return_counts=True)
+    if len(classes) < 2:
+        return None  # one-class agent → not discriminative
+    min_count = int(counts.min())
+    if min_count < _MIN_PER_CLASS:
+        return None
+    base = LogisticRegression(max_iter=1000, class_weight="balanced")
+    cv = min(_MAX_CV, min_count)
+    clf = CalibratedClassifierCV(base, cv=cv)
+    clf.fit(X, y)
+    return clf
+
+
+def _compute_centroids(X: np.ndarray, y: np.ndarray) -> Dict[str, np.ndarray]:
+    """L2-renormalized mean training vector per class — used for the
+    centroid_similarity corroborating signal at predict time."""
+    centroids: Dict[str, np.ndarray] = {}
+    for label in np.unique(y):
+        rows = X[y == label]
+        c = rows.mean(axis=0)
+        norm = np.linalg.norm(c)
+        centroids[str(label)] = (c / norm) if norm > 0 else c
+    return centroids
+
+
+def train(agent_host: str, vectors: List[List[float]], labels: List[str],
+         risk_categories: Optional[Dict[str, str]] = None) -> Dict[str, object]:
+    """Train (or retrain) the agent's multi-class model from examples.
+
+    labels: one fine-grained intent string per vector (may include
+    "__other__"/"__background__"). risk_categories maps intent -> risk
+    category (delete/edit/create/fetch_pii/fetch_generic), static per-intent
+    metadata carried on the model for predict() to return alongside a match.
+    Atomic swap on success; on too-few-samples the previous model (if any) is
+    left untouched and trained=False is returned.
+    """
+    if not vectors or len(vectors) != len(labels):
+        return {"trained": False, "reason": "empty_or_mismatched"}
+    X = np.asarray(vectors, dtype=np.float32)
+    y = np.asarray(labels, dtype=object)
+    X, y, dropped = _filter_min_class(X, y)
+    if dropped:
+        logger.info(f"[intent_clf] agent={agent_host} dropping sparse classes "
+                   f"(< {_MIN_PER_CLASS} samples): {dropped}")
+    clf = _fit(X, y)
+    if clf is None:
+        return {"trained": False, "reason": "insufficient_per_class_samples",
+                "n_samples": int(len(y)), "dropped_classes": dropped}
+    centroids = _compute_centroids(X, y)
+    model = _AgentModel(clf, centroids, dict(risk_categories or {}), int(len(y)))
+    with _lock:
+        _models[agent_host] = model  # atomic swap
+    classes = sorted(set(y.tolist()))
+    logger.info(f"[intent_clf] trained agent={agent_host} n={len(y)} classes={classes}"
+               + (f" dropped={dropped}" if dropped else ""))
+    return {"trained": True, "n_samples": int(len(y)), "classes": classes, "dropped_classes": dropped}
+
+
+def has_model(agent_host: str) -> bool:
+    return agent_host in _models
+
+
+def predict(agent_host: str, vectors: List[List[float]]) -> List[Optional[Dict[str, Any]]]:
+    """Return {intent, confidence, margin, centroid_similarity, risk_category}
+    per vector, or None per vector when no model exists (cold start)."""
+    model = _models.get(agent_host)
+    if model is None or not vectors:
+        return [None] * len(vectors)
+    try:
+        X = np.asarray(vectors, dtype=np.float32)
+        return model.predict_intent(X)
+    except Exception as exc:
+        logger.warning(f"[intent_clf] predict failed agent={agent_host}: {exc}")
+        return [None] * len(vectors)
+
+
+def serialize_model(agent_host: str) -> Optional[bytes]:
+    """Pickle the agent's fitted model for the Redis L2 cache. None if the
+    agent has no model (nothing to persist)."""
+    model = _models.get(agent_host)
+    if model is None:
+        return None
+    return pickle.dumps(model)
+
+
+def load_model(agent_host: str, blob: bytes) -> bool:
+    """Unpickle a model blob (fetched from Redis by the caller on an L1 miss)
+    and install it atomically. Fail-open: a corrupt/incompatible blob (e.g.
+    pickled by a different container version) just leaves the agent cold —
+    never raises."""
+    try:
+        model = pickle.loads(blob)
+    except Exception as exc:
+        logger.warning(f"[intent_clf] failed to load model blob agent={agent_host}: {exc}")
+        return False
+    if not isinstance(model, _AgentModel):
+        logger.warning(f"[intent_clf] model blob agent={agent_host} is not an _AgentModel")
+        return False
+    with _lock:
+        _models[agent_host] = model
+    return True
+
+
+def stats() -> Dict[str, int]:
+    return {"agents_with_models": len(_models)}

apps/agent-guard/python-service/intent-classifier-container/src/intent_classifier_service.py

@@ -0,0 +1,100 @@
+"""Per-customer intent-classifier service.
+
+Deployed one instance per tenant (colocated with or reachable from that
+customer's mcp-endpoint-shield), so /classify and /train never see another
+customer's traffic — statefulness (the in-process `_models` registry in
+intent_classifier.py) is safe here in a way it wouldn't be on the shared,
+multi-tenant embedder-container.
+
+Contract:
+  POST /train    {agent_host, vectors, labels, risk_categories} -> fits and
+                  installs the model locally (L1), and returns a pickled,
+                  base64-encoded blob of it so the caller (Go) can persist it
+                  to Redis (L2) for other replicas of this same service.
+  POST /classify {agent_host, vectors, model_cache_key} -> on an L1 miss,
+                  fetches model_cache_key from Redis (L2) and installs it
+                  before predicting. Never errors on a cold agent — returns an
+                  all-None result, matching today's fail-open-to-ESCALATE
+                  contract.
+"""
+
+import base64
+import logging
+from typing import Dict, List, Optional
+
+from fastapi import FastAPI
+from pydantic import BaseModel
+
+import intent_classifier
+import redis_client
+
+logger = logging.getLogger(__name__)
+
+app = FastAPI(title="Akto Agent Guard Intent Classifier", version="1.0.0")
+
+
+class ClassifyRequest(BaseModel):
+    agent_host: str
+    vectors: List[List[float]]
+    # Redis key mcp-endpoint-shield wrote the trained model artifact under
+    # (see mcp/intentguard/store.go) — used only on an L1 (`_models`) miss.
+    model_cache_key: str = ""
+
+
+class ClassifyResult(BaseModel):
+    intent: Optional[str] = None
+    confidence: Optional[float] = None
+    margin: Optional[float] = None
+    centroid_similarity: Optional[float] = None
+    risk_category: Optional[str] = None
+
+
+class ClassifyResponse(BaseModel):
+    results: List[ClassifyResult]
+
+
+class TrainRequest(BaseModel):
+    agent_host: str
+    vectors: List[List[float]]
+    labels: List[str]  # fine-grained intent per vector, may include "__other__"/"__background__"
+    risk_categories: Dict[str, str] = {}  # intent -> delete|edit|create|fetch_pii|fetch_generic
+
+
+class TrainResponse(BaseModel):
+    trained: bool
+    reason: Optional[str] = None
+    n_samples: int = 0
+    classes: List[str] = []
+    dropped_classes: List[str] = []
+    # base64(pickle(model)) — present only when trained=True. The caller
+    # persists this to Redis; this service never writes it there itself.
+    model_blob: Optional[str] = None
+
+
+@app.get("/health")
+def health():
+    return {"ok": True, **intent_classifier.stats()}
+
+
+@app.post("/classify", response_model=ClassifyResponse)
+def classify(req: ClassifyRequest):
+    if not intent_classifier.has_model(req.agent_host) and req.model_cache_key:
+        blob = redis_client.get(req.model_cache_key)
+        if blob is not None:
+            intent_classifier.load_model(req.agent_host, blob)
+
+    results = intent_classifier.predict(req.agent_host, req.vectors)
+    return ClassifyResponse(results=[
+        ClassifyResult(**r) if r is not None else ClassifyResult() for r in results
+    ])
+
+
+@app.post("/train", response_model=TrainResponse)
+def train(req: TrainRequest):
+    result = intent_classifier.train(req.agent_host, req.vectors, req.labels, req.risk_categories)
+    model_blob = None
+    if result.get("trained"):
+        blob = intent_classifier.serialize_model(req.agent_host)
+        if blob is not None:
+            model_blob = base64.b64encode(blob).decode("ascii")
+    return TrainResponse(**result, model_blob=model_blob)

apps/agent-guard/python-service/intent-classifier-container/src/redis_client.py

@@ -0,0 +1,51 @@
+"""Thin read-only Redis client for the model-artifact L2 cache.
+
+mcp-endpoint-shield (Go) is the single writer of model blobs — it persists
+whatever intent_classifier.train() returns to Redis under a key it constructs
+itself (see mcp/intentguard/store.go). This service never writes to Redis; it
+only reads a blob on an L1 (`_models`) miss, so key-format ownership stays in
+one place (Go).
+
+Fail-open throughout: if REDIS_URL is unset or Redis is unreachable, get()
+returns None and the caller treats the agent as cold (safe — falls through to
+ESCALATE upstream), never raises into the request path.
+"""
+
+import logging
+import os
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+_client = None
+_client_initialized = False
+
+
+def _get_client():
+    global _client, _client_initialized
+    if _client_initialized:
+        return _client
+    _client_initialized = True
+    url = os.environ.get("REDIS_URL", "").strip()
+    if not url:
+        return None
+    try:
+        import redis
+        _client = redis.Redis.from_url(url)
+    except Exception as exc:
+        logger.warning(f"[intent_classifier] redis init failed: {exc}")
+        _client = None
+    return _client
+
+
+def get(key: str) -> Optional[bytes]:
+    if not key:
+        return None
+    client = _get_client()
+    if client is None:
+        return None
+    try:
+        return client.get(key)
+    except Exception as exc:
+        logger.warning(f"[intent_classifier] redis GET failed key={key!r}: {exc}")
+        return None