From d79e30cd7af304e63319f439ffacb4196a15e8a3 Mon Sep 17 00:00:00 2001
From: LXW <1264174498@qq.com>
Date: Fri, 24 Nov 2023 17:52:30 +0800
Subject: [PATCH] (improvement)(knowledge) ModelSchema add tag type element
 (#419)

Co-authored-by: jolunoluo
---
 .../supersonic/chat/api/pojo/ModelSchema.java |  2 ++
 .../chat/api/pojo/SchemaElementType.java      |  3 ++-
 .../semantic/ModelSchemaBuilder.java          | 21 ++++++++++++++++---
 3 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/chat/api/src/main/java/com/tencent/supersonic/chat/api/pojo/ModelSchema.java b/chat/api/src/main/java/com/tencent/supersonic/chat/api/pojo/ModelSchema.java
index ec5e2e268..de0f8b1dd 100644
--- a/chat/api/src/main/java/com/tencent/supersonic/chat/api/pojo/ModelSchema.java
+++ b/chat/api/src/main/java/com/tencent/supersonic/chat/api/pojo/ModelSchema.java
@@ -13,6 +13,8 @@ public class ModelSchema {
     private Set<SchemaElement> metrics = new HashSet<>();
     private Set<SchemaElement> dimensions = new HashSet<>();
     private Set<SchemaElement> dimensionValues = new HashSet<>();
+    private Set<SchemaElement> tags = new HashSet<>();
+    @Deprecated
     private SchemaElement entity = new SchemaElement();
 
     public SchemaElement getElement(SchemaElementType elementType, long elementID) {
diff --git a/chat/api/src/main/java/com/tencent/supersonic/chat/api/pojo/SchemaElementType.java b/chat/api/src/main/java/com/tencent/supersonic/chat/api/pojo/SchemaElementType.java
index 9eaf598b3..891769c07 100644
--- a/chat/api/src/main/java/com/tencent/supersonic/chat/api/pojo/SchemaElementType.java
+++ b/chat/api/src/main/java/com/tencent/supersonic/chat/api/pojo/SchemaElementType.java
@@ -7,5 +7,6 @@ public enum SchemaElementType {
     VALUE,
     ENTITY,
     ID,
-    DATE
+    DATE,
+    TAG
 }
diff --git a/chat/knowledge/src/main/java/com/tencent/supersonic/knowledge/semantic/ModelSchemaBuilder.java b/chat/knowledge/src/main/java/com/tencent/supersonic/knowledge/semantic/ModelSchemaBuilder.java
index e8b608c80..62a33490d 100644
--- a/chat/knowledge/src/main/java/com/tencent/supersonic/knowledge/semantic/ModelSchemaBuilder.java
+++ b/chat/knowledge/src/main/java/com/tencent/supersonic/knowledge/semantic/ModelSchemaBuilder.java
@@ -13,6 +13,9 @@ import com.tencent.supersonic.semantic.api.model.pojo.SchemaItem;
 import com.tencent.supersonic.semantic.api.model.response.DimSchemaResp;
 import com.tencent.supersonic.semantic.api.model.response.MetricSchemaResp;
 import com.tencent.supersonic.semantic.api.model.response.ModelSchemaResp;
+import org.apache.logging.log4j.util.Strings;
+import org.springframework.beans.BeanUtils;
+import org.springframework.util.CollectionUtils;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashSet;
@@ -21,9 +24,6 @@ import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
 import java.util.stream.Collectors;
-import org.apache.logging.log4j.util.Strings;
-import org.springframework.beans.BeanUtils;
-import org.springframework.util.CollectionUtils;
 
 public class ModelSchemaBuilder {
 
@@ -62,6 +62,7 @@ public class ModelSchemaBuilder {
 
         Set<SchemaElement> dimensions = new HashSet<>();
         Set<SchemaElement> dimensionValues = new HashSet<>();
+        Set<SchemaElement> tags = new HashSet<>();
         for (DimSchemaResp dim : resp.getDimensions()) {
 
             List<String> alias = SchemaItem.getAliasList(dim.getAlias());
@@ -105,9 +106,23 @@ public class ModelSchemaBuilder {
                     .alias(new ArrayList<>(Arrays.asList(dimValueAlias.toArray(new String[0]))))
                     .build();
             dimensionValues.add(dimValueToAdd);
+            if (dim.getIsTag() == 1) {
+                SchemaElement tagToAdd = SchemaElement.builder()
+                        .model(resp.getId())
+                        .id(dim.getId())
+                        .name(dim.getName())
+                        .bizName(dim.getBizName())
+                        .type(SchemaElementType.TAG)
+                        .useCnt(dim.getUseCnt())
+                        .alias(alias)
+                        .schemaValueMaps(schemaValueMaps)
+                        .build();
+                tags.add(tagToAdd);
+            }
         }
         modelSchema.getDimensions().addAll(dimensions);
         modelSchema.getDimensionValues().addAll(dimensionValues);
+        modelSchema.getTags().addAll(tags);
 
         Entity entity = resp.getEntity();
         if (Objects.nonNull(entity)) {