Merge pull request #318 from aghassel/feat/topic-beam-search

feat(retrieval): add TopicBeamSearch retriever (TopicGraphRAG)

Author: mykola-pereyma

docs-site/src/content/docs/lexical-graph/traversal-based-search.mdx

@@ -89,13 +89,35 @@ An entity network context consists of a filtered and ranked network of entities
 
 ### Retrievers
 
-Traversal-based search provides three different retrievers:
+Traversal-based search provides four different retrievers:
 
   - The `ChunkBasedSearch` retriever uses a vector similarity search to find information that is similar to the original query. The retriever first finds relevant chunks using vector similarity search. From these chunks, the retriever traverses topics, statements, and facts. Chunk-based search tends to return a narrowly-scoped set of results based on the statement and fact neighbourhoods of chunks that match the original query.
 	- The `EntityBasedSearch` retriever uses as its starting points the entities in an entity network context. From these entities, the retriever traverses facts, statements and topics. Entity-based search tends to return a broadly-scoped set of results, based on the neighbourhoods of individual entities and the facts that connect entities.
 	- The `EntityNetworkSearch` retriever uses textual transcriptions of an entity network context to drive vector searches for information that is dissimilar to the original query but nonetheless structurally relevant for creating an accurate and full response. These vector searches return chunks that are similar to 'something different from the question being asked'. From these chunks, the retriever traverses topics, statements, and facts to explore the structurally relevant space of dissimilar content.
-	
-By default, the traversal-based search is configured to use a combination of `ChunkBasedSearch` and `EntityNetworkSearch`. Together, these two retrievers provide access to content that is similar to the question being asked, plus content that is similar to 'something different from the question being asked'.
+	- The `TopicBeamSearch` retriever starts from the most relevant topics (found by a vector similarity search over the topic index) and then performs a beam search across the topic graph to gather additional related topics, before expanding the winning topics into their statements. Because it is topic-first, it tends to assemble a focused, topic-coherent set of statements and works well on multi-hop, legal, and temporal questions.
+
+By default, the traversal-based search is configured to use a combination of `ChunkBasedSearch` and `EntityNetworkSearch`. Together, these two retrievers provide access to content that is similar to the question being asked, plus content that is similar to 'something different from the question being asked'. `TopicBeamSearch` is used as a standalone, topic-first alternative.
+
+#### TopicBeamSearch options
+
+`TopicBeamSearch` accepts a notable option:
+
+  - `topic_reranker` (default `'none'`): optionally reranks the gathered topics by relevance to the query before their statements are collected, keeping the top `max_topics`. It mirrors the statement `reranker` option and accepts `'tfidf'` or `'bedrock'` (Cohere). Topic reranking is an opt-in, per-corpus tuning lever: it can help when answers are concentrated in a few topics, but is not enabled by default.
+
+The beam's neighbour-edge strategy and other tuning are configured through `ProcessorArgs`
+(mirroring the `chunk_beam_*` parameters), so they can be passed to
+`for_traversal_based_search(...)`:
+
+  - `use_same_chunk_neighbors` (default `True`) and `use_adjacent_chunk_neighbors`
+    (default `True`): expand the beam to topics co-occurring in the same chunk and in the
+    adjacent (next) chunk. Adjacent-chunk expansion is on by default because it improved
+    accuracy across datasets with no regressions in our benchmarks.
+  - `use_entity_neighbors` (default `False`): expand to topics linked through shared
+    entities. This has a large fan-out and tends to add noise, so it is off by default;
+    when enabled, `max_entity_neighbors` (default `100`) caps the neighbours per topic,
+    ranked by entity-connection strength.
+  - `topic_beam_width` (default `100`), `topic_beam_max_depth` (default `6`),
+    `topic_top_k` (default `50`): beam breadth, walk depth, and number of seed topics.
 
 ### Search results
 

lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/post_processors/bedrock_context_format.py

@@ -88,7 +88,16 @@ def _postprocess_nodes(
         if not nodes:
             return [NodeWithScore(node=TextNode(text='No relevant context'))]
 
-        # Group nodes by source
+        # Traversal-based retrievers (e.g. TopicBeamSearch, TopicBasedSearch) return one
+        # node per SearchResult, carrying the full source->topics->statements tree in
+        # metadata['result']. Format those as nested source/topic/statement XML so the
+        # topic grouping is preserved, with statements kept in document order within
+        # each topic.
+        result_nodes = [n for n in nodes if n.node.metadata.get('result')]
+        if result_nodes:
+            return self._format_search_result_nodes(result_nodes)
+
+        # Legacy per-statement nodes: group by source, emit per-statement XML.
         sources: Dict[str, List[NodeWithScore]] = {}
         for node in nodes:
             source_id = node.node.metadata['source']['sourceId']
@@ -124,3 +133,51 @@ def _postprocess_nodes(
             formatted_sources.append("\n".join(source_output))
 
         return [NodeWithScore(node=TextNode(text=formatted_source)) for formatted_source in formatted_sources]
+
+    @staticmethod
+    def _doc_order_key(statement: dict):
+        """Document-order proxy for a statement: chunk sequence, then statement id."""
+        return (str(statement.get('chunkId') or ''), str(statement.get('statementId') or ''))
+
+    @staticmethod
+    def _format_statement_text(statement: dict) -> str:
+        text = statement.get('statement', '') or ''
+        details = statement.get('details')
+        if details:
+            details = details.strip().replace('\n', ', ')
+            return f"{text} (details: {details})"
+        return text
+
+    def _format_search_result_nodes(self, nodes: List[NodeWithScore]) -> List[NodeWithScore]:
+        """Format traversal SearchResult nodes as nested source/topic/statement XML.
+
+        One node = one source. Topic grouping is preserved; statements within a topic
+        are ordered by document position (chunkId, statementId) so the generator reads
+        them in their original narrative order.
+        """
+        formatted_sources = []
+        for source_count, node in enumerate(nodes, 1):
+            result = node.node.metadata.get('result') or {}
+            source = result.get('source', {}) or {}
+            out = [f"<source_{source_count}>"]
+
+            metadata = source.get('metadata', {}) or {}
+            if metadata:
+                out.append(f"<source_{source_count}_metadata>")
+                for key, value in sorted(metadata.items()):
+                    out.append(f"\t<{key}>{value}</{key}>")
+                out.append(f"</source_{source_count}_metadata>")
+
+            for topic_count, topic in enumerate(result.get('topics', []), 1):
+                topic_name = (topic.get('topic') or '').replace('"', "'")
+                out.append(f'<topic_{source_count}.{topic_count} name="{topic_name}">')
+                statements = sorted(topic.get('statements', []), key=self._doc_order_key)
+                for st_count, statement in enumerate(statements, 1):
+                    tag = f"statement_{source_count}.{topic_count}.{st_count}"
+                    out.append(f"<{tag}>{self._format_statement_text(statement)}</{tag}>")
+                out.append(f"</topic_{source_count}.{topic_count}>")
+
+            out.append(f"</source_{source_count}>")
+            formatted_sources.append("\n".join(out))
+
+        return [NodeWithScore(node=TextNode(text=fs)) for fs in formatted_sources]

lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/processors/__init__.py

@@ -15,6 +15,7 @@
 from .prune_statements import PruneStatements
 from .remove_versioning_metadata import RemoveVersioningMetadata
 from .rerank_statements import RerankStatements
+from .rerank_topics import RerankTopics
 from .rescore_results import RescoreResults
 from .simplify_single_topic_results import SimplifySingleTopicResults
 from .sort_results import SortResults

lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/processors/processor_args.py

@@ -88,6 +88,28 @@ def __init__(self, **kwargs):
         self.chunk_cosine_top_k = kwargs.get('chunk_cosine_top_k', 50)
         self.chunk_beam_width = kwargs.get('chunk_beam_width', 10)
         self.chunk_beam_max_depth = kwargs.get('chunk_beam_max_depth', 3)
+        # TopicBeamSearch tuning (mirrors the chunk_beam_* parameters above).
+        self.topic_beam_width = kwargs.get('topic_beam_width', 100)
+        self.topic_beam_max_depth = kwargs.get('topic_beam_max_depth', 6)
+        self.topic_top_k = kwargs.get('topic_top_k', 50)
+        # Beam neighbour-edge strategies. adjacent-chunk co-occurrence is enabled by
+        # default: benchmarks showed it improves accuracy across datasets with no
+        # regressions, whereas entity-overlap (off by default) tends to add noise.
+        self.use_entity_neighbors = kwargs.get('use_entity_neighbors', False)
+        self.use_same_chunk_neighbors = kwargs.get('use_same_chunk_neighbors', True)
+        self.use_adjacent_chunk_neighbors = kwargs.get('use_adjacent_chunk_neighbors', True)
+        # Cap on entity-overlap neighbours per topic (ranked by connection strength) to
+        # bound the fan-out when use_entity_neighbors is enabled.
+        self.max_entity_neighbors = kwargs.get('max_entity_neighbors', 100)
+        # Defensive cap on chunk-based neighbour strategies (same-chunk + adjacent-chunk),
+        # mirroring max_entity_neighbors. Bounds the neighbour set per topic in case a chunk
+        # is referenced by an unusually large number of topics.
+        self.max_chunk_neighbors = kwargs.get('max_chunk_neighbors', 100)
+        # Topic-level reranking (mirrors `reranker`): 'none' | 'tfidf' | 'bedrock'.
+        # When set, RerankTopics scores each topic (name + statements) against the query
+        # and keeps the top `max_topics` before statement selection.
+        self.topic_reranker = kwargs.get('topic_reranker', 'none')
+        self.max_topics = kwargs.get('max_topics', 40)
         self.no_cache = kwargs.get('no_cache', None)
 
 

lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/processors/rerank_topics.py

@@ -0,0 +1,123 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+RerankTopics processor.
+
+Follows the same convention as RerankStatements: the reranking strategy is selected by a
+string arg (``ProcessorArgs.topic_reranker`` -> 'tfidf' | 'bedrock' | 'none') and
+dispatched here, rather than via ad-hoc flags. Where RerankStatements reranks statements
+*within* a topic, this processor reranks whole topics by scoring each topic's
+``name + all its statements`` against the query, keeps the top ``max_topics``, and drops the
+rest before the token-budget truncation runs.
+
+It also propagates each surviving topic's relevance score down to any statement that does not
+already carry a score. This lets the statement reranker be turned off (``reranker='none'``)
+so that statement selection is driven purely by topic relevance -- the ablation that answers
+"if a topic ranks high by the reranker, is statement-level tfidf still needed?".
+"""
+import logging
+import time
+from typing import List, Dict
+
+import boto3
+
+from graphrag_toolkit.lexical_graph import GraphRAGConfig
+from graphrag_toolkit.lexical_graph.retrieval.processors.processor_base import ProcessorBase
+from graphrag_toolkit.lexical_graph.retrieval.processors.processor_args import ProcessorArgs
+from graphrag_toolkit.lexical_graph.retrieval.model import Topic
+from graphrag_toolkit.lexical_graph.utils.reranker_utils import score_values_with_tfidf
+from graphrag_toolkit.lexical_graph.metadata import FilterConfig
+
+from llama_index.core.schema import QueryBundle
+
+logger = logging.getLogger(__name__)
+
+
+class RerankTopics(ProcessorBase):
+    """Reranks and prunes whole topics by query relevance, using the strategy named in
+    ``ProcessorArgs.topic_reranker``. No-op when topic_reranker is unset/'none'."""
+
+    def __init__(self, args: ProcessorArgs, filter_config: FilterConfig):
+        super().__init__(args, filter_config)
+
+    def _topic_text(self, topic: Topic) -> str:
+        stmts = ' '.join([s.statement_str or '' for s in topic.statements])
+        return (f'{topic.topic}\n{stmts}')[:4000]  # cap per-doc length for the reranker
+
+    def _score_with_tfidf(self, texts: List[str], query: QueryBundle) -> List[float]:
+        scored = score_values_with_tfidf(texts, [query.query_str], len(texts))
+        return [scored.get(t, 0.0) for t in texts]
+
+    def _score_with_bedrock(self, texts: List[str], query: QueryBundle) -> List[float]:
+        region = boto3.Session().region_name or getattr(GraphRAGConfig, 'aws_region', None) or 'us-east-1'
+        client = boto3.client('bedrock-agent-runtime', region_name=region)
+        model_arn = f'arn:aws:bedrock:{region}::foundation-model/{GraphRAGConfig.bedrock_reranking_model}'
+        sources = [{'type': 'INLINE', 'inlineDocumentSource': {'type': 'TEXT', 'textDocument': {'text': t}}}
+                   for t in texts]
+        resp = client.rerank(
+            queries=[{'type': 'TEXT', 'textQuery': {'text': query.query_str}}],
+            sources=sources,
+            rerankingConfiguration={
+                'type': 'BEDROCK_RERANKING_MODEL',
+                'bedrockRerankingConfiguration': {
+                    'numberOfResults': len(texts),
+                    'modelConfiguration': {'modelArn': model_arn},
+                },
+            },
+        )
+        scores = [0.0] * len(texts)
+        for r in resp['results']:
+            scores[r['index']] = r['relevanceScore']
+        return scores
+
+    def _process_results(self, search_results, query: QueryBundle):
+        mode = (self.args.topic_reranker or 'none').lower()
+        if mode == 'none':
+            return search_results
+
+        # Flatten (result_index, topic) pairs and build one document per topic.
+        pairs = []
+        texts = []
+        for ri, sr in enumerate(search_results.results):
+            for topic in sr.topics:
+                pairs.append((ri, topic))
+                texts.append(self._topic_text(topic))
+        if not texts:
+            return search_results
+
+        start = time.time()
+        if mode == 'tfidf':
+            scores = self._score_with_tfidf(texts, query)
+        elif mode == 'bedrock':
+            scores = self._score_with_bedrock(texts, query)
+        else:
+            logger.warning(f'Unknown topic_reranker "{mode}", skipping topic rerank')
+            return search_results
+        logger.debug(f'Topic rerank ({mode}) of {len(texts)} topics: {(time.time()-start)*1000:.0f}ms')
+
+        # Rank topics globally; keep the top max_topics.
+        order = sorted(range(len(pairs)), key=lambda i: scores[i], reverse=True)
+        keep = set(order[: self.args.max_topics])
+
+        kept_topics_by_result: Dict[int, List[Topic]] = {}
+        for i, (ri, topic) in enumerate(pairs):
+            if i not in keep:
+                continue
+            # Propagate topic relevance to statements that have no score yet, so selection can
+            # run on topic relevance alone when the statement reranker is disabled.
+            for stmt in topic.statements:
+                if stmt.score is None or stmt.score == 0.0:
+                    stmt.score = scores[i]
+            kept_topics_by_result.setdefault(ri, []).append(topic)
+
+        # Rebuild results, dropping topics (and now-empty sources) that did not survive.
+        new_results = []
+        for ri, sr in enumerate(search_results.results):
+            kept = kept_topics_by_result.get(ri)
+            if not kept:
+                continue
+            sr.topics = kept
+            new_results.append(sr)
+        search_results = search_results.with_new_results(results=new_results)
+        return search_results

lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/retrievers/__init__.py

@@ -14,6 +14,8 @@
 from .query_mode_retriever import QueryModeRetriever
 from .chunk_cosine_search import ChunkCosineSimilaritySearch
 from .semantic_chunk_beam_search import SemanticChunkBeamGraphSearch
+from .beam_search_base import BeamSearch
+from .topic_beam_search import TopicBeamSearch
 
 # Mapping of deprecated class names to their module paths within the deprecated sub-package
 _DEPRECATED_NAMES = {