A framework for managing and executing a network of nodes using a language model. This class enables the creation of a graph structure for processing and analyzing data. Each node in the graph can perform specific operations, allowing for complex data processing workflows. The grap
| 14 | from KVCOMM.utils.log import logger |
| 15 | |
| 16 | class Graph(ABC): |
| 17 | """ |
| 18 | A framework for managing and executing a network of nodes using a language model. |
| 19 | |
| 20 | This class enables the creation of a graph structure for processing and analyzing data. Each node |
| 21 | in the graph can perform specific operations, allowing for complex data processing workflows. |
| 22 | The graph supports integration with language models, making it suitable for tasks that require |
| 23 | natural language processing capabilities. |
| 24 | |
| 25 | The communication of the node depends on the node.spatial_predecessors and node.spatial_successors. |
| 26 | |
| 27 | Attributes: |
| 28 | domain (str): The domain for which this graph is used. |
| 29 | llm_name (str): The name of the llm that used for processing within the nodes. |
| 30 | nodes (dict): A collection of nodes, each identified by a unique UUID. |
| 31 | |
| 32 | Methods: |
| 33 | build_graph(): Method to be implemented for constructing the graph structure. |
| 34 | add_node(node): Adds a new node to the graph with a unique identifier. |
| 35 | run(inputs, num_rounds=1, max_tries=3, max_time=600): Executes the graph for a specified number of rounds, processing provided inputs. |
| 36 | arun(input, num_rounds=1, max_tries=3, max_time=600, *, mode="default", **kwargs): Asynchronously executes the graph for a specified number of rounds, processing provided inputs. |
| 37 | update_memory(message=None): Propagates memory update across nodes; supports multi-request keys. |
| 38 | check_cycle(new_node, target_nodes): Detects if adding edges would create a cycle starting at `new_node`. |
| 39 | update_masks(): Returns current spatial and temporal mask parameters. |
| 40 | __getstate__(): Returns the state of the graph. |
| 41 | __setstate__(state): Sets the state of the graph. |
| 42 | __deepcopy__(memo): Creates a deep copy of the graph. |
| 43 | """ |
| 44 | |
| 45 | def __init__(self, |
| 46 | domain: str, |
| 47 | llm_name: Optional[str], |
| 48 | agent_names: List[str], |
| 49 | decision_method: str = None, |
| 50 | fixed_spatial_masks:List[List[int]] = None, |
| 51 | fixed_temporal_masks:List[List[int]] = None, |
| 52 | node_kwargs:List[Dict] = None, |
| 53 | kv_config: KVCommConfig | None = None, |
| 54 | ): |
| 55 | |
| 56 | num_agents = len(agent_names) |
| 57 | if fixed_spatial_masks is None: |
| 58 | fixed_spatial_masks = [[1 if i!=j else 0 for j in range(num_agents)] for i in range(num_agents)] |
| 59 | if fixed_temporal_masks is None: |
| 60 | fixed_temporal_masks = [[1 for _ in range(num_agents)] for _ in range(num_agents)] |
| 61 | spatial_mask_tensor = torch.as_tensor(fixed_spatial_masks, dtype=torch.float32).view(num_agents, num_agents) |
| 62 | temporal_mask_tensor = torch.as_tensor(fixed_temporal_masks, dtype=torch.float32).view(num_agents, num_agents) |
| 63 | assert spatial_mask_tensor.numel() == num_agents * num_agents, "The fixed_spatial_masks doesn't match the number of agents" |
| 64 | assert temporal_mask_tensor.numel() == num_agents * num_agents, "The fixed_temporal_masks doesn't match the number of agents" |
| 65 | self.kv_config = kv_config or KVCommConfig.from_env() |
| 66 | self.id:str = shortuuid.ShortUUID().random(length=4) |
| 67 | self.domain:str = domain |
| 68 | self.llm_name:str = llm_name |
| 69 | self.agent_names:List[str] = agent_names |
| 70 | self.decision_node:Node = AgentRegistry.get( |
| 71 | decision_method, |
| 72 | **{"domain": self.domain, "llm_name": self.llm_name, "llm_config": self.kv_config}, |
| 73 | ) if decision_method is not None else None |