src/graph/workflow.py - working - Yan_Yu/Nereus

src/graph/workflow.py
123 строки3 KB
yan
init: scaffold Nereus AI tutoring platform with multi-agent orchestration
02 авг 2026, 18:25
02 авг 2026, 18:2513d60da
"""LangGraph workflow assembly for the Nereus agent system.

Constructs the StateGraph with all nodes and edges, then compiles
it into a runnable CompiledGraph.
"""

import logging
from typing import Any, Dict

from langgraph.graph import StateGraph, END

from src.graph.state import AgentState
from src.graph.conditions import choose_action
from src.graph.nodes import (
    coach_node,
    tutor_node,
    examiner_node,
    deep_dive_node,
    finalize_node,
)

logger = logging.getLogger(__name__)


def build_graph() -> StateGraph:
    """Build the LangGraph StateGraph with all nodes and edges.

    Graph structure:
        START -> coach -> tutor -> examiner -> {choose_action}
            -> tutor (next block) / finalize / deep_dive -> examiner (loop)

    Returns:
        Uncompiled StateGraph instance.
    """
    # Create the StateGraph
    graph = StateGraph(AgentState)

    # Add all nodes
    graph.add_node("coach", coach_node)
    graph.add_node("tutor", tutor_node)
    graph.add_node("examiner", examiner_node)
    graph.add_node("deep_dive", deep_dive_node)
    graph.add_node("finalize", finalize_node)

    # Define the entry point
    graph.set_entry_point("coach")

    # coach -> tutor (always)
    graph.add_edge("coach", "tutor")

    # tutor -> examiner (always)
    graph.add_edge("tutor", "examiner")

    # examiner -> conditional edge based on choose_action
    graph.add_conditional_edges(
        "examiner",
        choose_action,
        {
            "tutor": "advance_and_tutor",  # Move to next block (via helper)
            "finalize": "finalize",        # All blocks completed
            "deep_dive": "deep_dive",      # Remedial study needed
        },
    )

    # Helper node that increments index and then goes to tutor
    def _advance_to_next_block(state: AgentState) -> Dict[str, Any]:
        """Increment block index and pass through to tutor."""
        current_index = state.get("current_block_index", 0)
        return {"current_block_index": current_index + 1}

    graph.add_node("advance_and_tutor", _advance_to_next_block)
    graph.add_edge("advance_and_tutor", "tutor")

    # deep_dive -> examiner (re-assessment after additional study)
    graph.add_edge("deep_dive", "examiner")

    # finalize -> END
    graph.add_edge("finalize", END)

    logger.info("Graph built successfully with all nodes and edges.")
    return graph


def compile_graph() -> Any:
    """Build and compile the LangGraph workflow.

    Returns:
        Compiled LangGraph application (Runnable).
    """
    graph = build_graph()
    compiled = graph.compile()
    logger.info("Graph compiled successfully.")
    return compiled


# Module-level compiled graph (created on import)
_app = compile_graph()


def run_graph(initial_state: Dict[str, Any]) -> Dict[str, Any]:
    """Run the workflow with the given initial state.

    This is the main entry point for executing the agent workflow.

    Args:
        initial_state: Dictionary with initial state values (user_name,
                       target_skill, current_level, etc.)

    Returns:
        Final state dictionary after the graph has completed execution.
    """
    result = _app.invoke(initial_state)
    logger.info("Graph execution completed.")
    return result


def get_graph():
    """Get the compiled graph instance.

    Returns:
        Compiled LangGraph application.
    """
    return _app