feat - for global_search, add the support for tracking max_token usage / use it to stuff as many reports in one LLM call / sort the keypoints by score

Author: ksachdeva

examples/simple-app/app/query.py

@@ -72,6 +72,7 @@ def global_search(
         community_level=cast(CommunityLevel, level),
         weight_calculator=CommunityWeightCalculator(),
         artifacts=artifacts,
+        token_counter=TiktokenCounter(),
     )
 
     kp_generator = KeyPointsGenerator(
@@ -83,7 +84,9 @@ def global_search(
     kp_aggregator = KeyPointsAggregator(
         llm=make_llm_instance(llm_type, llm_model, cache_dir),
         prompt_builder=KeyPointsAggregatorPromptBuilder(),
-        context_builder=KeyPointsContextBuilder(),
+        context_builder=KeyPointsContextBuilder(
+            token_counter=TiktokenCounter(),
+        ),
     )
 
     global_search = GlobalSearch(

src/langchain_graphrag/query/global_search/key_points_aggregator/context_builder.py

@@ -1,8 +1,11 @@
+import logging
+
 from langchain_core.documents import Document
 
 from langchain_graphrag.query.global_search.key_points_generator.utils import (
     KeyPointsResult,
 )
+from langchain_graphrag.utils.token_counter import TokenCounter
 
 _REPORT_TEMPLATE = """
 --- {analyst} ---
@@ -13,17 +16,48 @@
 
 """
 
+_LOGGER = logging.getLogger(__name__)
+
 
 class KeyPointsContextBuilder:
+    def __init__(
+        self,
+        token_counter: TokenCounter,
+        max_tokens: int = 8000,
+    ):
+        self._token_counter = token_counter
+        self._max_tokens = max_tokens
+
     def __call__(self, key_points: dict[str, KeyPointsResult]) -> list[Document]:
         documents: list[Document] = []
+        total_tokens = 0
+        max_token_limit_reached = False
         for k, v in key_points.items():
+            if max_token_limit_reached:
+                break
             for p in v.points:
                 report = _REPORT_TEMPLATE.format(
                     analyst=k,
                     score=p.score,
                     content=p.description,
                 )
-                documents.append(Document(page_content=report))
+                report_token = self._token_counter.count_tokens(report)
+                if total_tokens + report_token > self._max_tokens:
+                    _LOGGER.warning("Reached max tokens for key points aggregation ...")
+                    max_token_limit_reached = True
+                    break
+                total_tokens += report_token
+                documents.append(
+                    Document(
+                        page_content=report,
+                        metadata={"score": p.score, "analyst": k},
+                    )
+                )
 
-        return documents
+        # we now sort the documents based on the
+        # importance score of the key points
+        return sorted(
+            documents,
+            key=lambda x: x.metadata["score"],
+            reverse=True,
+        )

src/langchain_graphrag/query/global_search/key_points_generator/context_builder.py

@@ -1,3 +1,5 @@
+import logging
+
 from langchain_core.documents import Document
 
 from langchain_graphrag.indexing.artifacts import IndexerArtifacts
@@ -6,6 +8,7 @@
     CommunityWeightCalculator,
 )
 from langchain_graphrag.types.graphs.community import CommunityId, CommunityLevel
+from langchain_graphrag.utils.token_counter import TokenCounter
 
 _REPORT_TEMPLATE = """
 --- Report {report_id} ---
@@ -19,17 +22,23 @@
 
 """
 
+_LOGGER = logging.getLogger(__name__)
+
 
 class CommunityReportContextBuilder:
     def __init__(
         self,
         community_level: CommunityLevel,
         weight_calculator: CommunityWeightCalculator,
         artifacts: IndexerArtifacts,
+        token_counter: TokenCounter,
+        max_tokens: int = 8000,
     ):
         self._community_level = community_level
         self._weight_calculator = weight_calculator
         self._artifacts = artifacts
+        self._token_counter = token_counter
+        self._max_tokens = max_tokens
 
     def _filter_communities(self) -> list[CommunityReport]:
         df_entities = self._artifacts.entities
@@ -61,17 +70,37 @@ def __call__(self) -> list[Document]:
         reports = self._filter_communities()
 
         documents: list[Document] = []
+        report_str_accumulated: list[str] = []
+        token_count = 0
         for report in reports:
-            documents.append(  # noqa: PERF401
-                Document(
-                    page_content=_REPORT_TEMPLATE.format(
-                        report_id=report.id,
-                        title=report.title,
-                        weight=report.weight,
-                        rank=report.rank,
-                        content=report.content,
+            # we would try to combine multiple
+            # reports into a single document
+            # as long as we do not exceed the token limit
+
+            report_str = _REPORT_TEMPLATE.format(
+                report_id=report.id,
+                title=report.title,
+                weight=report.weight,
+                rank=report.rank,
+                content=report.content,
+            )
+
+            report_str_token_count = self._token_counter.count_tokens(report_str)
+
+            if token_count + report_str_token_count > self._max_tokens:
+                _LOGGER.debug("Reached max tokens for a community report call ...")
+                # we cut a new document here
+                documents.append(
+                    Document(
+                        page_content="\n\n".join(report_str_accumulated),
+                        metadata={"token_count": token_count},
                     )
                 )
-            )
+                # reset the token count and the accumulated string
+                token_count = 0
+                report_str_accumulated = []
+            else:
+                token_count += report_str_token_count
+                report_str_accumulated.append(report_str)
 
         return documents

src/langchain_graphrag/query/global_search/search.py

@@ -1,9 +1,13 @@
+import logging
 from typing import Iterator
 
-from langchain_core.runnables import Runnable, RunnablePassthrough
-
 from .key_points_aggregator import KeyPointsAggregator
 from .key_points_generator import KeyPointsGenerator
+from .key_points_generator.utils import (
+    KeyPointsResult,
+)
+
+_LOGGER = logging.getLogger(__name__)
 
 
 class GlobalSearch:
@@ -15,21 +19,27 @@ def __init__(
         self._kp_generator = kp_generator
         self._kp_aggregator = kp_aggregator
 
-    def invoke(self, query: str) -> str:
+    def _get_key_points(self, query: str) -> dict[str, KeyPointsResult]:
         generation_chain = self._kp_generator()
-        aggregation_chain = self._kp_aggregator()
-
         response = generation_chain.invoke(query)
 
+        if _LOGGER.getEffectiveLevel() == logging.INFO:
+            for k, v in response.items():
+                _LOGGER.info(f"{k} - {len(v.points)}")
+
+        return response
+
+    def invoke(self, query: str) -> str:
+        aggregation_chain = self._kp_aggregator()
+        response = self._get_key_points(query)
+
         return aggregation_chain.invoke(
             input=dict(report_data=response, global_query=query)
         )
 
     def stream(self, query: str) -> Iterator:
-        generation_chain = self._kp_generator()
         aggregation_chain = self._kp_aggregator()
-
-        response = generation_chain.invoke(query)
+        response = self._get_key_points(query)
 
         return aggregation_chain.stream(
             input=dict(report_data=response, global_query=query)