[Fix] QueryRRF Response (#223)

arakoodev · Sep 12, 2023 · f786cad · f786cad
1 parent 9fa2052
commit f786cad
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 3 deletions.
diff --git a/...pring/edgechain-app/src/main/java/com/edgechain/lib/index/client/impl/PostgresClient.java b/...pring/edgechain-app/src/main/java/com/edgechain/lib/index/client/impl/PostgresClient.java
@@ -284,6 +284,19 @@ public EdgeChain<List<PostgresWordEmbeddings>> queryRRF(PostgresEndpoint postgre
                   val.setRawText(
                       Objects.nonNull(row.get("raw_text")) ? (String) row.get("raw_text") : null);
 
+                    val.setFilename(
+                            Objects.nonNull(row.get("filename"))
+                                    ? (String) row.get("filename")
+                                    : null);
+                    val.setTimestamp(
+                            Objects.nonNull(row.get("timestamp"))
+                                    ? ((Timestamp) row.get("timestamp")).toLocalDateTime()
+                                    : null);
+                    val.setNamespace(
+                            Objects.nonNull(row.get("namespace"))
+                                    ? (String) row.get("namespace")
+                                    : null);
+
                   BigDecimal bigDecimal =
                       Objects.nonNull(row.get("rrf_score"))
                           ? (BigDecimal) row.get("rrf_score")

diff --git a/...hain-app/src/main/java/com/edgechain/lib/index/repositories/PostgresClientRepository.java b/...hain-app/src/main/java/com/edgechain/lib/index/repositories/PostgresClientRepository.java
@@ -258,7 +258,7 @@ public List<Map<String, Object>> queryRRF(
 
     StringBuilder query = new StringBuilder();
     query
-        .append("SELECT id, raw_text, document_date, metadata,\n")
+        .append("SELECT id, raw_text, document_date, metadata, namespace, filename, timestamp, \n")
         .append(
             String.format(
                 "%s / (ROW_NUMBER() OVER (ORDER BY text_rank DESC) + %s) + \n",
@@ -272,7 +272,7 @@ public List<Map<String, Object>> queryRRF(
                 "%s / (ROW_NUMBER() OVER (ORDER BY date_rank DESC) + %s) AS rrf_score\n",
                 dateWeight.getBaseWeight().getValue(), dateWeight.getFineTuneWeight()))
         .append("FROM ( ")
-        .append("SELECT sv.id, sv.raw_text, svtm.document_date, svtm.metadata, ")
+        .append("SELECT sv.id, sv.raw_text, sv.namespace, sv.filename, sv.timestamp, svtm.document_date, svtm.metadata, ")
         .append(
             String.format(
                 "ts_rank_cd(sv.tsv, plainto_tsquery('%s', '%s')) AS text_rank, ",
@@ -318,7 +318,6 @@ public List<Map<String, Object>> queryRRF(
     }
 
     query.append(" LIMIT ").append(topK).append(";");
-
     return jdbcTemplate.queryForList(query.toString());
   }