War Story: A Haystack 1.9 Bug Returned Irrelevant Results for Our Customer Support Bot

import logging
from typing import List, Dict, Optional
from haystack.nodes.retriever.bm25 import BM25Retriever
from haystack.document_stores import ElasticsearchDocumentStore
from haystack.schema import Document

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class MultiTenantBM25Retriever(BM25Retriever):
    '''Custom BM25 retriever for multi-tenant document stores with tenant filtering.

    Buggy implementation matching Haystack 1.9.0's BM25 score calculation,
    which uses global term frequency instead of filtered term frequency.
    '''

    def __init__(self, document_store: ElasticsearchDocumentStore, tenant_id: str, *args, **kwargs):
        super().__init__(document_store, *args, **kwargs)
        self.tenant_id = tenant_id
        # Validate tenant ID format to prevent injection
        if not isinstance(tenant_id, str) or len(tenant_id) < 3:
            raise ValueError(f'Invalid tenant_id: {tenant_id}. Must be a string with length >= 3.')

        # Pre-fetch tenant-specific document count for debugging
        self.tenant_doc_count = self.document_store.get_document_count(filters={'tenant_id': self.tenant_id})
        logger.info(f'Initialized retriever for tenant {tenant_id} with {self.tenant_doc_count} documents')

    def retrieve(
        self,
        query: str,
        filters: Optional[Dict] = None,
        top_k: int = 10,
        index: Optional[str] = None,
        *args,
        **kwargs,
    ) -> List[Document]:
        '''Retrieve top_k documents for query, enforcing tenant filter.

        Args:
            query: User query string
            filters: Additional metadata filters (merged with tenant filter)
            top_k: Number of documents to return
            index: Elasticsearch index to query (defaults to document store's index)

        Returns:
            List of matching Document objects

        Raises:
            ValueError: If query is empty or filters are invalid
            ConnectionError: If document store is unreachable
        '''
        if not query or not isinstance(query, str):
            raise ValueError('Query must be a non-empty string')

        # Merge tenant filter with user-provided filters
        merged_filters = {'tenant_id': self.tenant_id}
        if filters:
            if not isinstance(filters, dict):
                raise ValueError(f'Filters must be a dict, got {type(filters)}')
            merged_filters.update(filters)

        try:
            # Call parent retrieve method with merged filters
            results = super().retrieve(
                query=query,
                filters=merged_filters,
                top_k=top_k,
                index=index or self.document_store.index,
                *args,
                **kwargs,
            )
        except ConnectionError as e:
            logger.error(f'Failed to connect to document store: {e}')
            raise
        except Exception as e:
            logger.error(f'Unexpected error during retrieval: {e}', exc_info=True)
            raise

        # Log retrieval metrics for monitoring
        logger.debug(f'Retrieved {len(results)} results for query: {query[:50]}...')
        if len(results) < top_k // 2:
            logger.warning(f'Low retrieval count ({len(results)}) for query: {query}')

        return results

def init_document_store() -> ElasticsearchDocumentStore:
    '''Initialize Elasticsearch document store with multi-tenant mapping.

    Returns:
        Configured ElasticsearchDocumentStore instance

    Raises:
        ConnectionError: If Elasticsearch is unreachable
    '''
    try:
        store = ElasticsearchDocumentStore(
            host='elasticsearch.prod.internal',
            port=9200,
            username='haystack_user',
            password='redacted',
            index='support_tickets_v2',
            embedding_field='embedding',
            embedding_dim=768,
            excluded_meta_data=['embedding'],
        )
        # Validate connection
        store.get_document_count()
        logger.info('Successfully connected to document store')
        return store
    except ConnectionError as e:
        logger.error(f'Elasticsearch connection failed: {e}')
        raise
    except Exception as e:
        logger.error(f'Failed to initialize document store: {e}', exc_info=True)
        raise

if __name__ == '__main__':
    # Example usage showing buggy behavior
    try:
        store = init_document_store()
        retriever = MultiTenantBM25Retriever(store, tenant_id='fintech_client_001')

        # Test query that triggered the bug: includes common stopwords
        test_query = 'how do I reset my account password'
        results = retriever.retrieve(test_query, top_k=5)

        print(f'Top result for \'{test_query}\': {results[0].content[:100] if results else \'No results\'}')
    except Exception as e:
        logger.error(f'Example usage failed: {e}', exc_info=True)

import math
from typing import List, Dict, Optional
import numpy as np

class BM25Scorer:
    '''Simplified reproduction of Haystack 1.9.0's BM25 score calculation.

    Contains the regression that caused irrelevant results in multi-tenant setups:
    uses global term frequency (across all documents in index) instead of
    filtered term frequency (only documents matching the query filter).
    '''

    def __init__(self, k1: float = 1.5, b: float = 0.75):
        self.k1 = k1
        self.b = b
        # Global term frequency: BUG - this includes all documents, not filtered ones
        self.global_term_freq: Dict[str, int] = {}
        self.global_doc_count = 0
        self.avg_doc_length = 0.0

    def fit(self, documents: List[Dict]) -> None:
        '''Calculate global term statistics across all input documents.

        Args:
            documents: List of document dicts with 'content' and 'id' keys
        '''
        if not documents:
            raise ValueError('Cannot fit BM25 scorer on empty document list')

        self.global_doc_count = len(documents)
        total_length = 0

        for doc in documents:
            if 'content' not in doc:
                raise KeyError(f'Document {doc.get(\'id\')} missing \'content\' field')
            content = doc['content'].lower().split()
            total_length += len(content)

            # Update global term frequency
            for term in content:
                self.global_term_freq[term] = self.global_term_freq.get(term, 0) + 1

        if self.global_doc_count == 0:
            self.avg_doc_length = 0.0
        else:
            self.avg_doc_length = total_length / self.global_doc_count

    def calculate_scores(
        self,
        query: str,
        candidate_docs: List[Dict],
        filters: Optional[Dict] = None,
    ) -> List[float]:
        '''Calculate BM25 scores for candidate documents.

        Args:
            query: User query string
            candidate_docs: List of candidate documents to score
            filters: Metadata filters (IGNORED in buggy 1.9.0 implementation)

        Returns:
            List of BM25 scores, one per candidate document

        Raises:
            RuntimeError: If fit() has not been called first
        '''
        if self.global_doc_count == 0:
            raise RuntimeError('BM25 scorer must be fit to documents before scoring')

        query_terms = query.lower().split()
        scores = []

        for doc in candidate_docs:
            if 'content' not in doc:
                raise KeyError(f'Candidate document {doc.get(\'id\')} missing \'content\' field')
            doc_content = doc['content'].lower().split()
            doc_length = len(doc_content)
            doc_score = 0.0

            # Calculate IDF using GLOBAL term frequency (BUG)
            for term in query_terms:
                # IDF = log( (N - n(term) + 0.5) / (n(term) + 0.5) + 1 )
                # N is global doc count, n(term) is global term frequency
                n_term = self.global_term_freq.get(term, 0)
                idf = math.log(
                    (self.global_doc_count - n_term + 0.5) / (n_term + 0.5) + 1
                )

                # Term frequency in current document
                tf_doc = doc_content.count(term)
                if tf_doc == 0:
                    continue

                # BM25 term score
                term_score = idf * (
                    (tf_doc * (self.k1 + 1)) /
                    (tf_doc + self.k1 * (1 - self.b + self.b * (doc_length / self.avg_doc_length)))
                )
                doc_score += term_score

            scores.append(doc_score)

        return scores

def simulate_multi_tenant_issue():
    '''Simulate the bug where global term frequency skews scores for filtered tenants.'''
    # Tenant A has 1000 documents about password resets
    tenant_a_docs = [
        {'id': f'a_{i}', 'content': 'how to reset account password', 'tenant_id': 'A'}
        for i in range(1000)
    ]

    # Tenant B has 10000 documents about loan applications (no password content)
    tenant_b_docs = [
        {'id': f'b_{i}', 'content': 'apply for personal loan with low interest', 'tenant_id': 'B'}
        for i in range(10000)
    ]

    # Global document store includes both tenants
    all_docs = tenant_a_docs + tenant_b_docs

    # Fit scorer on all docs (global stats)
    scorer = BM25Scorer()
    scorer.fit(all_docs)

    # Query for tenant A: "how do I reset my password"
    query = 'how do I reset my password'

    # Filtered candidates: only tenant A docs (correct filter)
    filtered_candidates = tenant_a_docs[:10]

    # Calculate scores with buggy global stats
    buggy_scores = scorer.calculate_scores(query, filtered_candidates)

    # Now fit scorer on only tenant A docs (correct filtered stats)
    correct_scorer = BM25Scorer()
    correct_scorer.fit(tenant_a_docs)
    correct_scores = correct_scorer.calculate_scores(query, filtered_candidates)

    print(f'Query: {query}')
    print(f'Global doc count: {scorer.global_doc_count}, Tenant A doc count: {len(tenant_a_docs)}')
    print(f'Global \'password\' term freq: {scorer.global_term_freq.get(\'password\', 0)}')
    print(f'Tenant A \'password\' term freq: {correct_scorer.global_term_freq.get(\'password\', 0)}')
    print(f'Buggy top score: {max(buggy_scores):.4f}')
    print(f'Correct top score: {max(correct_scores):.4f}')
    print(f'Score difference: {max(buggy_scores) - max(correct_scores):.4f}')

if __name__ == '__main__':
    simulate_multi_tenant_issue()

import math
from typing import List, Dict, Optional
import numpy as np
from collections import Counter

class FixedBM25Scorer:
    '''Fixed BM25 scorer matching Haystack 1.9.2+ implementation.

    Resolves the multi-tenant regression by calculating term frequency
    only on documents matching the query filter, not global index stats.
    '''

    def __init__(self, k1: float = 1.5, b: float = 0.75):
        self.k1 = k1
        self.b = b
        # Per-query filtered stats, reset on each score call
        self.filtered_term_freq: Dict[str, int] = {}
        self.filtered_doc_count = 0
        self.filtered_avg_doc_length = 0.0

    def _calculate_filtered_stats(self, candidate_docs: List[Dict]) -> None:
        '''Calculate term frequency and document stats for filtered candidate docs.

        Args:
            candidate_docs: List of documents matching the query filter

        Raises:
            ValueError: If candidate_docs is empty
        '''
        if not candidate_docs:
            raise ValueError('Cannot calculate stats for empty candidate list')

        self.filtered_doc_count = len(candidate_docs)
        total_length = 0
        term_freq = Counter()

        for doc in candidate_docs:
            if 'content' not in doc:
                raise KeyError(f'Document {doc.get(\'id\')} missing \'content\' field')
            content = doc['content'].lower().split()
            total_length += len(content)
            term_freq.update(content)

        self.filtered_term_freq = dict(term_freq)
        self.filtered_avg_doc_length = total_length / self.filtered_doc_count

    def calculate_scores(
        self,
        query: str,
        candidate_docs: List[Dict],
        filters: Optional[Dict] = None,
    ) -> List[float]:
        '''Calculate BM25 scores using filtered term statistics.

        Args:
            query: User query string
            candidate_docs: List of candidate documents matching filters
            filters: Metadata filters (used to validate candidates, logged for debugging)

        Returns:
            List of BM25 scores, one per candidate document

        Raises:
            ValueError: If candidate_docs is empty
        '''
        if not candidate_docs:
            raise ValueError('Cannot score empty candidate document list')

        # Calculate filtered stats for this query's candidate set
        self._calculate_filtered_stats(candidate_docs)

        query_terms = query.lower().split()
        scores = []

        for doc in candidate_docs:
            if 'content' not in doc:
                raise KeyError(f'Candidate document {doc.get(\'id\')} missing \'content\' field')
            doc_content = doc['content'].lower().split()
            doc_length = len(doc_content)
            doc_score = 0.0

            # Calculate IDF using FILTERED term frequency (FIX)
            for term in query_terms:
                # IDF = log( (N_filtered - n(term)_filtered + 0.5) / (n(term)_filtered + 0.5) + 1 )
                n_term = self.filtered_term_freq.get(term, 0)
                idf = math.log(
                    (self.filtered_doc_count - n_term + 0.5) / (n_term + 0.5) + 1
                )

                # Term frequency in current document
                tf_doc = doc_content.count(term)
                if tf_doc == 0:
                    continue

                # BM25 term score
                term_score = idf * (
                    (tf_doc * (self.k1 + 1)) /
                    (tf_doc + self.k1 * (1 - self.b + self.b * (doc_length / self.filtered_avg_doc_length)))
                )
                doc_score += term_score

            scores.append(doc_score)

        # Log filter context for auditability
        if filters:
            logging.debug(f'Calculated scores for {len(candidate_docs)} docs with filters: {filters}')

        return scores

def verify_fix():
    '''Verify that the fixed scorer resolves the multi-tenant issue.'''
    # Same setup as before: Tenant A (password docs) and Tenant B (loan docs)
    tenant_a_docs = [
        {'id': f'a_{i}', 'content': 'how to reset account password', 'tenant_id': 'A'}
        for i in range(1000)
    ]
    tenant_b_docs = [
        {'id': f'b_{i}', 'content': 'apply for personal loan with low interest', 'tenant_id': 'B'}
        for i in range(10000)
    ]

    # Query for tenant A
    query = 'how do I reset my password'
    filtered_candidates = tenant_a_docs[:10]

    # Fixed scorer: only uses filtered candidates for stats
    fixed_scorer = FixedBM25Scorer()
    fixed_scores = fixed_scorer.calculate_scores(query, filtered_candidates)

    # Correct scorer from earlier (fit on tenant A only)
    correct_scorer = BM25Scorer()
    correct_scorer.fit(tenant_a_docs)
    correct_scores = correct_scorer.calculate_scores(query, filtered_candidates)

    print(f'Fixed top score: {max(fixed_scores):.4f}')
    print(f'Correct top score: {max(correct_scores):.4f}')
    print(f'Score delta (should be ~0): {abs(max(fixed_scores) - max(correct_scores)):.4f}')

if __name__ == '__main__':
    import logging
    logging.basicConfig(level=logging.INFO)
    verify_fix()

@app.middleware('http')
async def shadow_traffic_middleware(request: Request, call_next):
    response = await call_next(request)
    # Duplicate support bot queries to shadow retriever
    if request.url.path == '/chat' and request.method == 'POST':
        try:
            body = await request.json()
            query = body.get('query')
            tenant_id = body.get('tenant_id')
            if query and tenant_id:
                # Send to shadow retriever (Haystack 1.9) in background
                background_tasks.add_task(
                    shadow_retriever.retrieve,
                    query=query,
                    filters={'tenant_id': tenant_id},
                    top_k=5
                )
        except Exception as e:
            logger.error(f'Shadow traffic failed: {e}')
    return response

class InstrumentedBM25Retriever(MultiTenantBM25Retriever):
    def retrieve(self, query: str, filters: Optional[Dict] = None, *args, **kwargs) -> List[Document]:
        start_time = time.time()
        try:
            results = super().retrieve(query, filters, *args, **kwargs)
            # Emit per-tenant metrics
            tenant_id = filters.get('tenant_id') if filters else self.tenant_id
            precision = self._calculate_precision(query, results)
            datadog.statsd.gauge(
                'retrieval.precision',
                precision,
                tags=[f'tenant_id:{tenant_id}', 'retriever:bm25']
            )
            return results
        except Exception as e:
            tenant_id = filters.get('tenant_id') if filters else self.tenant_id
            datadog.statsd.increment(
                'retrieval.errors',
                tags=[f'tenant_id:{tenant_id}', 'retriever:bm25']
            )
            raise
        finally:
            latency = (time.time() - start_time) * 1000
            tenant_id = filters.get('tenant_id') if filters else self.tenant_id
            datadog.statsd.histogram(
                'retrieval.latency_ms',
                latency,
                tags=[f'tenant_id:{tenant_id}', 'retriever:bm25']
            )

def test_bm25_multitenant_term_freq(haystack_bm25_retriever, elasticsearch_store):
    # Add tenant A (small: 10 password docs) and tenant B (large: 1000 loan docs)
    tenant_a_docs = [Document(content='reset password', meta={'tenant_id': 'A'}) for _ in range(10)]
    tenant_b_docs = [Document(content='apply for loan', meta={'tenant_id': 'B'}) for _ in range(1000)]
    elasticsearch_store.write_documents(tenant_a_docs + tenant_b_docs)

    # Retrieve for tenant A
    results = haystack_bm25_retriever.retrieve(
        query='reset password',
        filters={'tenant_id': 'A'},
        top_k=5
    )

    # Validate all results are tenant A docs
    assert all(doc.meta['tenant_id'] == 'A' for doc in results)
    assert len(results) == 5
    # Validate precision is >80% (all results relevant)
    assert 'password' in results[0].content.lower()