Move the tests into MoonshotChatModelFunctionCallingIT

spring-projects · Jan 2, 2025 · 4f4c2d0 · 4f4c2d0
1 parent b61309f
commit 4f4c2d0
Show file tree

Hide file tree

Showing 3 changed files with 77 additions and 177 deletions.
diff --git a/...s/spring-ai-moonshot/src/main/java/org/springframework/ai/moonshot/MoonshotChatModel.java b/...s/spring-ai-moonshot/src/main/java/org/springframework/ai/moonshot/MoonshotChatModel.java
@@ -345,10 +345,10 @@ private ChatResponseMetadata from(ChatCompletion result) {
 	private ChatResponseMetadata from(ChatCompletion result, Usage usage) {
 		Assert.notNull(result, "Moonshot ChatCompletionResult must not be null");
 		return ChatResponseMetadata.builder()
-			.withId(result.id() != null ? result.id() : "")
-			.withUsage(usage)
-			.withModel(result.model() != null ? result.model() : "")
-			.withKeyValue("created", result.created() != null ? result.created() : 0L)
+			.id(result.id() != null ? result.id() : "")
+			.usage(usage)
+			.model(result.model() != null ? result.model() : "")
+			.keyValue("created", result.created() != null ? result.created() : 0L)
 			.build();
 	}
 

diff --git a/...spring-ai-moonshot/src/test/java/org/springframework/ai/moonshot/MoonShotChatModelIT.java b/...spring-ai-moonshot/src/test/java/org/springframework/ai/moonshot/MoonShotChatModelIT.java
diff --git a/...rc/test/java/org/springframework/ai/moonshot/chat/MoonshotChatModelFunctionCallingIT.java b/...rc/test/java/org/springframework/ai/moonshot/chat/MoonshotChatModelFunctionCallingIT.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2023-2024 the original author or authors.
+ * Copyright 2023-2025 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -17,6 +17,7 @@
 package org.springframework.ai.moonshot.chat;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 import java.util.Objects;
 import java.util.stream.Collectors;
@@ -53,6 +54,33 @@ class MoonshotChatModelFunctionCallingIT {
 	@Autowired
 	ChatModel chatModel;
 
+	private static final MoonshotApi.FunctionTool FUNCTION_TOOL = new MoonshotApi.FunctionTool(
+			MoonshotApi.FunctionTool.Type.FUNCTION, new MoonshotApi.FunctionTool.Function(
+					"Get the weather in location. Return temperature in 30°F or 30°C format.", "getCurrentWeather", """
+							{
+								"type": "object",
+								"properties": {
+									"location": {
+										"type": "string",
+										"description": "The city and state e.g. San Francisco, CA"
+									},
+									"lat": {
+										"type": "number",
+										"description": "The city latitude"
+									},
+									"lon": {
+										"type": "number",
+										"description": "The city longitude"
+									},
+									"unit": {
+										"type": "string",
+										"enum": ["C", "F"]
+									}
+								},
+								"required": ["location", "lat", "lon", "unit"]
+							}
+							"""));
+
 	@Test
 	void functionCallTest() {
 
@@ -89,6 +117,7 @@ void streamFunctionCallTest() {
 			.functionCallbacks(List.of(FunctionCallback.builder()
 				.function("getCurrentWeather", new MockWeatherService())
 				.description("Get the weather in location")
+				.inputType(MockWeatherService.Request.class)
 				.build()))
 			.build();
 
@@ -108,4 +137,47 @@ void streamFunctionCallTest() {
 		assertThat(content).contains("30", "10", "15");
 	}
 
+	@Test
+	public void toolFunctionCallWithUsage() {
+		var promptOptions = MoonshotChatOptions.builder()
+			.model(MoonshotApi.ChatModel.MOONSHOT_V1_8K.getValue())
+			.tools(Arrays.asList(FUNCTION_TOOL))
+			.functionCallbacks(List.of(FunctionCallback.builder()
+				.function("getCurrentWeather", new MockWeatherService())
+				.description("Get the weather in location. Return temperature in 36°F or 36°C format.")
+				.inputType(MockWeatherService.Request.class)
+				.build()))
+			.build();
+		Prompt prompt = new Prompt("What's the weather like in San Francisco? Return the temperature in Celsius.",
+				promptOptions);
+
+		ChatResponse chatResponse = this.chatModel.call(prompt);
+		assertThat(chatResponse).isNotNull();
+		assertThat(chatResponse.getResult().getOutput());
+		assertThat(chatResponse.getResult().getOutput().getText()).contains("San Francisco");
+		assertThat(chatResponse.getResult().getOutput().getText()).contains("30.0");
+		assertThat(chatResponse.getMetadata().getUsage().getTotalTokens()).isLessThan(450).isGreaterThan(280);
+	}
+
+	@Test
+	public void testStreamFunctionCallUsage() {
+		var promptOptions = MoonshotChatOptions.builder()
+			.model(MoonshotApi.ChatModel.MOONSHOT_V1_8K.getValue())
+			.tools(Arrays.asList(FUNCTION_TOOL))
+			.functionCallbacks(List.of(FunctionCallback.builder()
+				.function("getCurrentWeather", new MockWeatherService())
+				.description("Get the weather in location. Return temperature in 36°F or 36°C format.")
+				.inputType(MockWeatherService.Request.class)
+				.build()))
+			.build();
+		Prompt prompt = new Prompt("What's the weather like in San Francisco? Return the temperature in Celsius.",
+				promptOptions);
+
+		ChatResponse chatResponse = this.chatModel.stream(prompt).blockLast();
+		assertThat(chatResponse).isNotNull();
+		assertThat(chatResponse.getMetadata()).isNotNull();
+		assertThat(chatResponse.getMetadata().getUsage()).isNotNull();
+		assertThat(chatResponse.getMetadata().getUsage().getTotalTokens()).isLessThan(450).isGreaterThan(280);
+	}
+
 }