improved retrieval logic

This commit is contained in:
Ralph Soika 2025-06-14 13:01:37 +02:00
parent 38859b9b84
commit 41181227c6
4 changed files with 320 additions and 47 deletions

View file

@ -1,12 +1,14 @@
package ai.dynamixs.rag;
import java.net.InetSocketAddress;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.UUID;
import java.util.logging.Level;
import java.util.logging.Logger;
import java.util.stream.Collectors;
import org.apache.commons.compress.archivers.ArchiveException;
import org.eclipse.microprofile.config.inject.ConfigProperty;
@ -52,6 +54,7 @@ import jakarta.inject.Inject;
public class ClusterService {
public static final String KEYSPACE_REGEX = "^[a-z_]*[^-]$";
public static final int SEARCH_LIMIT_MAX = 100;
// mandatory environment settings
public static final String ENV_EMBEDDINGS_CLUSTER_CONTACTPOINTS = "EMBEDDINGS_CLUSTER_CONTACTPOINTS";
@ -127,7 +130,10 @@ public class ClusterService {
private CqlSession session;
private PreparedStatement insertVectorStmt;
// private PreparedStatement searchVectorStmt;
private final Map<String, PreparedStatement> searchVectorStatements = new HashMap<>();
private PreparedStatement selectVectorStmt;
private PreparedStatement removeVectorStmt;
@Resource
private TimerService timerService;
@ -342,6 +348,16 @@ public class ClusterService {
logger.info("│ ├── ✅ database schema created.");
}
/**
* Inserts an embedding together with a content chunk into the database. The
* uniqueID is the reference to the corresponding workitem holding the full meta
* data.
*
* @param uniqueID
* @param content
* @param vector
* @throws ClusterException
*/
public void insertVector(String uniqueID, String content, List<Float> vector) throws ClusterException {
// Prepare statement?
@ -367,44 +383,144 @@ public class ClusterService {
}
}
public List<String> searchEmbeddings(List<Float> textEmbeddings) throws ClusterException {
List<String> result = new ArrayList();
// Prepare statement?
/**
* Removes all embeddings for a given uniqueID. The method returns true if
* entries have existed.
*
* @param uniqueID
* @param content
* @param vector
* @throws ClusterException
*/
public boolean removeAllEmbeddings(String uniqueID) throws ClusterException {
// Prepare statement if not already done
if (removeVectorStmt == null) {
String deleteQuery = "DELETE FROM document_vectors WHERE id = ?";
removeVectorStmt = session.prepare(deleteQuery);
}
if (selectVectorStmt == null) {
String insertQuery = "SELECT id,content_chunk,similarity_cosine(content_vector, ?) " + //
" FROM document_vectors " + //
" ORDER BY content_vector ANN OF ? " + //
" LIMIT 5";
selectVectorStmt = session.prepare(insertQuery);
String selectQuery = "SELECT COUNT(*) FROM document_vectors WHERE id = ?";
selectVectorStmt = session.prepare(selectQuery);
}
try {
// CqlVector verwenden
CqlVector<Float> cqlVector = CqlVector.newInstance(textEmbeddings);
// test if entries exist
BoundStatement selectBoundStmt = selectVectorStmt.bind(uniqueID);
ResultSet resultSet = session.execute(selectBoundStmt);
Row row = resultSet.one();
ResultSet resultRows = session.execute(selectVectorStmt.bind(
if (row != null) {
long count = row.getLong(0);
if (count > 0) {
// remove all
BoundStatement deleteBoundStmt = removeVectorStmt.bind(uniqueID);
session.execute(deleteBoundStmt);
logger.info("│ ├── ✅ Removed " + count + " entries for uniqueID: " + uniqueID);
return true;
} else {
// logger.info("│ ├── ℹ️ No entries found for uniqueID: " + uniqueID);
}
} else {
// logger.info("│ ├── ℹ️ No entries found for uniqueID: " + uniqueID);
}
} catch (Exception e) {
throw new ClusterException(ClusterException.CLUSTER_ERROR,
"Failed to remove entries for uniqueID '" + uniqueID + "': " + e.getMessage(), e);
}
return false;
}
/**
* Performs a semantic search on the document vector database using cosine
* similarity.
* <p>
* This method searches for the most relevant document chunks based on the
* provided query embedding. If multiple chunks exist for the same document ID,
* only the chunk with the highest similarity score is returned. The results are
* sorted by relevance score in descending order.
*
* @param queryEmbedding The embedding vector of the search query as a list of
* float values. The dimensionality must match the
* DIMENSIONS constant used in the database schema.
* @return A list of {@link RetrievalResult} objects containing the most
* relevant documents, limited to a maximum of 5 results. Each result
* includes the document ID, the best matching content chunk, and the
* similarity score. The list is sorted by similarity score in
* descending order (most relevant first).
* @throws ClusterException if the database query fails or if there's an error
* processing the embedding vector
*
* @see RetrievalResult
* @since 1.0
*/
public List<RetrievalResult> searchEmbeddings(List<Float> embedding, int maxResults) throws ClusterException {
if (maxResults <= 0) {
throw new ClusterException(ClusterException.CLUSTER_ERROR,
"maxResults must be greater than 0, but was: " + maxResults);
}
// Calculate a reasonable DB limit based on maxResults
// Factor 3-4 times to have enough chunks for deduplication.
int dbLimit = Math.min(maxResults * 4, SEARCH_LIMIT_MAX); // Max 100
// Prepare statement if not already done (dynamisches Limit erfordert neue
// Statements)
String statementKey = "search_" + dbLimit;
PreparedStatement searchVectorStmt = searchVectorStatements.get(statementKey);
// Prepare statement if not already done
if (searchVectorStmt == null) {
String searchQuery = "SELECT id,content_chunk,similarity_cosine(content_vector, ?) " + //
" FROM document_vectors " + //
" ORDER BY content_vector ANN OF ? " + //
" LIMIT " + dbLimit;
searchVectorStmt = session.prepare(searchQuery);
}
try {
// use CqlVector
CqlVector<Float> cqlVector = CqlVector.newInstance(embedding);
ResultSet resultRows = session.execute(searchVectorStmt.bind(
cqlVector,
cqlVector));
// track best Similarity per uniqueID
Map<String, RetrievalResult> bestMatches = new HashMap<>();
for (Row row : resultRows) {
String id = row.getString(0);
result.add(id);
String contentChunk = row.getString(1);
Float similarity = row.getFloat(2);
System.out.println("--------------");
System.out.println("ID: " + id);
System.out.println("");
System.out.println(contentChunk + " " + similarity);
System.out.println("");
System.out.println("");
// test if we already have a match
RetrievalResult existingMatch = bestMatches.get(id);
if (existingMatch == null || similarity > existingMatch.getScore()) {
// better match found!
bestMatches.put(id, new RetrievalResult(id, contentChunk, similarity));
logger.info("│ ├── Found/Updated document: " + id + " with score: " + similarity);
}
}
// Sortiere nach Score (absteigend) und gebe Top 5 zurück
List<RetrievalResult> result = bestMatches.values().stream()
.sorted((a, b) -> Float.compare(b.getScore(), a.getScore()))
.limit(maxResults)
.peek(match -> {
System.out.println("--------------");
System.out.println("- $UniqueID: " + match.getUniqueId());
System.out.println("- Relevance Score: " + match.getScore());
System.out.println("");
System.out.println("- Best Matching Content: ");
System.out.println(match.getContent());
System.out.println("");
})
.collect(Collectors.toList());
logger.info("│ ├── ✅ Retrieved " + result.size() + " relevant documents");
return result;
} catch (Exception e) {
throw new ClusterException(ClusterException.CLUSTER_ERROR,
"Failed to select embeddings from keyspace: " + e.getMessage(), e);
}
return result;
}
}

View file

@ -32,9 +32,9 @@ import java.util.List;
import java.util.Optional;
import java.util.logging.Level;
import java.util.logging.Logger;
import java.util.stream.Collectors;
import org.eclipse.microprofile.config.inject.ConfigProperty;
import org.imixs.ai.workflow.ImixsAIPromptEvent;
import org.imixs.ai.workflow.OpenAIAPIService;
import org.imixs.workflow.ItemCollection;
import org.imixs.workflow.SignalAdapter;
@ -42,7 +42,6 @@ import org.imixs.workflow.engine.WorkflowService;
import org.imixs.workflow.exceptions.AdapterException;
import org.imixs.workflow.exceptions.PluginException;
import jakarta.enterprise.event.Event;
import jakarta.inject.Inject;
/**
@ -120,9 +119,6 @@ public class RAGAdapter implements SignalAdapter {
@Inject
private ClusterService clusterService;
@Inject
private Event<ImixsAIPromptEvent> llmPromptEventObservers = null;
/**
* Default Constructor
*/
@ -204,13 +200,10 @@ public class RAGAdapter implements SignalAdapter {
}
logger.info("├── post RAG index request: " + llmAPIEndpoint);
String promptTemplate = llmService.loadPromptTemplate(event);
logger.info("│ ├── PromptTemplate: ");
logger.info(promptTemplate);
String llmPrompt = llmService.buildPrompt(promptTemplate, workitem);
if (llmPrompt.isBlank()) {
throw new PluginException(
RAGAdapter.class.getSimpleName(), API_ERROR,
@ -223,19 +216,20 @@ public class RAGAdapter implements SignalAdapter {
logger.info(llmPrompt);
}
// remove old embeddings
clusterService.removeAllEmbeddings(workitem.getUniqueID());
// chunk text....
List<String> chunk_list = RAGUtil.chunkMarkupDocument(llmPrompt);
for (String chunk : chunk_list) {
if (llmAPIDebug) {
logger.info("├── 🔸 chunk: ");
logger.info("│ ├── 🔸 chunk: ");
logger.info(chunk);
}
List<Float> indexResult = llmService.postEmbedding(chunk, llmAPIEndpoint, llmAPIDebug);
// write to cassandra
clusterService.insertVector(workitem.getUniqueID(), chunk, indexResult);
if (llmAPIDebug) {
logger.info("├── ⇨ " + indexResult.size() + " floats stored in RAG db.");
logger.info("│ ├── ⇨ " + indexResult.size() + " floats stored in RAG db.");
}
}
@ -279,19 +273,15 @@ public class RAGAdapter implements SignalAdapter {
}
logger.info("├── post RAG Retrievalf request: " + llmAPIEndpoint);
String promptTemplate = llmService.loadPromptTemplate(event);
logger.info("│ ├── PromptTemplate: ");
logger.info(promptTemplate);
String llmPrompt = llmService.buildPrompt(promptTemplate, workitem);
if (llmPrompt.isBlank()) {
throw new PluginException(
RAGAdapter.class.getSimpleName(), API_ERROR,
"Unable to parse prompt for RAG indexing");
}
if (llmAPIDebug) {
logger.info("│ ├── Total Prompt Length = " + llmPrompt.length());
logger.info("│ ├── Prompt: ");
@ -299,21 +289,19 @@ public class RAGAdapter implements SignalAdapter {
}
// retrieve prompt....
List<Float> indexResult = llmService.postEmbedding(llmPrompt, llmAPIEndpoint, llmAPIDebug);
if (llmAPIDebug) {
logger.info("├── ⇨ " + indexResult.size() + " floats stored in RAG db.");
}
// search cassandra
List<String> ids = clusterService.searchEmbeddings(indexResult);
workitem.setItemValue("product.ref", ids);
List<RetrievalResult> retrievalResultList = clusterService.searchEmbeddings(indexResult, 5);
List<String> listOfIds = retrievalResultList.stream()
.map(RetrievalResult::getUniqueId)
.collect(Collectors.toList());
workitem.setItemValue("product.ref", listOfIds);
if (llmAPIDebug) {
logger.info(
"├── found " + ids.size() + " matches");
"├── found " + retrievalResultList.size() + " matches");
logger.info(
"├── ✅ Total processing time: " + (System.currentTimeMillis() - processingTime) + "ms");
}
@ -360,13 +348,10 @@ public class RAGAdapter implements SignalAdapter {
// adapt text...
llmAPIEndpoint = workflowService.adaptText(llmAPIEndpoint, null);
if (!llmAPIEndpoint.endsWith("/")) {
llmAPIEndpoint = llmAPIEndpoint + "/";
}
return llmAPIEndpoint;
}
}

View file

@ -0,0 +1,70 @@
package ai.dynamixs.rag;
/**
* Represents a single result from a semantic search operation in the RAG
* (Retrieval-Augmented Generation) system.
*
* This immutable data class encapsulates the information about a document that
* was retrieved based on semantic similarity to a query. Each result contains
* the document's unique identifier, the content of the best matching chunk, and
* the similarity score that indicates how well the document matches the search
* query.
*
* <p>
* The similarity score is typically a cosine similarity value between the query
* embedding and the document's content embedding, where higher values indicate
* better matches.
* </p>
*
* <p>
* <strong>Usage Example:</strong>
* </p>
*
* <pre>{@code
* List<RetrievalResult> results = searchEmbeddings(queryVector);
* for (RetrievalResult result : results) {
* System.out.println("Document: " + result.getUniqueId() +
* " (Score: " + result.getScore() + ")");
* System.out.println("Content: " + result.getContent());
* }
* }</pre>
*
* @author Your Name
* @version 1.0
* @since 1.0
* @see #getUniqueId()
* @see #getContent()
* @see #getScore()
*/
public class RetrievalResult {
private final String uniqueId;
private final String content;
private final Float score;
public RetrievalResult(String uniqueID, String content, Float score) {
this.uniqueId = uniqueID;
this.content = content;
this.score = score;
}
public String getUniqueId() {
return uniqueId;
}
public String getContent() {
return content;
}
public Float getScore() {
return score;
}
@Override
public String toString() {
return "RetrievalResult{" +
"$uniqueId='" + uniqueId + '\'' +
", score=" + score +
", content='" + content.substring(0, Math.min(50, content.length())) + "...'" +
'}';
}
}

View file

@ -72,6 +72,7 @@
<bpmn2:flowNodeRef>dataObject_axJA7Q</bpmn2:flowNodeRef>
<bpmn2:flowNodeRef>event_y0smQw</bpmn2:flowNodeRef>
<bpmn2:flowNodeRef>event_qb0xXA</bpmn2:flowNodeRef>
<bpmn2:flowNodeRef>event_jBz77w</bpmn2:flowNodeRef>
</bpmn2:lane>
</bpmn2:laneSet>
<bpmn2:startEvent id="StartEvent_1" name="Start">
@ -207,6 +208,7 @@ Stock: <itemvalue>_stock</itemvalue>
<bpmn2:incoming>SequenceFlow_42</bpmn2:incoming>
<bpmn2:incoming>sequenceFlow_N5KIRQ</bpmn2:incoming>
<bpmn2:outgoing>sequenceFlow_HCxYvw</bpmn2:outgoing>
<bpmn2:incoming>sequenceFlow_kvvOjA</bpmn2:incoming>
</bpmn2:task>
<bpmn2:task id="Task_4" imixs:processid="1400" name="Bestellt">
<bpmn2:extensionElements>
@ -1201,6 +1203,88 @@ result.isValid=true;
<bpmn2:association id="association_BseA8A" sourceRef="DataObject_2" targetRef="Task_3">
<bpmn2:documentation id="documentation_waNoXQ"/>
</bpmn2:association>
<bpmn2:intermediateCatchEvent id="event_jBz77w" imixs:activityid="50" name="(RAG)">
<bpmn2:extensionElements>
<imixs:item name="keylogtimeformat" type="xs:string">
<imixs:value><![CDATA[2]]></imixs:value>
</imixs:item>
<imixs:item name="keyarchive" type="xs:string">
<imixs:value><![CDATA[0]]></imixs:value>
</imixs:item>
<imixs:item name="txtmailsubject" type="xs:string">
<imixs:value/>
</imixs:item>
<imixs:item name="keyaccessmode" type="xs:string">
<imixs:value><![CDATA[0]]></imixs:value>
</imixs:item>
<imixs:item name="keyscheduledactivity" type="xs:string">
<imixs:value><![CDATA[0]]></imixs:value>
</imixs:item>
<imixs:item name="rtfmailbody" type="xs:string">
<imixs:value/>
</imixs:item>
<imixs:item name="keyversion" type="xs:string">
<imixs:value><![CDATA[0]]></imixs:value>
</imixs:item>
<imixs:item name="keyaddwritefields" type="xs:string">
<imixs:value><![CDATA[namCreator]]></imixs:value>
<imixs:value><![CDATA[namCurrentEditor]]></imixs:value>
<imixs:value><![CDATA[namprocessmanager]]></imixs:value>
<imixs:value><![CDATA[namprocessteam]]></imixs:value>
</imixs:item>
<imixs:item name="numnextactivityid" type="xs:int">
<imixs:value><![CDATA[0]]></imixs:value>
</imixs:item>
<imixs:item name="keyfollowup" type="xs:string">
<imixs:value><![CDATA[2]]></imixs:value>
</imixs:item>
<imixs:item name="keypublicresult" type="xs:string">
<imixs:value><![CDATA[1]]></imixs:value>
</imixs:item>
<imixs:item name="keylogdateformat" type="xs:string">
<imixs:value><![CDATA[2]]></imixs:value>
</imixs:item>
<imixs:item name="keyownershipfields" type="xs:string">
<imixs:value><![CDATA[namCurrentEditor]]></imixs:value>
<imixs:value><![CDATA[namCreator]]></imixs:value>
</imixs:item>
<imixs:item name="numnextid" type="xs:int">
<imixs:value><![CDATA[5000]]></imixs:value>
</imixs:item>
<imixs:item name="txtnextprocesstree" type="xs:string">
<imixs:value/>
</imixs:item>
<imixs:item name="txtactivityresult" type="xs:string">
<imixs:value><![CDATA[<imixs-ai name="INDEX">
<endpoint>https://llama.cpp2.imixs.com/</endpoint>
<debug>true</debug>
</imixs-ai>
]]></imixs:value>
</imixs:item>
<imixs:item name="txtname" type="xs:string">
<imixs:value><![CDATA[Speichern]]></imixs:value>
</imixs:item>
<imixs:item name="keyownershipmode" type="xs:string">
<imixs:value><![CDATA[0]]></imixs:value>
</imixs:item>
<imixs:item name="rtfresultlog" type="xs:string">
<imixs:value/>
</imixs:item>
<imixs:item name="keyupdateacl" type="xs:boolean">
<imixs:value>false</imixs:value>
</imixs:item>
</bpmn2:extensionElements>
<bpmn2:documentation id="documentation_2LV0ug"><![CDATA[RAG Test]]></bpmn2:documentation>
<bpmn2:signalEventDefinition id="signalEventDefinition_0qr0og" signalRef="signal_4"/>
<bpmn2:incoming>sequenceFlow_crpgGw</bpmn2:incoming>
<bpmn2:outgoing>sequenceFlow_kvvOjA</bpmn2:outgoing>
</bpmn2:intermediateCatchEvent>
<bpmn2:sequenceFlow id="sequenceFlow_kvvOjA" sourceRef="event_jBz77w" targetRef="Task_3">
<bpmn2:documentation id="documentation_Ebp1MA"/>
</bpmn2:sequenceFlow>
<bpmn2:association id="association_SjKymg" sourceRef="dataObject_axJA7Q" targetRef="event_jBz77w">
<bpmn2:documentation id="documentation_EaPADA"/>
</bpmn2:association>
</bpmn2:process>
<bpmn2:process id="process_2" name="Default Process" processType="Public">
<bpmn2:documentation id="documentation_imRbaw"/>
@ -1501,6 +1585,24 @@ result.isValid=true;
<di:waypoint x="728.0" y="205.0"/>
<di:waypoint x="728.0" y="310.0"/>
</bpmndi:BPMNEdge>
<bpmndi:BPMNShape bpmnElement="event_jBz77w" id="BPMNShape_MQkY4w">
<dc:Bounds height="36.0" width="36.0" x="797.0" y="427.0"/>
<bpmndi:BPMNLabel id="BPMNLabel_WgIqjg">
<dc:Bounds height="20.0" width="100.0" x="765.0" y="466.0"/>
</bpmndi:BPMNLabel>
</bpmndi:BPMNShape>
<bpmndi:BPMNEdge bpmnElement="sequenceFlow_kvvOjA" id="BPMNEdge_St0T1w" sourceElement="BPMNShape_MQkY4w" targetElement="BPMNShape_Task_9">
<di:waypoint x="815.0" y="427.0"/>
<di:waypoint x="815.0" y="394.0"/>
<di:waypoint x="761.0" y="394.0"/>
<di:waypoint x="761.0" y="360.0"/>
</bpmndi:BPMNEdge>
<bpmndi:BPMNEdge bpmnElement="association_SjKymg" id="BPMNEdge_mSixYw" sourceElement="BPMNShape_3PbNFQ" targetElement="BPMNShape_MQkY4w">
<di:waypoint x="548.0" y="470.0"/>
<di:waypoint x="548.0" y="533.0"/>
<di:waypoint x="815.0" y="533.0"/>
<di:waypoint x="815.0" y="463.0"/>
</bpmndi:BPMNEdge>
</bpmndi:BPMNPlane>
<bpmndi:BPMNLabelStyle id="BPMNLabelStyle_1">
<dc:Font name="arial" size="9.0"/>