Optimize the response of AI agent APIs (#373) (#374)

* Optimize the response of executing AI agents Signed-off-by: gaobinlong <[email protected]> * Add change log Signed-off-by: gaobinlong <[email protected]> * Optimize the code Signed-off-by: gaobinlong <[email protected]> * Optimize the response Signed-off-by: gaobinlong <[email protected]> * Optimize the code Signed-off-by: gaobinlong <[email protected]> --------- Signed-off-by: gaobinlong <[email protected]> (cherry picked from commit 581a7ca) Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> # Conflicts: # CHANGELOG.md Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
opensearch-project · Dec 3, 2024 · b40dc4e · b40dc4e
1 parent 32a8cee
commit b40dc4e
Show file tree

Hide file tree

Showing 3 changed files with 132 additions and 25 deletions.
diff --git a/server/routes/agent_routes.ts b/server/routes/agent_routes.ts
@@ -39,7 +39,20 @@ export function registerAgentRoutes(router: IRouter, assistantService: Assistant
         );
         return res.ok({ body: response });
       } catch (e) {
-        return res.badRequest();
+        context.assistant_plugin.logger.error('Execute agent failed!', e);
+        if (e.statusCode >= 400 && e.statusCode <= 499) {
+          return res.customError({
+            body: e.body,
+            statusCode: e.statusCode,
+            headers: e.headers,
+          });
+        } else {
+          return res.customError({
+            body: 'Execute agent failed!',
+            statusCode: 500,
+            headers: e.headers,
+          });
+        }
       }
     })
   );

diff --git a/server/routes/summary_routes.ts b/server/routes/summary_routes.ts
@@ -47,14 +47,33 @@ export function registerSummaryAssistantRoutes(
         req.body.index && req.body.dsl && req.body.topNLogPatternData
           ? LOG_PATTERN_SUMMARY_AGENT_CONFIG_ID
           : SUMMARY_AGENT_CONFIG_ID;
-      const response = await assistantClient.executeAgentByConfigName(agentConfigId, {
-        context: req.body.context,
-        question: req.body.question,
-        index: req.body.index,
-        input: req.body.dsl,
-        topNLogPatternData: req.body.topNLogPatternData,
-      });
-      let summary;
+
+      let response;
+      try {
+        response = await assistantClient.executeAgentByConfigName(agentConfigId, {
+          context: req.body.context,
+          question: req.body.question,
+          index: req.body.index,
+          input: req.body.dsl,
+          topNLogPatternData: req.body.topNLogPatternData,
+        });
+      } catch (e) {
+        context.assistant_plugin.logger.error('Execute agent failed!', e);
+        if (e.statusCode >= 400 && e.statusCode <= 499) {
+          return res.customError({
+            body: e.body,
+            statusCode: e.statusCode,
+            headers: e.headers,
+          });
+        } else {
+          return res.customError({
+            body: 'Execute agent failed!',
+            statusCode: 500,
+            headers: e.headers,
+          });
+        }
+      }
+
       let insightAgentIdExists = false;
       try {
         if (req.body.insightType) {
@@ -65,19 +84,24 @@ export function registerSummaryAssistantRoutes(
           ));
         }
       } catch (e) {
-        context.assistant_plugin.logger.debug(
+        context.assistant_plugin.logger.error(
           `Cannot find insight agent for ${req.body.insightType}`,
           e
         );
       }
-      try {
-        summary = response.body.inference_results[0].output[0].result;
+
+      const summary = response.body.inference_results[0]?.output[0]?.result;
+      if (summary) {
         return res.ok({ body: { summary, insightAgentIdExists } });
-      } catch (e) {
-        return res.badRequest({ body: e });
+      } else {
+        return res.customError({
+          body: 'Execute agent failed with empty response!',
+          statusCode: 500,
+        });
       }
     })
   );
+
   router.post(
     {
       path: SUMMARY_ASSISTANT_API.INSIGHT,
@@ -105,15 +129,38 @@ export function registerSummaryAssistantRoutes(
         req.body.summaryType,
         client
       );
-      const response = await assistantClient.executeAgent(insightAgentId, {
-        context: req.body.context,
-        summary: req.body.summary,
-        question: req.body.question,
-      });
+
       try {
-        return res.ok({ body: response.body.inference_results[0].output[0].result });
+        const response = await assistantClient.executeAgent(insightAgentId, {
+          context: req.body.context,
+          summary: req.body.summary,
+          question: req.body.question,
+        });
+
+        const insight = response.body.inference_results[0]?.output[0]?.result;
+        if (insight) {
+          return res.ok({ body: { insight } });
+        } else {
+          return res.customError({
+            body: 'Execute agent failed with empty response!',
+            statusCode: 500,
+          });
+        }
       } catch (e) {
-        return res.badRequest({ body: e });
+        context.assistant_plugin.logger.error('Execute agent failed!', e);
+        if (e.statusCode >= 400 && e.statusCode <= 499) {
+          return res.customError({
+            body: e.body,
+            statusCode: e.statusCode,
+            headers: e.headers,
+          });
+        } else {
+          return res.customError({
+            body: 'Execute agent failed!',
+            statusCode: 500,
+            headers: e.headers,
+          });
+        }
       }
     })
   );
@@ -167,10 +214,31 @@ export function registerData2SummaryRoutes(
             question: req.body.question,
           }
         );
+
         const result = response.body.inference_results[0].output[0].result;
-        return res.ok({ body: result });
+        if (result) {
+          return res.ok({ body: result });
+        } else {
+          return res.customError({
+            body: 'Execute agent failed with empty response!',
+            statusCode: 500,
+          });
+        }
       } catch (e) {
-        return res.badRequest({ body: e });
+        context.assistant_plugin.logger.error('Execute agent failed!', e);
+        if (e.statusCode >= 400 && e.statusCode <= 499) {
+          return res.customError({
+            body: e.body,
+            statusCode: e.statusCode,
+            headers: e.headers,
+          });
+        } else {
+          return res.customError({
+            body: 'Execute agent failed!',
+            statusCode: 500,
+            headers: e.headers,
+          });
+        }
       }
     })
   );

diff --git a/server/routes/text2viz_routes.ts b/server/routes/text2viz_routes.ts
@@ -86,7 +86,20 @@ export function registerText2VizRoutes(router: IRouter, assistantService: Assist
         }
         return res.badRequest();
       } catch (e) {
-        return res.badRequest();
+        context.assistant_plugin.logger.error('Execute agent failed!', e);
+        if (e.statusCode >= 400 && e.statusCode <= 499) {
+          return res.customError({
+            body: e.body,
+            statusCode: e.statusCode,
+            headers: e.headers,
+          });
+        } else {
+          return res.customError({
+            body: 'Execute agent failed!',
+            statusCode: 500,
+            headers: e.headers,
+          });
+        }
       }
     })
   );
@@ -115,7 +128,20 @@ export function registerText2VizRoutes(router: IRouter, assistantService: Assist
         const result = JSON.parse(response.body.inference_results[0].output[0].result);
         return res.ok({ body: result });
       } catch (e) {
-        return res.badRequest();
+        context.assistant_plugin.logger.error('Execute agent failed!', e);
+        if (e.statusCode >= 400 && e.statusCode <= 499) {
+          return res.customError({
+            body: e.body,
+            statusCode: e.statusCode,
+            headers: e.headers,
+          });
+        } else {
+          return res.customError({
+            body: 'Execute agent failed!',
+            statusCode: 500,
+            headers: e.headers,
+          });
+        }
       }
     })
   );