Function build_subgraph

src/WorkFlow.py:152–246

(node_map: Dict[str, NodeData], llm)

        return "False"

def build_subgraph(node_map: Dict[str, NodeData], llm) -> StateGraph:
    # Define the state machine
    subgraph = StateGraph(PipelineState)

    # Start node, only one start point
    start_node = find_nodes_by_type(node_map, "START")[0]
    logger(f"Start root ID: {start_node.uniq_id}")

    # Step nodes
    step_nodes = find_nodes_by_type(node_map, "STEP")
    for current_node in step_nodes:
        if current_node.tool:
            tool_info = tool_info_registry[current_node.tool]
            prompt_template = f"""
            history: {{history}}
            {current_node.description}
            Available tool: {tool_info}
            Based on Available tool, arguments in the json format:
            "function": "<func_name>", "args": [<arg1>, <arg2>, ...]

            next stage directly parse then run <func_name>(<arg1>,<arg2>, ...) make sure syntax is right json and align function siganture
            """
            subgraph.add_node(
                current_node.uniq_id, 
                lambda state, template=prompt_template, llm=llm, name=current_node.name : execute_tool(name, state, template, llm)
            )
        else:
            prompt_template=f"""
            history: {{history}}
            {current_node.description}
            you reply in the json format
            """
            subgraph.add_node(
                current_node.uniq_id, 
                lambda state, template=prompt_template, llm=llm, name=current_node.name: execute_step(name, state, template, llm)
            )

    # Add INFO nodes
    info_nodes = find_nodes_by_type(node_map, "INFO")
    for info_node in info_nodes:
        # INFO nodes just append predefined information to the state history
        subgraph.add_node(
            info_node.uniq_id, 
            lambda state, template=info_node.description, llm=llm, name=info_node.name: info_add(name, state, template, llm)
        )
    
    # Add SUBGRAPH nodes
    subgraph_nodes = find_nodes_by_type(node_map, "SUBGRAPH")
    for sg_node in subgraph_nodes:
        subgraph.add_node(
            sg_node.uniq_id, 
            lambda state, llm=llm, name=sg_node.name, sg_name=sg_node.name: sg_add(name, state, sg_name)
        )

    # Edges
    # Find all next nodes from start_node
    next_node_ids = start_node.nexts
    next_nodes = [node_map[next_id] for next_id in next_node_ids]