Spaces:
Sleeping
Sleeping
| from langgraph.graph import StateGraph, START, END | |
| from .func import State | |
| from langgraph.graph.state import CompiledStateGraph | |
| from langgraph.store.memory import InMemoryStore | |
| class PrimaryChatBot: | |
| def __init__(self): | |
| pass | |
| def should_continue(state: State): | |
| messages = state["messages"] | |
| last_message = messages[-1] | |
| if not last_message.tool_calls: | |
| return "end" | |
| else: | |
| return "continue" | |
| def node(self, graph: StateGraph): | |
| return graph | |
| def edge(self, graph: StateGraph): | |
| return graph | |
| def __call__(self, checkpointer=InMemoryStore()) -> CompiledStateGraph: | |
| graph = StateGraph(State) | |
| graph: StateGraph = self.node(graph) | |
| graph: StateGraph = self.edge(graph) | |
| return graph.compile(checkpointer=checkpointer) | |