Graph of Thoughts: Solving Complex Problems with Large Language Models

import time # Assuming time is needed for timestamps later

class Prompter:
    def __init__(self, templates):
        """
        Initialize with a dictionary of prompt templates for different operations.
        
        :param templates: Dict mapping operation types (e.g., 'generate', 'merge') 
                          to f-string style templates. Templates should anticipate 
                          placeholders like {node_list}, {node_count}, {context}.
        """
        self.templates = templates
    
    def create_prompt(self, operation_type, active_nodes, context=None):
        """
        Builds a prompt string for the LLM based on the operation and graph state.
        
        :param operation_type: The intended action (e.g., "generate", "merge", "refine").
        :param active_nodes: List of node dictionaries (containing 'content', 'id', etc.) 
                             relevant to this operation.
        :param context: Optional string providing broader context or specific instructions.
        :return: A formatted prompt string ready for the LLM.
        """
        if operation_type not in self.templates:
            raise ValueError(f"No template defined for operation: {operation_type}")
            
        template = self.templates[operation_type]
        
        # Format node content clearly for the LLM
        node_text = ""
        for idx, node in enumerate(active_nodes):
            # Include ID for potential reference by the LLM, if needed
            node_text += f"--- Thought Node {node.get('id', idx)} ---\n" 
            node_text += f"{node.get('content', 'N/A')}\n"
            # Conditionally add metadata useful for the operation
            if 'score' in node and operation_type in ['refine', 'validate']:
                 node_text += f"  (Current Score: {node['score']:.2f})\n"
            node_text += "---\n"
        
        # Format context cleanly
        context_text = f"\nAdditional Context:\n{context}\n" if context else ""
        
        # Populate the template
        try:
            prompt = template.format(
                node_list=node_text,
                node_count=len(active_nodes),
                context=context_text
            )
        except KeyError as e:
            print(f"Warning: Template placeholder {e} not found or not needed for operation '{operation_type}'.")
            # Attempt to format with available keys, might fail if essential keys missing
            prompt = template.format_map(default_factory(str, 
                node_list=node_text, 
                node_count=len(active_nodes), 
                context=context_text
            ))

        return prompt

# Helper for potentially missing keys in format
from collections import defaultdict
class default_factory(defaultdict):
    def __missing__(self, key):
        return key # Return the key itself if missing

prompt_templates = {
    "generate": "Given the following thought(s), generate {k} distinct and novel extensions or alternative approaches:\n{node_list}\n{context}\nPresent each new idea clearly, perhaps numbered.",
    
    "merge": "Synthesize the core ideas from the following {node_count} thoughts into a single, more comprehensive solution. Resolve contradictions and combine strengths:\n{node_list}\n{context}\nProvide the unified solution:",
    
    "refine": "Critically evaluate and improve the following thought. Focus on addressing weaknesses, increasing precision, or enhancing its applicability:\n{node_list}\n{context}\nProvide the refined thought:",
    
    "validate": "Assess the validity and quality of the following thought. Identify logical flaws, missing information, or questionable assumptions. Provide a rating (0-10, 10 is best) and a brief justification:\n{node_list}\n{context}\nRating: [Your Rating]\nJustification: [Your Justification]"
}

import re
import json
import time

class Parser:
    def __init__(self, operation_parsers=None):
        """
        Initialize with optional specialized parsers for specific operation types.
        
        :param operation_parsers: Dict mapping operation types (e.g., 'generate') 
                                  to custom parsing functions. If None, uses defaults.
        """
        self.default_parser = self._basic_paragraph_parse
        self.operation_parsers = operation_parsers or {}
        # Add a specific parser for JSON if the LLM is instructed to provide it
        self.operation_parsers.setdefault('parse_json', self._parse_json_output)

    def parse_llm_output(self, llm_output, operation_type):
        """
        Parses the LLM's raw text output based on the operation that generated it.
        
        :param llm_output: The raw string response from the LLM.
        :param operation_type: String indicating the operation type (helps select the right parser).
        :return: A list of thought dictionaries, ready to be added as nodes.
                 Each dict should ideally have 'content' and 'created_at'.
        """
        # Use a specialized parser if one exists for this operation type
        parser_func = self.operation_parsers.get(operation_type, self.default_parser)
        
        # Attempt parsing, fall back to basic if specialized fails
        try:
            parsed_thoughts = parser_func(llm_output)
        except Exception as e:
            print(f"Parser for '{operation_type}' failed: {e}. Falling back to basic paragraph parsing.")
            parsed_thoughts = self._basic_paragraph_parse(llm_output)
            
        # Ensure basic structure and timestamp
        structured_thoughts = []
        for thought_content in parsed_thoughts:
            if isinstance(thought_content, dict) and 'content' in thought_content:
                thought = thought_content
            elif isinstance(thought_content, str):
                thought = {'content': thought_content}
            else:
                print(f"Warning: Skipping malformed parsed item: {thought_content}")
                continue # Skip malformed items

            if 'created_at' not in thought:
                 thought['created_at'] = time.time()
            structured_thoughts.append(thought)

        return structured_thoughts

    def _basic_paragraph_parse(self, llm_output):
        """Default parser: Treats non-empty paragraphs separated by double newlines as thoughts."""
        if not isinstance(llm_output, str): # Basic type check
             return []
        paragraphs = [p.strip() for p in llm_output.split("\n\n") if p.strip()]
        # Return list of strings directly, structuring happens in parse_llm_output
        return paragraphs

    def _parse_json_output(self, llm_output):
        """Attempts to parse structured JSON output, potentially embedded in text."""
        if not isinstance(llm_output, str):
            return []
        try:
            # Try finding JSON within ```json ... ``` blocks first
            json_match = re.search(r'```json\n(.*?)\n```', llm_output, re.DOTALL | re.IGNORECASE)
            if json_match:
                json_str = json_match.group(1)
            else:
                # If no block, try finding the first plausible JSON object/array
                json_match = re.search(r'($$.*$$|\{.*\})', llm_output, re.DOTALL)
                if not json_match:
                    raise ValueError("No JSON structure found")
                json_str = json_match.group(0)

            # Clean up potential artifacts before parsing
            json_str = json_str.strip()
            parsed_data = json.loads(json_str)

            # Expecting a list of thoughts (dicts) or a single thought (dict)
            if isinstance(parsed_data, list):
                # Assume list contains thought dictionaries
                return parsed_data 
            elif isinstance(parsed_data, dict):
                # Wrap single dictionary in a list
                return [parsed_data]
            else:
                raise ValueError(f"Parsed JSON is not a list or dict: {type(parsed_data)}")

        except (json.JSONDecodeError, ValueError, AttributeError) as e:
            print(f"JSON parsing failed: {e}. Input: '{llm_output[:100]}...'")
            # Important: Reraise or return empty to signal failure to the caller
            raise e # Or return [] if fallback is desired in the caller

import time

class ScoringValidator:
    def __init__(self, scoring_functions=None, domain_validators=None, llm_client=None, validation_template=None):
        """
        Initializes the scorer/validator.
        
        :param scoring_functions: Dict mapping metric names (e.g., 'clarity') to functions 
                                  that take a thought dict and return a score.
        :param domain_validators: Dict mapping domain names (e.g., 'math') to functions 
                                  that take a thought dict and return (bool_valid, details_dict).
        :param llm_client: An optional function/client to call an LLM for validation.
        :param validation_template: An f-string template for LLM-based validation prompts, 
                                    expecting a {thought_content} placeholder.
        """
        # Default scoring functions (examples)
        self.scoring_functions = scoring_functions or {
            "length": lambda t: len(t.get("content", "").split()),
            "novelty": lambda t, graph: self._calculate_novelty(t, graph) # Example needing graph access
        }
        
        self.domain_validators = domain_validators or {}
        self.llm_client = llm_client
        self.validation_template = validation_template

    def score_and_validate(self, thought, graph_state, metrics=None, domain=None):
        """
        Scores a thought using specified metrics and validates it based on domain or LLM.
        
        :param thought: The thought dictionary to evaluate.
        :param graph_state: The current graph dictionary (needed for context-aware scoring like novelty).
        :param metrics: List of metric names (keys in self.scoring_functions) to compute. 
                        If None, computes all available metrics.
        :param domain: Optional string specifying the problem domain for specialized validation.
        :return: The input thought dictionary, updated with 'scores' dict and 'validation' dict.
        """
        thought.setdefault('scores', {})
        thought.setdefault('validation', {'is_valid': None, 'details': None}) # Initialize validation structure

        metrics_to_run = metrics or list(self.scoring_functions.keys())
        
        # Calculate scores
        for metric in metrics_to_run:
            if metric in self.scoring_functions:
                try:
                    # Pass graph state if function expects it (e.g., for novelty)
                    score_func = self.scoring_functions[metric]
                    sig = inspect.signature(score_func)
                    if 'graph' in sig.parameters or 'graph_state' in sig.parameters:
                         thought['scores'][metric] = score_func(thought, graph_state)
                    else:
                         thought['scores'][metric] = score_func(thought)
                except Exception as e:
                    print(f"Error scoring metric '{metric}': {e}")
                    thought['scores'][metric] = None # Indicate scoring error

        # Calculate overall score (example: average, handle potential Nones)
        valid_scores = [s for s in thought['scores'].values() if isinstance(s, (int, float))]
        if valid_scores:
             thought['scores']['overall'] = sum(valid_scores) / len(valid_scores)
        else:
             thought['scores']['overall'] = 0 # Default if no valid scores

        # Perform validation
        valid = None
        details = {}
        try:
            if domain and domain in self.domain_validators:
                valid, details = self.domain_validators[domain](thought)
            elif self.llm_client and self.validation_template:
                valid, details = self._validate_with_llm(thought)
            else:
                # Fallback to basic validation if no specific method applies
                valid, details = self._basic_content_validation(thought)
        except Exception as e:
             print(f"Error during validation: {e}")
             valid = False # Default to invalid on error
             details = {'error': str(e)}

        thought['validation']['is_valid'] = bool(valid) # Ensure boolean
        thought['validation']['details'] = details
            
        return thought
        
    def _calculate_novelty(self, thought, graph):
        """Example context-aware scorer: checks similarity to existing nodes."""
        # Placeholder: Requires a similarity function (e.g., embedding distance)
        # compare thought['content'] against content of nodes in graph['nodes']
        # Return a score (e.g., 1 - max_similarity)
        return 0.8 # Dummy value

    def _basic_content_validation(self, thought):
        """Simple default validation: checks for minimal content."""
        content = thought.get("content", "")
        is_valid = isinstance(content, str) and len(content.strip()) > 5 # Arbitrary min length
        details = {'reason': 'Basic content check' if is_valid else 'Content too short or missing'}
        return is_valid, details

    def _validate_with_llm(self, thought):
        """Uses an LLM to perform validation based on a template."""
        if not self.llm_client or not self.validation_template:
             return False, {'error': 'LLM validator not configured'}
             
        try:
            prompt = self.validation_template.format(thought_content=thought.get("content", ""))
            response = self.llm_client(prompt)
            
            # Heuristic parsing of LLM validation response (needs improvement)
            lower_response = response.lower()
            is_valid = "valid" in lower_response or "correct" in lower_response or "rating: [8-9]|10" in lower_response
            if "invalid" in lower_response or "incorrect" in lower_response or "flawed" in lower_response or "rating: [0-5]" in lower_response:
                 is_valid = False
            
            return is_valid, {'llm_response': response}
        except Exception as e:
             print(f"LLM validation failed: {e}")
             return False, {'error': f"LLM call failed: {e}"}

# --- Example Domain Validators (outside the class) ---
import inspect # Needed for signature inspection

def create_math_validator():
    """Creates a validator function for mathematical thoughts."""
    def validate_math_thought(thought):
        content = thought.get("content", "")
        details = {}
        # Basic checks (can be much more sophisticated)
        has_equation = re.search(r'=|\+|-|\*|/', content) is not None
        has_final_answer = re.search(r'answer is|result is|final value is', content, re.IGNORECASE) is not None
        no_obvious_error = "error" not in content.lower() and "cannot compute" not in content.lower()
        
        is_valid = has_equation and has_final_answer and no_obvious_error
        details = {'has_equation': has_equation, 'has_final_answer': has_final_answer, 'no_obvious_error': no_obvious_error}
        return is_valid, details
    return validate_math_thought

def create_code_validator(language="python"):
    """Creates a validator function for code snippets."""
    def validate_code_thought(thought):
        content = thought.get("content", "")
        details = {'language': language}
        is_valid = False
        
        # Extract code blocks if necessary
        code_match = re.search(r'```(?:\w*\n)?(.*?)```', content, re.DOTALL)
        code_to_validate = code_match.group(1).strip() if code_match else content.strip()

        if not code_to_validate:
             details['error'] = "No code content found"
             return False, details

        if language == "python":
            try:
                import ast
                ast.parse(code_to_validate)
                details["syntax_check"] = "Passed"
                is_valid = True # Basic syntax validity
            except SyntaxError as e:
                details["syntax_check"] = f"Failed: {e}"
                is_valid = False
            # Add more checks: linting, simple execution tests?
        else:
            details["syntax_check"] = f"Validation for {language} not implemented"
            is_valid = True # Assume valid if no validator exists

        return is_valid, details
    return validate_code_thought

# Example usage
domain_validators = {
    "math": create_math_validator(),
    "python_code": create_code_validator("python")
}
# scorer = ScoringValidator(domain_validators=domain_validators) # Initialize with these

import time

class Controller:
    def __init__(self, graph_of_operations, scoring_validator, llm_client, prompter, parser):
        """
        Initializes the Controller.
        
        :param graph_of_operations: A list of operation dictionaries defining the strategy.
        :param scoring_validator: An instance of ScoringValidator.
        :param llm_client: Function/client to interact with the LLM.
        :param prompter: An instance of Prompter.
        :param parser: An instance of Parser.
        """
        self.go_ops = graph_of_operations
        self.scoring_validator = scoring_validator
        self.llm_client = llm_client # Renamed from llm_api for clarity
        self.prompter = prompter
        self.parser = parser
        
        # Graph Reasoning State (GRS)
        self.graph = {
            "nodes": [],  # List of thought dictionaries {id, content, scores, validation, ...}
            "edges": [],  # List of tuples (from_node_id, to_node_id, {'operation': type, ...})
            "metadata": {"created_at": time.time(), "status": "initialized"}
        }
        self._active_node_ids = set() # Use a set for efficient management
        self._node_counter = 0

    def _get_next_node_id(self):
        """Generates unique node IDs."""
        nid = self._node_counter
        self._node_counter += 1
        return nid

    def initialize_graph(self, initial_content):
        """Sets up the graph with one or more initial thoughts."""
        initial_thoughts = []
        if isinstance(initial_content, str):
            initial_thoughts = [{'content': initial_content}]
        elif isinstance(initial_content, list):
            initial_thoughts = [{'content': c} for c in initial_content if isinstance(c, str)]
        else:
             raise ValueError("initial_content must be a string or list of strings")

        if not initial_thoughts:
             raise ValueError("No valid initial thoughts provided")

        initial_node_ids = []
        for thought in initial_thoughts:
            node_id = self._add_node(thought) # _add_node now handles scoring/validation
            initial_node_ids.append(node_id)
            
        self._active_node_ids = set(initial_node_ids)
        self.graph["metadata"]["status"] = "running"
        print(f"Graph initialized with {len(self._active_node_ids)} node(s).")

    def run(self):
        """Executes the sequence of operations defined in the GoO."""
        if self.graph["metadata"]["status"] == "initialized":
             raise RuntimeError("Graph not initialized. Call initialize_graph() first.")
             
        total_ops = len(self.go_ops)
        for op_idx, op_config in enumerate(self.go_ops):
            op_type = op_config.get("type")
            print(f"\n--- Executing Op {op_idx+1}/{total_ops}: {op_type} ---")
            
            if not self._active_node_ids and op_type not in ["select"]: # Allow select to potentially recover active nodes
                 print("Warning: No active nodes to operate on. Skipping operation.")
                 continue

            # Dynamically call the appropriate execution method
            exec_method_name = f"_execute_{op_type}"
            if hasattr(self, exec_method_name) and callable(getattr(self, exec_method_name)):
                try:
                    getattr(self, exec_method_name)(op_config)
                except Exception as e:
                     print(f"Error executing operation '{op_type}': {e}")
                     # Decide whether to halt or continue
                     # self.graph["metadata"]["status"] = "error"
                     # return None # Example: Halt on error
                     print("Attempting to continue...")
            else:
                print(f"Warning: Unknown operation type '{op_type}'. Skipping.")

            # Optional: Add graph validation/cleanup steps here

            if not self._active_node_ids and op_idx < total_ops - 1:
                 print("Warning: No active nodes remaining after operation. Subsequent operations might fail.")

        self.graph["metadata"]["status"] = "completed"
        return self._get_final_solution() # Return the best node based on some criteria

    def _add_node(self, thought_content_dict):
        """Adds a node to the graph, assigns ID, scores, validates, and returns ID."""
        node_id = self._get_next_node_id()
        thought_content_dict['id'] = node_id
        
        # Ensure basic structure even if parser failed partially
        thought_content_dict.setdefault('content', '')
        thought_content_dict.setdefault('created_at', time.time())
        
        # Score and validate the thought using the current graph state
        scored_thought = self.scoring_validator.score_and_validate(
            thought_content_dict, 
            self.graph, # Pass the whole graph for context
            domain=thought_content_dict.get('domain') # Allow domain hint
        )
        
        self.graph["nodes"].append(scored_thought)
        # print(f"Added Node {node_id}: Score={scored_thought['scores'].get('overall', 'N/A'):.2f}, Valid={scored_thought['validation'].get('is_valid', 'N/A')}")
        return node_id

    def _add_edge(self, from_id, to_id, operation_details):
        """Adds a directed edge with metadata about the operation."""
        # Ensure valid IDs
        if self._get_node_by_id(from_id) is None or self._get_node_by_id(to_id) is None:
             print(f"Warning: Attempting to add edge with invalid node ID ({from_id} -> {to_id}). Skipping.")
             return
        edge = (from_id, to_id, operation_details)
        self.graph["edges"].append(edge)

    def _get_nodes_by_ids(self, node_ids):
        """Retrieves full node dictionaries for a set of IDs."""
        # More efficient than repeated _get_node_by_id calls
        nodes_dict = {node['id']: node for node in self.graph['nodes']}
        return [nodes_dict[nid] for nid in node_ids if nid in nodes_dict]

    def _get_active_nodes(self):
        """Returns the list of currently active node dictionaries."""
        return self._get_nodes_by_ids(self._active_node_ids)

    def _execute_generate(self, op_config):
        """Generates new thoughts from selected source nodes."""
        k = op_config.get("k", 1) # Default to generating 1 per source
        context = op_config.get("context")
        source_node_ids = op_config.get("source_nodes", list(self._active_node_ids)) # Default to active
        
        newly_generated_ids = set()
        
        source_nodes = self._get_nodes_by_ids(source_node_ids)
        if not source_nodes:
            print("Generate: No valid source nodes found.")
            return

        for source_node in source_nodes:
            prompt = self.prompter.create_prompt(
                "generate", 
                [source_node], # Generate from one source at a time
                context=context.format(k=k) if context and '{k}' in context else context # Allow k in context
            )
            llm_output = self.llm_client(prompt)
            new_thoughts = self.parser.parse_llm_output(llm_output, "generate")
            
            added_count = 0
            for thought in new_thoughts:
                if added_count >= k: break # Respect k limit
                new_node_id = self._add_node(thought)
                self._add_edge(source_node["id"], new_node_id, {"operation": "generate"})
                newly_generated_ids.add(new_node_id)
                added_count += 1
        
        # Update active nodes: replace sources with generated, or add to active set
        if op_config.get("replace_active", True):
            self._active_node_ids = newly_generated_ids
        else:
            self._active_node_ids.update(newly_generated_ids)
        print(f"Generate: Created {len(newly_generated_ids)} new nodes. Active: {len(self._active_node_ids)}")

    def _execute_merge(self, op_config):
        """Merges multiple thoughts into one or more new thoughts."""
        input_selector = op_config.get("input_selector", "active") # active, top_k, specific_ids
        repeat = op_config.get("repeat", 1) # How many merge attempts
        context = op_config.get("context")

        nodes_to_merge_ids = self._select_node_ids(op_config) # Use helper for selection logic
        
        if len(nodes_to_merge_ids) < 2:
             print(f"Merge: Need at least 2 nodes to merge, found {len(nodes_to_merge_ids)}. Skipping.")
             return

        nodes_to_merge = self._get_nodes_by_ids(nodes_to_merge_ids)
        newly_merged_ids = set()

        for _ in range(repeat):
            prompt = self.prompter.create_prompt("merge", nodes_to_merge, context=context)
            llm_output = self.llm_client(prompt)
            # Assume merge typically produces one coherent thought, but parser handles lists
            merged_thoughts = self.parser.parse_llm_output(llm_output, "merge") 
            
            for thought in merged_thoughts:
                new_node_id = self._add_node(thought)
                for source_node in nodes_to_merge: # Add edges from all parents
                    self._add_edge(source_node["id"], new_node_id, {"operation": "merge"})
                newly_merged_ids.add(new_node_id)
                if len(merged_thoughts) == 1: break # Often only expect one merge result

        if op_config.get("update_active", True):
            # Replace merged nodes with the new merged node(s)
            self._active_node_ids = (self._active_node_ids - set(nodes_to_merge_ids)) | newly_merged_ids
        else:
             self._active_node_ids.update(newly_merged_ids) # Just add merged nodes
        print(f"Merge: Created {len(newly_merged_ids)} merged nodes from {len(nodes_to_merge_ids)} inputs. Active: {len(self._active_node_ids)}")

    def _execute_refine(self, op_config):
        """Refines selected thoughts, creating new versions."""
        context = op_config.get("context")
        nodes_to_refine_ids = self._select_node_ids(op_config) # Use helper

        if not nodes_to_refine_ids:
             print("Refine: No nodes selected for refinement. Skipping.")
             return

        nodes_to_refine = self._get_nodes_by_ids(nodes_to_refine_ids)
        newly_refined_ids = set()

        for node in nodes_to_refine:
            prompt = self.prompter.create_prompt("refine", [node], context=context)
            llm_output = self.llm_client(prompt)
            # Assume refinement produces one improved thought
            refined_thoughts = self.parser.parse_llm_output(llm_output, "refine") 
            
            for thought in refined_thoughts: # Usually just one
                new_node_id = self._add_node(thought)
                self._add_edge(node["id"], new_node_id, {"operation": "refine"})
                newly_refined_ids.add(new_node_id)
                if len(refined_thoughts) == 1: break

        if op_config.get("update_active", True):
             # Replace refined nodes with their new versions
             self._active_node_ids = (self._active_node_ids - set(nodes_to_refine_ids)) | newly_refined_ids
        else:
             self._active_node_ids.update(newly_refined_ids) # Add refined nodes
        print(f"Refine: Created {len(newly_refined_ids)} refined nodes from {len(nodes_to_refine_ids)} inputs. Active: {len(self._active_node_ids)}")


    def _select_node_ids(self, op_config):
        """Helper function to select node IDs based on criteria in op_config."""
        selector_type = op_config.get("input_selector", "active") # Default source is active nodes
        source_pool_ids = set(op_config.get("source_pool", list(self._active_node_ids))) # Pool to select from

        if not source_pool_ids:
             return [] # Return empty list if the source pool is empty

        source_pool_nodes = self._get_nodes_by_ids(source_pool_ids)

        if selector_type == "active":
            return list(self._active_node_ids) # Return current active nodes directly
        
        elif selector_type == "top_k":
            k = op_config.get("k", 1)
            metric = op_config.get("metric", "scores.overall") # Support nested keys
            
            # Helper to get nested metric values safely
            def get_metric_value(node, metric_key):
                keys = metric_key.split('.')
                value = node
                try:
                    for key in keys:
                        value = value[key]
                    return value if isinstance(value, (int, float)) else 0
                except (KeyError, TypeError):
                    return 0 # Default for missing or non-numeric scores
            
            sorted_nodes = sorted(
                source_pool_nodes,
                key=lambda n: get_metric_value(n, metric),
                reverse=True # Higher score is better
            )
            return [node["id"] for node in sorted_nodes[:k]]
            
        elif selector_type == "valid_only":
             return [
                 node["id"] for node in source_pool_nodes
                 if node.get("validation", {}).get("is_valid", False)
             ]
             
        elif selector_type == "threshold":
            metric = op_config.get("metric", "scores.overall")
            threshold = op_config.get("threshold", 0.5)
            return [
                node["id"] for node in source_pool_nodes
                if get_metric_value(node, metric) >= threshold
            ]

        elif selector_type == "latest":
             k = op_config.get("k", 1)
             sorted_nodes = sorted(
                 source_pool_nodes,
                 key=lambda n: n.get("created_at", 0),
                 reverse=True
             )
             return [node["id"] for node in sorted_nodes[:k]]
             
        elif selector_type == "specific_ids":
             # Directly use provided IDs, filtering by those actually in the pool
             specified_ids = set(op_config.get("node_ids", []))
             return list(source_pool_ids.intersection(specified_ids))

        elif selector_type == "custom":
             custom_selector = op_config.get("selector_function")
             if custom_selector and callable(custom_selector):
                 # Pass necessary info to custom function
                 return custom_selector(self.graph, list(source_pool_ids)) 
             else:
                 print("Warning: Custom selector specified but no valid function provided.")
                 return list(source_pool_ids) # Fallback

        else:
            print(f"Warning: Unknown selector type '{selector_type}'. Using 'active'.")
            return list(self._active_node_ids)

    def _execute_select(self, op_config):
        """Selects a subset of nodes to become the new active set."""
        selected_ids = self._select_node_ids(op_config) # Reuse the selection logic

        # Ensure at least one node remains active if possible, unless explicitly allowed empty
        if not selected_ids and self._active_node_ids and not op_config.get("allow_empty_selection", False):
            print("Select: Selection resulted in zero nodes. Retaining previous active set.")
            # Optionally select the best single node as fallback?
            # best_fallback = self._select_node_ids({'input_selector': 'top_k', 'k': 1})
            # self._active_node_ids = set(best_fallback) if best_fallback else self._active_node_ids
            pass # Keep current active nodes
        else:
            self._active_node_ids = set(selected_ids)
            
        print(f"Select: New active node set size: {len(self._active_node_ids)}")
        # This operation primarily modifies the active set, doesn't usually add nodes/edges.

    def _get_node_by_id(self, node_id):
        """Efficiently retrieves a node by its ID."""
        # Consider indexing nodes by ID in a dict for large graphs if performance matters
        for node in self.graph["nodes"]:
            if node["id"] == node_id:
                return node
        return None # Return None if not found

    def _get_final_solution(self):
        """Selects the 'best' node from the graph as the final answer."""
        
        # Prioritize valid nodes first
        valid_nodes = [n for n in self.graph["nodes"] if n.get("validation", {}).get("is_valid", False)]
        
        candidate_pool = valid_nodes if valid_nodes else self.graph["nodes"] # Fallback to all nodes if no valid ones
        
        if not candidate_pool:
             print("Error: No nodes in the graph to select a final solution from.")
             return None

        # Select the node with the highest overall score from the candidate pool
        # Use the same safe metric getter as in _select_node_ids
        def get_metric_value(node, metric_key="scores.overall"):
             keys = metric_key.split('.')
             value = node
             try:
                 for key in keys:
                     value = value[key]
                 return value if isinstance(value, (int, float)) else 0
             except (KeyError, TypeError):
                 return 0

        best_node = max(candidate_pool, key=lambda n: get_metric_value(n))
        
        print(f"\n--- Final Solution (Node {best_node['id']}) ---")
        print(f"Score: {get_metric_value(best_node):.2f}")
        print(f"Content: {best_node.get('content')[:200]}...") # Show snippet
        return best_node

# Example GoO: Solve a complex design problem
graph_of_operations = [
    # 1. Initial Brainstorm: Generate diverse starting points
    {"type": "generate", "k": 5, "context": "Generate 5 distinct high-level approaches to solve [Problem Description]."},
    
    # 2. Initial Filter: Select the top 3 promising approaches
    {"type": "select", "input_selector": "top_k", "k": 3, "metric": "scores.overall"}, 
    
    # 3. Elaborate Approaches: Develop each selected approach further
    {"type": "generate", "k": 2, "replace_active": False, # Add elaborations, don't replace parents yet
     "context": "For the given approach, detail the first 2-3 key steps or components."},
     
    # 4. Focused Refinement/Critique: Improve the elaborated steps
    {"type": "refine", "input_selector": "latest", "k": 6, # Refine the 6 new nodes from step 3
     "update_active": True, # Replace originals with refined versions
     "context": "Critique this step/component. Identify potential flaws and suggest improvements."},
     
    # 5. Synthesis Attempt: Merge the refined components from the best original approach
    # (Requires more complex selection logic not fully shown here)
    # {"type": "merge", "input_selector": "custom", "selector_function": select_best_approach_components, "repeat": 1},
    
    # 6. Final Selection: Pick the best synthesized solution
    {"type": "select", "input_selector": "top_k", "k": 1, "metric": "scores.overall"},
]

# Simplified conceptual function (actual call uses Controller methods)
def generate_thought(source_node_content, llm, k=1):
    prompt = f"Based on: '{source_node_content}'. Generate {k} distinct next steps or ideas."
    response = llm(prompt)
    # Parsing logic omitted
    new_contents = parse_response(response, k) 
    return [{'content': c} for c in new_contents] 

# Simplified conceptual function
def refine_thought(thought_to_refine, llm, critique=None):
    context = f"Critique: {critique}" if critique else ""
    prompt = f"Improve the following thought: '{thought_to_refine['content']}'. {context}"
    response = llm(prompt)
    return {'content': response, 'refined_from': thought_to_refine['id']}

# Simplified conceptual function
def merge_thoughts(thoughts_to_merge, llm):
    contents = [f"Idea {i+1}: {t['content']}" for i, t in enumerate(thoughts_to_merge)]
    prompt = f"Combine the following distinct ideas into a single coherent synthesis:\n\n" + "\n".join(contents) + "\n\nSynthesized Idea:"
    response = llm(prompt)
    return {
        'content': response,
        'merged_from': [t['id'] for t in thoughts_to_merge]
    }

# Conceptual GoO for NDA Merging
nda_merger_operations = [
    # 1. Deconstruct: Extract key clauses/sections from each source NDA.
    #    - Input: Each original NDA text node.
    #    - Output: Multiple nodes per NDA, each representing a key clause (Definition, Term, Confidentiality Obligations, etc.).
    {"type": "generate", "k": 'all', # Special value meaning extract all relevant parts
     "context": "Identify and extract key clauses (like Definitions, Term, Confidentiality Obligations, Exclusions, Remedies) from this NDA document. Present each clause separately.",
     "source_nodes": ["NDA1_raw", "NDA2_raw", "NDA3_raw"]}, # Assuming initial nodes have these IDs

    # 2. Cluster: Group similar clauses across all source documents.
    #    - Input: All extracted clause nodes from step 1.
    #    - Output: Nodes representing *groups* of related clauses (e.g., "Confidentiality Definition Group", "Term Group").
    {"type": "merge", # Using merge conceptually for grouping/clustering
     "input_selector": "custom", "selector_function": find_clause_nodes, # Selects all clause nodes
     "repeat": 'auto', # Dynamically create groups based on similarity
     "context": "Analyze these extracted clauses. Group clauses covering the same topic (e.g., Definition of Confidential Information, Permitted Disclosure). For each group, summarize the core concept and list the source clause IDs."},

    # 3. Synthesize: Draft merged clauses for each group, resolving differences.
    #    - Input: Each clause group node from step 2.
    #    - Output: Nodes containing proposed merged text for each clause type.
    {"type": "generate", "k": 1, # Generate one best merged version per group
     "input_selector": "custom", "selector_function": find_group_nodes,
     "context": "Based on the clauses in this group, draft a single, comprehensive clause that covers the essential points from all sources. Prioritize clarity and consistency. Note any significant variations that couldn't be reconciled."},

    # 4. Assemble: Combine the drafted merged clauses into a full draft document.
    #    - Input: All drafted merged clause nodes from step 3.
    #    - Output: One or more full draft NDA nodes.
    {"type": "merge", # Using merge to assemble the document
     "input_selector": "custom", "selector_function": find_draft_clause_nodes,
     "repeat": 1, # Create one primary draft
     "context": "Assemble these drafted clauses into a coherent, well-structured NDA document. Ensure logical flow and consistent definitions."},

    # 5. Review & Refine: Improve the full draft for legal soundness, clarity, and completeness.
    #    - Input: The full draft NDA node(s).
    #    - Output: Final polished NDA node.
    {"type": "refine", "input_selector": "active", # Refine the assembled draft
     "context": "Review this draft NDA for legal clarity, completeness, internal consistency, and grammatical correctness. Identify any ambiguities or potential issues and provide the improved final version."}
]

# Conceptual GoO for Math
math_operations = [
    {"type": "generate", "k": 3, "context": "Outline 3 different methods to solve this problem."},
    {"type": "generate", "k": 1, "replace_active": False, # Elaborate each method
     "context": "Work through the steps of this method."},
    {"type": "refine", "input_selector": "latest", "k": 3, # Refine the worked steps
     "context": "Check each calculation step for errors and clarity."},
    {"type": "validate", "domain": "math", # Use math-specific validator
     "input_selector": "active"},
    {"type": "select", "criteria": "valid_only"},
    {"type": "merge", "context": "Combine the valid solutions/methods into a final answer, explaining each approach."}
]

# Conceptual GoO for Code Gen
code_operations = [
    {"type": "generate", "k": 2, "context": "Generate two distinct implementations for this function specification."},
    {"type": "validate", "domain": "python_code", # Use code validator (syntax, maybe tests)
     "input_selector": "active"},
    {"type": "select", "criteria": "valid_only"}, # Discard implementations that fail basic checks
    {"type": "refine", "input_selector": "active",
     "context": "Refactor this code for improved efficiency and readability based on standard practices."},
    {"type": "validate", "domain": "python_code", # Re-validate after refactoring
     "input_selector": "active"},
    {"type": "select", "criteria": "top_k", "k": 1, "metric": "scores.efficiency"}, # Pick best performing valid one
    {"type": "refine", "input_selector": "active",
     "context": "Add clear documentation (docstrings, comments) to the final code."}
]

# Conceptual GoO for Writing
writing_operations = [
    {"type": "generate", "k": 3, "context": "Generate 3 different outlines for an article on [Topic]."},
    {"type": "select", "criteria": "top_k", "k": 1}, # Choose the best outline
    {"type": "generate", "k": 1, # Draft based on the selected outline
     "context": "Expand this outline into a first draft, focusing on covering all key points."},
    {"type": "refine", "input_selector": "active",
     "context": "Improve the clarity, flow, and structure of this draft."},
    {"type": "refine", "input_selector": "active",
     "context": "Enhance the writing style, engagement, and persuasiveness."},
    {"type": "refine", "input_selector": "active",
     "context": "Perform a final proofread for grammar, spelling, and punctuation errors."}
]