brain_server/langgraph/graph_definition.py

"""????? ???? LangGraph ?? ???? ????."""
import logging
from typing import Dict, Any
from agents.planner_agent import PlannerAgent
from agents.reasoning_agent import ReasoningAgent
from agents.tool_selector_agent import ToolSelectorAgent
from agents.supervisor_agent import SupervisorAgent
from agents.auto_heal_agent import AutoHealAgent
from agents.memory_agent import MemoryAgent
from api.deps import get_logger

logger = get_logger("kapo.langgraph")


class SimpleGraph:
    """???? ???? ?? LangGraph ??? ???? ???????."""

    def run(self, user_input: str, context: Dict[str, Any]):
        planner = PlannerAgent()
        reasoning = ReasoningAgent()
        tool_selector = ToolSelectorAgent()
        supervisor = SupervisorAgent()
        auto_heal = AutoHealAgent()
        memory = MemoryAgent()

        plan = planner.run(user_input, context)
        rationale = reasoning.run(user_input, plan)
        memory.write_short_term("last_plan", {"plan": plan, "rationale": rationale})
        return {
            "plan": plan,
            "rationale": rationale,
            "tool_selector": tool_selector,
            "supervisor": supervisor,
            "auto_heal": auto_heal,
        }


def get_graph():
    """????? ??????? LangGraph ?? ????? ???? ?????? SimpleGraph."""
    try:
        import langgraph  # noqa: F401
        logger.info("Using langgraph")
        return SimpleGraph()
    except Exception:
        logger.warning("LangGraph not available; using SimpleGraph")
        return SimpleGraph()