|
1 | 1 | package io.quarkiverse.langchain4j.watsonx.deployment; |
2 | 2 |
|
| 3 | +import static com.github.tomakehurst.wiremock.client.WireMock.equalTo; |
| 4 | +import static com.github.tomakehurst.wiremock.client.WireMock.matchingJsonPath; |
3 | 5 | import static org.junit.jupiter.api.Assertions.assertEquals; |
4 | 6 | import static org.junit.jupiter.api.Assertions.assertThrows; |
5 | 7 |
|
| 8 | +import java.util.Date; |
| 9 | + |
6 | 10 | import jakarta.inject.Inject; |
7 | 11 | import jakarta.inject.Singleton; |
8 | 12 |
|
|
11 | 15 | import org.junit.jupiter.api.Test; |
12 | 16 | import org.junit.jupiter.api.extension.RegisterExtension; |
13 | 17 |
|
| 18 | +import dev.langchain4j.service.SystemMessage; |
14 | 19 | import dev.langchain4j.service.UserMessage; |
15 | 20 | import dev.langchain4j.service.V; |
16 | 21 | import io.quarkiverse.langchain4j.RegisterAiService; |
17 | 22 | import io.quarkus.test.QuarkusUnitTest; |
18 | 23 |
|
19 | | -public class ResponseSchemaOffTest { |
| 24 | +public class ResponseSchemaOffTest extends WireMockAbstract { |
20 | 25 |
|
21 | 26 | @RegisterExtension |
22 | 27 | static QuarkusUnitTest unitTest = new QuarkusUnitTest() |
23 | 28 | .overrideRuntimeConfigKey("quarkus.langchain4j.watsonx.base-url", WireMockUtil.URL_WATSONX_SERVER) |
24 | 29 | .overrideRuntimeConfigKey("quarkus.langchain4j.watsonx.iam.base-url", WireMockUtil.URL_IAM_SERVER) |
25 | 30 | .overrideRuntimeConfigKey("quarkus.langchain4j.watsonx.api-key", WireMockUtil.API_KEY) |
26 | 31 | .overrideRuntimeConfigKey("quarkus.langchain4j.watsonx.project-id", WireMockUtil.PROJECT_ID) |
| 32 | + .overrideConfigKey("quarkus.langchain4j.watsonx.mode", "generation") |
27 | 33 | .overrideConfigKey("quarkus.langchain4j.response-schema", "false") |
28 | 34 | .setArchiveProducer(() -> ShrinkWrap.create(JavaArchive.class).addClass(WireMockUtil.class)); |
29 | 35 |
|
30 | | - @RegisterAiService |
| 36 | + @Override |
| 37 | + void handlerBeforeEach() { |
| 38 | + mockServers.mockIAMBuilder(200) |
| 39 | + .grantType(langchain4jWatsonConfig.defaultConfig().iam().grantType()) |
| 40 | + .response(WireMockUtil.BEARER_TOKEN, new Date()) |
| 41 | + .build(); |
| 42 | + } |
| 43 | + |
| 44 | + @RegisterAiService(chatMemoryProviderSupplier = RegisterAiService.NoChatMemoryProviderSupplier.class) |
31 | 45 | @Singleton |
32 | 46 | interface OnMethodAIService { |
33 | | - String poem(@UserMessage String message, @V("topic") String topic); |
| 47 | + String poem1(@UserMessage String message, @V("topic") String topic); |
| 48 | + |
| 49 | + Poem poem2(@UserMessage String message); |
| 50 | + |
| 51 | + @UserMessage("{message}") |
| 52 | + Poem poem3(String message); |
| 53 | + |
| 54 | + @SystemMessage("SystemMessage") |
| 55 | + @UserMessage("{message}") |
| 56 | + Poem poem4(String message); |
| 57 | + |
| 58 | + public record Poem(String text) { |
| 59 | + }; |
34 | 60 | } |
35 | 61 |
|
36 | 62 | @Inject |
37 | 63 | OnMethodAIService onMethodAIService; |
38 | 64 |
|
| 65 | + static String POEM_RESPONSE = """ |
| 66 | + { |
| 67 | + "model_id": "mistralai/mistral-large", |
| 68 | + "created_at": "2024-01-21T17:06:14.052Z", |
| 69 | + "results": [ |
| 70 | + { |
| 71 | + "generated_text": "{ \\\"text\\\": \\\"Poem\\\" }", |
| 72 | + "generated_token_count": 5, |
| 73 | + "input_token_count": 50, |
| 74 | + "stop_reason": "eos_token", |
| 75 | + "seed": 2123876088 |
| 76 | + } |
| 77 | + ] |
| 78 | + } |
| 79 | + """; |
| 80 | + |
39 | 81 | @Test |
40 | | - void on_method_ai_service() throws Exception { |
| 82 | + void test_poem_1() throws Exception { |
41 | 83 | var ex = assertThrows(RuntimeException.class, |
42 | | - () -> onMethodAIService.poem("{response_schema} Generate a poem about {topic}", "dog")); |
| 84 | + () -> onMethodAIService.poem1("{response_schema} Generate a poem about {topic}", "dog")); |
43 | 85 | assertEquals( |
44 | 86 | "The {response_schema} placeholder cannot be used if the property quarkus.langchain4j.response-schema is set to false. Found in: io.quarkiverse.langchain4j.watsonx.deployment.ResponseSchemaOffTest$OnMethodAIService", |
45 | 87 | ex.getMessage()); |
| 88 | + |
| 89 | + mockServers.mockWatsonxBuilder(WireMockUtil.URL_WATSONX_GENERATION_API, 200) |
| 90 | + .body(matchingJsonPath("$.input", equalTo("Generate a poem about dog"))) |
| 91 | + .response(WireMockUtil.RESPONSE_WATSONX_GENERATION_API) |
| 92 | + .build(); |
| 93 | + |
| 94 | + assertEquals("AI Response", onMethodAIService.poem1("Generate a poem about {topic}", "dog")); |
46 | 95 | } |
47 | 96 |
|
| 97 | + @Test |
| 98 | + void test_poem_2() { |
| 99 | + mockServers.mockWatsonxBuilder(WireMockUtil.URL_WATSONX_GENERATION_API, 200) |
| 100 | + .body(matchingJsonPath("$.input", equalTo("Generate a poem about dog"))) |
| 101 | + .response(POEM_RESPONSE) |
| 102 | + .build(); |
| 103 | + |
| 104 | + assertEquals("Poem", onMethodAIService.poem2("Generate a poem about dog").text); |
| 105 | + } |
| 106 | + |
| 107 | + @Test |
| 108 | + void test_poem_3() { |
| 109 | + mockServers.mockWatsonxBuilder(WireMockUtil.URL_WATSONX_GENERATION_API, 200) |
| 110 | + .body(matchingJsonPath("$.input", equalTo("Generate a poem about dog"))) |
| 111 | + .response(POEM_RESPONSE) |
| 112 | + .build(); |
| 113 | + |
| 114 | + assertEquals("Poem", onMethodAIService.poem3("Generate a poem about dog").text); |
| 115 | + } |
| 116 | + |
| 117 | + @Test |
| 118 | + void test_poem_4() { |
| 119 | + mockServers.mockWatsonxBuilder(WireMockUtil.URL_WATSONX_GENERATION_API, 200) |
| 120 | + .body(matchingJsonPath("$.input", equalTo("SystemMessage\nGenerate a poem about dog"))) |
| 121 | + .response(POEM_RESPONSE) |
| 122 | + .build(); |
| 123 | + |
| 124 | + assertEquals("Poem", onMethodAIService.poem4("Generate a poem about dog").text); |
| 125 | + } |
48 | 126 | } |
0 commit comments