Spaces:
Running
Running
File size: 4,616 Bytes
b7f63db 156dd84 b7f63db 0f123dc b7f63db bc3dd51 b7f63db 156dd84 b7f63db 0f123dc b7f63db bc3dd51 b7f63db 156dd84 b7f63db 0f123dc b7f63db bc3dd51 b7f63db 156dd84 b7f63db 0f123dc b7f63db bc3dd51 b7f63db 156dd84 b7f63db 0f123dc b7f63db bc3dd51 b7f63db | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | """
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
|