refactor - in local & global search api, stick to callable objects that return langchain chains to keep the API surface minimal & consistent

Author: ksachdeva

README.md

@@ -105,13 +105,16 @@ retriever = LocalSearchRetriever(
     artifacts=artifacts,
 )
 
-# Get a langchain chain to do local search
-search_chain = make_local_search_chain(
+# Build the LocalSearch object
+local_search = LocalSearch(
     prompt_builder=LocalSearchPromptBuilder(),
     llm=make_llm_instance(llm_type, llm_model, cache_dir),
     retriever=retriever,
 )
 
+# it's a callable that returns the chain
+search_chain = local_search()
+
 # you could invoke
 # print(search_chain.invoke(query))
 

examples/simple-app/app/query.py

@@ -39,9 +39,9 @@
     KeyPointsGeneratorPromptBuilder,
 )
 from langchain_graphrag.query.local_search import (
+    LocalSearch,
     LocalSearchPromptBuilder,
     LocalSearchRetriever,
-    make_local_search_chain,
 )
 from langchain_graphrag.query.local_search.context_builders import (
     ContextBuilder,
@@ -163,13 +163,15 @@ def local_search(
         artifacts=artifacts,
     )
 
-    # Get a langchain chain to do local search
-    search_chain = make_local_search_chain(
+    local_search = LocalSearch(
         prompt_builder=LocalSearchPromptBuilder(),
         llm=make_llm_instance(llm_type, llm_model, cache_dir),
         retriever=retriever,
     )
 
+    # get the chain
+    search_chain = local_search()
+
     # you could invoke
     # print(search_chain.invoke(query))
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "langchain-graphrag"
-version = "0.0.2-beta.9"
+version = "0.0.2-beta.10"
 description = "Implementation of GraphRAG (https://arxiv.org/pdf/2404.16130)"
 authors = [{ name = "Kapil Sachdeva", email = "notan@email.com" }]
 dependencies = [

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

@@ -1,11 +1,10 @@
 """KeyPointsAggregator module."""
 
-from .aggregator import KeyPointsAggregator, make_key_points_aggregator_chain
+from .aggregator import KeyPointsAggregator
 from .context_builder import KeyPointsContextBuilder
 from .prompt_builder import KeyPointsAggregatorPromptBuilder
 
 __all__ = [
-    "make_key_points_aggregator_chain",
     "KeyPointsAggregatorPromptBuilder",
     "KeyPointsContextBuilder",
     "KeyPointsAggregator",

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

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 import operator
 from functools import partial
 
@@ -28,30 +26,6 @@ def _kp_result_to_docs(
     return context_builder(key_points)
 
 
-def make_key_points_aggregator_chain(
-    llm: BaseLLM,
-    prompt_builder: PromptBuilder,
-    context_builder: KeyPointsContextBuilder,
-) -> Runnable:
-    kp_lambda = partial(_kp_result_to_docs, context_builder=context_builder)
-
-    prompt, output_parser = prompt_builder.build()
-
-    search_chain: Runnable = (
-        {
-            "report_data": operator.itemgetter("report_data")
-            | RunnableLambda(kp_lambda)
-            | _format_docs,
-            "global_query": operator.itemgetter("global_query"),
-        }
-        | prompt
-        | llm
-        | output_parser
-    )
-
-    return search_chain
-
-
 class KeyPointsAggregator:
     def __init__(
         self,
@@ -64,8 +38,19 @@ def __init__(
         self._context_builder = context_builder
 
     def __call__(self) -> Runnable:
-        return make_key_points_aggregator_chain(
-            llm=self._llm,
-            prompt_builder=self._prompt_builder,
+        kp_lambda = partial(
+            _kp_result_to_docs,
             context_builder=self._context_builder,
         )
+
+        prompt, output_parser = self._prompt_builder.build()
+        base_chain = prompt | self._llm | output_parser
+
+        search_chain: Runnable = {
+            "report_data": operator.itemgetter("report_data")
+            | RunnableLambda(kp_lambda)
+            | _format_docs,
+            "global_query": operator.itemgetter("global_query"),
+        } | base_chain
+
+        return search_chain

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

@@ -1,11 +1,10 @@
 """Key Points generator module."""
 
 from .context_builder import CommunityReportContextBuilder
-from .generator import KeyPointsGenerator, make_key_points_generator_chain
+from .generator import KeyPointsGenerator
 from .prompt_builder import KeyPointsGeneratorPromptBuilder
 
 __all__ = [
-    "make_key_points_generator_chain",
     "KeyPointsGeneratorPromptBuilder",
     "CommunityReportContextBuilder",
     "KeyPointsGenerator",

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

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 from langchain_core.documents import Document
 from langchain_core.language_models import BaseLLM
 from langchain_core.runnables import Runnable, RunnableParallel
@@ -15,28 +13,6 @@ def _format_docs(documents: list[Document]) -> str:
     return context_data_str
 
 
-def make_key_points_generator_chain(
-    llm: BaseLLM,
-    prompt_builder: PromptBuilder,
-    context_builder: CommunityReportContextBuilder,
-) -> Runnable:
-    prompt, output_parser = prompt_builder.build()
-
-    documents = context_builder()
-
-    chains: list[Runnable] = []
-
-    for d in documents:
-        d_context_data = _format_docs([d])
-        d_prompt = prompt.partial(context_data=d_context_data)
-        generator_chain: Runnable = d_prompt | llm | output_parser
-        chains.append(generator_chain)
-
-    analysts = [f"Analayst-{i}" for i in range(1, len(chains) + 1)]
-
-    return RunnableParallel(dict(zip(analysts, chains, strict=True)))
-
-
 class KeyPointsGenerator:
     def __init__(
         self,
@@ -49,8 +25,18 @@ def __init__(
         self._context_builder = context_builder
 
     def __call__(self) -> Runnable:
-        return make_key_points_generator_chain(
-            llm=self._llm,
-            prompt_builder=self._prompt_builder,
-            context_builder=self._context_builder,
-        )
+        prompt, output_parser = self._prompt_builder.build()
+
+        documents = self._context_builder()
+
+        chains: list[Runnable] = []
+
+        for d in documents:
+            d_context_data = _format_docs([d])
+            d_prompt = prompt.partial(context_data=d_context_data)
+            generator_chain: Runnable = d_prompt | self._llm | output_parser
+            chains.append(generator_chain)
+
+        analysts = [f"Analayst-{i}" for i in range(1, len(chains) + 1)]
+
+        return RunnableParallel(dict(zip(analysts, chains, strict=True)))

src/langchain_graphrag/query/local_search/__init__.py

@@ -2,10 +2,10 @@
 
 from .prompt_builder import LocalSearchPromptBuilder
 from .retriever import LocalSearchRetriever
-from .search import make_local_search_chain
+from .search import LocalSearch
 
 __all__ = [
-    "make_local_search_chain",
+    "LocalSearch",
     "LocalSearchPromptBuilder",
     "LocalSearchRetriever",
 ]

src/langchain_graphrag/query/local_search/search.py

@@ -12,21 +12,25 @@ def _format_docs(documents: list[Document]) -> str:
     return context_data_str
 
 
-def make_local_search_chain(
-    llm: BaseLLM,
-    prompt_builder: PromptBuilder,
-    retriever: BaseRetriever,
-) -> Runnable:
-    prompt, output_parser = prompt_builder.build()
-
-    search_chain: Runnable = (
-        {
-            "context_data": retriever | _format_docs,
+class LocalSearch:
+    def __init__(
+        self,
+        llm: BaseLLM,
+        prompt_builder: PromptBuilder,
+        retriever: BaseRetriever,
+    ):
+        self._llm = llm
+        self._prompt_builder = prompt_builder
+        self._retriever = retriever
+
+    def __call__(self) -> Runnable:
+        prompt, output_parser = self._prompt_builder.build()
+
+        base_chain = prompt | self._llm | output_parser
+
+        search_chain: Runnable = {
+            "context_data": self._retriever | _format_docs,
             "local_query": RunnablePassthrough(),
-        }
-        | prompt
-        | llm
-        | output_parser
-    )
+        } | base_chain
 
-    return search_chain
+        return search_chain