Multi-tenancy step 2: per-tenant knowledge-base isolation

.env.example

@@ -45,9 +45,11 @@ CHUNK_OVERLAP=150
 REDACT_PII=true
 
 # Multi-tenancy — tenant id comes from the X-Tenant-ID header and namespaces
-# graph threads + the approval queue. Missing header falls back to
-# DEFAULT_TENANT (single-tenant deploys need nothing here). Set REQUIRE_TENANT
-# to reject requests without a valid tenant id. (Knowledge base is still shared.)
+# graph threads, the approval queue, and the knowledge base (chunks are tagged
+# at ingest, retrieval is filtered per tenant). Missing header falls back to
+# DEFAULT_TENANT (single-tenant deploys need nothing here; legacy untagged
+# chunks are auto-backfilled to DEFAULT_TENANT on startup). Set REQUIRE_TENANT
+# to reject requests without a valid tenant id.
 DEFAULT_TENANT=default
 REQUIRE_TENANT=false
 

README.md

@@ -166,6 +166,9 @@ Chat UI + admin views: live traces, eval scores, approval queue. The full-stack
 - ✅ Qdrant + alternative vector-store adapters
 - ✅ Additional reference examples (insurance, support)
 - Multi-tenancy
+  - ✅ Per-tenant thread + approval isolation
+  - ✅ Per-tenant knowledge-base isolation
+  - Pluggable auth (replace the trusted X-Tenant-ID header)
 
 ---
 

agentforge/agents/nodes.py

@@ -83,7 +83,7 @@ def supervisor_node(state: AgentState) -> dict:
 
 
 def retrieve_node(state: AgentState) -> dict:
-    result = retrieve(state["redacted_question"])
+    result = retrieve(state["redacted_question"], state.get("tenant_id"))
     return {
         "context": result.context_block(),
         "citations": [c.__dict__ for c in result.citations],

agentforge/agents/state.py

@@ -20,6 +20,10 @@ class AgentState(TypedDict, total=False):
     redacted_question: str
     pii_found: list[str]
 
+    # Tenant whose knowledge base retrieval is scoped to. Set by the API from
+    # the resolved tenant; falls back to the default tenant when absent.
+    tenant_id: str
+
     # Supervisor routing decision: "knowledge" (answer from docs) | "action"
     # (perform a sensitive operation). Set by ``supervisor_node``.
     route: str

agentforge/api/main.py

@@ -50,6 +50,18 @@
 @asynccontextmanager
 async def lifespan(_app: FastAPI):
     setup_observability()
+    # Stamp any pre-multitenancy (untagged) chunks as the default tenant so an
+    # existing single-tenant corpus keeps answering with zero re-ingest.
+    try:
+        from agentforge.rag.store import backfill_tenant
+
+        stamped = backfill_tenant(settings.default_tenant)
+        if stamped:
+            logger.info(
+                "Backfilled %d untagged chunks to tenant %r", stamped, settings.default_tenant
+            )
+    except Exception:
+        logger.warning("Tenant backfill skipped (store not ready)", exc_info=True)
     if settings.auto_ingest:
         try:
             from agentforge.rag.ingest import ingest_if_empty
@@ -152,12 +164,12 @@ def evals() -> EvalReport | None:
 
 
 @app.get("/documents", response_model=list[DocumentSummaryItem])
-def documents() -> list[DocumentSummaryItem]:
+def documents(tenant: str = Depends(resolve_tenant)) -> list[DocumentSummaryItem]:
     from agentforge.rag.catalog import list_documents
 
     return [
         DocumentSummaryItem(source=d.source, title=d.title, chunks=d.chunks)
-        for d in list_documents()
+        for d in list_documents(tenant)
     ]
 
 
@@ -181,7 +193,8 @@ def chat(req: ChatRequest, tenant: str = Depends(resolve_tenant)) -> ChatRespons
     thread_id = validate_thread_id(req.thread_id) if req.thread_id else str(uuid.uuid4())
     graph = get_compiled_graph()
     result = graph.invoke(
-        {"question": req.message}, config=_run_config(scoped_thread(tenant, thread_id))
+        {"question": req.message, "tenant_id": tenant},
+        config=_run_config(scoped_thread(tenant, thread_id)),
     )
     resp = _to_response(thread_id, result)
     _record_domain_metrics(resp)
@@ -220,7 +233,7 @@ async def event_generator():
         yield {"event": "thread", "data": thread_id}
         # Stream LLM tokens as the answer/act_agent specialists produce them.
         async for event in graph.astream_events(
-            {"question": req.message}, config=config, version="v2"
+            {"question": req.message, "tenant_id": tenant}, config=config, version="v2"
         ):
             if event["event"] == "on_chat_model_stream":
                 if event["metadata"].get("langgraph_node") not in _streaming_nodes:

agentforge/api/tenancy.py

@@ -1,12 +1,11 @@
 """Tenant identity + per-tenant scoping for graph threads and approvals.
 
-Step 1 of multi-tenancy: install the seam. The tenant is taken from the
-``X-Tenant-ID`` header (trusted for now — authentication arrives in a later
-step) and used to namespace checkpoint threads and the approval registry, so
-one tenant can never see or resume another's runs.
-
-What this does *not* yet do: isolate the knowledge base. All tenants still
-share one corpus at this stage — that is a separate, later step.
+The tenant is taken from the ``X-Tenant-ID`` header (trusted for now —
+authentication arrives in a later step) and used to namespace checkpoint
+threads and the approval registry, so one tenant can never see or resume
+another's runs. The knowledge base is scoped separately (chunks are tagged with
+tenant_id at ingest and retrieval is filtered to the caller's tenant); see
+``agentforge.rag.store.tenant_filter``.
 """
 
 from __future__ import annotations

agentforge/cli.py

@@ -2,6 +2,7 @@
 
     agentforge ingest [CORPUS_DIR]   # load a corpus into the vector store
     agentforge ask "QUESTION"        # one-shot query against the agent
+    agentforge backfill-tenant       # stamp untagged chunks with the default tenant
     agentforge serve                 # run the FastAPI gateway
 """
 
@@ -11,18 +12,34 @@
 import uuid
 
 
+def _tenant(args: argparse.Namespace) -> str:
+    from agentforge.config import get_settings
+
+    return args.tenant or get_settings().default_tenant
+
+
 def _cmd_ingest(args: argparse.Namespace) -> None:
     from agentforge.rag.ingest import ingest
 
-    count = ingest(args.corpus_dir)
-    print(f"Ingested {count} chunks from {args.corpus_dir}")
+    tenant = _tenant(args)
+    count = ingest(args.corpus_dir, tenant)
+    print(f"Ingested {count} chunks from {args.corpus_dir} for tenant {tenant!r}")
+
+
+def _cmd_backfill_tenant(args: argparse.Namespace) -> None:
+    from agentforge.rag.store import backfill_tenant
+
+    tenant = _tenant(args)
+    count = backfill_tenant(tenant)
+    print(f"Backfilled {count} untagged chunks to tenant {tenant!r}")
 
 
 def _cmd_ask(args: argparse.Namespace) -> None:
     from agentforge.agents import get_compiled_graph
 
     config = {"configurable": {"thread_id": str(uuid.uuid4())}}
-    result = get_compiled_graph().invoke({"question": args.question}, config=config)
+    invoke_input = {"question": args.question, "tenant_id": _tenant(args)}
+    result = get_compiled_graph().invoke(invoke_input, config=config)
     if result.get("__interrupt__"):
         print("[approval required]", result["__interrupt__"][0].value)
     else:
@@ -49,10 +66,18 @@ def main() -> None:
     p_ingest.add_argument(
         "corpus_dir", nargs="?", default="examples/banking-compliance/corpus"
     )
+    p_ingest.add_argument("--tenant", help="Tag chunks for this tenant (default: DEFAULT_TENANT)")
     p_ingest.set_defaults(func=_cmd_ingest)
 
+    p_backfill = sub.add_parser(
+        "backfill-tenant", help="Stamp untagged chunks with the default (or given) tenant"
+    )
+    p_backfill.add_argument("--tenant", help="Tenant id to stamp (default: DEFAULT_TENANT)")
+    p_backfill.set_defaults(func=_cmd_backfill_tenant)
+
     p_ask = sub.add_parser("ask", help="One-shot query against the agent")
     p_ask.add_argument("question")
+    p_ask.add_argument("--tenant", help="Scope retrieval to this tenant (default: DEFAULT_TENANT)")
     p_ask.set_defaults(func=_cmd_ask)
 
     p_serve = sub.add_parser("serve", help="Run the FastAPI gateway")

agentforge/config.py

@@ -72,11 +72,12 @@ class Settings(BaseSettings):
 
     # --- Multi-tenancy ---------------------------------------------------
     # Tenant identity comes from the X-Tenant-ID header and namespaces graph
-    # threads + the approval queue so tenants can't see/resume each other's
-    # runs. With require_tenant False (default), a request without the header
-    # falls back to default_tenant, so single-tenant deploys need no config.
-    # Set require_tenant True to reject requests that omit a valid tenant id.
-    # NOTE: the knowledge base is still shared across tenants at this stage.
+    # threads, the approval queue, and the knowledge base (chunks are tagged with
+    # tenant_id at ingest and retrieval is filtered to the caller's tenant) so
+    # tenants can't see/resume each other's runs or documents. With
+    # require_tenant False (default), a request without the header falls back to
+    # default_tenant, so single-tenant deploys need no config. Set require_tenant
+    # True to reject requests that omit a valid tenant id.
     default_tenant: str = Field(default="default")
     require_tenant: bool = Field(default=False)
 

agentforge/rag/catalog.py

@@ -13,7 +13,8 @@
 from agentforge.config import Settings, get_settings, libpq_url
 
 # Group the collection's chunks by source document. Joins the embedding rows to
-# their collection by name so we only count this app's collection.
+# their collection by name so we only count this app's collection, and scopes to
+# one tenant's documents via the tenant_id metadata tag.
 _SQL = """
 SELECT
     e.cmetadata ->> 'source'  AS source,
@@ -22,6 +23,7 @@
 FROM langchain_pg_embedding e
 JOIN langchain_pg_collection c ON c.uuid = e.collection_id
 WHERE c.name = %s
+  AND e.cmetadata ->> 'tenant_id' = %s
 GROUP BY e.cmetadata ->> 'source'
 ORDER BY source
 """
@@ -34,28 +36,39 @@ class DocumentSummary:
     chunks: int
 
 
-def _from_pgvector(settings: Settings) -> list[DocumentSummary]:
+def _from_pgvector(settings: Settings, tenant_id: str) -> list[DocumentSummary]:
     import psycopg
 
     with psycopg.connect(libpq_url(settings.database_url), connect_timeout=3) as conn:
-        rows = conn.execute(_SQL, (settings.collection_name,)).fetchall()
+        rows = conn.execute(_SQL, (settings.collection_name, tenant_id)).fetchall()
     return [
         DocumentSummary(source=row[0] or "(unknown)", title=row[1] or "", chunks=row[2])
         for row in rows
     ]
 
 
-def _from_qdrant(settings: Settings) -> list[DocumentSummary]:
-    from qdrant_client import QdrantClient
+def _from_qdrant(settings: Settings, tenant_id: str) -> list[DocumentSummary]:
+    from qdrant_client import QdrantClient, models
 
     client = QdrantClient(url=settings.qdrant_url, api_key=settings.qdrant_api_key)
     titles: dict[str, str] = {}
     chunks: dict[str, int] = defaultdict(int)
 
+    scope = models.Filter(
+        must=[
+            models.FieldCondition(
+                key="metadata.tenant_id", match=models.MatchValue(value=tenant_id)
+            )
+        ]
+    )
     offset = None
     while True:
         points, offset = client.scroll(
-            settings.collection_name, with_payload=True, limit=256, offset=offset
+            settings.collection_name,
+            scroll_filter=scope,
+            with_payload=True,
+            limit=256,
+            offset=offset,
         )
         for point in points:
             meta = (point.payload or {}).get("metadata", {})
@@ -71,13 +84,14 @@ def _from_qdrant(settings: Settings) -> list[DocumentSummary]:
     ]
 
 
-def list_documents() -> list[DocumentSummary]:
-    """One row per ingested source document, or ``[]`` if the store is unreachable."""
+def list_documents(tenant_id: str | None = None) -> list[DocumentSummary]:
+    """One row per ingested source document for the tenant, or ``[]`` if unreachable."""
     settings = get_settings()
+    tenant = tenant_id or settings.default_tenant
     try:
         if settings.vector_store_backend.lower() == "qdrant":
-            return _from_qdrant(settings)
-        return _from_pgvector(settings)
+            return _from_qdrant(settings, tenant)
+        return _from_pgvector(settings, tenant)
     except Exception:
         # Store not provisioned yet / unreachable / extra missing — empty catalog.
         return []

agentforge/rag/ingest.py

@@ -11,8 +11,9 @@
 
 from langchain_core.documents import Document
 
+from agentforge.config import get_settings
 from agentforge.rag.chunking import split_documents
-from agentforge.rag.store import get_vector_store
+from agentforge.rag.store import TENANT_FIELD, get_vector_store
 
 SUPPORTED_SUFFIXES = {".md", ".txt", ".markdown"}
 
@@ -37,25 +38,33 @@ def load_corpus(corpus_dir: str | Path) -> list[Document]:
     return documents
 
 
-def ingest(corpus_dir: str | Path) -> int:
-    """Load → chunk → embed → store. Returns the number of chunks written."""
+def ingest(corpus_dir: str | Path, tenant_id: str | None = None) -> int:
+    """Load → chunk → tag with tenant → embed → store. Returns chunks written.
+
+    Every chunk is stamped with ``tenant_id`` so retrieval can scope to one
+    tenant's knowledge base. Defaults to the configured default tenant.
+    """
+    tenant = tenant_id or get_settings().default_tenant
     documents = load_corpus(corpus_dir)
     chunks = split_documents(documents)
+    for chunk in chunks:
+        chunk.metadata[TENANT_FIELD] = tenant
     if chunks:
         get_vector_store().add_documents(chunks)
     return len(chunks)
 
 
-def ingest_if_empty(corpus_dir: str | Path) -> int:
-    """Ingest only when the store is empty — idempotent across restarts.
+def ingest_if_empty(corpus_dir: str | Path, tenant_id: str | None = None) -> int:
+    """Ingest only when the tenant's store is empty — idempotent across restarts.
 
     Returns the number of chunks written (0 if already populated).
     """
     from agentforge.rag.store import collection_is_empty
 
-    if not collection_is_empty():
+    tenant = tenant_id or get_settings().default_tenant
+    if not collection_is_empty(tenant):
         return 0
-    return ingest(corpus_dir)
+    return ingest(corpus_dir, tenant)
 
 
 def main() -> None:

agentforge/rag/retriever.py

@@ -15,7 +15,7 @@
 from langchain_core.documents import Document
 
 from agentforge.config import get_settings
-from agentforge.rag.store import get_vector_store
+from agentforge.rag.store import get_vector_store, tenant_filter
 
 
 @dataclass
@@ -50,12 +50,18 @@ def context_block(self) -> str:
         return "\n\n".join(parts)
 
 
-def retrieve(query: str) -> RetrievalResult:
+def retrieve(query: str, tenant_id: str | None = None) -> RetrievalResult:
     settings = get_settings()
     store = get_vector_store()
 
+    # Scope retrieval to the tenant's chunks; default tenant when unset so
+    # single-tenant callers keep working unchanged.
+    tenant = tenant_id or settings.default_tenant
+
     # Normalized relevance in [0, 1]; higher means more similar.
-    scored = store.similarity_search_with_relevance_scores(query, k=settings.retrieval_k)
+    scored = store.similarity_search_with_relevance_scores(
+        query, k=settings.retrieval_k, filter=tenant_filter(tenant)
+    )
 
     result = RetrievalResult()
     for doc, relevance in scored: