io.weaviate.client.v1.classifications.api.Scheduler Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of client Show documentation
A java client for Weaviate Vector Search Engine
There is a newer version: 4.9.0
package io.weaviate.client.v1.classifications.api;

import io.weaviate.client.v1.classifications.model.Classification;
import io.weaviate.client.v1.classifications.model.ClassificationFilters;
import org.apache.commons.lang3.ObjectUtils;
import io.weaviate.client.Config;
import io.weaviate.client.base.BaseClient;
import io.weaviate.client.base.ClientResult;
import io.weaviate.client.base.Response;
import io.weaviate.client.base.Result;
import io.weaviate.client.base.http.HttpClient;
import io.weaviate.client.v1.filters.WhereFilter;

public class Scheduler extends BaseClient implements ClientResult {

  private String classificationType;
  private String className;
  private String[] classifyProperties;
  private String[] basedOnProperties;
  private WhereFilter sourceWhereFilter;
  private WhereFilter trainingSetWhereFilter;
  private WhereFilter targetWhereFilter;
  private Boolean waitForCompletion;
  private Object settings;

  private Getter getter;

  public Scheduler(HttpClient httpClient, Config config) {
    super(httpClient, config);
    this.getter = new Getter(httpClient, config);
    this.waitForCompletion = false;
  }

  public Scheduler withType(String classificationType) {
    this.classificationType = classificationType;
    return this;
  }

  public Scheduler withClassName(String className) {
    this.className = className;
    return this;
  }

  public Scheduler withClassifyProperties(String[] classifyProperties) {
    this.classifyProperties = classifyProperties;
    return this;
  }

  public Scheduler withBasedOnProperties(String[] basedOnProperties) {
    this.basedOnProperties = basedOnProperties;
    return this;
  }

  public Scheduler withSourceWhereFilter(WhereFilter whereFilter) {
    this.sourceWhereFilter = whereFilter;
    return this;
  }

  public Scheduler withTrainingSetWhereFilter(WhereFilter whereFilter) {
    this.trainingSetWhereFilter = whereFilter;
    return this;
  }

  public Scheduler withTargetWhereFilter(WhereFilter whereFilter) {
    this.targetWhereFilter = whereFilter;
    return this;
  }

  public Scheduler withSettings(Object settings) {
    this.settings = settings;
    return this;
  }

  public Scheduler withWaitForCompletion() {
    this.waitForCompletion = true;
    return this;
  }

  private Classification waitForCompletion(String id) throws InterruptedException {
    for (; ; ) {
      Result result = getter.withID(id).run();
      if (result == null || result.getResult() == null) {
        return null;
      }
      Classification runningClassification = result.getResult();
      if (runningClassification.getStatus() == "running") {
        Thread.sleep(2000);
      } else {
        return runningClassification;
      }
    }
  }

  @Override
  public Result run() {
    Classification config = Classification.builder()
      .basedOnProperties(basedOnProperties)
      .className(className)
      .classifyProperties(classifyProperties)
      .type(classificationType)
      .settings(settings)
      .filters(getClassificationFilters(sourceWhereFilter, targetWhereFilter, trainingSetWhereFilter))
      .build();
    Response resp = sendPostRequest("/classifications", config, Classification.class);
    if (resp.getStatusCode() == 201) {
      if (waitForCompletion) {
        try {
          Classification c = waitForCompletion(resp.getBody().getId());
          return new Result<>(resp.getStatusCode(), c, null);
        } catch (InterruptedException e) {
          return new Result<>(resp);
        }
      }
      return new Result<>(resp);
    }
    return new Result<>(resp);
  }

  private ClassificationFilters getClassificationFilters(WhereFilter sourceWhere, WhereFilter targetWhere, WhereFilter trainingSetWhere) {
    if (ObjectUtils.anyNotNull(sourceWhere, targetWhere, trainingSetWhere)) {
      return ClassificationFilters.builder()
        .sourceWhere(sourceWhere)
        .targetWhere(targetWhere)
        .trainingSetWhere(trainingSetWhere)
        .build();
    }
    return null;
  }
}