mirror of
https://github.com/CopilotKit/CopilotKit.git
synced 2026-09-14 16:26:20 +08:00
Alem Tuzlak
69 lines
2.2 KiB
Python
69 lines
2.2 KiB
Python
"""
|
|
This is the main entry point for the AI.
|
|
It defines the workflow graph and the entry point for the agent.
|
|
"""
|
|
|
|
# pylint: disable=line-too-long, unused-import
|
|
from typing import cast
|
|
|
|
from langchain_core.messages import AIMessage, ToolMessage
|
|
from langgraph.checkpoint.memory import MemorySaver
|
|
from langgraph.graph import END, START, StateGraph
|
|
|
|
from src.chat import chat_node
|
|
from src.search import search_node
|
|
from src.state import AgentState
|
|
from src.trips import perform_trips_node, trips_node
|
|
|
|
|
|
# Route is responsible for determing the next node based on the last message. This
|
|
# is needed because LangGraph does not automatically route to nodes, instead that
|
|
# is handled through code.
|
|
def route(state: AgentState):
|
|
"""Route after the chat node."""
|
|
messages = state.get("messages", [])
|
|
if messages and isinstance(messages[-1], AIMessage):
|
|
ai_message = cast(AIMessage, messages[-1])
|
|
|
|
# If the last AI message has tool calls we need to determine to route to the
|
|
# trips_node or search_node based on the tool name.
|
|
if ai_message.tool_calls:
|
|
tool_name = ai_message.tool_calls[0]["name"]
|
|
if tool_name in [
|
|
"add_trips",
|
|
"update_trips",
|
|
"delete_trips",
|
|
"select_trip",
|
|
]:
|
|
return "trips_node"
|
|
if tool_name in ["search_for_places"]:
|
|
return "search_node"
|
|
return "chat_node"
|
|
|
|
if messages and isinstance(messages[-1], ToolMessage):
|
|
return "chat_node"
|
|
|
|
return END
|
|
|
|
|
|
graph_builder = StateGraph(AgentState)
|
|
|
|
graph_builder.add_node("chat_node", chat_node)
|
|
graph_builder.add_node("trips_node", trips_node)
|
|
graph_builder.add_node("search_node", search_node)
|
|
graph_builder.add_node("perform_trips_node", perform_trips_node)
|
|
|
|
graph_builder.add_conditional_edges(
|
|
"chat_node", route, ["search_node", "chat_node", "trips_node", END]
|
|
)
|
|
|
|
graph_builder.add_edge(START, "chat_node")
|
|
graph_builder.add_edge("search_node", "chat_node")
|
|
graph_builder.add_edge("perform_trips_node", "chat_node")
|
|
graph_builder.add_edge("trips_node", "perform_trips_node")
|
|
|
|
graph = graph_builder.compile(
|
|
interrupt_after=["trips_node"],
|
|
checkpointer=MemorySaver(),
|
|
)
|