Skip to content

Commit

Permalink
fix: import
Browse files Browse the repository at this point in the history
  • Loading branch information
@Angular2Guy
Angular2Guy committed Jan 20, 2024
1 parent 6a080bd commit 560d661
Show file tree
Hide file tree
Showing 2 changed files with 80 additions and 13 deletions.
88 changes: 77 additions & 11 deletions backend/src/main/java/ch/xxx/aidoclibchat/adapter/repository/DocumentVSRepositoryBean.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,8 +12,13 @@
*/
package ch.xxx.aidoclibchat.adapter.repository;

import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.List;
import java.util.Map;
import java.util.stream.IntStream;

import org.postgresql.util.PGobject;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.vectorstore.PgVectorStore;
Expand All @@ -23,19 +28,33 @@
import org.springframework.ai.vectorstore.filter.Filter.ExpressionType;
import org.springframework.ai.vectorstore.filter.Filter.Key;
import org.springframework.ai.vectorstore.filter.Filter.Value;
import org.springframework.ai.vectorstore.filter.converter.FilterExpressionConverter;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.jdbc.core.RowMapper;
import org.springframework.stereotype.Repository;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.pgvector.PGvector;

import ch.xxx.aidoclibchat.domain.common.MetaData;
import ch.xxx.aidoclibchat.domain.common.MetaData.DataType;
import ch.xxx.aidoclibchat.domain.model.entity.DocumentVsRepository;

@Repository
public class DocumentVSRepositoryBean implements DocumentVsRepository {
private final String vectorTableName;
private final VectorStore vectorStore;
private final JdbcTemplate jdbcTemplate;
private final ObjectMapper objectMapper;
private final FilterExpressionConverter filterExpressionConverter;

public DocumentVSRepositoryBean(JdbcTemplate jdbcTemplate, EmbeddingClient embeddingClient) {
public DocumentVSRepositoryBean(JdbcTemplate jdbcTemplate, EmbeddingClient embeddingClient, ObjectMapper objectMapper) {
this.jdbcTemplate = jdbcTemplate;
this.objectMapper = objectMapper;
this.vectorStore = new PgVectorStore(jdbcTemplate, embeddingClient);
this.filterExpressionConverter = ((PgVectorStore) this.vectorStore).filterExpressionConverter;
this.vectorTableName = PgVectorStore.VECTOR_TABLE_NAME;
}

@Override
Expand Down Expand Up @@ -67,20 +86,67 @@ 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.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()))))));
String nativeFilterExpression = this.filterExpressionConverter.convertExpression(new Filter.Expression(ExpressionType.NE,
new Key(MetaData.DATATYPE), new Value(DataType.DOCUMENT.toString())));

String jsonPathFilter = " WHERE metadata::jsonb @@ '" + nativeFilterExpression + "'::jsonpath ";

return this.jdbcTemplate.query(
String.format("SELECT * FROM %s %s LIMIT ? ", this.vectorTableName, jsonPathFilter),
new DocumentRowMapper(this.objectMapper), 100000);
}

@Override
public void deleteByIds(List<String> ids) {
this.vectorStore.delete(ids);
}

private static class DocumentRowMapper implements RowMapper<Document> {

private static final String COLUMN_EMBEDDING = "embedding";

private static final String COLUMN_METADATA = "metadata";

private static final String COLUMN_ID = "id";

private static final String COLUMN_CONTENT = "content";

private ObjectMapper objectMapper;

public DocumentRowMapper(ObjectMapper objectMapper) {
this.objectMapper = objectMapper;
}

@Override
public Document mapRow(ResultSet rs, int rowNum) throws SQLException {
String id = rs.getString(COLUMN_ID);
String content = rs.getString(COLUMN_CONTENT);
PGobject pgMetadata = rs.getObject(COLUMN_METADATA, PGobject.class);
PGobject embedding = rs.getObject(COLUMN_EMBEDDING, PGobject.class);

Map<String, Object> metadata = toMap(pgMetadata);

Document document = new Document(id, content, metadata);
document.setEmbedding(toDoubleList(embedding));

return document;
}

private List<Double> toDoubleList(PGobject embedding) throws SQLException {
float[] floatArray = new PGvector(embedding.getValue()).toArray();
return IntStream.range(0, floatArray.length).mapToDouble(i -> floatArray[i]).boxed().toList();
}

private Map<String, Object> toMap(PGobject pgObject) {

String source = pgObject.getValue();
try {
return (Map<String, Object>) objectMapper.readValue(source, Map.class);
}
catch (JsonProcessingException e) {
throw new RuntimeException(e);
}
}

}
}
5 changes: 3 additions & 2 deletions backend/src/main/java/ch/xxx/aidoclibchat/usecase/service/TableService.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -76,8 +76,9 @@ public void searchTables(SearchDto searchDto) {
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: {}, content: {}, distance: {}", myDoc.getMetadata().get(MetaData.DATANAME), myDoc.getContent(),
myDoc.getMetadata().get(MetaData.DISTANCE)));
rowDocuments.forEach(
myDoc -> LOGGER.info("name: {}, content: {}, distance: {}", myDoc.getMetadata().get(MetaData.DATANAME),
myDoc.getContent(), myDoc.getMetadata().get(MetaData.DISTANCE)));
}

@Async
Expand Down

0 comments on commit 560d661

Please sign in to comment.