Skip to content

Commit 1298a28

Browse files
jitokimilayaperumalg
authored andcommitted
add OllamaChatModelTests
Signed-off-by: jitokim <pigberger70@gmail.com>
1 parent 08b0c30 commit 1298a28

File tree

1 file changed

+80
-0
lines changed

1 file changed

+80
-0
lines changed
Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
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

Comments
 (0)