gistfile1.txt

from typing import List, Dict, Optional


class ChatNode:
    """Represents a node in the chat flow tree."""

    def __init__(self, text: str, color: Optional[str] = None):
        self.text = text
        self.color = color
        self.children: List[ChatNode] = []  # List of child nodes

    def add_child(self, text: str, color: Optional[str] = None) -> "ChatNode":
        """Adds a new child node to this node."""
        new_node = ChatNode(text, color)
        self.children.append(new_node)
        return new_node  # Return the newly created node for potential further use

    def __str__(self):
        return (
            f"ChatNode(text='{self.text}', color={self.color}, "
            f"children={[str(child) for child in self.children]})"
        )

    def __repr__(self):
        return self.__str__()

    def to_dict(self):
        """Converts the node and its children to a dictionary."""
        return {
            "text": self.text,
            "color": self.color,
            "children": [child.to_dict() for child in self.children],
        }


class ChatFlow:
    """Represents a chat flow as a tree structure."""

    def __init__(self, name: str):
        self.name = name
        self.root_nodes: List[ChatNode] = []  # List of root nodes (for a forest-like structure)

    def add_root_node(self, text: str, color: Optional[str] = None) -> ChatNode:
        """Adds a new root node to the chat flow."""
        new_node = ChatNode(text, color)
        self.root_nodes.append(new_node)
        return new_node # return new node to potentially add children right after

    def add_message_to_node(
        self, parent_node: ChatNode, text: str, color: Optional[str] = None
    ) -> ChatNode:
        """Adds a new message as the last child of the specified parent node.

        Args:
            parent_node: The existing node to add the new message to.
            text: The text content of the new message.
            color: An optional CSS color string.

        Returns:
            The new created ChatNode
        """
        return parent_node.add_child(text, color)
    
    def find_node(self, text: str) -> Optional[ChatNode]:
      """
      Helper method to find node by text content.

      It performes Breadth-First Search (BFS) to traverse the chat flow tree.

      Args:
          text: the content of the node we are looking for

      Returns:
          ChatNode: the found node or None
      """
      queue = self.root_nodes.copy()

      while queue:
        current_node = queue.pop(0)
        if current_node.text == text:
          return current_node
        queue.extend(current_node.children)
      return None
    

    def __str__(self):
        return f"ChatFlow(name='{self.name}', root_nodes={self.root_nodes})"

    def __repr__(self):
        return self.__str__()

    def to_dict(self):
        """Converts the entire chat flow to a dictionary."""
        return {
            "name": self.name,
            "root_nodes": [node.to_dict() for node in self.root_nodes],
        }
    
    def add_message_to_last_node(self, text: str, color: Optional[str] = None) -> Optional[ChatNode]:
      """
      Adds a new message to the last created leaf node.

      Args:
          text: The text of the new message.
          color: Optional CSS color for the message.

      Returns:
           ChatNode: created message, or None if chat flow is empty
      """
      if not self.root_nodes:
          return None
      
      # Find the last added leaf node using a Breadth-First approach, but prioritized by depth
      queue = self.root_nodes.copy()
      last_node = None
      while queue:
        current_node = queue.pop(0)
        if not current_node.children: # it's a leaf node, so we can stop here and update it's children.
            last_node = current_node
        else: # not the latest leaf, need to keep traversing the children
             # Add the current_node children to the start of the queue, for prioritizing the last added node.
             queue = current_node.children + queue

      if last_node:
          return self.add_message_to_node(last_node, text, color)

      return None # this should not be reached


# --- Example Usage ---
if __name__ == "__main__":
    # Create a chat flow
    negotiation_flow = ChatFlow("Price Negotiation")

    # Add initial root nodes (starting messages)
    root1 = negotiation_flow.add_root_node(
        "Hello! I'm interested in buying this item."
    )
    root2 = negotiation_flow.add_root_node("What is your best price?", "blue")

    # Add children to the first root node
    child1 = negotiation_flow.add_message_to_node(root1, "I saw it listed for $100.")
    grandchild1 = negotiation_flow.add_message_to_node(child1, "Is that correct?")

    # Add a child to the second root node
    child2 = negotiation_flow.add_message_to_node(root2, "I can offer $80.", "green")

    # Example using add_message_to_last_node
    negotiation_flow.add_message_to_last_node("Ok, deal!", "orange")

    # Find a specific node (e.g., the node with "I can offer $80.")
    found_node = negotiation_flow.find_node("I can offer $80.")
    if found_node:
        print(f"\nFound node: {found_node}")
        # Add a child to the found node
        negotiation_flow.add_message_to_node(found_node, "Sounds good!")
    else:
        print("\nNode not found.")

    # Print the chat flow (using __str__)
    print("\n",negotiation_flow)

    # Convert to dictionary (for frontend)
    print("\nChatFlow as Dictionary:")
    print(negotiation_flow.to_dict())

    # --- Example of adding more nodes ---
    empty_flow = ChatFlow("Empty Flow")
    empty_flow.add_message_to_last_node("This should not be added.")
    print("\nEmpty Flow")
    print(empty_flow) # expect to be empty

    new_flow = ChatFlow("New Flow")
    first_root = new_flow.add_root_node("First root message")
    first_child = new_flow.add_message_to_last_node("Added to the first root node.") # Added to the first root node.
    second_root = new_flow.add_root_node("Second root")
    second_root_child = new_flow.add_message_to_last_node("Added to the second root") #Added to the second root
    another_child_to_first = new_flow.add_message_to_node(first_root, "Another one") # another child of the first root.
    new_flow.add_message_to_last_node("And, one more.") #added to "Another one"

    print("\nNew Flow: ")
    print(new_flow.to_dict())