use new graph state

lang_agent/graphs/routing.py

@@ -14,6 +14,8 @@ import glob
 from lang_agent.config import KeyConfig
 from lang_agent.tool_manager import ToolManager, ToolManagerConfig
 from lang_agent.base import GraphBase
+from lang_agent.graphs.graph_state import State
+from lang_agent.graphs.tool_nodes import AnnotatedToolNode, ToolNode
 
 from langchain.chat_models import init_chat_model
 from langchain_core.messages import SystemMessage, HumanMessage, BaseMessage
@@ -60,16 +62,6 @@ class Route(BaseModel):
     )
 
 
-class State(TypedDict):
-    inp: Tuple[Dict[str, List[SystemMessage | HumanMessage]], 
-               Dict[str, Dict[str, str|int]]]
-    messages: List[SystemMessage | HumanMessage]
-    decision: str
-    subgraph_states: Dict[str, Any]   # NOTE: Naively assuming subgraphs
-                                      #       won't be so complicated
-
-
-
 class RoutingGraph(GraphBase):
     def __init__(self, config: RoutingConfig):
         self.config = config