spring-ai/ollama-local/src/test/java/com/ankurm/ollamalocal/LocalChatAndEmbeddingTest.java

package com.ankurm.ollamalocal;

import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.time.Duration;

import com.ankurm.ollamalocal.support.Transcript;
import org.junit.jupiter.api.MethodOrderer;
import org.junit.jupiter.api.Order;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.TestMethodOrder;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.testcontainers.ollama.OllamaContainer;
import org.testcontainers.utility.DockerImageName;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.metadata.ChatResponseMetadata;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.ollama.api.OllamaChatOptions;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.test.context.DynamicPropertyRegistry;
import org.springframework.test.context.DynamicPropertySource;

import static org.assertj.core.api.Assertions.assertThat;

/**
 * One real, local Ollama server for the whole class, started from a baked image that already
 * has {@code qwen2.5:0.5b} (chat) and {@code all-minilm} (embeddings) pulled -- see
 * {@code scripts/bake-image.sh}. No network call happens during these tests: the container
 * starts from a local image and both models are already on disk inside it.
 *
 * <p>Every other module in this series drives {@code ChatClient} against a hand-written
 * {@code ScriptedChatModel} so the tests are deterministic and need no live model. This module
 * is the deliberate exception: the whole point of "run it locally" is a real model answering a
 * real prompt, so these tests assert only what is genuinely deterministic about a live small
 * model (a non-blank response, a confirmed-empty model registry right after an explicit unload)
 * and leave exact wording, and exact timings, alone.
 */
@Testcontainers
@SpringBootTest
@TestMethodOrder(MethodOrderer.OrderAnnotation.class)
class LocalChatAndEmbeddingTest {

	@Container
	static final OllamaContainer OLLAMA = new OllamaContainer(
			DockerImageName.parse("ollama-baked-qwen05b-minilm:local").asCompatibleSubstituteFor("ollama/ollama"));

	@DynamicPropertySource
	static void ollamaProperties(DynamicPropertyRegistry registry) {
		registry.add("spring.ai.ollama.base-url", OLLAMA::getEndpoint);
	}

	private static final HttpClient HTTP = HttpClient.newHttpClient();

	@Autowired
	private ChatClient chatClient;

	@Autowired
	private EmbeddingModel embeddingModel;

	@Test
	@Order(1)
	void reloadAfterAConfirmedUnload() throws Exception {
		// keepAlive("0") tells Ollama to unload the model from memory as soon as this call
		// finishes -- the same knob you'd reach for in production to free RAM/VRAM between
		// bursts of traffic. Passed through ChatClient's per-call options() exactly like any
		// other OllamaChatOptions field.
		this.chatClient.prompt()
			.user("OK")
			.options(OllamaChatOptions.builder().model("qwen2.5:0.5b").keepAlive("0"))
			.call()
			.content();

		// Don't take the unload on faith -- ask Ollama's own /api/ps, which lists every
		// currently loaded model, and confirm the registry is really empty before measuring
		// what a reload costs.
		String modelsAfterUnload = getModelRegistry();
		assertThat(modelsAfterUnload).contains("\"models\":[]");

		ChatClient.CallResponseSpec response = this.chatClient.prompt()
			.user("Reply with a single short sentence: why do developers like small local models?")
			.call();

		String content = response.content();
		ChatResponseMetadata metadata = response.chatResponse().getMetadata();

		assertThat(content).isNotBlank();
		// Loading is real work -- reading weights off disk and initializing the runtime -- so
		// even a fast reload can't be negative or missing.
		Duration loadDuration = metadata.get("load-duration");
		assertThat(loadDuration).isNotNull();
		assertThat(loadDuration.isNegative()).isFalse();

		Transcript.write("01-chat-reload-after-unload",
				"/api/ps immediately after the unload call: " + modelsAfterUnload + "\n\n"
						+ "prompt: \"Reply with a single short sentence: why do developers like small local models?\"\n"
						+ "response: " + content + "\n\n" + formatOllamaMetadata(metadata));
	}

	@Test
	@Order(2)
	void backToBackCallReusesTheAlreadyLoadedModel() {
		ChatClient.CallResponseSpec response = this.chatClient.prompt()
			.user("Reply with a single short sentence: what is Testcontainers for?")
			.call();

		String content = response.content();
		ChatResponseMetadata metadata = response.chatResponse().getMetadata();

		assertThat(content).isNotBlank();

		Transcript.write("02-chat-back-to-back-call",
				"prompt: \"Reply with a single short sentence: what is Testcontainers for?\"\n"
						+ "response: " + content + "\n\n" + formatOllamaMetadata(metadata));
	}

	@Test
	@Order(3)
	void embeddingsAreDeterministicallySized() {
		float[] vector = this.embeddingModel.embed("Spring AI can run entirely against a local Ollama server.");

		assertThat(vector).hasSize(384);
		assertThat(this.embeddingModel.dimensions()).isEqualTo(384);

		StringBuilder prefix = new StringBuilder();
		for (int i = 0; i < 8; i++) {
			if (i > 0) {
				prefix.append(", ");
			}
			prefix.append(String.format("%.5f", vector[i]));
		}

		Transcript.write("03-embedding-dimensions",
				"model: all-minilm\n" + "dimensions(): " + this.embeddingModel.dimensions() + "\n"
						+ "vector.length: " + vector.length + "\n" + "first 8 values: [" + prefix + ", ...]");
	}

	private static String getModelRegistry() throws Exception {
		HttpRequest request = HttpRequest.newBuilder(URI.create(OLLAMA.getEndpoint() + "/api/ps")).GET().build();
		return HTTP.send(request, HttpResponse.BodyHandlers.ofString()).body();
	}

	private static String formatOllamaMetadata(ChatResponseMetadata metadata) {
		Duration total = metadata.get("total-duration");
		Duration load = metadata.get("load-duration");
		Duration promptEval = metadata.get("prompt-eval-duration");
		Duration eval = metadata.get("eval-duration");
		Integer promptEvalCount = metadata.get("prompt-eval-count");
		Integer evalCount = metadata.get("eval-count");

		StringBuilder sb = new StringBuilder();
		sb.append("total-duration: ").append(total.toMillis()).append("ms\n");
		sb.append("load-duration: ").append(load.toMillis()).append("ms\n");
		sb.append("prompt-eval-count: ").append(promptEvalCount);
		sb.append(", prompt-eval-duration: ").append(promptEval.toMillis()).append("ms\n");
		sb.append("eval-count: ").append(evalCount);
		sb.append(", eval-duration: ").append(eval.toMillis()).append("ms");

		if (evalCount != null && eval != null && eval.toNanos() > 0) {
			double tokensPerSecond = evalCount / (eval.toNanos() / 1_000_000_000.0);
			sb.append(String.format("%n%.2f tokens/sec (eval-count / eval-duration)", tokensPerSecond));
		}
		return sb.toString();
	}

}