graph: check nil choices from final response (#747)

Author: WineChord

graph/checkpoint_test.go

@@ -1112,6 +1112,22 @@ func (e *emptyModel) GenerateContent(ctx context.Context, req *model.Request) (<
 }
 func (e *emptyModel) Info() model.Info { return model.Info{Name: "empty"} }
 
+type emptyChoicesModel struct{}
+
+func (e *emptyChoicesModel) GenerateContent(
+	ctx context.Context,
+	req *model.Request,
+) (<-chan *model.Response, error) {
+	ch := make(chan *model.Response, 1)
+	ch <- &model.Response{}
+	close(ch)
+	return ch, nil
+}
+
+func (e *emptyChoicesModel) Info() model.Info {
+	return model.Info{Name: "empty_choices"}
+}
+
 func TestExecuteModelWithEvents_NoResponseError(t *testing.T) {
 	tracer := oteltrace.NewNoopTracerProvider().Tracer("t")
 	_, span := tracer.Start(context.Background(), "s")
@@ -1130,6 +1146,29 @@ func TestExecuteModelWithEvents_NoResponseError(t *testing.T) {
 	require.Error(t, err)
 }
 
+func TestExecuteModelWithEvents_NoChoicesError(t *testing.T) {
+	tracer := oteltrace.NewNoopTracerProvider().Tracer("t")
+	_, span := tracer.Start(context.Background(), "s")
+	_, err := executeModelWithEvents(context.Background(),
+		modelExecutionConfig{
+			ModelCallbacks: nil,
+			LLMModel:       &emptyChoicesModel{},
+			Request: &model.Request{
+				Messages: []model.Message{
+					model.NewUserMessage("hi"),
+				},
+			},
+			EventChan:    make(chan *event.Event, 1),
+			InvocationID: "inv",
+			SessionID:    "sid",
+			AppName:      "app",
+			UserID:       "user",
+			Span:         span,
+			NodeID:       "n",
+		})
+	require.Error(t, err)
+}
+
 func TestRunModel_BeforeModelCustomResponse(t *testing.T) {
 	// Callbacks return custom response, dummy model should not be called
 	cbs := model.NewCallbacks().RegisterBeforeModel(func(ctx context.Context, req *model.Request) (*model.Response, error) {

graph/state_graph.go

@@ -1483,6 +1483,11 @@ type modelExecutionConfig struct {
 	Span           oteltrace.Span
 }
 
+const (
+	errMsgNoModelResponse = "no response received from model"
+	errMsgNoModelChoices  = "model returned no choices"
+)
+
 // executeModelWithEvents executes the model with event processing.
 func executeModelWithEvents(ctx context.Context, config modelExecutionConfig) (any, error) {
 	ctx, responseChan, err := runModel(ctx, config.ModelCallbacks, config.LLMModel, config.Request)
@@ -1544,8 +1549,18 @@ func executeModelWithEvents(ctx context.Context, config modelExecutionConfig) (a
 		finalResponse = response
 	}
 	if finalResponse == nil {
-		config.Span.SetAttributes(attribute.String("trpc.go.agent.error", "no response received from model"))
-		return nil, errors.New("no response received from model")
+		config.Span.SetAttributes(attribute.String(
+			"trpc.go.agent.error",
+			errMsgNoModelResponse,
+		))
+		return nil, errors.New(errMsgNoModelResponse)
+	}
+	if len(finalResponse.Choices) == 0 {
+		config.Span.SetAttributes(attribute.String(
+			"trpc.go.agent.error",
+			errMsgNoModelChoices,
+		))
+		return nil, errors.New(errMsgNoModelChoices)
 	}
 	if len(finalResponse.Choices[0].Message.ToolCalls) < len(toolCalls) {
 		finalResponse.Choices[0].Message.ToolCalls = toolCalls