feat: use row embeddings

Angular2Guy · Jan 20, 2024 · 085a6a2 · 085a6a2
1 parent 8469390
commit 085a6a2
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 12 deletions.
diff --git a/backend/src/main/java/ch/xxx/aidoclibchat/adapter/repository/DocumentVSRepositoryBean.java b/backend/src/main/java/ch/xxx/aidoclibchat/adapter/repository/DocumentVSRepositoryBean.java
@@ -45,17 +45,17 @@ public void add(List<Document> documents) {
 
 	@Override
 	public List<Document> retrieve(String query, DataType dataType, int k, double threshold) {
-		return this.vectorStore.similaritySearch(SearchRequest.query(query)
-				.withFilterExpression(
-						new Filter.Expression(ExpressionType.EQ, new Key(MetaData.DATATYPE), new Value(dataType.toString())))
-				.withTopK(k).withSimilarityThreshold(threshold));
+		return this.vectorStore
+				.similaritySearch(SearchRequest
+						.query(query).withFilterExpression(new Filter.Expression(ExpressionType.EQ,
+								new Key(MetaData.DATATYPE), new Value(dataType.toString())))
+						.withTopK(k).withSimilarityThreshold(threshold));
 	}
 
 	@Override
 	public List<Document> retrieve(String query, DataType dataType, int k) {
-		return this.vectorStore.similaritySearch(SearchRequest.query(query)
-				.withFilterExpression(
-						new Filter.Expression(ExpressionType.EQ, new Key(MetaData.DATATYPE), new Value(dataType.toString())))
+		return this.vectorStore.similaritySearch(SearchRequest.query(query).withFilterExpression(
+				new Filter.Expression(ExpressionType.EQ, new Key(MetaData.DATATYPE), new Value(dataType.toString())))
 				.withTopK(k));
 	}
 
@@ -67,12 +67,18 @@ public List<Document> retrieve(String query, DataType dataType) {
 
 	@Override
 	public List<Document> findAllTableDocuments() {
-		return this.vectorStore.similaritySearch(SearchRequest.defaults().withSimilarityThresholdAll().withTopK(Integer.MAX_VALUE).withFilterExpression(new Filter.Expression(
-				ExpressionType.OR,
-				new Filter.Expression(ExpressionType.EQ, new Key(MetaData.DATATYPE), new Value(DataType.COLUMN.toString())),
-				new Filter.Expression(ExpressionType.EQ, new Key(MetaData.DATATYPE), new Value(DataType.TABLE.toString())))));
+		return this.vectorStore
+				.similaritySearch(SearchRequest.defaults().withSimilarityThresholdAll().withTopK(Integer.MAX_VALUE)
+						.withFilterExpression(new Filter.Expression(ExpressionType.OR,
+								new Filter.Expression(ExpressionType.EQ, new Key(MetaData.DATATYPE),
+										new Value(DataType.COLUMN.toString())),
+								new Filter.Expression(ExpressionType.OR,
+										new Filter.Expression(ExpressionType.EQ, new Key(MetaData.DATATYPE),
+												new Value(DataType.TABLE.toString())),
+										new Filter.Expression(ExpressionType.EQ, new Key(MetaData.DATATYPE),
+												new Value(DataType.ROW.toString()))))));
 	}
-	
+
 	@Override
 	public void deleteByIds(List<String> ids) {
 		this.vectorStore.delete(ids);

diff --git a/backend/src/main/java/ch/xxx/aidoclibchat/usecase/service/TableService.java b/backend/src/main/java/ch/xxx/aidoclibchat/usecase/service/TableService.java
@@ -67,10 +67,17 @@ public void searchTables(SearchDto searchDto) {
 				searchDto.getResultAmount());
 		var columnDocuments = this.documentVsRepository.retrieve(searchDto.getSearchString(), MetaData.DataType.COLUMN,
 				searchDto.getResultAmount());
+		var rowDocuments = this.documentVsRepository.retrieve(searchDto.getSearchString(), MetaData.DataType.ROW,
+				searchDto.getResultAmount());
+		LOGGER.info("Table: ");
 		tableDocuments.forEach(myDoc -> LOGGER.info("name: {}, distance: {}",
 				myDoc.getMetadata().get(MetaData.DATANAME), myDoc.getMetadata().get(MetaData.DISTANCE)));
+		LOGGER.info("Column: ");
 		columnDocuments.forEach(myDoc -> LOGGER.info("name: {}, distance: {}",
 				myDoc.getMetadata().get(MetaData.DATANAME), myDoc.getMetadata().get(MetaData.DISTANCE)));
+		LOGGER.info("Row: ");
+		rowDocuments.forEach(myDoc -> LOGGER.info("name: {}, distance: {}", myDoc.getMetadata().get(MetaData.DATANAME),
+				myDoc.getMetadata().get(MetaData.DISTANCE)));
 	}
 
 	@Async