pkg/aflow: handle ISE errors from LLM as temporal

2 files changed, 12 insertions, 0 deletions

diff --git a/pkg/aflow/llm_agent.go b/pkg/aflow/llm_agent.go
index 0997b9f09..d06e73642 100644
--- a/pkg/aflow/llm_agent.go
+++ b/pkg/aflow/llm_agent.go
@@ -440,6 +440,10 @@ func parseLLMError(resp *genai.GenerateContentResponse, err error, model string,
 		strings.Contains(apiErr.Message, "The input token count exceeds the maximum") {
 		return 0, &tokenOverflowError{err}
 	}
+	if apiErr.Code == http.StatusInternalServerError && try < maxLLMRetryIters {
+		// Let's assume ISE is just something temporal on the server side.
+		return time.Second, nil
+	}
 	return 0, err
 }
 
diff --git a/pkg/aflow/llm_agent_test.go b/pkg/aflow/llm_agent_test.go
index adf2a1278..ac3187884 100644
--- a/pkg/aflow/llm_agent_test.go
+++ b/pkg/aflow/llm_agent_test.go
@@ -59,6 +59,14 @@ func TestParseLLMError(t *testing.T) {
 				Message: `The input token count exceeds the maximum number of tokens allowed 1048576.`,
 			}},
 		},
+		{
+			resp: nil,
+			inputErr: genai.APIError{
+				Code:    500,
+				Message: `Internal error encountered.`,
+			},
+			retry: time.Second,
+		},
 	}
 	for i, test := range tests {
 		t.Run(fmt.Sprint(i), func(t *testing.T) {