|
17 | 17 | package org.springframework.ai.chat.cache.semantic;
|
18 | 18 |
|
19 | 19 | import com.redis.testcontainers.RedisStackContainer;
|
| 20 | +import io.micrometer.observation.ObservationRegistry; |
| 21 | +import io.micrometer.observation.tck.TestObservationRegistry; |
20 | 22 | import org.junit.jupiter.api.AfterEach;
|
21 | 23 | import org.junit.jupiter.api.Test;
|
22 | 24 | import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
|
|
26 | 28 | import org.springframework.ai.chat.model.ChatResponse;
|
27 | 29 | import org.springframework.ai.chat.model.Generation;
|
28 | 30 | import org.springframework.ai.embedding.EmbeddingModel;
|
| 31 | +import org.springframework.ai.model.tool.ToolCallingManager; |
29 | 32 | import org.springframework.ai.openai.OpenAiChatModel;
|
30 | 33 | import org.springframework.ai.openai.OpenAiChatOptions;
|
31 | 34 | import org.springframework.ai.openai.OpenAiEmbeddingModel;
|
|
42 | 45 | import org.springframework.boot.test.context.runner.ApplicationContextRunner;
|
43 | 46 | import org.springframework.context.annotation.Bean;
|
44 | 47 | import org.springframework.data.redis.connection.jedis.JedisConnectionFactory;
|
| 48 | +import org.springframework.retry.support.RetryTemplate; |
| 49 | + |
45 | 50 | import org.testcontainers.junit.jupiter.Container;
|
46 | 51 | import org.testcontainers.junit.jupiter.Testcontainers;
|
47 | 52 | import redis.clients.jedis.JedisPooled;
|
@@ -207,18 +212,24 @@ public SemanticCache semanticCache(EmbeddingModel embeddingModel,
|
207 | 212 |
|
208 | 213 | @Bean(name = "openAiEmbeddingModel")
|
209 | 214 | public EmbeddingModel embeddingModel() {
|
210 |
| - return new OpenAiEmbeddingModel(new OpenAiApi(System.getenv("OPENAI_API_KEY"))); |
| 215 | + return new OpenAiEmbeddingModel(OpenAiApi.builder().apiKey(System.getenv("OPENAI_API_KEY")).build()); |
| 216 | + } |
| 217 | + |
| 218 | + @Bean |
| 219 | + public TestObservationRegistry observationRegistry() { |
| 220 | + return TestObservationRegistry.create(); |
211 | 221 | }
|
212 | 222 |
|
213 | 223 | @Bean(name = "openAiChatModel")
|
214 |
| - public OpenAiChatModel openAiChatModel() { |
215 |
| - var openAiApi = new OpenAiApi(System.getenv("OPENAI_API_KEY")); |
| 224 | + public OpenAiChatModel openAiChatModel(ObservationRegistry observationRegistry) { |
| 225 | + var openAiApi = OpenAiApi.builder().apiKey(System.getenv("OPENAI_API_KEY")).build(); |
216 | 226 | var openAiChatOptions = OpenAiChatOptions.builder()
|
217 | 227 | .model("gpt-3.5-turbo")
|
218 | 228 | .temperature(0.4)
|
219 | 229 | .maxTokens(200)
|
220 | 230 | .build();
|
221 |
| - return new OpenAiChatModel(openAiApi, openAiChatOptions); |
| 231 | + return new OpenAiChatModel(openAiApi, openAiChatOptions, ToolCallingManager.builder().build(), |
| 232 | + RetryTemplate.defaultInstance(), observationRegistry); |
222 | 233 | }
|
223 | 234 |
|
224 | 235 | }
|
|
0 commit comments