Spaces:
Sleeping
Sleeping
| """ | |
| StateGraph construction and compilation. | |
| Wires together all nodes and conditional edges into a single compiled graph. | |
| """ | |
| from __future__ import annotations | |
| import logging | |
| from langgraph.graph import END, StateGraph | |
| from src.graphs.state import AgentState | |
| from src.graphs.nodes import ( | |
| entry_node, | |
| semantic_router_node, | |
| llm_router_node, | |
| error_node, | |
| swap_agent_node, | |
| lending_agent_node, | |
| staking_agent_node, | |
| liquidity_agent_node, | |
| dca_agent_node, | |
| strategy_agent_node, | |
| crypto_agent_node, | |
| search_agent_node, | |
| default_agent_node, | |
| database_agent_node, | |
| portfolio_advisor_node, | |
| ) | |
| from src.graphs.edges import decide_route, after_llm_router | |
| from src.agents.formatter.node import formatter_node | |
| logger = logging.getLogger(__name__) | |
| # All agent nodes that feed into the formatter | |
| _AGENT_NODES = [ | |
| "swap_agent_node", | |
| "lending_agent_node", | |
| "staking_agent_node", | |
| "liquidity_agent_node", | |
| "dca_agent_node", | |
| "strategy_agent_node", | |
| "crypto_agent_node", | |
| "search_agent_node", | |
| "default_agent_node", | |
| "database_agent_node", | |
| "portfolio_advisor_node", | |
| ] | |
| def build_graph() -> StateGraph: | |
| """ | |
| Construct and compile the Zico agent StateGraph. | |
| Flow: | |
| entry → semantic_router → {decide_route} → agent / llm_router / error | |
| llm_router → {after_llm_router} → agent | |
| agent → formatter → END | |
| error → END | |
| """ | |
| graph = StateGraph(AgentState) | |
| # --- Register nodes --- | |
| graph.add_node("entry_node", entry_node) | |
| graph.add_node("semantic_router_node", semantic_router_node) | |
| graph.add_node("llm_router_node", llm_router_node) | |
| graph.add_node("error_node", error_node) | |
| graph.add_node("formatter_node", formatter_node) | |
| graph.add_node("swap_agent_node", swap_agent_node) | |
| graph.add_node("lending_agent_node", lending_agent_node) | |
| graph.add_node("staking_agent_node", staking_agent_node) | |
| graph.add_node("liquidity_agent_node", liquidity_agent_node) | |
| graph.add_node("dca_agent_node", dca_agent_node) | |
| graph.add_node("strategy_agent_node", strategy_agent_node) | |
| graph.add_node("crypto_agent_node", crypto_agent_node) | |
| graph.add_node("search_agent_node", search_agent_node) | |
| graph.add_node("default_agent_node", default_agent_node) | |
| graph.add_node("database_agent_node", database_agent_node) | |
| graph.add_node("portfolio_advisor_node", portfolio_advisor_node) | |
| # --- Entry point --- | |
| graph.set_entry_point("entry_node") | |
| # --- Linear edges --- | |
| graph.add_edge("entry_node", "semantic_router_node") | |
| # --- Conditional: after semantic router --- | |
| graph.add_conditional_edges( | |
| "semantic_router_node", | |
| decide_route, | |
| { | |
| "error_node": "error_node", | |
| "llm_router_node": "llm_router_node", | |
| "swap_agent_node": "swap_agent_node", | |
| "lending_agent_node": "lending_agent_node", | |
| "staking_agent_node": "staking_agent_node", | |
| "liquidity_agent_node": "liquidity_agent_node", | |
| "dca_agent_node": "dca_agent_node", | |
| "strategy_agent_node": "strategy_agent_node", | |
| "crypto_agent_node": "crypto_agent_node", | |
| "search_agent_node": "search_agent_node", | |
| "default_agent_node": "default_agent_node", | |
| "database_agent_node": "database_agent_node", | |
| "portfolio_advisor_node": "portfolio_advisor_node", | |
| }, | |
| ) | |
| # --- Conditional: after LLM router --- | |
| graph.add_conditional_edges( | |
| "llm_router_node", | |
| after_llm_router, | |
| { | |
| "swap_agent_node": "swap_agent_node", | |
| "lending_agent_node": "lending_agent_node", | |
| "staking_agent_node": "staking_agent_node", | |
| "liquidity_agent_node": "liquidity_agent_node", | |
| "dca_agent_node": "dca_agent_node", | |
| "strategy_agent_node": "strategy_agent_node", | |
| "crypto_agent_node": "crypto_agent_node", | |
| "search_agent_node": "search_agent_node", | |
| "default_agent_node": "default_agent_node", | |
| "database_agent_node": "database_agent_node", | |
| "portfolio_advisor_node": "portfolio_advisor_node", | |
| }, | |
| ) | |
| # --- All agent nodes → formatter → END --- | |
| for node_name in _AGENT_NODES: | |
| graph.add_edge(node_name, "formatter_node") | |
| graph.add_edge("formatter_node", END) | |
| # --- Error → END --- | |
| graph.add_edge("error_node", END) | |
| compiled = graph.compile() | |
| logger.info("StateGraph compiled: %d nodes", len(_AGENT_NODES) + 5) | |
| return compiled | |