Change JsonOutputParser to MapOutputParser

spring-projects · Aug 19, 2023 · f2f81cd · f2f81cd
1 parent 8e276ca
commit f2f81cd
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 18 deletions.
diff --git a/...ore/src/main/java/org/springframework/ai/parser/AbstractMessageConverterOutputParser.java b/...ore/src/main/java/org/springframework/ai/parser/AbstractMessageConverterOutputParser.java
@@ -2,7 +2,7 @@
 
 import org.springframework.messaging.converter.MessageConverter;
 
-public abstract class AbstractMessageConverterOutputParser implements OutputParser<Object> {
+public abstract class AbstractMessageConverterOutputParser<T> implements OutputParser<T> {
 
 	private MessageConverter messageConverter;
 

diff --git a/...framework/ai/parser/JsonOutputParser.java → ...gframework/ai/parser/MapOutputParser.java b/...framework/ai/parser/JsonOutputParser.java → ...gframework/ai/parser/MapOutputParser.java
@@ -1,28 +1,27 @@
 package org.springframework.ai.parser;
 
+import com.fasterxml.jackson.databind.JsonNode;
 import org.springframework.messaging.Message;
 import org.springframework.messaging.converter.MappingJackson2MessageConverter;
-import org.springframework.messaging.converter.MessageConverter;
 import org.springframework.messaging.support.MessageBuilder;
 
 import java.nio.charset.StandardCharsets;
+import java.util.HashMap;
+import java.util.Map;
 
 /**
  * Uses Jackson
  */
-public class JsonOutputParser extends AbstractMessageConverterOutputParser {
+public class MapOutputParser extends AbstractMessageConverterOutputParser<Map<String, Object>> {
 
-	private Class dataType;
-
-	public JsonOutputParser(Class dataType) {
+	public MapOutputParser() {
 		super(new MappingJackson2MessageConverter());
-		this.dataType = dataType;
 	}
 
 	@Override
-	public Object parse(String text) {
+	public Map<String, Object> parse(String text) {
 		Message<?> message = MessageBuilder.withPayload(text.getBytes(StandardCharsets.UTF_8)).build();
-		return getMessageConverter().fromMessage(message, dataType);
+		return (Map) getMessageConverter().fromMessage(message, HashMap.class);
 	}
 
 	@Override
@@ -32,7 +31,7 @@ public String getFormat() {
 				The data structure for the JSON should match this Java class: %s
 				Do not include any explanations, only provide a RFC8259 compliant JSON response following this format without deviation.
 				 """;
-		return String.format(raw, dataType.getCanonicalName());
+		return String.format(raw, "java.util.HashMap");
 	}
 
 }
diff --git a/...-ai-openai/src/test/java/org/springframework/ai/openai/client/ClientIntegrationTests.java b/...-ai-openai/src/test/java/org/springframework/ai/openai/client/ClientIntegrationTests.java
@@ -4,8 +4,8 @@
 import org.springframework.ai.client.AiResponse;
 import org.springframework.ai.client.Generation;
 import org.springframework.ai.openai.testutils.AbstractIntegrationTest;
-import org.springframework.ai.parser.JsonOutputParser;
 import org.springframework.ai.parser.ListOutputParser;
+import org.springframework.ai.parser.MapOutputParser;
 import org.springframework.ai.prompt.Prompt;
 import org.springframework.ai.prompt.PromptTemplate;
 import org.springframework.ai.prompt.SystemPromptTemplate;
@@ -17,7 +17,6 @@
 import org.springframework.core.io.Resource;
 
 import java.util.Arrays;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
@@ -64,8 +63,8 @@ void outputParser() {
 	}
 
 	@Test
-	void jsonOutputParser() {
-		JsonOutputParser outputParser = new JsonOutputParser(HashMap.class);
+	void mapOutputParser() {
+		MapOutputParser outputParser = new MapOutputParser();
 
 		String format = outputParser.getFormat();
 		String template = """
@@ -77,10 +76,8 @@ void jsonOutputParser() {
 		Prompt prompt = new Prompt(promptTemplate.createMessage());
 		Generation generation = openAiClient.generate(prompt).getGeneration();
 
-		Object result = outputParser.parse(generation.getText());
-		System.out.println(result);
-		assertThat(result).isNotNull();
-		assertThat(((Map) result).get("numbers")).isEqualTo(Arrays.asList(1, 2, 3, 4, 5, 6, 7, 8, 9));
+		Map<String, Object> result = outputParser.parse(generation.getText());
+		assertThat(result.get("numbers")).isEqualTo(Arrays.asList(1, 2, 3, 4, 5, 6, 7, 8, 9));
 
 	}