feat - in local & global search, add the support to output raw chunk …

…(instead of processing via associated OutputParser) [default is False]
ksachdeva · Oct 8, 2024 · 2c5210b · 2c5210b
1 parent 69d9aa9
commit 2c5210b
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 2 deletions.
diff --git a/examples/simple-app/app/query.py b/examples/simple-app/app/query.py
@@ -76,6 +76,7 @@ def global_search(
         True,  # noqa: FBT003
         help="Repeat instructions in the prompt",
     ),
+    output_raw: bool = typer.Option(False, help="Output raw response"),  # noqa: FBT001, FBT003
     enable_langsmith: bool = typer.Option(False, help="Enable Langsmith"),  # noqa: FBT001, FBT003
 ):
     if enable_langsmith:
@@ -99,6 +100,7 @@ def global_search(
             ],
             ["Show References", str(show_references)],
             ["Repeat Instructions In Prompt", str(repeat_instructions)],
+            ["Output Raw", str(output_raw)],
         ]
     )
 
@@ -133,6 +135,7 @@ def global_search(
         context_builder=KeyPointsContextBuilder(
             token_counter=TiktokenCounter(),
         ),
+        output_raw=output_raw,
     )
 
     global_search = GlobalSearch(
@@ -169,6 +172,7 @@ def local_search(
         True,  # noqa: FBT003
         help="Repeat instructions in the prompt",
     ),
+    output_raw: bool = typer.Option(False, help="Output raw response"),  # noqa: FBT001, FBT003
     enable_langsmith: bool = typer.Option(False, help="Enable Langsmith"),  # noqa: FBT001, FBT003
 ):
     if enable_langsmith:
@@ -196,6 +200,7 @@ def local_search(
             ],
             ["Show References", str(show_references)],
             ["Repeat Instructions In Prompt", str(repeat_instructions)],
+            ["Output Raw", str(output_raw)],
         ]
     )
 
@@ -255,6 +260,7 @@ def local_search(
             ollama_num_context=ollama_num_context,
         ),
         retriever=retriever,
+        output_raw=output_raw,
     )
 
     # get the chain

diff --git a/src/langchain_graphrag/query/global_search/key_points_aggregator/aggregator.py b/src/langchain_graphrag/query/global_search/key_points_aggregator/aggregator.py
@@ -32,10 +32,13 @@ def __init__(
         llm: BaseLLM,
         prompt_builder: PromptBuilder,
         context_builder: KeyPointsContextBuilder,
+        *,
+        output_raw: bool = False,
     ):
         self._llm = llm
         self._prompt_builder = prompt_builder
         self._context_builder = context_builder
+        self._output_raw = output_raw
 
     def __call__(self) -> Runnable:
         kp_lambda = partial(
@@ -44,7 +47,10 @@ def __call__(self) -> Runnable:
         )
 
         prompt, output_parser = self._prompt_builder.build()
-        base_chain = prompt | self._llm | output_parser
+        base_chain = prompt | self._llm
+
+        if not self._output_raw:
+            base_chain = base_chain | output_parser
 
         search_chain: Runnable = {
             "report_data": operator.itemgetter("report_data")

diff --git a/src/langchain_graphrag/query/local_search/search.py b/src/langchain_graphrag/query/local_search/search.py
@@ -18,15 +18,21 @@ def __init__(
         llm: BaseLLM,
         prompt_builder: PromptBuilder,
         retriever: BaseRetriever,
+        *,
+        output_raw: bool = False,
     ):
         self._llm = llm
         self._prompt_builder = prompt_builder
         self._retriever = retriever
+        self._output_raw = output_raw
 
     def __call__(self) -> Runnable:
         prompt, output_parser = self._prompt_builder.build()
 
-        base_chain = prompt | self._llm | output_parser
+        base_chain = prompt | self._llm
+
+        if not self._output_raw:
+            base_chain = base_chain | output_parser
 
         search_chain: Runnable = {
             "context_data": self._retriever | _format_docs,