From 1f71793afe203bce07e39327b7da0c9b6283ec93 Mon Sep 17 00:00:00 2001
From: Songkan Tang <songkant@amazon.com>
Date: Mon, 5 Aug 2024 11:40:32 +0800
Subject: [PATCH] Change the embedding vector target

Signed-off-by: Songkan Tang <songkant@amazon.com>
---
 .../org/opensearch/agent/tools/RCATool.java    | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/src/main/java/org/opensearch/agent/tools/RCATool.java b/src/main/java/org/opensearch/agent/tools/RCATool.java
index 6bc286fd..955ccd6a 100644
--- a/src/main/java/org/opensearch/agent/tools/RCATool.java
+++ b/src/main/java/org/opensearch/agent/tools/RCATool.java
@@ -117,22 +117,26 @@ public <T> void runOption1(Map<String, ?> knowledgeBase, ActionListener<T> liste
 
     @SuppressWarnings("unchecked")
     public <T> void runOption2(Map<String, ?> knowledgeBase, ActionListener<T> listener) {
-        // input phenomenon's embedded dense vector
         String phenomenon = (String) knowledgeBase.get("phenomenon");
-        List<RealVector> inputVector = getEmbeddedVector(Collections.singletonList(phenomenon));
 
-        // api response embedded dense vectors
+        // API response embedded vectors
         List<Map<String, String>> causes = (List<Map<String, String>>) knowledgeBase.get("causes");
-        List<String> apiResponses = causes.stream()
+        List<String> responses = causes.stream()
             .map(cause -> cause.get("response"))
             .collect(Collectors.toList());
-        List<RealVector> rootCauseVectors = getEmbeddedVector(apiResponses);
+        List<RealVector> responseVectors = getEmbeddedVector(responses);
 
-        Map<String, Double> dotProductMap = IntStream.range(0, rootCauseVectors.size())
+        // expected API response embedded vectors
+        List<String> expectedResponses = causes.stream()
+            .map(cause -> cause.get("expected_response"))
+            .collect(Collectors.toList());
+        List<RealVector> expectedResponseVectors = getEmbeddedVector(expectedResponses);
+
+        Map<String, Double> dotProductMap = IntStream.range(0, causes.size())
             .boxed()
             .collect(Collectors.toMap(
                 i -> causes.get(i).get("reason"),
-                i -> inputVector.get(0).dotProduct(rootCauseVectors.get(i))
+                i -> responseVectors.get(i).dotProduct(expectedResponseVectors.get(i))
             ));
 
         Optional<Map.Entry<String, Double>> mapEntry =