WalkXR-AI / src / walkxr_ai / core / routing.py
routing.py
Raw

"""
routing.py

Routing helpers for the WalkXR LangGraph workflow.

These functions keep conditional graph logic separate from graph construction,
so branching decisions are easy to read and extend later.
"""

from __future__ import annotations
from walkxr_ai.core.state import WalkState

MEMORY_UPDATE_INTERVAL = 10


def route_after_call_llm(state: WalkState) -> str:
    """
    Routes the graph after the call_llm node.

    Returns:
    - "error" if the LLM step failed
    - "call_supervisor" if the response was generated successfully
    """
    if state.get("error"):
        return "error"

    return "call_supervisor"

def route_after_supervisor_score(state: WalkState) -> str:
    """
    Routes the graph after the reflection node.

    Returns:
    - "exit" if the supervisor scored the llm response as "good"
        or maximum response generation attempts have been reached
    - "loop" if the supervisor scored the llm response as "bad"
    """
    if state["supervisor_score"] == "good":
        return "exit"
    elif state["supervisor_retries"] >= state["max_retries_allowed"]:
        return "fail-safe"
    return "loop"


def route_after_update_conversation_history(state: WalkState) -> str:
    """
    Routes the graph after the conversation history has been updated.

    Returns:
    - "update_memory" when the current completed turn hits the memory interval
    - "end" otherwise
    """
    completed_turn = state.get("turn_index", 0) + 1

    # Update memory every MEMORY_UPDATE_INTERVAL turns
    if completed_turn % MEMORY_UPDATE_INTERVAL == 0:
        return "update_memory"

    return "end"