wehub-resource-sync
850 lines
32 KiB
Python
850 lines
32 KiB
Python
import dspy
|
|
import numpy as np
|
|
import re
|
|
import threading
|
|
from typing import Set, Dict, List, Optional, Union, Tuple
|
|
|
|
from .encoder import Encoder
|
|
from .interface import Information
|
|
|
|
|
|
class ConversationTurn:
|
|
"""
|
|
A class to represent a turn in a conversation.
|
|
|
|
Attributes:
|
|
role (str): A short phrase of the role of the speaker for the current conversation turn.
|
|
raw_utterance (str): The response generated by the LM model without polished style and tone.
|
|
utterance_type (str): The type of utterance (e.g., statement, question).
|
|
claim_to_make (Optional[str]): The point that this utterance tries to make. Should be empty if the utterance type is questioning.
|
|
utterance (Optional[str]): The response generated by the model with polished style and tone. Defaults to raw_utterance if not provided.
|
|
queries (List[str]): The queries used to gather information to have a grounded answer.
|
|
raw_retrieved_info (List['Information']): A list of Information type that is retrieved.
|
|
cited_info (Dict[int, 'Information']): A dictionary where the key is the citation index and the value is Information type.
|
|
role_description (Optional[str]): A few sentences description of the role. Defaults to an empty string if not provided.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
role: str,
|
|
raw_utterance: str,
|
|
utterance_type: str,
|
|
claim_to_make: Optional[str] = None,
|
|
utterance: Optional[str] = None,
|
|
queries: Optional[List[str]] = None,
|
|
raw_retrieved_info: Optional[List[Information]] = None,
|
|
cited_info: Optional[List[Information]] = None,
|
|
):
|
|
self.utterance = utterance if utterance is not None else raw_utterance
|
|
self.raw_utterance = raw_utterance
|
|
self.role = role if ":" not in role else role.split(":")[0]
|
|
self.role_description = "" if ":" not in role else role.split(":")[1]
|
|
self.queries = queries if queries is not None else []
|
|
self.raw_retrieved_info = (
|
|
raw_retrieved_info if raw_retrieved_info is not None else []
|
|
)
|
|
self.cited_info = cited_info if cited_info is not None else {}
|
|
self.utterance_type = utterance_type
|
|
self.claim_to_make = claim_to_make if claim_to_make is not None else ""
|
|
|
|
def get_all_citation_index(self):
|
|
citation_pattern = re.compile(r"\[(\d+)\]")
|
|
return list(map(int, citation_pattern.findall(self.utterance)))
|
|
|
|
def to_dict(self):
|
|
raw_retrieved_info = [info.to_dict() for info in self.raw_retrieved_info]
|
|
return {
|
|
"utterance": self.utterance,
|
|
"raw_utterance": self.raw_utterance,
|
|
"role": self.role,
|
|
"role_description": self.role_description,
|
|
"queries": self.queries,
|
|
"utterance_type": self.utterance_type,
|
|
"claim_to_make": self.claim_to_make,
|
|
"raw_retrieved_info": raw_retrieved_info,
|
|
"cited_info": None,
|
|
}
|
|
|
|
@classmethod
|
|
def from_dict(cls, conv_turn_dict: Dict):
|
|
raw_retrieved_info = [
|
|
Information.from_dict(info) for info in conv_turn_dict["raw_retrieved_info"]
|
|
]
|
|
|
|
return cls(
|
|
utterance=conv_turn_dict["utterance"],
|
|
raw_utterance=conv_turn_dict["raw_utterance"],
|
|
role=f"{conv_turn_dict['role']}: {conv_turn_dict['role_description']}",
|
|
queries=conv_turn_dict["queries"],
|
|
raw_retrieved_info=raw_retrieved_info,
|
|
cited_info=None,
|
|
utterance_type=conv_turn_dict["utterance_type"],
|
|
claim_to_make=conv_turn_dict["claim_to_make"],
|
|
)
|
|
|
|
|
|
class KnowledgeNode:
|
|
"""
|
|
Class representing a node in the knowledge base.
|
|
|
|
Attributes:
|
|
name (str): The name of the node.
|
|
content (list): A list of Information instances.
|
|
children (list): A list of child KnowledgeNode instances.
|
|
parent (KnowledgeNode): The parent node of the current node.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
name: str,
|
|
content: Optional[str] = None,
|
|
parent: Optional["KnowledgeNode"] = None,
|
|
children: Optional[List["KnowledgeNode"]] = None,
|
|
synthesize_output: Optional[str] = None,
|
|
need_regenerate_synthesize_output: bool = True,
|
|
):
|
|
"""
|
|
Initializes a KnowledgeNode instance.
|
|
|
|
Args:
|
|
name (str): The name of the node.
|
|
content (list, optional): A list of information uuid. Defaults to None.
|
|
parent (KnowledgeNode, optional): The parent node of the current node. Defaults to None.
|
|
"""
|
|
self.name = name
|
|
self.content: Set[int] = set(content) if content is not None else set()
|
|
self.children = [] if children is None else children
|
|
self.parent = parent
|
|
self.synthesize_output = synthesize_output
|
|
self.need_regenerate_synthesize_output = need_regenerate_synthesize_output
|
|
|
|
def collect_all_content(self):
|
|
"""
|
|
Collects all content from the current node and its descendants.
|
|
|
|
Returns:
|
|
Set[int]: A set containing all content from the current node and its descendants.
|
|
"""
|
|
all_content = set(self.content)
|
|
for child in self.children:
|
|
all_content.update(child.collect_all_content())
|
|
return all_content
|
|
|
|
def has_child(self, child_node_name: str):
|
|
"""
|
|
Check if the node has the child of given name.
|
|
"""
|
|
return child_node_name in [child.name for child in self.children]
|
|
|
|
def add_child(self, child_node_name: str, duplicate_handling: str = "skip"):
|
|
"""
|
|
Adds a child node to the current node.
|
|
duplicate_handling (str): How to handle duplicate nodes. Options are "skip", "none", and "raise error".
|
|
"""
|
|
if self.has_child(child_node_name):
|
|
if duplicate_handling == "skip":
|
|
for child in self.children:
|
|
if child.name == child_node_name:
|
|
return child
|
|
elif duplicate_handling == "raise error":
|
|
raise Exception(
|
|
f"Insert node error. Node {child_node_name} already exists under its parent node {self.name}."
|
|
)
|
|
child_node = KnowledgeNode(name=child_node_name, parent=self)
|
|
self.children.append(child_node)
|
|
return child_node
|
|
|
|
def get_parent(self):
|
|
"""
|
|
Returns the parent node of the current node.
|
|
|
|
Returns:
|
|
KnowledgeNode: The parent node of the current node.
|
|
"""
|
|
return self.parent
|
|
|
|
def get_children(self):
|
|
"""
|
|
Returns the children of the current node.
|
|
|
|
Returns:
|
|
list: A list of child KnowledgeNode instances.
|
|
"""
|
|
return self.children
|
|
|
|
def get_children_names(self):
|
|
"""
|
|
Returns a list of children names.
|
|
"""
|
|
return [child.name for child in self.children]
|
|
|
|
def __repr__(self):
|
|
"""
|
|
Returns a string representation of the KnowledgeNode instance.
|
|
|
|
Returns:
|
|
str: String representation of the KnowledgeNode instance.
|
|
"""
|
|
return f"KnowledgeNode(name={self.name}, content={self.content}, children={len(self.children)})"
|
|
|
|
def get_path_from_root(self, root: Optional["KnowledgeNode"] = None):
|
|
"""
|
|
Get a list of names from the root to this node.
|
|
|
|
Returns:
|
|
List[str]: A list of node names from the root to this node.
|
|
"""
|
|
path = []
|
|
current_node = self
|
|
while current_node:
|
|
path.append(current_node.name)
|
|
if root is not None and current_node.name == root.name:
|
|
break
|
|
current_node = current_node.parent
|
|
return path[::-1]
|
|
|
|
def insert_information(self, information_index: int):
|
|
if information_index not in self.content:
|
|
self.need_regenerate_synthesize_output = True
|
|
self.content.add(information_index)
|
|
|
|
def get_all_descendents(self) -> List["KnowledgeNode"]:
|
|
"""
|
|
Get a list of all descendant nodes.
|
|
|
|
Returns:
|
|
List[KnowledgeNode]: A list of all descendant nodes.
|
|
"""
|
|
descendents = []
|
|
|
|
def collect_descendents(node):
|
|
for child in node.children:
|
|
descendents.append(child)
|
|
collect_descendents(child)
|
|
|
|
collect_descendents(self)
|
|
return descendents
|
|
|
|
def get_all_predecessors(self) -> List["KnowledgeNode"]:
|
|
"""
|
|
Get a list of all predecessor nodes (from current node to root).
|
|
|
|
Returns:
|
|
List[KnowledgeNode]: A list of all predecessor nodes.
|
|
"""
|
|
predecessors = []
|
|
current_node = self.parent
|
|
while current_node is not None:
|
|
predecessors.append(current_node)
|
|
current_node = current_node.parent
|
|
return predecessors
|
|
|
|
def to_dict(self):
|
|
"""
|
|
Converts the KnowledgeNode instance to a dictionary representation.
|
|
|
|
Returns:
|
|
dict: The dictionary representation of the KnowledgeNode.
|
|
"""
|
|
return {
|
|
"name": self.name,
|
|
"content": list(self.content),
|
|
"children": [child.to_dict() for child in self.children],
|
|
"parent": self.parent.name if self.parent else None,
|
|
"synthesize_output": self.synthesize_output,
|
|
"need_regenerate_synthesize_output": self.need_regenerate_synthesize_output,
|
|
}
|
|
|
|
@classmethod
|
|
def from_dict(cls, data):
|
|
"""
|
|
Constructs a KnowledgeNode instance from a dictionary representation.
|
|
|
|
Args:
|
|
data (dict): The dictionary representation of the KnowledgeNode.
|
|
|
|
Returns:
|
|
KnowledgeNode: The constructed KnowledgeNode instance.
|
|
"""
|
|
|
|
def helper(cls, data, parent_node=None):
|
|
if parent_node is not None:
|
|
assert data["parent"] is not None and data["parent"] == parent_node.name
|
|
node = cls(
|
|
name=data["name"],
|
|
content=data["content"],
|
|
parent=parent_node,
|
|
children=None,
|
|
synthesize_output=data.get("synthesize_output", None),
|
|
need_regenerate_synthesize_output=data.get(
|
|
"need_regenerate_synthesize_output", True
|
|
),
|
|
)
|
|
for child_data in data["children"]:
|
|
child_node = helper(cls, child_data, parent_node=node)
|
|
node.children.append(child_node)
|
|
return node
|
|
|
|
return helper(cls, data)
|
|
|
|
|
|
class KnowledgeBase:
|
|
"""
|
|
Represents the dynamic, hierarchical mind map used in Co-STORM to track and organize discourse.
|
|
|
|
The knowledge base serves as a shared conceptual space between the user and the system, allowing for effective collaboration by reducing the user's cognitive load and ensuring that the discourse is easy to follow.
|
|
|
|
The knowledge base is structured as a tree (or mind map) that dynamically organizes collected information and concepts as the conversation progresses.
|
|
|
|
The mind map consists of concepts (nodes) and edges that represent parent-child relationships among topics. Each concept is linked to retrieved information,
|
|
which is placed under the most appropriate concept based on its associated question and semantic similarity.
|
|
|
|
For more details, please refer to Section 3.2 of Co-STORM paper: https://www.arxiv.org/pdf/2408.15232
|
|
Attributes:
|
|
root (KnowledgeNode): The root node of the hierarchical knowledge base, representing the top-level concept.
|
|
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
topic: str,
|
|
knowledge_base_lm: Union[dspy.dsp.LM, dspy.dsp.HFModel],
|
|
node_expansion_trigger_count: int,
|
|
encoder: Encoder,
|
|
):
|
|
"""
|
|
Initializes a KnowledgeBase instance.
|
|
|
|
Args:
|
|
topic (str): The topic of the knowledge base
|
|
expand_node_module (dspy.Module): The module that organize knowledge base in place.
|
|
The module should accept knowledge base as param. E.g. expand_node_module(self)
|
|
article_generation_module (dspy.Module): The module that generate report from knowledge base.
|
|
The module should return string. E.g. report = article_generation_module(self)
|
|
"""
|
|
from .collaborative_storm.modules.article_generation import (
|
|
ArticleGenerationModule,
|
|
)
|
|
from .collaborative_storm.modules.information_insertion_module import (
|
|
InsertInformationModule,
|
|
ExpandNodeModule,
|
|
)
|
|
from .collaborative_storm.modules.knowledge_base_summary import (
|
|
KnowledgeBaseSummaryModule,
|
|
)
|
|
|
|
self.topic: str = topic
|
|
self.encoder: Encoder = encoder
|
|
|
|
self.information_insert_module = InsertInformationModule(
|
|
engine=knowledge_base_lm, encoder=self.encoder
|
|
)
|
|
self.expand_node_module = ExpandNodeModule(
|
|
engine=knowledge_base_lm,
|
|
information_insert_module=self.information_insert_module,
|
|
node_expansion_trigger_count=node_expansion_trigger_count,
|
|
)
|
|
self.article_generation_module = ArticleGenerationModule(
|
|
engine=knowledge_base_lm
|
|
)
|
|
self.gen_summary_module = KnowledgeBaseSummaryModule(engine=knowledge_base_lm)
|
|
|
|
self.root: KnowledgeNode = KnowledgeNode(name="root")
|
|
self.kb_embedding = {
|
|
"hash": hash(""),
|
|
"encoded_structure": np.array([[]]),
|
|
"structure_string": "",
|
|
}
|
|
self.info_uuid_to_info_dict: Dict[int, Information] = {}
|
|
self.info_hash_to_uuid_dict: Dict[int, int] = {}
|
|
self._lock = threading.Lock()
|
|
|
|
def to_dict(self):
|
|
info_uuid_to_info_dict = {
|
|
key: value.to_dict() for key, value in self.info_uuid_to_info_dict.items()
|
|
}
|
|
return {
|
|
"topic": self.topic,
|
|
"tree": self.root.to_dict(),
|
|
"info_uuid_to_info_dict": info_uuid_to_info_dict,
|
|
"info_hash_to_uuid_dict": self.info_hash_to_uuid_dict,
|
|
}
|
|
|
|
@classmethod
|
|
def from_dict(
|
|
cls,
|
|
data: Dict,
|
|
knowledge_base_lm: Union[dspy.dsp.LM, dspy.dsp.HFModel],
|
|
node_expansion_trigger_count: int,
|
|
encoder: Encoder,
|
|
):
|
|
knowledge_base = cls(
|
|
topic=data["topic"],
|
|
knowledge_base_lm=knowledge_base_lm,
|
|
node_expansion_trigger_count=node_expansion_trigger_count,
|
|
encoder=encoder,
|
|
)
|
|
knowledge_base.root = KnowledgeNode.from_dict(data["tree"])
|
|
knowledge_base.info_hash_to_uuid_dict = {
|
|
int(key): int(value)
|
|
for key, value in data["info_hash_to_uuid_dict"].items()
|
|
}
|
|
info_uuid_to_info_dict = {
|
|
int(key): Information.from_dict(value)
|
|
for key, value in data["info_uuid_to_info_dict"].items()
|
|
}
|
|
knowledge_base.info_uuid_to_info_dict = info_uuid_to_info_dict
|
|
return knowledge_base
|
|
|
|
def get_knowledge_base_structure_embedding(
|
|
self, root: Optional[KnowledgeNode] = None
|
|
) -> Tuple[np.ndarray, List[str]]:
|
|
outline_string = self.get_node_hierarchy_string(
|
|
include_indent=False,
|
|
include_full_path=True,
|
|
include_hash_tag=False,
|
|
root=root,
|
|
)
|
|
outline_string_hash = hash(outline_string)
|
|
if outline_string_hash != self.kb_embedding["hash"]:
|
|
outline_strings: List[str] = outline_string.split("\n")
|
|
cleaned_outline_strings = [
|
|
outline.replace(" -> ", ", ") for outline in outline_strings
|
|
]
|
|
encoded_outline = self.encoder.encode(cleaned_outline_strings)
|
|
self.kb_embedding = {
|
|
"hash": outline_string_hash,
|
|
"encoded_structure": encoded_outline,
|
|
"structure_string": outline_strings,
|
|
}
|
|
return (
|
|
self.kb_embedding["encoded_structure"],
|
|
self.kb_embedding["structure_string"],
|
|
)
|
|
|
|
def traverse_down(self, node):
|
|
"""
|
|
Traverses the tree downward from the given node.
|
|
|
|
Args:
|
|
node (KnowledgeNode): The node to start the traversal from.
|
|
|
|
Returns:
|
|
list: A list of KnowledgeNode instances in the order they were visited.
|
|
"""
|
|
nodes = []
|
|
|
|
def _traverse(current_node):
|
|
nodes.append(current_node)
|
|
for child in current_node.get_children():
|
|
_traverse(child)
|
|
|
|
_traverse(node)
|
|
return nodes
|
|
|
|
def traverse_up(self, node):
|
|
"""
|
|
Traverses the tree upward from the given node.
|
|
|
|
Args:
|
|
node (KnowledgeNode): The node to start the traversal from.
|
|
|
|
Returns:
|
|
list: A list of KnowledgeNode instances in the order they were visited.
|
|
"""
|
|
nodes = []
|
|
while node is not None:
|
|
nodes.append(node)
|
|
node = node.get_parent()
|
|
return nodes
|
|
|
|
def collect_all_nodes(self):
|
|
nodes = []
|
|
|
|
def _collect(node):
|
|
nodes.append(node)
|
|
for child in node.children:
|
|
_collect(child)
|
|
|
|
_collect(self.root)
|
|
return nodes
|
|
|
|
def insert_node(
|
|
self,
|
|
new_node_name,
|
|
parent_node: Optional[KnowledgeNode] = None,
|
|
duplicate_handling="skip",
|
|
):
|
|
"""
|
|
Inserts a new node into the knowledge base under the specified parent node.
|
|
|
|
Args:
|
|
new_node_name (str): The name of the new node.
|
|
parent_node_name (str): The name of the parent node. If None, the new node is inserted under the root.
|
|
duplicate_handling (str): How to handle duplicate nodes. Options are "skip", "none", and "raise error".
|
|
"""
|
|
if parent_node is None:
|
|
return self.root.add_child(
|
|
new_node_name, duplicate_handling=duplicate_handling
|
|
)
|
|
else:
|
|
return parent_node.add_child(
|
|
new_node_name, duplicate_handling=duplicate_handling
|
|
)
|
|
|
|
def find_node(self, current_node, node_name):
|
|
"""
|
|
Finds a node by name in the knowledge base.
|
|
|
|
Args:
|
|
current_node (KnowledgeNode): The node to start the search from.
|
|
node_name (str): The name of the node to find.
|
|
|
|
Returns:
|
|
KnowledgeNode: The node with the specified name, or None if not found.
|
|
"""
|
|
if current_node.name == node_name:
|
|
return current_node
|
|
for child in current_node.get_children():
|
|
result = self.find_node(child, node_name)
|
|
if result is not None:
|
|
return result
|
|
return None
|
|
|
|
def insert_from_outline_string(self, outline_string, duplicate_handling="skip"):
|
|
"""
|
|
Creates and inserts nodes into the knowledge base from a string outline.
|
|
|
|
Args:
|
|
outline_string (str): The outline string where each line starts with '#' denoting the level.
|
|
duplicate_handling (str): How to handle duplicate nodes. Options are "skip", "none", and "raise error".
|
|
"""
|
|
last_node_at_level = {}
|
|
for line in outline_string.split("\n"):
|
|
level = line.count("#")
|
|
if level > 0:
|
|
title = line.strip("# ").strip()
|
|
if title.lower() in ["overview", "summary", "introduction"]:
|
|
continue
|
|
parent_node = None if level == 1 else last_node_at_level.get(level - 1)
|
|
new_node = self.insert_node(
|
|
new_node_name=title,
|
|
parent_node=parent_node,
|
|
duplicate_handling=duplicate_handling,
|
|
)
|
|
last_node_at_level[level] = new_node
|
|
for deeper_level in list(last_node_at_level.keys()):
|
|
if deeper_level > level:
|
|
del last_node_at_level[deeper_level]
|
|
|
|
def get_node_hierarchy_string(
|
|
self,
|
|
include_indent=False,
|
|
include_full_path=False,
|
|
include_hash_tag=True,
|
|
include_node_content_count=False,
|
|
cited_indices: Optional[List[int]] = None,
|
|
root: Optional[KnowledgeNode] = None,
|
|
) -> str:
|
|
def find_node_contain_index(node, index):
|
|
"""
|
|
Traverses the tree downward from the given node.
|
|
|
|
Args:
|
|
node (KnowledgeNode): The node to start the traversal from.
|
|
|
|
Returns:
|
|
list: A list of KnowledgeNode instances in the order they were visited.
|
|
"""
|
|
nodes = []
|
|
|
|
def _traverse(current_node):
|
|
if current_node is not None and index in current_node.content:
|
|
nodes.append(current_node)
|
|
for child in current_node.get_children():
|
|
_traverse(child)
|
|
|
|
_traverse(node)
|
|
return nodes
|
|
|
|
paths_to_highlight = set()
|
|
nodes_to_include = set()
|
|
if cited_indices is not None:
|
|
for index in cited_indices:
|
|
for cur_node in find_node_contain_index(self.root, index):
|
|
paths_to_highlight.add(" -> ".join(cur_node.get_path_from_root()))
|
|
nodes_to_include.add(cur_node)
|
|
nodes_to_include.update(cur_node.get_all_descendents())
|
|
predecessors = cur_node.get_all_predecessors()
|
|
for predecessor in predecessors:
|
|
nodes_to_include.update(predecessor.children)
|
|
nodes_to_include.update(predecessors)
|
|
|
|
def should_include_node(node):
|
|
if cited_indices is None:
|
|
return True
|
|
return node in nodes_to_include
|
|
|
|
def should_omit_child_nodes(node):
|
|
if cited_indices is None:
|
|
return False
|
|
for child in node.children:
|
|
if should_include_node(child):
|
|
return False
|
|
return True
|
|
|
|
def helper(cur_root, level):
|
|
to_return = []
|
|
if cur_root is not None:
|
|
should_include_current_node = should_include_node(cur_root)
|
|
|
|
indent = "" if not include_indent else "\t" * (level - 1)
|
|
full_path = " -> ".join(cur_root.get_path_from_root(root=root))
|
|
node_info = cur_root.name if not include_full_path else full_path
|
|
hash_tag = "#" * level + " " if include_hash_tag else ""
|
|
content_count = (
|
|
f" ({len(cur_root.content)})" if include_node_content_count else ""
|
|
)
|
|
special_note = (
|
|
""
|
|
if cited_indices is None or full_path not in paths_to_highlight
|
|
else " ⭐"
|
|
)
|
|
|
|
if should_include_current_node:
|
|
to_return.append(
|
|
f"{indent}{hash_tag}{node_info}{content_count}{special_note}"
|
|
)
|
|
if should_omit_child_nodes(cur_root):
|
|
if len(cur_root.children) > 0:
|
|
child_indent = indent = (
|
|
"" if not include_indent else "\t" * (level)
|
|
)
|
|
to_return.append(f"{child_indent}...")
|
|
else:
|
|
for child in cur_root.children:
|
|
to_return.extend(helper(child, level + 1))
|
|
return to_return
|
|
|
|
to_return = []
|
|
if root is None and self.root is not None:
|
|
for child in self.root.children:
|
|
to_return.extend(helper(child, level=1))
|
|
else:
|
|
to_return.extend(helper(root, level=1))
|
|
|
|
return "\n".join(to_return)
|
|
|
|
def find_node_by_path(
|
|
self,
|
|
path: str,
|
|
missing_node_handling="abort",
|
|
root: Optional[KnowledgeNode] = None,
|
|
):
|
|
"""
|
|
Returns the target node given a path string.
|
|
|
|
Args:
|
|
path (str): The path to the node, with node names connected by " -> ".
|
|
missing_node_handling (str): How to handle missing nodes. Options are "abort", "create", and "raise error".
|
|
|
|
Returns:
|
|
KnowledgeNode: The target node.
|
|
"""
|
|
node_names = path.split(" -> ")
|
|
current_node = self.root if root is None else root
|
|
|
|
for name in node_names[1:]:
|
|
found_node = next(
|
|
(child for child in current_node.children if child.name == name), None
|
|
)
|
|
if found_node is None:
|
|
if missing_node_handling == "abort":
|
|
return
|
|
elif missing_node_handling == "create":
|
|
new_node = current_node.add_child(child_node_name=name)
|
|
current_node = new_node
|
|
elif missing_node_handling == "raise error":
|
|
structure = self.get_node_hierarchy_string(
|
|
include_indent=True,
|
|
include_full_path=False,
|
|
include_hash_tag=True,
|
|
)
|
|
raise Exception(
|
|
f"Insert information error. Unable to find node {{{name}}} under {{{current_node.name}}}\n{structure}"
|
|
)
|
|
else:
|
|
current_node = found_node
|
|
return current_node
|
|
|
|
def insert_information(
|
|
self,
|
|
path: str,
|
|
information: Information,
|
|
missing_node_handling="abort",
|
|
root: Optional[KnowledgeNode] = None,
|
|
):
|
|
"""
|
|
Inserts information into the knowledge base at the specified path.
|
|
|
|
Args:
|
|
path (str): The placement path string, connected by " -> " linking the name of nodes.
|
|
information (Information): The information to insert.
|
|
missing_node_handling (str): How to handle missing nodes. Options are "abort", "create", and "raise error".
|
|
Return:
|
|
uuid of insertion information
|
|
"""
|
|
with self._lock:
|
|
target_node: KnowledgeNode = self.find_node_by_path(
|
|
path=path, missing_node_handling=missing_node_handling, root=root
|
|
)
|
|
information_hash = hash(information)
|
|
if information.citation_uuid == -1:
|
|
info_citation_uuid = self.info_hash_to_uuid_dict.get(
|
|
information_hash, len(self.info_hash_to_uuid_dict) + 1
|
|
)
|
|
information.citation_uuid = info_citation_uuid
|
|
self.info_hash_to_uuid_dict[information_hash] = info_citation_uuid
|
|
self.info_uuid_to_info_dict[info_citation_uuid] = information
|
|
if target_node is not None:
|
|
self.info_uuid_to_info_dict[information.citation_uuid].meta[
|
|
"placement"
|
|
] = " -> ".join(target_node.get_path_from_root())
|
|
target_node.insert_information(information.citation_uuid)
|
|
|
|
def trim_empty_leaf_nodes(self):
|
|
"""
|
|
Trims all leaf nodes that do not have any content. Iteratively does it until all leaf nodes have at least one content.
|
|
"""
|
|
|
|
def trim_node(node):
|
|
if not node.children and not node.content:
|
|
return True
|
|
node.children = [child for child in node.children if not trim_node(child)]
|
|
return not node.children and not node.content
|
|
|
|
# Start the trimming process from the root
|
|
while True:
|
|
before_trim = len(self.get_all_leaf_nodes())
|
|
trim_node(self.root)
|
|
after_trim = len(self.get_all_leaf_nodes())
|
|
if before_trim == after_trim:
|
|
break
|
|
|
|
def get_all_leaf_nodes(self):
|
|
"""
|
|
Helper function to get all leaf nodes.
|
|
|
|
Returns:
|
|
List[KnowledgeNode]: A list of all leaf nodes in the knowledge base.
|
|
"""
|
|
leaf_nodes = []
|
|
|
|
def find_leaf_nodes(node):
|
|
if not node.children:
|
|
leaf_nodes.append(node)
|
|
for child in node.children:
|
|
find_leaf_nodes(child)
|
|
|
|
find_leaf_nodes(self.root)
|
|
return leaf_nodes
|
|
|
|
def merge_single_child_nodes(self):
|
|
"""
|
|
Merges content of a node with its single child and removes the child node.
|
|
Iteratively does this from leaf nodes back to the root.
|
|
"""
|
|
|
|
def merge_node(node):
|
|
# Recursively merge children first
|
|
for child in node.children:
|
|
merge_node(child)
|
|
|
|
# If the node has exactly one child, merge its content with the child and remove the child
|
|
if len(node.children) == 1:
|
|
single_child = node.children[0]
|
|
node.content.update(single_child.content)
|
|
node.children = single_child.children
|
|
for grandchild in node.children:
|
|
grandchild.parent = node
|
|
|
|
merge_node(self.root)
|
|
|
|
def update_all_info_path(self):
|
|
def _helper(node):
|
|
for citation_idx in node.content:
|
|
self.info_uuid_to_info_dict[citation_idx].meta["placement"] = (
|
|
" -> ".join(node.get_path_from_root())
|
|
)
|
|
for child in node.children:
|
|
_helper(child)
|
|
|
|
_helper(self.root)
|
|
|
|
def update_from_conv_turn(
|
|
self,
|
|
conv_turn: ConversationTurn,
|
|
allow_create_new_node: bool = False,
|
|
insert_under_root: bool = False,
|
|
):
|
|
if conv_turn is None:
|
|
return
|
|
info_to_insert = list(conv_turn.cited_info.values())
|
|
if insert_under_root:
|
|
for info in info_to_insert:
|
|
self.insert_information(path=self.root.name, information=info)
|
|
else:
|
|
self.information_insert_module(
|
|
knowledge_base=self,
|
|
information=info_to_insert,
|
|
allow_create_new_node=allow_create_new_node,
|
|
)
|
|
old_to_new_citation_idx_mapping = {
|
|
old_idx: info.citation_uuid
|
|
for old_idx, info in conv_turn.cited_info.items()
|
|
}
|
|
|
|
for old_idx, new_idx in old_to_new_citation_idx_mapping.items():
|
|
conv_turn.utterance = conv_turn.utterance.replace(
|
|
f"[{old_idx}]", f"[_{new_idx}_]"
|
|
)
|
|
conv_turn.raw_utterance = conv_turn.raw_utterance.replace(
|
|
f"[{old_idx}]", f"[_{new_idx}_]"
|
|
)
|
|
for _, new_idx in old_to_new_citation_idx_mapping.items():
|
|
conv_turn.utterance = conv_turn.utterance.replace(
|
|
f"[_{new_idx}_]", f"[{new_idx}]"
|
|
)
|
|
conv_turn.utterance.replace("[-1]", "")
|
|
conv_turn.raw_utterance = conv_turn.raw_utterance.replace(
|
|
f"[_{new_idx}_]", f"[{new_idx}]"
|
|
)
|
|
conv_turn.raw_utterance.replace("[-1]", "")
|
|
conv_turn.cited_info = None
|
|
|
|
def get_knowledge_base_summary(self):
|
|
return self.gen_summary_module(self)
|
|
|
|
def reorganize(self):
|
|
"""
|
|
Reorganizes the knowledge base through two main processes: top-down expansion and bottom-up cleaning.
|
|
|
|
The reorganization process ensures that the knowledge base remains well-structured and relevant as new information is added. It consists of the following steps:
|
|
1.Top-Down Expansion: Expands nodes that have accumulated significant amounts of information by creating subtopics,
|
|
ensuring that each concept remains specific and manageable.
|
|
2.Bottom-Up Cleaning: Cleans the knowledge base by removing empty leaf nodes (nodes with no supporting information)
|
|
and merging nodes that have only a single child, simplifying the structure and maintaining clarity.
|
|
"""
|
|
# pre-processing
|
|
self.trim_empty_leaf_nodes()
|
|
self.merge_single_child_nodes()
|
|
# expand nodes
|
|
self.expand_node_module(knowledge_base=self)
|
|
# clean up
|
|
self.trim_empty_leaf_nodes()
|
|
self.merge_single_child_nodes()
|
|
self.update_all_info_path()
|
|
|
|
def to_report(self):
|
|
return self.article_generation_module(knowledge_base=self)
|