From 16da4d4e6b48da7c8b8c1f885b41db0453030246 Mon Sep 17 00:00:00 2001 From: Johnny Huynh <27847622+johnnyhuy@users.noreply.github.com> Date: Wed, 11 Nov 2026 02:05:02 +0900 Subject: [PATCH] fix(rag): apply channel scope before retrieval limit Co-authored-by: OpenAI Codex --- backend/app/rag.py | 3 +- backend/app/vector_store.py | 8 +- backend/tests/test_agent_runtime.py | 1 + backend/tests/test_retrieval.py | 110 ++++++++++++++++++++++++++++ backend/tests/test_rules.py | 2 +- mise.toml | 2 +- 6 files changed, 120 insertions(+), 6 deletions(-) create mode 100644 backend/tests/test_retrieval.py diff --git a/backend/app/rag.py b/backend/app/rag.py index eda94fc..e75a858 100644 --- a/backend/app/rag.py +++ b/backend/app/rag.py @@ -93,7 +93,8 @@ def _store_conversation_turn(self, conversation_id: str, user_id: str, channel_i def _agent_payload(self, question: str, channel_id=None, conversation_id=None, request_id=None, slack=None): - results = vector_store.similarity_search(question) + search_options = {"channel_id": channel_id} if channel_id else {} + results = vector_store.similarity_search(question, **search_options) if channel_id: results = [r for r in results if r.get("metadata", {}).get("channel_id") == channel_id] identifier = request_id or agent_runtime.request_id() diff --git a/backend/app/vector_store.py b/backend/app/vector_store.py index 2d3fb20..e27ac65 100644 --- a/backend/app/vector_store.py +++ b/backend/app/vector_store.py @@ -67,14 +67,16 @@ def add_documents(self, texts: List[str], metadatas: List[Dict[str, Any]] = None return self.vector_store.add_texts(texts=texts, metadatas=metadatas) - def similarity_search(self, query: str, k: int = None) -> List[Dict[str, Any]]: - """Search for similar documents""" + def similarity_search(self, query: str, k: int = None, *, + channel_id: str = None) -> List[Dict[str, Any]]: + """Search within a channel before limiting results, or globally if unscoped.""" self._ensure_initialized() if not self.vector_store: raise RuntimeError("Vector store not initialized") k = k or settings.RETRIEVAL_TOP_K - results = self.vector_store.similarity_search(query, k=k) + search_options = {"filter": {"channel_id": channel_id}} if channel_id else {} + results = self.vector_store.similarity_search(query, k=k, **search_options) return [ { diff --git a/backend/tests/test_agent_runtime.py b/backend/tests/test_agent_runtime.py index e94517f..82d2e70 100644 --- a/backend/tests/test_agent_runtime.py +++ b/backend/tests/test_agent_runtime.py @@ -39,6 +39,7 @@ def test_rag_routes_retrieved_context_and_stable_slack_identity(monkeypatch): submit = Mock(return_value={"id": "job-1", "status": "pending"}) monkeypatch.setattr(agent_runtime, "submit", submit) module.rag_engine.submit_slack("Question", "C1", "team:thread:C1:1.2", "event-1", "1.2") + vector.similarity_search.assert_called_once_with("Question", channel_id="C1") payload = submit.call_args.args[0] assert payload["request_id"] == "event-1" assert payload["conversation_id"] == "team:thread:C1:1.2" diff --git a/backend/tests/test_retrieval.py b/backend/tests/test_retrieval.py new file mode 100644 index 0000000..ece3b0f --- /dev/null +++ b/backend/tests/test_retrieval.py @@ -0,0 +1,110 @@ +"""Channel retrieval regressions using local Chroma and fixed embeddings.""" +import importlib.util +from pathlib import Path +import sys +from types import SimpleNamespace +from unittest.mock import Mock + +import chromadb +from chromadb.config import Settings as ChromaSettings +from langchain_community.vectorstores import Chroma +from langchain_core.embeddings import Embeddings +import pytest +from chromadb.telemetry.product.posthog import Posthog + +from app.config import settings +from app.rules import RuleEngine +from app.vector_store import VectorStore + + +class FixtureEmbeddings(Embeddings): + def embed_documents(self, texts): + distances = { + "Other channel first": 0.0, + "Other channel second": 0.1, + "Shared document": 0.2, + "Target channel first": 1.0, + "Target channel second": 2.0, + } + return [[distances[text], 1.0] for text in texts] + + def embed_query(self, text): + return [0.0, 1.0] + + +@pytest.fixture +def vector(tmp_path, monkeypatch): + monkeypatch.setattr(settings, "RETRIEVAL_TOP_K", 2) + monkeypatch.setattr(Posthog, "capture", lambda *args, **kwargs: None) + client = chromadb.PersistentClient( + path=str(tmp_path), settings=ChromaSettings(anonymized_telemetry=False), + ) + store = VectorStore() + store.vector_store = Chroma( + client=client, collection_name="retrieval-fixture", + embedding_function=FixtureEmbeddings(), + ) + store._initialized = True + store.vector_store.add_texts( + ["Other channel first", "Other channel second", "Shared document", + "Target channel first", "Target channel second"], + [{"source": "slack", "channel_id": "COTHER"}, + {"source": "slack", "channel_id": "COTHER"}, + {"source": "api"}, + {"source": "slack", "channel_id": "CTARGET"}, + {"source": "slack", "channel_id": "CTARGET"}], + ) + return store + + +def test_channel_filter_runs_before_top_k(vector): + results = vector.similarity_search("How do I connect?", channel_id="CTARGET") + assert [item["content"] for item in results] == [ + "Target channel first", "Target channel second", + ] + + +def test_unscoped_search_preserves_global_ranking_and_explicit_k(vector): + results = vector.similarity_search("How do I connect?", k=3) + assert [item["content"] for item in results] == [ + "Other channel first", "Other channel second", "Shared document", + ] + + +def test_empty_channel_does_not_fall_back_to_global_results(vector): + assert vector.similarity_search("How do I connect?", channel_id="CEMPTY") == [] + + +@pytest.mark.parametrize("path", ["api", "slack"]) +def test_rag_payload_uses_scoped_matches(vector, monkeypatch, path): + monkeypatch.setitem(sys.modules, "app.vector_store", SimpleNamespace(vector_store=vector)) + spec = importlib.util.spec_from_file_location( + "retrieval_rag_under_test", Path(__file__).parents[1] / "app/rag.py", + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + monkeypatch.setattr(module.rules, "rule_engine", RuleEngine([])) + generate = Mock(return_value={"answer": "fixture answer"}) + submit = Mock(return_value={"id": "fixture-job", "status": "pending"}) + monkeypatch.setattr(module.agent_runtime, "generate", generate) + monkeypatch.setattr(module.agent_runtime, "submit", submit) + + if path == "api": + module.rag_engine.generate_response( + "How do I connect?", channel_id="CTARGET", request_id="fixture-api", + ) + payload = generate.call_args.args[0] + else: + module.rag_engine.submit_slack( + "How do I connect?", "CTARGET", "fixture:thread", "fixture-slack", "1.2", + ) + payload = submit.call_args.args[0] + assert payload["slack"] == {"channel": "CTARGET", "thread_ts": "1.2"} + + assert payload["context"] == ( + "From slack:\nTarget channel first\n\nFrom slack:\nTarget channel second" + ) + assert payload["sources"] == [ + {"source": "slack", "channel_id": "CTARGET"}, + {"source": "slack", "channel_id": "CTARGET"}, + ] diff --git a/backend/tests/test_rules.py b/backend/tests/test_rules.py index 5b7ff30..41749e1 100644 --- a/backend/tests/test_rules.py +++ b/backend/tests/test_rules.py @@ -195,7 +195,7 @@ def test_unmatched_slack_question_uses_retrieval_without_pinned_answer(self, mon self._install_engine(monkeypatch, engine(rule("vpn", ["reset vpn"]))) monkeypatch.setattr( rag_module.vector_store, "similarity_search", - lambda _: [{"content": "Doc", "metadata": {"source": "slack", "channel_id": "C1"}}], + lambda _, **kwargs: [{"content": "Doc", "metadata": {"source": "slack", "channel_id": "C1"}}], ) submitted = {} def fake_submit(payload): diff --git a/mise.toml b/mise.toml index 0af49c2..93f4bc3 100644 --- a/mise.toml +++ b/mise.toml @@ -51,7 +51,7 @@ run = "docker compose ps" [tasks.test-backend] description = "Run backend tests" dir = "backend" -run = "uv run pytest tests/test_file_extraction_units.py tests/test_logging_config.py tests/test_logging.py tests/test_agent_runtime.py tests/test_rules.py tests/test_conversations.py tests/test_documents.py -v" +run = "uv run pytest tests/test_file_extraction_units.py tests/test_logging_config.py tests/test_logging.py tests/test_agent_runtime.py tests/test_rules.py tests/test_conversations.py tests/test_documents.py tests/test_retrieval.py -v" [tasks.migrate] description = "Run database migrations (upgrade to latest)"