52 lines
1.5 KiB
TypeScript
52 lines
1.5 KiB
TypeScript
|
import { StateGraph, START, END } from "@langchain/langgraph";
|
||
|
|
import { TripPlannerAnnotation, TripPlannerState } from "./types";
|
||
|
|
import { extraction } from "./nodes/extraction";
|
||
|
|
import { callTools } from "./nodes/tools";
|
||
|
|
import { classify } from "./nodes/classify";
|
||
|
|
|
||
|
|
function routeStart(state: TripPlannerState): "classify" | "extraction" {
|
||
|
|
if (!state.tripDetails) {
|
||
|
|
return "extraction";
|
||
|
|
}
|
||
|
|
|
||
|
|
return "classify";
|
||
|
|
}
|
||
|
|
|
||
|
|
function routeAfterClassifying(
|
||
|
|
state: TripPlannerState,
|
||
|
|
): "callTools" | "extraction" {
|
||
|
|
// if `tripDetails` is undefined, this means they are not relevant to the conversation
|
||
|
|
if (!state.tripDetails) {
|
||
|
|
return "extraction";
|
||
|
|
}
|
||
|
|
|
||
|
|
// otherwise, they are relevant, and we should route to callTools
|
||
|
|
return "callTools";
|
||
|
|
}
|
||
|
|
|
||
|
|
function routeAfterExtraction(
|
||
|
|
state: TripPlannerState,
|
||
|
|
): "callTools" | typeof END {
|
||
|
|
// if `tripDetails` is undefined, this means they're missing some fields.
|
||
|
|
if (!state.tripDetails) {
|
||
|
|
return END;
|
||
|
|
}
|
||
|
|
|
||
|
|
return "callTools";
|
||
|
|
}
|
||
|
|
|
||
|
|
const builder = new StateGraph(TripPlannerAnnotation)
|
||
|
|
.addNode("classify", classify)
|
||
|
|
.addNode("extraction", extraction)
|
||
|
|
.addNode("callTools", callTools)
|
||
|
|
.addConditionalEdges(START, routeStart, ["classify", "extraction"])
|
||
|
|
.addConditionalEdges("classify", routeAfterClassifying, [
|
||
|
|
"callTools",
|
||
|
|
"extraction",
|
||
|
|
])
|
||
|
|
.addConditionalEdges("extraction", routeAfterExtraction, ["callTools", END])
|
||
|
|
.addEdge("callTools", END);
|
||
|
|
|
||
|
|
export const tripPlannerGraph = builder.compile();
|
||
|
|
tripPlannerGraph.name = "Trip Planner";
|