Add File

agents-flex-store/agents-flex-store-pgvector/src/main/java/com/agentsflex/store/pgvector/PgvectorVectorStore.java

@@ -0,0 +1,231 @@
+/*
+ *  Copyright (c) 2023-2025, Agents-Flex (fuhai999@gmail.com).
+ *  <p>
+ *  Licensed under the Apache License, Version 2.0 (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *  <p>
+ *  http://www.apache.org/licenses/LICENSE-2.0
+ *  <p>
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+package com.agentsflex.store.pgvector;
+
+import com.agentsflex.core.document.Document;
+import com.agentsflex.core.store.DocumentStore;
+import com.agentsflex.core.store.SearchWrapper;
+import com.agentsflex.core.store.StoreOptions;
+import com.agentsflex.core.store.StoreResult;
+import com.alibaba.fastjson.JSON;
+import com.alibaba.fastjson.JSONObject;
+import org.postgresql.ds.PGSimpleDataSource;
+import org.postgresql.util.PGobject;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.sql.*;
+import java.util.*;
+
+public class PgvectorVectorStore extends DocumentStore {
+    private static final Logger logger = LoggerFactory.getLogger(PgvectorVectorStore.class);
+    public static final double DEFAULT_SIMILARITY_THRESHOLD = 0.3;
+    private final PGSimpleDataSource dataSource;
+    private final String defaultCollectionName;
+    private final PgvectorVectorStoreConfig config;
+
+
+    public PgvectorVectorStore(PgvectorVectorStoreConfig config) {
+        dataSource = new PGSimpleDataSource();
+        dataSource.setServerNames(new String[]{config.getHost() + ":" + config.getPort()});
+        dataSource.setUser(config.getUsername());
+        dataSource.setPassword(config.getPassword());
+        dataSource.setDatabaseName(config.getDatabaseName());
+        if (!config.getProperties().isEmpty()) {
+            config.getProperties().forEach((k, v) -> {
+                try {
+                    dataSource.setProperty(k, v);
+                } catch (SQLException e) {
+                    logger.error("set pg property error", e);
+                }
+            });
+        }
+
+        this.defaultCollectionName = config.getDefaultCollectionName();
+        this.config = config;
+
+        // 异步初始化数据库
+        new Thread(this::initDb).start();
+    }
+
+    public void initDb() {
+        // 启动的时候初始化向量表, 需要数据库支持pgvector插件
+        // pg管理员需要在对应的库上执行 CREATE EXTENSION IF NOT EXISTS vector;
+        if (config.isAutoCreateCollection()) {
+            createCollection(defaultCollectionName);
+        }
+    }
+
+    private Connection getConnection() throws SQLException {
+        Connection connection = dataSource.getConnection();
+        connection.setAutoCommit(false);
+        return connection;
+    }
+
+    @Override
+    public StoreResult storeInternal(List<Document> documents, StoreOptions options) {
+
+        // 表名
+        String collectionName = options.getCollectionNameOrDefault(defaultCollectionName);
+
+        try (Connection connection = getConnection()) {
+            PreparedStatement pstmt = connection.prepareStatement("insert into " + collectionName + " (id, content, vector, metadata) values (?, ?, ?, ?::jsonb)");
+            for (Document doc : documents) {
+                Map<String, Object> metadatas = doc.getMetadataMap();
+                JSONObject jsonObject = JSON.parseObject(JSON.toJSONBytes(metadatas == null ? Collections.EMPTY_MAP : metadatas));
+                pstmt.setString(1, String.valueOf(doc.getId()));
+                pstmt.setString(2, doc.getContent());
+                pstmt.setObject(3, PgvectorUtil.toPgVector(doc.getVector()));
+                pstmt.setString(4, jsonObject.toString());
+                pstmt.addBatch();
+            }
+
+            pstmt.executeBatch();
+            connection.commit();
+        } catch (SQLException e) {
+            logger.error("store vector error", e);
+            return StoreResult.fail();
+        }
+        return StoreResult.successWithIds(documents);
+    }
+
+    private Boolean createCollection(String collectionName) {
+        try (Connection connection = getConnection()) {
+            try (CallableStatement statement = connection.prepareCall("CREATE TABLE IF NOT EXISTS " + collectionName +
+                " (id varchar(100) PRIMARY KEY, content text, vector vector(" + config.getVectorDimension() + "), metadata jsonb)")) {
+                statement.execute();
+            }
+
+            // 默认情况下，pgvector 执行精确的最近邻搜索，从而提供完美的召回率. 可以通过索引来修改 pgvector 的搜索方式，以获得更好的性能。
+            // By default, pgvector performs exact nearest neighbor search, which provides perfect recall.
+            if (config.isUseHnswIndex()) {
+                try (Statement stmt = connection.createStatement()) {
+                    stmt.execute("CREATE INDEX IF NOT EXISTS " + collectionName + "_vector_idx ON " + collectionName +
+                        " USING hnsw (vector vector_cosine_ops)");
+                }
+            }
+
+        } catch (SQLException e) {
+            logger.error("create collection error", e);
+            return false;
+        }
+
+        return true;
+    }
+
+    @Override
+    public StoreResult deleteInternal(Collection<?> ids, StoreOptions options) {
+        StringBuilder sql = new StringBuilder("DELETE FROM " + options.getCollectionNameOrDefault(defaultCollectionName) + " WHERE id IN (");
+        for (int i = 0; i < ids.size(); i++) {
+            sql.append("?");
+            if (i < ids.size() - 1) {
+                sql.append(",");
+            }
+        }
+        sql.append(")");
+
+        try (Connection connection = getConnection()) {
+            PreparedStatement pstmt = connection.prepareStatement(sql.toString());
+            ArrayList<?> list = new ArrayList<>(ids);
+            for (int i = 0; i < list.size(); i++) {
+                pstmt.setString(i + 1, (String) list.get(i));
+            }
+
+            pstmt.executeUpdate();
+            connection.commit();
+        } catch (Exception e) {
+            logger.error("delete document error: " + e, e);
+            return StoreResult.fail();
+        }
+
+        return StoreResult.success();
+
+    }
+
+    @Override
+    public List<Document> searchInternal(SearchWrapper searchWrapper, StoreOptions options) {
+        StringBuilder sql = new StringBuilder("select ");
+        if (searchWrapper.isOutputVector()) {
+            sql.append("id, vector, content, metadata");
+        } else {
+            sql.append("id,  content, metadata");
+        }
+
+        sql.append(" from ").append(options.getCollectionNameOrDefault(defaultCollectionName));
+        sql.append(" where vector <=> ? < ? order by vector <=> ? LIMIT ?");
+
+        try (Connection connection = getConnection()){
+            // 使用余弦距离计算最相似的文档
+            PreparedStatement stmt = connection.prepareStatement(sql.toString());
+
+            PGobject vector = PgvectorUtil.toPgVector(searchWrapper.getVector());
+            stmt.setObject(1, vector);
+            stmt.setObject(2, Optional.ofNullable(searchWrapper.getMinScore()).orElse(DEFAULT_SIMILARITY_THRESHOLD));
+            stmt.setObject(3, vector);
+            stmt.setObject(4, searchWrapper.getMaxResults());
+
+            ResultSet resultSet = stmt.executeQuery();
+            List<Document> documents = new ArrayList<>();
+            while (resultSet.next()) {
+                Document doc = new Document();
+                doc.setId(resultSet.getString("id"));
+                doc.setContent(resultSet.getString("content"));
+                doc.addMetadata(JSON.parseObject(resultSet.getString("metadata")));
+
+                if (searchWrapper.isOutputVector()) {
+                    String vectorStr = resultSet.getString("vector");
+                    doc.setVector(PgvectorUtil.fromPgVector(vectorStr));
+                }
+
+                documents.add(doc);
+            }
+
+            return documents;
+        } catch (Exception e) {
+            logger.error("Error searching in pgvector", e);
+            return Collections.emptyList();
+        }
+    }
+
+    @Override
+    public StoreResult updateInternal(List<Document> documents, StoreOptions options) {
+        if (documents == null || documents.isEmpty()) {
+            return StoreResult.success();
+        }
+
+        StringBuilder sql = new StringBuilder("UPDATE " + options.getCollectionNameOrDefault(defaultCollectionName) + " SET ");
+        sql.append("content = ?, vector = ?, metadata = ?::jsonb WHERE id = ?");
+        try (Connection connection = getConnection()) {
+            PreparedStatement pstmt = connection.prepareStatement(sql.toString());
+            for (Document doc : documents) {
+                Map<String, Object> metadatas = doc.getMetadataMap();
+                JSONObject metadataJson = JSON.parseObject(JSON.toJSONBytes(metadatas == null ? Collections.EMPTY_MAP : metadatas));
+                pstmt.setString(1, doc.getContent());
+                pstmt.setObject(2, PgvectorUtil.toPgVector(doc.getVector()));
+                pstmt.setString(3, metadataJson.toString());
+                pstmt.setString(4, String.valueOf(doc.getId()));
+                pstmt.addBatch();
+            }
+
+            pstmt.executeUpdate();
+            connection.commit();
+        } catch (Exception e) {
+            logger.error("Error update in pgvector", e);
+            return StoreResult.fail();
+        }
+        return StoreResult.successWithIds(documents);
+    }
+}