import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.ollama.OllamaChatModel;
import org.junit.jupiter.api.Test;
import org.testcontainers.containers.GenericContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;

@Testcontainers
class OllamaChatModelTest {

    /**
     * The first time you run this test, it will download a Docker image with Ollama and a model.
     * It might take a few minutes.
     * <p>
     * This test uses modified Ollama Docker images, which already contain models inside them.
     * All images with pre-packaged models are available here: https://hub.docker.com/repositories/langchain4j
     * <p>
     * However, you are not restricted to these images.
     * You can run any model from https://ollama.ai/library by following these steps:
     * 1. Run "docker run -d -v ollama:/root/.ollama -p 11434:11434 --name ollama ollama/ollama"
     * 2. Run "docker exec -it ollama ollama run mistral" <- specify the desired model here
     */

    static String MODEL_NAME = "orca-mini"; // try "mistral", "llama2", "codellama", "phi" or "tinyllama"

    @Container
    static GenericContainer<?> ollama = new GenericContainer<>("langchain4j/ollama-" + MODEL_NAME + ":latest")
            .withExposedPorts(11434);

    @Test
    void simple_example() {

        ChatLanguageModel model = OllamaChatModel.builder()
                .baseUrl(baseUrl())
                .modelName(MODEL_NAME)
                .build();

        String answer = model.generate("Provide 3 short bullet points explaining why Java is awesome");

        System.out.println(answer);
    }

    @Test
    void json_output_example() {

        ChatLanguageModel model = OllamaChatModel.builder()
                .baseUrl(baseUrl())
                .modelName(MODEL_NAME)
                .format("json")
                .build();

        String json = model.generate("Give me a JSON with 2 fields: name and age of a John Doe, 42");

        System.out.println(json);
    }

    static String baseUrl() {
        return String.format("http://%s:%d", ollama.getHost(), ollama.getFirstMappedPort());
    }
}

dev.langchain4j.model.ollama.OllamaChatModel.builder
[((1404, 1529), 'dev.langchain4j.model.ollama.OllamaChatModel.builder'), ((1404, 1504), 'dev.langchain4j.model.ollama.OllamaChatModel.builder'), ((1404, 1465), 'dev.langchain4j.model.ollama.OllamaChatModel.builder'), ((1758, 1915), 'dev.langchain4j.model.ollama.OllamaChatModel.builder'), ((1758, 1890), 'dev.langchain4j.model.ollama.OllamaChatModel.builder'), ((1758, 1858), 'dev.langchain4j.model.ollama.OllamaChatModel.builder'), ((1758, 1819), 'dev.langchain4j.model.ollama.OllamaChatModel.builder')]