Extract result mapper.

Author: jasperblues

src/main/kotlin/com/embabel/agent/rag/neo/drivine/ContentElementMapper.kt

@@ -0,0 +1,42 @@
+package com.embabel.agent.rag.neo.drivine
+
+import com.embabel.agent.rag.model.Chunk
+import com.embabel.agent.rag.model.ContentElement
+import com.embabel.agent.rag.model.MaterializedDocument
+import org.springframework.stereotype.Component
+
+@Component
+class ContentElementMapper {
+
+    fun rowToContentElement(row: Map<*, *>): ContentElement {
+        val metadata = mutableMapOf<String, Any>()
+        metadata["source"] = row["metadata_source"] ?: "unknown"
+        val labels = row["labels"] as? List<String> ?: error("Must have labels")
+        if (labels.contains("Chunk"))
+            return Chunk(
+                id = row["id"] as String,
+                text = row["text"] as String,
+                parentId = row["parentId"] as String,
+                metadata = metadata,
+            )
+        if (labels.contains("Document")) {
+            val ingestionDate = when (val rawDate = row["ingestionDate"]) {
+                is java.time.Instant -> rawDate
+                is java.time.ZonedDateTime -> rawDate.toInstant()
+                is Long -> java.time.Instant.ofEpochMilli(rawDate)
+                is String -> java.time.Instant.parse(rawDate)
+                null -> java.time.Instant.now()
+                else -> java.time.Instant.now()
+            }
+            return MaterializedDocument(
+                id = row["id"] as String,
+                title = row["id"] as String,
+                children = emptyList(),
+                metadata = metadata,
+                uri = row["uri"] as String,
+                ingestionTimestamp = ingestionDate,
+            )
+        }
+        throw RuntimeException("Don't know how to map: $labels")
+    }
+}

src/main/kotlin/com/embabel/agent/rag/neo/drivine/DrivineCypherSearch.kt

@@ -1,6 +1,10 @@
 package com.embabel.agent.rag.neo.drivine
 
 import com.embabel.agent.rag.model.*
+import com.embabel.agent.rag.service.Cluster
+import com.embabel.agent.rag.service.ClusterFinder
+import com.embabel.agent.rag.service.ClusterRetrievalRequest
+import com.embabel.agent.rag.service.TypedEntitySearch
 import com.embabel.common.core.types.SimilarityResult
 import com.embabel.common.core.types.SimpleSimilaritySearchResult
 import org.drivine.manager.PersistenceManager
@@ -9,6 +13,8 @@ import org.drivine.utils.ObjectUtils
 import org.slf4j.Logger
 import org.slf4j.LoggerFactory
 import org.springframework.stereotype.Service
+import org.springframework.transaction.annotation.Transactional
+import kotlin.text.get
 
 @Service
 class DrivineCypherSearch(
@@ -220,4 +226,45 @@ class DrivineCypherSearch(
             else -> 0
         }
     }
+
+//    @Transactional(readOnly = true)
+//    override fun <E> findClusters(opts: ClusterRetrievalRequest<E>): List<Cluster<E>> {
+//        val labels = opts.entitySearch?.labels?.toList() ?: error("Must specify labels in entity search for clustering")
+//        val desiredType = (opts.entitySearch as? TypedEntitySearch)?.entities?.first() ?: OgmMappedEntity::class.java
+//        val params = mapOf(
+//            "labels" to labels,
+//            "vectorIndex" to opts.vectorIndex,
+//            "similarityThreshold" to opts.similarityThreshold,
+//            "topK" to opts.topK,
+//        )
+//        val result = query(
+//            purpose = "cluster",
+//            query = "vector_cluster",
+//            params = params,
+//        )
+//        return result.map { row ->
+//            val anchor = row["anchor"] as E
+//            val similar = (row["similar"]) as Array<Map<String, Any>>
+//            val similarityResults = similar.mapNotNull { similarEntityMap ->
+//                val inode = similarEntityMap["match"] as InternalNode
+//                val matchId = (inode.get("id") as String)
+//                val score = similarEntityMap["score"] as Double
+//                val match = try {
+//                    currentSession().load(desiredType, matchId)
+//                } catch (e: Exception) {
+//                    logger.warn("Could not load entity of type $desiredType with id $matchId", e)
+//                    null
+//                }
+//                if (match == null) {
+//                    // Shouldn't happen...query is likely incorrect
+//                    logger.warn("Could not load match for $similarEntityMap, type=${desiredType}, id=$matchId")
+//                    null
+//                } else {
+//                    logger.debug("Found match: {} with score {}", match, "%.2f".format(score))
+//                    SimpleSimilaritySearchResult(match, score) as SimilarityResult<E>
+//                }
+//            }
+//            Cluster(anchor, similarityResults)
+//        }
+//    }
 }

src/main/kotlin/com/embabel/agent/rag/neo/drivine/DrivineStore.kt

@@ -31,6 +31,7 @@ class DrivineStore(
     override val enhancers: List<RetrievableEnhancer> = emptyList(),
     val properties: NeoRagServiceProperties,
     private val cypherSearch: CypherSearch,
+    private val contentElementMapper: ContentElementMapper,
     modelProvider: ModelProvider,
     platformTransactionManager: PlatformTransactionManager,
 ) : AbstractChunkingContentElementRepository(properties), ChunkingContentElementRepository, RagFacetProvider {
@@ -99,7 +100,7 @@ class DrivineStore(
                 .withStatement(statement)
                 .bind(parameters)
                 .transform(Map::class.java)
-                .map { rowToContentElement(it) }
+                .map { contentElementMapper.rowToContentElement(it) }
 
             val result = persistenceManager.maybeGetOne(spec)
             logger.debug("Root document with URI {} found: {}", uri, result != null)
@@ -158,7 +159,7 @@ class DrivineStore(
             .withStatement(statement)
             .bind(mapOf("ids" to chunkIds))
             .transform(Map::class.java)
-            .map({ rowToContentElement(it) })
+            .map({ contentElementMapper.rowToContentElement(it) })
             .filter { it is Chunk }
             .map { it as Chunk }
 
@@ -171,7 +172,7 @@ class DrivineStore(
             .withStatement(statement)
             .bind(mapOf("id" to id))
             .transform(Map::class.java)
-            .map({ rowToContentElement(it) })
+            .map({ contentElementMapper.rowToContentElement(it) })
         return persistenceManager.maybeGetOne(spec)
     }
 
@@ -260,38 +261,6 @@ class DrivineStore(
               } AS result
             """.trimIndent()
 
-    private fun rowToContentElement(row: Map<*, *>): ContentElement {
-        val metadata = mutableMapOf<String, Any>()
-        metadata["source"] = row["metadata_source"] ?: "unknown"
-        val labels = row["labels"] as? List<String> ?: error("Must have labels")
-        if (labels.contains("Chunk"))
-            return Chunk(
-                id = row["id"] as String,
-                text = row["text"] as String,
-                parentId = row["parentId"] as String,
-                metadata = metadata,
-            )
-        if (labels.contains("Document")) {
-            val ingestionDate = when (val rawDate = row["ingestionDate"]) {
-                is java.time.Instant -> rawDate
-                is java.time.ZonedDateTime -> rawDate.toInstant()
-                is Long -> java.time.Instant.ofEpochMilli(rawDate)
-                is String -> java.time.Instant.parse(rawDate)
-                null -> java.time.Instant.now()
-                else -> java.time.Instant.now()
-            }
-            return MaterializedDocument(
-                id = row["id"] as String,
-                title = row["id"] as String,
-                children = emptyList(),
-                metadata = metadata,
-                uri = row["uri"] as String,
-                ingestionTimestamp = ingestionDate,
-            )
-        }
-        throw RuntimeException("Don't know how to map: $labels")
-    }
-
     private val readonlyTransactionTemplate = TransactionTemplate(platformTransactionManager).apply {
         isReadOnly = true
         propagationBehavior = TransactionDefinition.PROPAGATION_REQUIRED

src/main/resources/cypher/vector_cluster.cypher

@@ -19,7 +19,10 @@ CALL apoc.cypher.parallel(
 "item"
 ) YIELD value
   WHERE size(value.similar) > 0
-RETURN value.anchorNode as anchor,
-       value.similar as similar,
-       size(value.similar) as similarCount
-  ORDER BY similarCount DESC
+
+RETURN {
+         anchor: properties(value.anchorNode),
+         similar: properties(value.similar),
+         similarCount: size(value.similar)
+       } AS result
+  ORDER BY result.count DESC