mirror of
https://github.com/imartinez/privateGPT.git
synced 2026-07-16 17:00:12 +00:00
92 lines
3.2 KiB
Python
92 lines
3.2 KiB
Python
from llama_index.core.bridge.pydantic import Field
|
|
from llama_index.core.postprocessor.types import BaseNodePostprocessor
|
|
from llama_index.core.schema import NodeWithScore, QueryBundle
|
|
|
|
from private_gpt.components.node_store.node_store_component import NodeStoreComponent
|
|
|
|
|
|
class WindowPrevNextReplacementPostProcessor(BaseNodePostprocessor):
|
|
"""Post-processor to return the previous or next nodes to the current node.
|
|
|
|
This post-processor is useful for generating a window context out of the LLM,
|
|
useful for better user understanding of the context.
|
|
|
|
Args:
|
|
docstore (BaseDocumentStore): The document store.
|
|
num_nodes (int): The number of nodes to return (default: 1)
|
|
mode (str): The mode of the post-processor.
|
|
Can be "previous", "next", or "both.
|
|
"""
|
|
|
|
node_component: NodeStoreComponent
|
|
collection: str
|
|
window_length: int = Field(default=1)
|
|
mode: str = Field(default="next")
|
|
|
|
@classmethod
|
|
def class_name(cls) -> str:
|
|
return "WindowPrevNextReplacementPostProcessor"
|
|
|
|
def get_sibling_text(
|
|
self,
|
|
node_with_score: NodeWithScore,
|
|
max_length: int,
|
|
forward: bool = True,
|
|
) -> list[str]:
|
|
current_length = 0
|
|
result_text = []
|
|
current_node = node_with_score.node
|
|
|
|
while current_length < max_length:
|
|
explored_node_info = (
|
|
current_node.next_node if forward else current_node.prev_node
|
|
)
|
|
if explored_node_info is None:
|
|
break
|
|
|
|
explored_node = self.node_component.get_node(
|
|
self.collection, explored_node_info.node_id
|
|
)
|
|
if explored_node is None:
|
|
break
|
|
|
|
node_content = explored_node.get_content()
|
|
|
|
if forward:
|
|
# Add the text to the end of the list
|
|
# If the text is too long, we need to truncate it from the end
|
|
result_text.append(node_content[: max_length - current_length])
|
|
else:
|
|
# Add the text to the beginning of the list
|
|
# If the text is too long, we need to truncate it from the start
|
|
result_text.insert(0, node_content[-(max_length - current_length) :])
|
|
|
|
current_length += len(result_text[-1]) if forward else len(result_text[0])
|
|
|
|
if current_length >= max_length:
|
|
break
|
|
|
|
current_node = explored_node
|
|
|
|
return result_text
|
|
|
|
def _postprocess_nodes(
|
|
self,
|
|
nodes: list[NodeWithScore],
|
|
query_bundle: QueryBundle | None = None,
|
|
) -> list[NodeWithScore]:
|
|
"""Postprocess nodes."""
|
|
for node in nodes:
|
|
if self.mode == "previous" or self.mode == "both":
|
|
node.metadata["previous_texts"] = self.get_sibling_text(
|
|
node, self.window_length, forward=False
|
|
)
|
|
node.node.excluded_llm_metadata_keys.append("previous_texts")
|
|
if self.mode == "next" or self.mode == "both":
|
|
node.metadata["next_texts"] = self.get_sibling_text(
|
|
node, self.window_length, forward=True
|
|
)
|
|
node.node.excluded_llm_metadata_keys.append("next_texts")
|
|
|
|
return nodes
|