diff --git a/.env.example b/.env.example index d5864f5..ef4ad31 100644 --- a/.env.example +++ b/.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 diff --git a/README.md b/README.md index 5aaa5c5..924a4a4 100644 --- a/README.md +++ b/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) --- diff --git a/agentforge/agents/nodes.py b/agentforge/agents/nodes.py index 6f44c51..d228589 100644 --- a/agentforge/agents/nodes.py +++ b/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], diff --git a/agentforge/agents/state.py b/agentforge/agents/state.py index 8f414d1..018f4ff 100644 --- a/agentforge/agents/state.py +++ b/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 diff --git a/agentforge/api/main.py b/agentforge/api/main.py index aa621ac..22474e2 100644 --- a/agentforge/api/main.py +++ b/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: diff --git a/agentforge/api/tenancy.py b/agentforge/api/tenancy.py index 0f14ba9..b2886cc 100644 --- a/agentforge/api/tenancy.py +++ b/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 diff --git a/agentforge/cli.py b/agentforge/cli.py index 404304a..eb329db 100644 --- a/agentforge/cli.py +++ b/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") diff --git a/agentforge/config.py b/agentforge/config.py index 154d9f8..cd119f1 100644 --- a/agentforge/config.py +++ b/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) diff --git a/agentforge/rag/catalog.py b/agentforge/rag/catalog.py index 5b54c40..5f05535 100644 --- a/agentforge/rag/catalog.py +++ b/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 [] diff --git a/agentforge/rag/ingest.py b/agentforge/rag/ingest.py index afe6fdf..10c2271 100644 --- a/agentforge/rag/ingest.py +++ b/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: diff --git a/agentforge/rag/retriever.py b/agentforge/rag/retriever.py index a7b827e..fbdb106 100644 --- a/agentforge/rag/retriever.py +++ b/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: diff --git a/agentforge/rag/store.py b/agentforge/rag/store.py index 77bfd8d..64c1f97 100644 --- a/agentforge/rag/store.py +++ b/agentforge/rag/store.py @@ -70,14 +70,105 @@ def get_vector_store(): ) -def collection_is_empty() -> bool: +# Metadata key every chunk is tagged with so retrieval can scope to one tenant. +TENANT_FIELD = "tenant_id" + + +def tenant_filter(tenant_id: str): + """A backend-appropriate metadata filter selecting one tenant's chunks. + + Both backends store the chunk's LangChain metadata, but their filter + dialects differ: pgvector takes a Mongo-style operator dict over the jsonb + column, while Qdrant needs a typed ``Filter`` over the ``metadata`` payload + key. Returning the right object here keeps ``retrieve`` backend-agnostic. + """ + backend = get_settings().vector_store_backend.lower() + if backend == "qdrant": + from qdrant_client import models + + return models.Filter( + must=[ + models.FieldCondition( + key=f"metadata.{TENANT_FIELD}", + match=models.MatchValue(value=tenant_id), + ) + ] + ) + # pgvector (langchain_postgres) operator-dict over the jsonb metadata. + return {TENANT_FIELD: {"$eq": tenant_id}} + + +def collection_is_empty(tenant_id: str | None = None) -> bool: """True if the store holds no documents (used to guard auto-ingest). ``similarity_search`` returns the top-k regardless of distance, so a single - hit means the collection is populated; an empty list means it isn't. + hit means the collection is populated; an empty list means it isn't. When + ``tenant_id`` is given, the check is scoped to that tenant's chunks. """ try: - return len(get_vector_store().similarity_search("ping", k=1)) == 0 + kwargs = {"filter": tenant_filter(tenant_id)} if tenant_id is not None else {} + return len(get_vector_store().similarity_search("ping", k=1, **kwargs)) == 0 except Exception: # Collection not created yet / store unreachable — treat as empty. return True + + +def backfill_tenant(tenant_id: str) -> int: + """Stamp untagged legacy chunks with ``tenant_id`` (metadata-only, no re-embed). + + Idempotent: only rows missing the ``tenant_id`` field are touched, so it is + safe to run on every startup. Lets pre-multitenancy corpora keep working with + zero re-ingest — they simply become the default tenant's knowledge base. + Returns the number of chunks updated. + """ + settings = get_settings() + if settings.vector_store_backend.lower() == "qdrant": + return _backfill_qdrant(settings, tenant_id) + return _backfill_pgvector(settings, tenant_id) + + +def _backfill_pgvector(settings: Settings, tenant_id: str) -> int: + import json + + import psycopg + + from agentforge.config import libpq_url + + # Add tenant_id only to this collection's rows that don't already have it. + sql = """ + UPDATE langchain_pg_embedding e + SET cmetadata = e.cmetadata || %s::jsonb + FROM langchain_pg_collection c + WHERE c.uuid = e.collection_id + AND c.name = %s + AND NOT (e.cmetadata ? 'tenant_id') + """ + patch = json.dumps({TENANT_FIELD: tenant_id}) + with psycopg.connect(libpq_url(settings.database_url), connect_timeout=5) as conn: + cur = conn.execute(sql, (patch, settings.collection_name)) + conn.commit() + return cur.rowcount + + +def _backfill_qdrant(settings: Settings, tenant_id: str) -> int: + from qdrant_client import QdrantClient, models + + client = QdrantClient(url=settings.qdrant_url, api_key=settings.qdrant_api_key) + if not client.collection_exists(settings.collection_name): + return 0 + # Points whose metadata.tenant_id is absent → set it to the default tenant. + untagged = models.Filter( + must=[models.IsEmptyCondition(is_empty=models.PayloadField(key=f"metadata.{TENANT_FIELD}"))] + ) + before = client.count(settings.collection_name, count_filter=untagged, exact=True).count + if before: + # ``key="metadata"`` merges into the nested metadata object rather than + # overwriting it, so source/title survive the stamp. + client.set_payload( + collection_name=settings.collection_name, + payload={TENANT_FIELD: tenant_id}, + key="metadata", + points=untagged, + wait=True, + ) + return before diff --git a/tests/test_graph_flow.py b/tests/test_graph_flow.py index 62330ca..15b2a26 100644 --- a/tests/test_graph_flow.py +++ b/tests/test_graph_flow.py @@ -14,7 +14,7 @@ def _config(thread_id: str) -> dict: def test_grounded_answer_includes_citations(fresh_graph, monkeypatch): monkeypatch.setattr(nodes, "get_fast_model", lambda: route_model("knowledge")) - monkeypatch.setattr(nodes, "retrieve", lambda q: grounded_retrieval()) + monkeypatch.setattr(nodes, "retrieve", lambda q, tenant_id=None: grounded_retrieval()) monkeypatch.setattr( nodes, "get_chat_model", @@ -33,7 +33,7 @@ def test_grounded_answer_includes_citations(fresh_graph, monkeypatch): def test_out_of_scope_question_is_refused(fresh_graph, monkeypatch): monkeypatch.setattr(nodes, "get_fast_model", lambda: route_model("knowledge")) - monkeypatch.setattr(nodes, "retrieve", lambda q: empty_retrieval()) + monkeypatch.setattr(nodes, "retrieve", lambda q, tenant_id=None: empty_retrieval()) monkeypatch.setattr( nodes, "get_chat_model", diff --git a/tests/test_kb_isolation.py b/tests/test_kb_isolation.py new file mode 100644 index 0000000..222fee8 --- /dev/null +++ b/tests/test_kb_isolation.py @@ -0,0 +1,103 @@ +"""Per-tenant knowledge-base isolation (step 2). + +Offline coverage of the tag-on-write / filter-on-read seam: ingest stamps every +chunk with the tenant, retrieval forwards a tenant-scoped filter to the store, +the graph threads the tenant through, and the pgvector filter has the shape the +backend expects. No real Postgres/Qdrant is needed — a recording fake store +stands in for the backend. +""" + +from __future__ import annotations + +from langchain_core.documents import Document + +from agentforge.agents import nodes +from agentforge.rag import ingest as ingest_mod +from agentforge.rag import retriever as retriever_mod +from agentforge.rag import store + + +class _RecordingStore: + """Captures add_documents chunks and the filter passed to retrieval.""" + + def __init__(self, scored=None): + self.added: list[Document] = [] + self.last_filter = "" + self._scored = scored or [] + + def add_documents(self, chunks): + self.added.extend(chunks) + + def similarity_search_with_relevance_scores(self, query, k, filter=None): + self.last_filter = filter + return self._scored + + +def test_ingest_tags_every_chunk_with_tenant(monkeypatch): + fake = _RecordingStore() + monkeypatch.setattr(ingest_mod, "get_vector_store", lambda: fake) + monkeypatch.setattr( + ingest_mod, + "load_corpus", + lambda _dir: [Document(page_content="hello world", metadata={"source": "s", "title": "t"})], + ) + + count = ingest_mod.ingest("ignored", tenant_id="acme") + + assert count == len(fake.added) >= 1 + assert {c.metadata["tenant_id"] for c in fake.added} == {"acme"} + + +def test_ingest_defaults_to_default_tenant(monkeypatch): + fake = _RecordingStore() + monkeypatch.setattr(ingest_mod, "get_vector_store", lambda: fake) + monkeypatch.setattr( + ingest_mod, "load_corpus", lambda _dir: [Document(page_content="x", metadata={})] + ) + + ingest_mod.ingest("ignored") # no tenant given + + from agentforge.config import get_settings + + assert {c.metadata["tenant_id"] for c in fake.added} == {get_settings().default_tenant} + + +def test_retrieve_forwards_tenant_filter(monkeypatch): + fake = _RecordingStore(scored=[]) + monkeypatch.setattr(retriever_mod, "get_vector_store", lambda: fake) + + retriever_mod.retrieve("q", tenant_id="acme") + + # pgvector is the default backend in CI -> Mongo-style operator dict. + assert fake.last_filter == {"tenant_id": {"$eq": "acme"}} + + +def test_retrieve_uses_default_tenant_when_unset(monkeypatch): + fake = _RecordingStore(scored=[]) + monkeypatch.setattr(retriever_mod, "get_vector_store", lambda: fake) + + retriever_mod.retrieve("q") # no tenant + + from agentforge.config import get_settings + + assert fake.last_filter == {"tenant_id": {"$eq": get_settings().default_tenant}} + + +def test_retrieve_node_threads_tenant_from_state(monkeypatch): + seen = {} + + def fake_retrieve(query, tenant_id=None): + seen["query"] = query + seen["tenant"] = tenant_id + return retriever_mod.RetrievalResult() + + monkeypatch.setattr(nodes, "retrieve", fake_retrieve) + + nodes.retrieve_node({"redacted_question": "q", "tenant_id": "globex"}) + + assert seen == {"query": "q", "tenant": "globex"} + + +def test_pgvector_tenant_filter_shape(): + # Default backend is pgvector; the filter is the jsonb operator dict. + assert store.tenant_filter("acme") == {"tenant_id": {"$eq": "acme"}} diff --git a/tests/test_pii_flow.py b/tests/test_pii_flow.py index b6ddb38..7507380 100644 --- a/tests/test_pii_flow.py +++ b/tests/test_pii_flow.py @@ -11,7 +11,7 @@ def test_pii_redacted_before_retrieval(fresh_graph, monkeypatch): seen: dict[str, str] = {} - def spy_retrieve(query: str): + def spy_retrieve(query: str, tenant_id=None): seen["query"] = query return empty_retrieval() diff --git a/tests/test_supervisor.py b/tests/test_supervisor.py index bcc7ad4..dc83076 100644 --- a/tests/test_supervisor.py +++ b/tests/test_supervisor.py @@ -17,7 +17,7 @@ def _config(thread_id: str) -> dict: return {"configurable": {"thread_id": thread_id}} -def _no_retrieve(query: str): +def _no_retrieve(query: str, tenant_id=None): raise AssertionError("knowledge path was taken on an action request") @@ -47,7 +47,7 @@ def test_action_request_routes_to_act_agent_without_retrieval(fresh_graph, monke def test_unrecognized_classification_defaults_to_knowledge(fresh_graph, monkeypatch): # A garbage one-word reply must fall back to the safe (tool-less) path. monkeypatch.setattr(nodes, "get_fast_model", lambda: FakeChatModel(responses=[AIMessage("???")])) - monkeypatch.setattr(nodes, "retrieve", lambda q: empty_retrieval()) + monkeypatch.setattr(nodes, "retrieve", lambda q, tenant_id=None: empty_retrieval()) monkeypatch.setattr( nodes, "get_chat_model", lambda: FakeChatModel(responses=[AIMessage("Out of scope.")]) )