feat(a2a): ship RAG retrieval in the shared image; fix the RAG that was inert in prod

agent-templates/a2a/Dockerfile

@@ -33,7 +33,8 @@ COPY contracts/a2a/v1/client/ /app/contract-client/
 RUN pip install --no-cache-dir /app/contract-client
 
 # The shared A2A server + its runtime deps within the templates tree:
-#   a2a/        the server itself (python -m a2a.runtime)
+#   a2a/        the server itself (python -m a2a.runtime), INCLUDING a2a/rag/
+#               (retrieval against the shared Chroma instance — no separate COPY)
 #   providers/  `from providers import get_provider` (Managed-Agents binding)
 #   sync/       common/driver/role_loader the anthropic provider delegates to
 COPY a2a/ /app/a2a/
@@ -56,7 +57,11 @@ USER 10001
 # (service.port), so image and chart stay consistent via one source. Same rule for the
 # ADVERTISED endpoint: a per-product pod declares `a2a.inClusterUrl` in that values file
 # (contract v1.2.0), never an env var.
-ENV HOST=0.0.0.0 \
+# RAG is OFF unless the chart says otherwise. a2a.rag.build_from_env() returns None
+# when A2A_RAG_ENABLED is unset, and RAISES when it is set but misconfigured — so a
+# RAG-enabled pod can never start believing it is connected to a store it is not.
+ENV A2A_RAG_ENABLED=0 \
+    HOST=0.0.0.0 \
     A2A_VALUES_FILE=/config/values.json \
     A2A_REPOS_DIR=/repos \
     AGENT_PROVIDER=anthropic \

agent-templates/a2a/rag/__init__.py

@@ -0,0 +1,35 @@
+"""Shared RAG retrieval for the A2A server image.
+
+One Chroma instance (FuzeInfra's), one database allocated to A2A, one collection
+per tenant. Retrieval only — see client.py for why indexing stays with whoever
+owns the documents.
+
+    from a2a.rag import build_from_env, RagUnavailable
+
+    retriever = build_from_env()      # None when RAG is deliberately disabled
+    if retriever:
+        chunks = retriever.search("FuzeAgent", "how does auth work?")
+"""
+
+from .client import Chunk, RagRetriever
+from .config import RagConfig, from_env as config_from_env
+from .errors import RagConfigError, RagError, RagUnavailable
+
+__all__ = [
+    "Chunk", "RagRetriever", "RagConfig", "RagConfigError", "RagError",
+    "RagUnavailable", "build_from_env", "config_from_env",
+]
+
+
+def build_from_env(env: dict | None = None) -> RagRetriever | None:
+    """Retriever from the environment, or None when RAG is disabled.
+
+    None is returned ONLY for a deliberate opt-out (A2A_RAG_ENABLED unset).
+    Enabled-but-misconfigured raises RagConfigError; enabled-but-unreachable
+    raises RagUnavailable. Neither is degraded into a working-looking retriever
+    that answers every question with nothing.
+    """
+    config = config_from_env(env)
+    if config is None:
+        return None
+    return RagRetriever(config)

agent-templates/a2a/rag/client.py

@@ -0,0 +1,141 @@
+"""Retrieval against the shared Chroma instance. RETRIEVAL ONLY — never indexing.
+
+The boundary matters. Indexing needs the original documents, and those live with
+whoever owns them (FuzeAgent's orchestrator keeps them on a filesystem under
+KNOWLEDGE_STORAGE_PATH). The A2A server is a multi-tenant pod that must not hold
+any tenant's document store, so it never writes: it embeds a query, asks Chroma,
+and returns the chunk text Chroma already carries.
+
+That last part is why no document database is needed here. The indexer stores the
+chunk TEXT alongside the vector (`collection.add(..., documents=text_chunks)`),
+and the query asks for it back (`include=["documents", ...]`). A second hop to
+Mongo or Postgres to re-fetch the source would be needed only if Chroma held
+vectors alone — it does not. The document store is for re-indexing and download,
+not for answering a query.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+
+from .config import RagConfig
+from .errors import RagUnavailable
+
+log = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class Chunk:
+    """One retrieved passage. `score` is cosine similarity in [0, 1]."""
+
+    document_id: str
+    chunk_index: int
+    content: str
+    score: float
+    metadata: dict
+
+
+class RagRetriever:
+    """Chroma-backed retrieval for one A2A deployment.
+
+    Connects EAGERLY at construction. A lazy connection would move the failure to
+    the first query, where it reads as "no results" — the exact confusion this
+    package exists to remove.
+    """
+
+    def __init__(self, config: RagConfig, *, client=None, embedder=None):
+        self.config = config
+        self._embedder = embedder
+        self._client = client if client is not None else self._connect()
+
+    def _connect(self):
+        try:
+            import chromadb
+        except ImportError as exc:  # pragma: no cover - dependency is declared
+            raise RagUnavailable(
+                "chromadb is not installed but RAG is enabled; the image is built "
+                "without its retrieval dependency"
+            ) from exc
+        c = self.config
+        try:
+            return chromadb.HttpClient(
+                host=c.host,
+                port=c.port,
+                ssl=c.ssl,
+                tenant=c.tenant,
+                database=c.database,
+                settings=c.chroma_settings(),
+            )
+        except Exception as exc:
+            # Re-raised, never swallowed. A pod that cannot reach its vector store
+            # should fail its readiness probe, not serve confident empty answers.
+            raise RagUnavailable(
+                f"cannot reach Chroma at {c.host}:{c.port} "
+                f"(database={c.database!r}): {exc}"
+            ) from exc
+
+    def _embed(self, query: str) -> list[float]:
+        if self._embedder is None:
+            try:
+                from sentence_transformers import SentenceTransformer
+            except ImportError as exc:  # pragma: no cover
+                raise RagUnavailable(
+                    "sentence-transformers is not installed but RAG is enabled"
+                ) from exc
+            self._embedder = SentenceTransformer(self.config.embedding_model)
+        return self._embedder.encode(query).tolist()
+
+    def search(self, tenant: str, query: str, *, limit: int = 5,
+               min_score: float = 0.0) -> list[Chunk]:
+        """Chunks for `tenant`, most similar first.
+
+        An empty list means the corpus had nothing above `min_score`. It never
+        means the store was unreachable — that raises RagUnavailable.
+        """
+        if not query or not query.strip():
+            return []
+        name = self.config.collection_for(tenant)
+        try:
+            collection = self._client.get_collection(name=name)
+        except Exception as exc:
+            # A tenant with no collection has simply indexed nothing. That IS an
+            # empty corpus, so it is an empty result rather than an error — but it
+            # is logged, because "this tenant has no collection" and "this tenant's
+            # collection is empty" are worth telling apart in a transcript.
+            log.info("no Chroma collection %r for tenant %r (%s); empty corpus",
+                     name, tenant, exc)
+            return []
+        try:
+            res = collection.query(
+                query_embeddings=[self._embed(query)],
+                n_results=limit,
+                include=["documents", "metadatas", "distances"],
+            )
+        except Exception as exc:
+            raise RagUnavailable(f"query against {name!r} failed: {exc}") from exc
+
+        docs = (res.get("documents") or [[]])[0]
+        metas = (res.get("metadatas") or [[]])[0]
+        dists = (res.get("distances") or [[]])[0]
+        out: list[Chunk] = []
+        for doc, meta, dist in zip(docs, metas, dists):
+            meta = meta or {}
+            score = 1.0 - float(dist)
+            if score < min_score:
+                continue
+            out.append(Chunk(
+                document_id=str(meta.get("document_id", "")),
+                chunk_index=int(meta.get("chunk_index", 0) or 0),
+                content=doc,
+                score=score,
+                metadata=dict(meta),
+            ))
+        return out
+
+    def healthcheck(self) -> None:
+        """Raise RagUnavailable unless the store answers. Used by readiness."""
+        try:
+            self._client.heartbeat()
+        except Exception as exc:
+            raise RagUnavailable(f"Chroma heartbeat failed: {exc}") from exc

agent-templates/a2a/rag/config.py

@@ -0,0 +1,159 @@
+"""Environment -> RagConfig, with no default that can be silently wrong.
+
+The rule this file exists to enforce: a value whose wrong default still lets the
+process start is worse than no default at all. `CHROMA_HOST` defaulting to
+"localhost" is exactly that — the pod starts, connects to nothing, and reports
+an empty corpus forever.
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+
+from .errors import RagConfigError
+
+# Chroma resolves this string with an import path, not a keyword. The orchestrator
+# passed the literal "basic", which resolves to no class at all — so the client
+# was constructed without the auth it appeared to be configuring. Naming the two
+# real providers here keeps that mistake from being re-typed.
+TOKEN_AUTH_PROVIDER = "chromadb.auth.token_authn.TokenAuthClientProvider"
+BASIC_AUTH_PROVIDER = "chromadb.auth.basic_authn.BasicAuthClientProvider"
+
+_TRUE = {"1", "true", "yes", "on"}
+
+
+def _flag(name: str, default: str = "0") -> bool:
+    return os.environ.get(name, default).strip().lower() in _TRUE
+
+
+@dataclass(frozen=True)
+class RagConfig:
+    """Everything the retriever needs, validated.
+
+    `database` is the Chroma *database* the family allocates to A2A on the shared
+    FuzeInfra instance — one instance, its own database, per-tenant collections
+    inside it. Sharing the instance without a dedicated database would put
+    every product's collections in one flat namespace where a name collision is
+    a cross-tenant data leak, not a mistake.
+    """
+
+    host: str
+    port: int
+    ssl: bool
+    tenant: str
+    database: str
+    collection_prefix: str
+    embedding_model: str
+    auth_provider: str | None
+    auth_credentials: str | None
+    auth_header: str
+
+    @property
+    def enabled(self) -> bool:
+        return True
+
+    def collection_for(self, tenant: str) -> str:
+        """Per-tenant collection name. Isolation is by collection, not by filter.
+
+        A `where` filter on a shared collection is one forgotten argument away
+        from returning another tenant's chunks. A separate collection cannot be
+        read by omission.
+        """
+        slug = "".join(c if c.isalnum() or c in "-_" else "-" for c in tenant.lower())
+        if not slug:
+            raise RagConfigError("tenant name is empty after slugification")
+        return f"{self.collection_prefix}{slug}"
+
+    def chroma_settings(self):
+        """`chromadb.config.Settings` for this config, or None when unauthenticated.
+
+        Kept here rather than in the client so the self-tests can assert the exact
+        provider string without importing chromadb.
+        """
+        from chromadb.config import Settings  # imported lazily; see client.py
+
+        if not self.auth_provider:
+            return Settings()
+        kwargs = {
+            "chroma_client_auth_provider": self.auth_provider,
+            "chroma_client_auth_credentials": self.auth_credentials,
+        }
+        if self.auth_provider == TOKEN_AUTH_PROVIDER:
+            kwargs["chroma_auth_token_transport_header"] = self.auth_header
+        return Settings(**kwargs)
+
+
+def from_env(env: dict | None = None) -> RagConfig | None:
+    """Build a config, or return None when RAG is deliberately off.
+
+    None means "not enabled" and is the ONLY silent outcome. Enabled-but-broken
+    raises: that is a deploy defect and must not present as a working pod.
+    """
+    env = os.environ if env is None else env
+
+    def get(name, default=None):
+        v = env.get(name, default)
+        return v.strip() if isinstance(v, str) else v
+
+    enabled = (get("A2A_RAG_ENABLED", "0") or "0").lower() in _TRUE
+    if not enabled:
+        return None
+
+    host = get("CHROMA_HOST")
+    if not host:
+        raise RagConfigError(
+            "A2A_RAG_ENABLED is set but CHROMA_HOST is not. Refusing to fall back to "
+            "localhost: that is how RAG runs 'successfully' against nothing. Set "
+            "CHROMA_HOST to the shared instance, or unset A2A_RAG_ENABLED."
+        )
+
+    try:
+        port = int(get("CHROMA_PORT", "8000"))
+    except (TypeError, ValueError) as exc:
+        raise RagConfigError(f"CHROMA_PORT is not an integer: {get('CHROMA_PORT')!r}") from exc
+
+    database = get("CHROMA_DATABASE")
+    if not database:
+        raise RagConfigError(
+            "A2A_RAG_ENABLED is set but CHROMA_DATABASE is not. The shared instance is "
+            "shared: without its own database, A2A's collections land in whatever "
+            "namespace the default happens to be, alongside every other product's."
+        )
+
+    token = get("CHROMA_AUTH_TOKEN")
+    basic = get("CHROMA_AUTH_BASIC")
+    if token and basic:
+        raise RagConfigError(
+            "Both CHROMA_AUTH_TOKEN and CHROMA_AUTH_BASIC are set; pick one."
+        )
+    if token:
+        provider, credentials = TOKEN_AUTH_PROVIDER, token
+    elif basic:
+        if ":" not in basic:
+            raise RagConfigError("CHROMA_AUTH_BASIC must be 'user:password'.")
+        provider, credentials = BASIC_AUTH_PROVIDER, basic
+    elif (get("CHROMA_ALLOW_UNAUTHENTICATED", "0") or "0").lower() in _TRUE:
+        # Local dev / CI against an ephemeral Chroma. Explicit, because an empty
+        # credential silently meaning "no auth" is what shipped last time.
+        provider, credentials = None, None
+    else:
+        raise RagConfigError(
+            "A2A_RAG_ENABLED is set but no Chroma credential is configured. Set "
+            "CHROMA_AUTH_TOKEN (preferred) or CHROMA_AUTH_BASIC, or set "
+            "CHROMA_ALLOW_UNAUTHENTICATED=1 to say out loud that this instance has "
+            "no auth. An empty credential is never treated as 'no auth wanted'."
+        )
+
+    return RagConfig(
+        host=host,
+        port=port,
+        ssl=(get("CHROMA_SSL", "0") or "0").lower() in _TRUE,
+        tenant=get("CHROMA_TENANT", "default_tenant"),
+        database=database,
+        collection_prefix=get("CHROMA_COLLECTION_PREFIX", "a2a-"),
+        embedding_model=get("RAG_EMBEDDING_MODEL", "all-MiniLM-L6-v2"),
+        auth_provider=provider,
+        auth_credentials=credentials,
+        auth_header=get("CHROMA_AUTH_HEADER", "Authorization"),
+    )

agent-templates/a2a/rag/errors.py

@@ -0,0 +1,31 @@
+"""RAG failure modes, kept distinct because conflating them is the defect.
+
+`services/orchestrator/rag_integration.py` wrapped connect-time failure in a
+broad `except Exception` that set `chroma_client = None` and carried on. Every
+later search then returned an empty result — indistinguishable, to the caller
+and to an operator reading a transcript, from "the corpus genuinely has nothing
+relevant". RAG was inert in production for as long as that was true, and the
+only trace was one log line at startup.
+
+So: a retriever that cannot reach its store RAISES. It never returns `[]`.
+"""
+
+
+class RagError(Exception):
+    """Base for everything in this package."""
+
+
+class RagConfigError(RagError):
+    """The configuration is incoherent — raised at construction, before serving.
+
+    Deliberately fatal rather than degraded. A RAG-enabled deployment missing
+    CHROMA_HOST is a deploy defect, and the useful moment to say so is the pod's
+    first seconds, not silently on every query afterwards.
+    """
+
+
+class RagUnavailable(RagError):
+    """The store is configured but unreachable, or the query failed.
+
+    Distinct from "no matching chunks", which is an ordinary empty list.
+    """