| 
 | 1 | +/*  | 
 | 2 | + * Copyright 2023-2024 the original author or authors.  | 
 | 3 | + *  | 
 | 4 | + * Licensed under the Apache License, Version 2.0 (the "License");  | 
 | 5 | + * you may not use this file except in compliance with the License.  | 
 | 6 | + * You may obtain a copy of the License at  | 
 | 7 | + *  | 
 | 8 | + *      https://www.apache.org/licenses/LICENSE-2.0  | 
 | 9 | + *  | 
 | 10 | + * Unless required by applicable law or agreed to in writing, software  | 
 | 11 | + * distributed under the License is distributed on an "AS IS" BASIS,  | 
 | 12 | + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  | 
 | 13 | + * See the License for the specific language governing permissions and  | 
 | 14 | + * limitations under the License.  | 
 | 15 | + */  | 
 | 16 | + | 
 | 17 | +package org.springframework.ai.ollama;  | 
 | 18 | + | 
 | 19 | +import static org.junit.jupiter.api.Assertions.assertEquals;  | 
 | 20 | +import static org.junit.jupiter.api.Assertions.assertThrows;  | 
 | 21 | +import static org.mockito.BDDMockito.given;  | 
 | 22 | + | 
 | 23 | +import java.time.Duration;  | 
 | 24 | + | 
 | 25 | +import org.junit.jupiter.api.Test;  | 
 | 26 | +import org.junit.jupiter.api.extension.ExtendWith;  | 
 | 27 | +import org.mockito.Mock;  | 
 | 28 | +import org.mockito.junit.jupiter.MockitoExtension;  | 
 | 29 | + | 
 | 30 | +import org.springframework.ai.chat.metadata.ChatResponseMetadata;  | 
 | 31 | +import org.springframework.ai.ollama.api.OllamaApi;  | 
 | 32 | +import org.springframework.ai.ollama.api.OllamaModel;  | 
 | 33 | +import org.springframework.ai.ollama.api.OllamaOptions;  | 
 | 34 | + | 
 | 35 | +/**  | 
 | 36 | + * @author Jihoon Kim  | 
 | 37 | + * @since 1.0.0  | 
 | 38 | + */  | 
 | 39 | +@ExtendWith(MockitoExtension.class)  | 
 | 40 | +public class OllamaChatModelTests {  | 
 | 41 | + | 
 | 42 | +	@Mock  | 
 | 43 | +	OllamaApi ollamaApi;  | 
 | 44 | + | 
 | 45 | +	@Mock  | 
 | 46 | +	OllamaApi.ChatResponse response;  | 
 | 47 | + | 
 | 48 | +	@Test  | 
 | 49 | +	public void buildOllamaChatModel() {  | 
 | 50 | +		Exception exception = assertThrows(IllegalArgumentException.class,  | 
 | 51 | +				() -> OllamaChatModel.builder()  | 
 | 52 | +					.withOllamaApi(ollamaApi)  | 
 | 53 | +					.withDefaultOptions(OllamaOptions.create().withModel(OllamaModel.LLAMA2))  | 
 | 54 | +					.withModelManagementOptions(null)  | 
 | 55 | +					.build());  | 
 | 56 | +		assertEquals("modelManagementOptions must not be null", exception.getMessage());  | 
 | 57 | +	}  | 
 | 58 | + | 
 | 59 | +	@Test  | 
 | 60 | +	public void buildChatResponseMetadata() {  | 
 | 61 | +		Duration evalDuration = Duration.ofSeconds(1);  | 
 | 62 | +		Integer evalCount = 101;  | 
 | 63 | + | 
 | 64 | +		Duration promptEvalDuration = Duration.ofSeconds(8);  | 
 | 65 | +		Integer promptEvalCount = 808;  | 
 | 66 | + | 
 | 67 | +		given(response.evalDuration()).willReturn(evalDuration);  | 
 | 68 | +		given(response.evalCount()).willReturn(evalCount);  | 
 | 69 | +		given(response.promptEvalDuration()).willReturn(promptEvalDuration);  | 
 | 70 | +		given(response.promptEvalCount()).willReturn(promptEvalCount);  | 
 | 71 | + | 
 | 72 | +		ChatResponseMetadata metadata = OllamaChatModel.from(response);  | 
 | 73 | + | 
 | 74 | +		assertEquals(evalDuration, metadata.get("eval-duration"));  | 
 | 75 | +		assertEquals(evalCount, metadata.get("eval-count"));  | 
 | 76 | +		assertEquals(promptEvalDuration, metadata.get("prompt-eval-duration"));  | 
 | 77 | +		assertEquals(promptEvalCount, metadata.get("prompt-eval-count"));  | 
 | 78 | +	}  | 
 | 79 | + | 
 | 80 | +}  | 
0 commit comments