diff --git a/apps/open-swe/src/graphs/programmer/index.ts b/apps/open-swe/src/graphs/programmer/index.ts index 4ef92469..4ec8b1bd 100644 --- a/apps/open-swe/src/graphs/programmer/index.ts +++ b/apps/open-swe/src/graphs/programmer/index.ts @@ -1,6 +1,7 @@ -import { END, Send, START, StateGraph } from "@langchain/langgraph"; +import { Command, END, Send, START, StateGraph } from "@langchain/langgraph"; import { GraphAnnotation, + GraphConfig, GraphConfiguration, GraphState, } from "@open-swe/shared/open-swe/types"; @@ -39,12 +40,12 @@ function lastMessagesMissingToolCalls( * Otherwise, it ends the process. * * @param {GraphState} state - The current graph state. - * @returns {"reviewer-subgraph" | "take-action" | "request-help" | "generate-action" | Send} The next node to execute, or END if the process should stop. + * @returns {"route-to-review-or-conclusion" | "take-action" | "request-help" | "generate-action" | Send} The next node to execute, or END if the process should stop. */ function routeGeneratedAction( state: GraphState, ): - | "reviewer-subgraph" + | "route-to-review-or-conclusion" | "take-action" | "request-help" | "generate-action" @@ -82,7 +83,7 @@ function routeGeneratedAction( } // No tool calls, route to reviewer subgraph - return "reviewer-subgraph"; + return "route-to-review-or-conclusion"; } /** @@ -101,6 +102,22 @@ function routeGenerateActionsOrEnd( return "generate-action"; } +function routeToReviewOrConclusion( + state: GraphState, + config: GraphConfig, +): Command { + const maxAllowedReviews = config.configurable?.maxReviewCount ?? 3; + if (state.reviewsCount >= maxAllowedReviews) { + return new Command({ + goto: "generate-conclusion", + }); + } + + return new Command({ + goto: "reviewer-subgraph", + }); +} + const workflow = new StateGraph(GraphAnnotation, GraphConfiguration) .addNode("initialize", initializeSandbox) .addNode("generate-action", generateAction) @@ -109,12 +126,19 @@ const workflow = new StateGraph(GraphAnnotation, GraphConfiguration) }) .addNode("update-plan", updatePlan) .addNode("progress-plan-step", progressPlanStep, { - ends: ["summarize-history", "generate-action", "reviewer-subgraph"], + ends: [ + "summarize-history", + "generate-action", + "route-to-review-or-conclusion", + ], }) .addNode("generate-conclusion", generateConclusion) .addNode("request-help", requestHelp, { ends: ["generate-action", END], }) + .addNode("route-to-review-or-conclusion", routeToReviewOrConclusion, { + ends: ["generate-conclusion", "route-to-review-or-conclusion"], + }) .addNode("reviewer-subgraph", reviewerGraph) .addNode("open-pr", openPullRequest) .addNode("diagnose-error", diagnoseError) @@ -124,7 +148,7 @@ const workflow = new StateGraph(GraphAnnotation, GraphConfiguration) .addConditionalEdges("generate-action", routeGeneratedAction, [ "take-action", "request-help", - "reviewer-subgraph", + "route-to-review-or-conclusion", "update-plan", "generate-action", ]) diff --git a/apps/open-swe/src/graphs/programmer/nodes/progress-plan-step.ts b/apps/open-swe/src/graphs/programmer/nodes/progress-plan-step.ts index 7354b7a2..92b505e2 100644 --- a/apps/open-swe/src/graphs/programmer/nodes/progress-plan-step.ts +++ b/apps/open-swe/src/graphs/programmer/nodes/progress-plan-step.ts @@ -187,7 +187,7 @@ Once you've determined the status of the current task, call either the \`mark_ta taskPlan: updatedPlanTasks, }; return new Command({ - goto: "reviewer-subgraph", + goto: "route-to-review-or-conclusion", update: commandUpdate, }); } diff --git a/apps/open-swe/src/graphs/reviewer/nodes/final-review.ts b/apps/open-swe/src/graphs/reviewer/nodes/final-review.ts index de812433..6d8b6bc9 100644 --- a/apps/open-swe/src/graphs/reviewer/nodes/final-review.ts +++ b/apps/open-swe/src/graphs/reviewer/nodes/final-review.ts @@ -157,5 +157,6 @@ export async function finalReview( taskPlan: updatedTaskPlan, messages: messagesUpdate, internalMessages: messagesUpdate, + reviewsCount: (state.reviewsCount || 0) + 1, }; } diff --git a/packages/shared/src/open-swe/reviewer/types.ts b/packages/shared/src/open-swe/reviewer/types.ts index 2a19051b..84e6545e 100644 --- a/packages/shared/src/open-swe/reviewer/types.ts +++ b/packages/shared/src/open-swe/reviewer/types.ts @@ -100,6 +100,16 @@ export const ReviewerGraphStateObj = MessagesZodState.extend({ fn: (_state, update) => update, }, }), + /** + * The number of times the reviewer subgraph has been executed. + */ + reviewsCount: withLangGraph(z.custom(), { + reducer: { + schema: z.custom(), + fn: (_state, update) => update, + }, + default: () => 0, + }), }); export type ReviewerGraphState = z.infer; diff --git a/packages/shared/src/open-swe/types.ts b/packages/shared/src/open-swe/types.ts index ce72eb70..377b0fd9 100644 --- a/packages/shared/src/open-swe/types.ts +++ b/packages/shared/src/open-swe/types.ts @@ -241,14 +241,14 @@ export const GraphAnnotation = MessagesZodState.extend({ }, }), /** - * The review generated by the reviewer subgraph + * The number of times the reviewer subgraph has been executed. */ - review: withLangGraph(z.custom(), { + reviewsCount: withLangGraph(z.custom(), { reducer: { - schema: z.custom(), + schema: z.custom(), fn: (_state, update) => update, }, - default: () => "", + default: () => 0, }), // ---NOT USED--- @@ -292,6 +292,11 @@ export const GraphConfigurationMetadata: { "Maximum number of review actions allowed during the review phase. An action consists of a tool call.", }, }, + maxReviewCount: { + x_open_swe_ui_config: { + type: "hidden", + }, + }, programmerModelName: { x_open_swe_ui_config: { type: "select", @@ -514,6 +519,12 @@ export const GraphConfiguration = z.object({ mcpServers: withLangGraph(z.string().optional(), { metadata: GraphConfigurationMetadata.mcpServers, }), + /** + * The maxium number of times the reviewer subgraph can be executed. + */ + maxReviewCount: withLangGraph(z.number().optional(), { + metadata: GraphConfigurationMetadata.maxReviewCount, + }), }); export type GraphConfig = LangGraphRunnableConfig<