Fix Ollama Kotlin function callback integration tests

sobychacko · markpollack · commit 5e1f681da25e · 2024-11-08T21:24:28.000-05:00
- Add default values to KotlinRequest data class to fix Jackson deserialization
- Replace hardcoded model name with OllamaModel.LLAMA3_2 for consistency
- Extract common test setup into companion objects for Ollama tests
- Replace var with val using when expression in weather service
- Clean up imports and remove unused annotations
- Standardize property access using Kotlin conventions
diff --git a/spring-ai-spring-boot-autoconfigure/src/test/kotlin/org/springframework/ai/autoconfigure/ollama/tool/FunctionCallbackContextKotlinIT.kt b/spring-ai-spring-boot-autoconfigure/src/test/kotlin/org/springframework/ai/autoconfigure/ollama/tool/FunctionCallbackContextKotlinIT.kt
@@ -14,43 +14,51 @@
  * limitations under the License.
  */
 
-
 package org.springframework.ai.autoconfigure.ollama.tool
 
 import org.assertj.core.api.Assertions.assertThat
+import org.junit.jupiter.api.BeforeAll
 import org.junit.jupiter.api.Test
-import org.junit.jupiter.api.condition.DisabledIf
 import org.slf4j.LoggerFactory
+
 import org.springframework.ai.autoconfigure.ollama.BaseOllamaIT
 import org.springframework.ai.autoconfigure.ollama.OllamaAutoConfiguration
 import org.springframework.ai.chat.messages.UserMessage
 import org.springframework.ai.chat.prompt.Prompt
 import org.springframework.ai.model.function.FunctionCallingOptions
 import org.springframework.ai.ollama.OllamaChatModel
+import org.springframework.ai.ollama.api.OllamaModel
 import org.springframework.ai.ollama.api.OllamaOptions
 import org.springframework.boot.autoconfigure.AutoConfigurations
 import org.springframework.boot.test.context.runner.ApplicationContextRunner
 import org.springframework.context.annotation.Bean
 import org.springframework.context.annotation.Configuration
 import org.springframework.context.annotation.Description
-import org.testcontainers.junit.jupiter.Testcontainers
 
 class FunctionCallbackContextKotlinIT : BaseOllamaIT() {
 
-	private val logger = LoggerFactory.getLogger(FunctionCallbackContextKotlinIT::class.java)
+	companion object {
+
+		private val MODEL_NAME = OllamaModel.LLAMA3_2.getName();
+
+		@JvmStatic
+		@BeforeAll
+		fun beforeAll() {
+			initializeOllama(MODEL_NAME)
+		}
+	}
 
-	private val MODEL_NAME = "qwen2.5:3b"
+	private val logger = LoggerFactory.getLogger(FunctionCallbackContextKotlinIT::class.java)
 
-	val contextRunner = buildOllamaApiWithModel(MODEL_NAME).let { baseUrl ->
-		ApplicationContextRunner().withPropertyValues(
-			"spring.ai.ollama.baseUrl=$baseUrl",
+	private val contextRunner = ApplicationContextRunner()
+		.withPropertyValues(
+			"spring.ai.ollama.baseUrl=${getBaseUrl()}",
 			"spring.ai.ollama.chat.options.model=$MODEL_NAME",
 			"spring.ai.ollama.chat.options.temperature=0.5",
 			"spring.ai.ollama.chat.options.topK=10"
 		)
-			.withConfiguration(AutoConfigurations.of(OllamaAutoConfiguration::class.java))
-			.withUserConfiguration(Config::class.java)
-	}
+		.withConfiguration(AutoConfigurations.of(OllamaAutoConfiguration::class.java))
+		.withUserConfiguration(Config::class.java)
 
 	@Test
 	fun functionCallTest() {
@@ -98,18 +106,13 @@ class FunctionCallbackContextKotlinIT : BaseOllamaIT() {
 		@Bean
 		@Description("Find the weather conditions, forecasts, and temperatures for a location, like a city or state.")
 		open fun weatherInfo(): (KotlinRequest) -> KotlinResponse = { request ->
-			var temperature = 10.0
-			if (request.location.contains("Paris")) {
-				temperature = 15.0
+			val temperature = when {
+				request.location.contains("Paris") -> 15.0
+				request.location.contains("Tokyo") -> 10.0
+				request.location.contains("San Francisco") -> 30.0
+				else -> 10.0
 			}
-			else if (request.location.contains("Tokyo")) {
-				temperature = 10.0
-			}
-			else if (request.location.contains("San Francisco")) {
-				temperature = 30.0
-			}
-			KotlinResponse(temperature, 15.0, 20.0, 2.0, 53, 45, Unit.C);
+			KotlinResponse(temperature, 15.0, 20.0, 2.0, 53, 45, Unit.C)
 		}
-
 	}
 }
diff --git a/spring-ai-spring-boot-autoconfigure/src/test/kotlin/org/springframework/ai/autoconfigure/ollama/tool/FunctionCallbackWrapperKotlinIT.kt b/spring-ai-spring-boot-autoconfigure/src/test/kotlin/org/springframework/ai/autoconfigure/ollama/tool/FunctionCallbackWrapperKotlinIT.kt
@@ -14,13 +14,13 @@
  * limitations under the License.
  */
 
-
 package org.springframework.ai.autoconfigure.ollama.tool
 
 import org.assertj.core.api.Assertions.assertThat
+import org.junit.jupiter.api.BeforeAll
 import org.junit.jupiter.api.Test
-import org.junit.jupiter.api.condition.DisabledIf
 import org.slf4j.LoggerFactory
+
 import org.springframework.ai.autoconfigure.ollama.BaseOllamaIT
 import org.springframework.ai.autoconfigure.ollama.OllamaAutoConfiguration
 import org.springframework.ai.chat.messages.UserMessage
@@ -29,29 +29,37 @@ import org.springframework.ai.model.function.FunctionCallback
 import org.springframework.ai.model.function.FunctionCallbackWrapper
 import org.springframework.ai.model.function.FunctionCallingOptions
 import org.springframework.ai.ollama.OllamaChatModel
+import org.springframework.ai.ollama.api.OllamaModel
 import org.springframework.ai.ollama.api.OllamaOptions
 import org.springframework.boot.autoconfigure.AutoConfigurations
 import org.springframework.boot.test.context.runner.ApplicationContextRunner
 import org.springframework.context.annotation.Bean
 import org.springframework.context.annotation.Configuration
-import org.testcontainers.junit.jupiter.Testcontainers
 
 class FunctionCallbackWrapperKotlinIT : BaseOllamaIT() {
 
-	private val logger = LoggerFactory.getLogger(FunctionCallbackWrapperKotlinIT::class.java)
+	companion object {
 
-	private val MODEL_NAME = "qwen2.5:3b"
+		private val MODEL_NAME = OllamaModel.LLAMA3_2.getName();
 
-	val contextRunner = buildOllamaApiWithModel(MODEL_NAME).let { baseUrl ->
-		ApplicationContextRunner().withPropertyValues(
-			"spring.ai.ollama.baseUrl=$baseUrl",
+		@JvmStatic
+		@BeforeAll
+		fun beforeAll() {
+			initializeOllama(MODEL_NAME)
+		}
+	}
+
+	private val logger = LoggerFactory.getLogger(FunctionCallbackWrapperKotlinIT::class.java)
+
+	private val contextRunner = ApplicationContextRunner()
+		.withPropertyValues(
+			"spring.ai.ollama.baseUrl=${getBaseUrl()}",
 			"spring.ai.ollama.chat.options.model=$MODEL_NAME",
 			"spring.ai.ollama.chat.options.temperature=0.5",
 			"spring.ai.ollama.chat.options.topK=10"
 		)
-			.withConfiguration(AutoConfigurations.of(OllamaAutoConfiguration::class.java))
-			.withUserConfiguration(Config::class.java)
-	}
+		.withConfiguration(AutoConfigurations.of(OllamaAutoConfiguration::class.java))
+		.withUserConfiguration(Config::class.java)
 
 	@Test
 	fun functionCallTest() {
diff --git a/spring-ai-spring-boot-autoconfigure/src/test/kotlin/org/springframework/ai/autoconfigure/ollama/tool/MockKotlinWeatherService.kt b/spring-ai-spring-boot-autoconfigure/src/test/kotlin/org/springframework/ai/autoconfigure/ollama/tool/MockKotlinWeatherService.kt
@@ -60,12 +60,24 @@ enum class Unit(val unitName: String) {
  */
 @JsonInclude(Include.NON_NULL)
 @JsonClassDescription("Weather API request")
-data class KotlinRequest(@get:JsonProperty(required = true, value = "location") @get:JsonPropertyDescription("The city and state e.g. San Francisco, CA") val location: String,
-						 @get:JsonProperty(required = true, value = "lat") @get:JsonPropertyDescription("The city latitude") val lat: Double,
-						 @get:JsonProperty(required = true, value = "lon") @get:JsonPropertyDescription("The city longitude") val lon: Double,
-						 @get:JsonProperty(required = true, value = "unit") @get:JsonPropertyDescription("Temperature unit") val unit: Unit) {
+data class KotlinRequest(
+	@get:JsonProperty(required = true, value = "location")
+	@get:JsonPropertyDescription("The city and state e.g. San Francisco, CA")
+	val location: String = "",
+
+	@get:JsonProperty(required = true, value = "lat")
+	@get:JsonPropertyDescription("The city latitude")
+	val lat: Double = 0.0,
+
+	@get:JsonProperty(required = true, value = "lon")
+	@get:JsonPropertyDescription("The city longitude")
+	val lon: Double = 0.0,
+
+	@get:JsonProperty(required = true, value = "unit")
+	@get:JsonPropertyDescription("Temperature unit")
+	val unit: Unit = Unit.C
+)
 
-}
 
 /**
  * Weather Function response.
