Fix failing OpenSearch IT and disable Mistral function calling IT

- PaymentStatusFunctionCallingIT in org.springframework.ai.mistralai.api.tool
  is failing.  Needs investigation.
  See https://github.com/spring-projects/spring-ai/issues/1853
- OpenSearchVectorStoreWithOllamaIT updated to pull model if not available.
This commit is contained in:
Mark Pollack
2024-12-02 16:45:02 -05:00
parent 771c4a3182
commit b405536bba
2 changed files with 24 additions and 0 deletions

View File

@@ -24,6 +24,7 @@ import java.util.function.Function;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.junit.jupiter.api.Disabled;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.slf4j.Logger;
@@ -52,6 +53,7 @@ import static org.assertj.core.api.Assertions.assertThat;
* @author Christian Tzolov
* @since 0.8.1
*/
@Disabled("See https://github.com/spring-projects/spring-ai/issues/1853")
@EnabledIfEnvironmentVariable(named = "MISTRAL_AI_API_KEY", matches = ".+")
public class PaymentStatusFunctionCallingIT {

View File

@@ -43,6 +43,9 @@ import org.springframework.ai.ollama.OllamaEmbeddingModel;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaModel;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.ollama.management.ModelManagementOptions;
import org.springframework.ai.ollama.management.OllamaModelManager;
import org.springframework.ai.ollama.management.PullModelStrategy;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
@@ -56,6 +59,12 @@ import static org.hamcrest.Matchers.hasSize;
@EnabledIfEnvironmentVariable(named = "OLLAMA_TESTS_ENABLED", matches = "true")
class OpenSearchVectorStoreWithOllamaIT {
private static final Duration DEFAULT_TIMEOUT = Duration.ofMinutes(10);
private static final int DEFAULT_MAX_RETRIES = 2;
private static final String OLLAMA_LOCAL_URL = "http://localhost:11434";
@Container
private static final OpensearchContainer<?> opensearchContainer = new OpensearchContainer<>(
OpenSearchImage.DEFAULT_IMAGE);
@@ -72,6 +81,19 @@ class OpenSearchVectorStoreWithOllamaIT {
Awaitility.setDefaultPollInterval(2, TimeUnit.SECONDS);
Awaitility.setDefaultPollDelay(Duration.ZERO);
Awaitility.setDefaultTimeout(Duration.ofMinutes(1));
// Ensure the model is pulled before running tests
ensureModelIsPresent(OllamaModel.MXBAI_EMBED_LARGE.getName());
}
private static void ensureModelIsPresent(final String model) {
final OllamaApi api = new OllamaApi(OLLAMA_LOCAL_URL);
final var modelManagementOptions = ModelManagementOptions.builder()
.withMaxRetries(DEFAULT_MAX_RETRIES)
.withTimeout(DEFAULT_TIMEOUT)
.build();
final var ollamaModelManager = new OllamaModelManager(api, modelManagementOptions);
ollamaModelManager.pullModel(model, PullModelStrategy.WHEN_MISSING);
}
private String getText(String uri) {