office-dynamixs-ai/dynamixs-ai-rag/src/main/java/ai/dynamixs/rag/ClusterService.java

package ai.dynamixs.rag;

import java.net.InetSocketAddress;
import java.time.Instant;
import java.util.Optional;
import java.util.UUID;
import java.util.logging.Level;
import java.util.logging.Logger;

import org.apache.commons.compress.archivers.ArchiveException;
import org.eclipse.microprofile.config.inject.ConfigProperty;

import com.datastax.oss.driver.api.core.CqlSession;
import com.datastax.oss.driver.api.core.CqlSessionBuilder;
import com.datastax.oss.driver.api.core.cql.BoundStatement;
import com.datastax.oss.driver.api.core.cql.PreparedStatement;
import com.datastax.oss.driver.api.core.servererrors.InvalidQueryException;

import jakarta.annotation.PostConstruct;
import jakarta.annotation.PreDestroy;
import jakarta.annotation.Resource;
import jakarta.ejb.Singleton;
import jakarta.ejb.Startup;
import jakarta.ejb.Timeout;
import jakarta.ejb.Timer;
import jakarta.ejb.TimerConfig;
import jakarta.ejb.TimerService;
import jakarta.inject.Inject;

/**
 * The ClusterService provides methods to persist the content of a Imixs
 * Document into a Cassandra keystore.
 * <p>
 * The service saves the content in XML format. The size of an XML
 * representation of a Imixs document is only slightly different in size from
 * the serialized map object. This is the reason why we do not store the
 * document map in a serialized object format.
 * <p>
 * The ClusterService creates a Core-KeySpace automatically which is used for
 * the internal management.
 * 
 * @author rsoika
 * 
 */
@Singleton
@Startup
public class ClusterService {

    public static final String KEYSPACE_REGEX = "^[a-z_]*[^-]$";

    // mandatory environment settings
    public static final String ENV_EMBEDDINGS_CLUSTER_CONTACTPOINTS = "EMBEDDINGS_CLUSTER_CONTACTPOINTS";
    public static final String ENV_EMBEDDINGS_CLUSTER_KEYSPACE = "EMBEDDINGS_CLUSTER_KEYSPACE";

    // optional environment settings
    public static final String ENV_EMBEDDINGS_CLUSTER_AUTH_USER = "EMBEDDINGS_CLUSTER_AUTH_USER";
    public static final String ENV_EMBEDDINGS_CLUSTER_AUTH_PASSWORD = "EMBEDDINGS_CLUSTER_AUTH_PASSWORD";
    public static final String ENV_EMBEDDINGS_CLUSTER_SSL = "EMBEDDINGS_CLUSTER_SSL";
    public static final String ENV_EMBEDDINGS_CLUSTER_SSL_TRUSTSTOREPATH = "EMBEDDINGS_CLUSTER_SSL_TRUSTSTOREPATH";
    public static final String ENV_EMBEDDINGS_CLUSTER_SSL_TRUSTSTOREPASSWORD = "EMBEDDINGS_CLUSTER_SSL_TRUSTSTOREPASSWORD";
    public static final String ENV_EMBEDDINGS_CLUSTER_SSL_KEYSTOREPATH = "EMBEDDINGS_CLUSTER_SSL_KEYSTOREPATH";
    public static final String ENV_EMBEDDINGS_CLUSTER_SSL_KEYSTOREPASSWORD = "EMBEDDINGS_CLUSTER_SSL_KEYSTOREPASSWORD";

    public static final String ENV_EMBEDDINGS_CLUSTER_REPLICATION_FACTOR = "EMBEDDINGS_CLUSTER_REPLICATION_FACTOR";
    public static final String ENV_EMBEDDINGS_CLUSTER_REPLICATION_CLASS = "EMBEDDINGS_CLUSTER_REPLICATION_CLASS";

    // workflow rest service endpoint
    public static final String ENV_WORKFLOW_SERVICE_ENDPOINT = "WORKFLOW_SERVICE_ENDPOINT";
    public static final String ENV_WORKFLOW_SERVICE_USER = "WORKFLOW_SERVICE_USER";
    public static final String ENV_WORKFLOW_SERVICE_PASSWORD = "WORKFLOW_SERVICE_PASSWORD";
    public static final String ENV_WORKFLOW_SERVICE_AUTHMETHOD = "WORKFLOW_SERVICE_AUTHMETHOD";

    private static Logger logger = Logger.getLogger(ClusterService.class.getName());

    @Inject
    @ConfigProperty(name = ENV_EMBEDDINGS_CLUSTER_REPLICATION_FACTOR, defaultValue = "1")
    String repFactor;

    @Inject
    @ConfigProperty(name = ENV_EMBEDDINGS_CLUSTER_REPLICATION_CLASS, defaultValue = "SimpleStrategy")
    String repClass;

    @Inject
    @ConfigProperty(name = ENV_EMBEDDINGS_CLUSTER_CONTACTPOINTS)
    Optional<String> contactPoint;

    @Inject
    @ConfigProperty(name = ENV_EMBEDDINGS_CLUSTER_KEYSPACE)
    Optional<String> keySpace;

    @Inject
    @ConfigProperty(name = ENV_EMBEDDINGS_CLUSTER_AUTH_USER)
    Optional<String> userid;

    @Inject
    @ConfigProperty(name = ENV_EMBEDDINGS_CLUSTER_AUTH_PASSWORD)
    Optional<String> password;

    @Inject
    @ConfigProperty(name = ENV_EMBEDDINGS_CLUSTER_SSL, defaultValue = "false")
    boolean bUseSSL;

    @Inject
    @ConfigProperty(name = ENV_EMBEDDINGS_CLUSTER_SSL_TRUSTSTOREPATH)
    Optional<String> truststorePath;

    @Inject
    @ConfigProperty(name = ENV_EMBEDDINGS_CLUSTER_SSL_TRUSTSTOREPASSWORD)
    Optional<String> truststorePwd;

    @Inject
    @ConfigProperty(name = ENV_EMBEDDINGS_CLUSTER_SSL_KEYSTOREPATH)
    Optional<String> keystorePath;

    @Inject
    @ConfigProperty(name = ENV_EMBEDDINGS_CLUSTER_SSL_KEYSTOREPASSWORD)
    Optional<String> keystorePwd;

    private CqlSession session;

    @Resource
    private TimerService timerService;

    @PostConstruct
    public void init() {
        scheduleClusterCheck();
    }

    /**
     * Versucht, die Session herzustellen. Bei Fehler wird ein neuer Timer geplant.
     */
    private void scheduleClusterCheck() {
        try {
            logger.log(Level.INFO, "├── Connecting Cassandra cluster...");
            getSession();
            logger.log(Level.INFO, "├── ✅ Successfully connected to Cassandra!");
        } catch (ClusterException e) {
            logger.log(Level.WARNING, "├── " + e.getMessage());
            logger.log(Level.WARNING, "├── ⚠️ Cluster not ready yet. Retrying in 10 seconds...");
            // retry in 10 sec
            timerService.createSingleActionTimer(10000, new TimerConfig());
        }
    }

    /**
     * called by timer. Tries to connect cassandra clauster
     */
    @Timeout
    public void retryClusterCheck(Timer timer) {
        scheduleClusterCheck(); // retry Check
    }

    @PreDestroy
    private void tearDown() {
        // close session and cluster object
        if (session != null) {
            session.close();
        }
    }

    /**
     * Returns the Cassandra Session object. If the session is not yet created, the
     * method tries to init the session the first time. If initialization failed the
     * method throws a ClusterException
     * <p>
     * Note: the initialization is only successful if the cassandra cluster is fully
     * started.
     * <p>
     * The method also verifies the keyspace and creates a table schema if not yet
     * available.
     * 
     * @return
     */
    public CqlSession getSession() throws ClusterException {
        if (session == null) {
            try {
                logger.info("├── initializing cluster...");
                session = initSession();
                if (session != null) {
                    logger.info("│   ├── keySpace connection status = OK");
                    initKeySpace();
                }
            } catch (Exception e) {
                throw new ClusterException(ClusterException.CLUSTER_ERROR,
                        "Failed to init cassandra session: " + e.getMessage(), e);
            }
        }
        return session;
    }

    /**
     * Returns a cassandra session for the archive KeySpace. The keyspace is defined
     * by the environmetn variable EMBEDDINGS__CLUSTER_KEYSPACE. If no keyspace with
     * this given keyspace name exists, the method creates the keyspace and table
     * schemas.
     * 
     * @throws ArchiveException
     */
    private void initKeySpace() throws ClusterException {
        String keySpacename = keySpace.orElse(null);
        if (keySpacename == null) {
            throw new ClusterException(ClusterException.INVALID_KEYSPACE, "EMBEDDINGS_CLUSTER_KEYSPACE not set!");
        }
        if (!isValidKeyspaceName(keySpacename)) {
            throw new ClusterException(ClusterException.INVALID_KEYSPACE,
                    "keyspace '" + keySpacename + "' name invalid.");
        }

        logger.info("│   ├── connecting keyspace='" + keySpacename + "'...");
        // try to open keySpace
        try {
            // session = cluster.connect(keySpace.get());
        } catch (InvalidQueryException e) {
            logger.warning("......connecting keySpace '" + keySpacename + "' failed: " + e.getMessage());
            createKeySpace(keySpacename);
        }

    }

    /**
     * This method creates a Cassandra Cluster object. The cluster is defined by
     * ContactPoints provided in the environmetn variable
     * 'EMBEDDINGS__CLUSTER_CONTACTPOINTS' or in the imixs.property
     * 'archive.cluster.contactpoints'
     * 
     * @return Cassandra Cluster instacne
     * @throws ArchiveException
     */
    protected CqlSession initSession() throws ClusterException {
        // Boolean used to check if at least one host could be resolved
        boolean found = false;

        if (!contactPoint.isPresent() || contactPoint.get().isEmpty()) {
            throw new ClusterException(ClusterException.MISSING_CONTACTPOINT,
                    "missing cluster contact points - verify configuration!");
        }

        logger.info("│   ├── cluster connecting: " + contactPoint.get());

        CqlSessionBuilder builder = CqlSession.builder();
        String[] hosts = contactPoint.get().split(",");
        for (String host : hosts) {
            try {
                logger.info("│   ├── adding host: " + host + ":9042");
                builder.addContactPoint(new InetSocketAddress(host, 9042));
                // One host could be resolved
                found = true;
            } catch (IllegalArgumentException e) {
                // This host could not be resolved so we log a message and keep going
                logger.warning("...the host '" + host + "' is unknown so it will be ignored");
            }
        }
        if (!found) {
            // No host could be resolved so we throw an exception
            throw new IllegalStateException("All provided hosts are unknown - check cluster status and configuration!");
        }

        // builder.withLoadBalancingPolicy(new RoundRobinPolicy());
        // builder.withRetryPolicy(DefaultRetryPolicy.INSTANCE);

        // set optional credentials...
        if (userid.isPresent() && !userid.get().isEmpty()) {
            builder = builder.withAuthCredentials(userid.get(), password.get());
        }

        // use SSL ?
        // if (bUseSSL) {
        // try {
        // // create ssl options based on environment settings...
        // SSLOptions options = createSSLOptions();
        // logger.info("...creating SSL cluster session...");
        // builder.withSSL(options);
        // } catch (KeyManagementException | UnrecoverableKeyException |
        // KeyStoreException | NoSuchAlgorithmException
        // | CertificateException | IOException e) {
        // logger.severe("Failed to connect withSSL: " + e.getMessage());
        // e.printStackTrace();
        // }
        // }

        builder.withLocalDatacenter("datacenter1");
        return builder.build();
        // cluster = builder.build();
        // cluster.init();
        // logger.info("...cluster conection status = OK");

        // return cluster;

    }

    /**
     * Helper Method to generate sslOpens object for cassandra dataStax driver
     * withSSL.
     * <p>
     * At least the truststorePath is needed to establish a one way SSL connection.
     * <p>
     * The keystorePath is only needed in case 'require_client_auth = true' is set
     * in the cassandra.yaml file
     * 
     * @return sslOptons
     * @throws KeyStoreException
     * @throws FileNotFoundException
     * @throws IOException
     * @throws NoSuchAlgorithmException
     * @throws KeyManagementException
     * @throws CertificateException
     * @throws UnrecoverableKeyException
     */
    // private SSLOptions createSSLOptions() throws KeyStoreException,
    // FileNotFoundException, IOException,
    // NoSuchAlgorithmException, KeyManagementException, CertificateException,
    // UnrecoverableKeyException {

    // TrustManagerFactory tmf = null;
    // if (truststorePath.isPresent() && !truststorePath.get().isEmpty()) {
    // KeyStore tks = KeyStore.getInstance("JKS");
    // tks.load((InputStream) new FileInputStream(new File(truststorePath.get())),
    // truststorePwd.get().toCharArray());
    // tmf =
    // TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm());
    // tmf.init(tks);
    // } else {
    // logger.info("SSLOptions without truststore...");
    // }

    // KeyManagerFactory kmf = null;
    // if (keystorePath.isPresent() && !keystorePath.get().isEmpty()) {
    // KeyStore kks = KeyStore.getInstance("JKS");
    // kks.load((InputStream) new FileInputStream(new File(keystorePath.get())),
    // keystorePwd.get().toCharArray());
    // kmf = KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm());
    // kmf.init(kks, keystorePwd.get().toCharArray());
    // } else {
    // logger.info("SSLOptions without keystore...");
    // }

    // SSLContext sslContext = SSLContext.getInstance("TLS");
    // sslContext.init(kmf != null ? kmf.getKeyManagers() : null, tmf != null ?
    // tmf.getTrustManagers() : null,
    // new SecureRandom());
    // RemoteEndpointAwareJdkSSLOptions sslOptions =
    // RemoteEndpointAwareJdkSSLOptions.builder()
    // .withSSLContext(sslContext).build();
    // return sslOptions;
    // }

    /**
     * Test if the keyspace name is valid.
     * 
     * @param keySpace
     * @return
     */
    public boolean isValidKeyspaceName(String keySpace) {
        if (keySpace == null || keySpace.isEmpty()) {
            return false;
        }

        return keySpace.matches(KEYSPACE_REGEX);

    }

    /**
     * This method creates a cassandra keySpace.
     * <p>
     * Depending on the KeyspaceType, the method creates the core table schema to
     * store configurations, or the extended Archive table schcema which is used to
     * store imixs documents.
     * 
     * @param cluster
     * @throws ArchiveException
     */
    protected void createKeySpace(String keySpace) throws ClusterException {
        logger.info("│   ├── creating new keyspace '" + keySpace + "'...");

        String statement = "CREATE KEYSPACE IF NOT EXISTS " + keySpace + " WITH replication = {'class': '" + repClass
                + "', 'replication_factor': " + repFactor + "};";

        session.execute(statement);
        logger.info("│   ├── keyspace created...");
        if (session != null) {
            logger.info("│   ├──  create schema...");
            // now create table schemas
            createArchiveTableSchema();
        }

    }

    /**
     * This helper method creates the ImixsArchive document table schema if not yet
     * exists
     */
    protected void createArchiveTableSchema() {

        logger.info("│   ├── Create Data Table...");

        String query = "CREATE TABLE IF NOT EXISTS embeddings.document_vectors (\n" + //
                "  record_id timeuuid,\n" + //
                "  id uuid,\n" + //
                "  content_chunk text,\n" + //
                "  content_vector VECTOR <FLOAT, 768>,\n" + //
                "  created_at timestamp,\n" + //
                "  PRIMARY KEY (id, created_at)\n" + //
                ")" + //
                "WITH CLUSTERING ORDER BY (created_at DESC);";
        session.execute(query);

        logger.info("│   ├── Create Index...");
        query = "CREATE INDEX IF NOT EXISTS edv_ann_index\n" + //
                "  ON embeddings.document_vectors(content_vector) USING 'sai';";

        session.execute(query);

    }

    public void insertVector(String content, float[] vector) {

        // Generiere UUID und aktuelle Zeit
        UUID id = UUID.randomUUID();
        Instant createdAt = Instant.now();

        // CQL-Query mit Platzhaltern
        String insertQuery = """
                INSERT INTO embeddings.document_vectors
                    (record_id, id, content_chunk, content_vector, created_at)
                VALUES (now(), ?, ?, ?, ?)
                """;

        // Vorbereiten und Ausführen der Query
        PreparedStatement preparedStmt = session.prepare(insertQuery);
        BoundStatement boundStmt = preparedStmt.bind(id, content, vector, createdAt);
        session.execute(boundStmt);

        // id = uuid.uuid4()
        // content_chunk = content
        // content_vector = vector
        // created_at = datetime.datetime.now()
        // insert_query = """
        // INSERT INTO embeddings.document_vectors (record_id, id, content_chunk,
        // content_vector, created_at)
        // VALUES (now(), ?, ?, ?, ?)
        // """
        // prepared_stmt = session.prepare(insert_query)
        // session.execute(prepared_stmt, [
        // id,
        // content_chunk,
        // content_vector,
        // created_at
        // ])
    }

}