Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

use AiInput/AiOuput in chain instead of raw hashmaps

Browse Source
This commit is contained in:
Mark Pollack
2023-08-18 17:40:52 -04:00
parent b229bf898e
commit 2b10843537
5 changed files with 60 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
spring-ai-core/src/main/java/org/springframework/ai/chain/AbstractChain.java
View File

@@ -26,27 +26,27 @@ public abstract class AbstractChain implements Chain {
// TODO validation of input/outputs
@Override
public Map<String, Object> apply(Map<String, Object> inputMap) {
Map<String, Object> inputMapToUse = preProcess(inputMap);
Map<String, Object> outputMap = doApply(inputMapToUse);
Map<String, Object> outputMapToUse = postProcess(inputMapToUse, outputMap);
return outputMapToUse;
public AiOutput apply(AiInput aiInput) {
AiInput aiInputToUse = preProcess(aiInput);
AiOutput aiOutput = doApply(aiInputToUse);
Map<String, Object> outputMapToUse = postProcess(aiInput, aiOutput);
return new AiOutput(outputMapToUse);
}
@Override
public Map<String, Object> preProcess(Map<String, Object> inputMap) {
validateInputs(inputMap);
return inputMap;
public AiInput preProcess(AiInput aiInput) {
validateInputs(aiInput.getInputData());
return aiInput;
}
protected abstract Map<String, Object> doApply(Map<String, Object> inputMap);
protected abstract AiOutput doApply(AiInput aiInput);
@Override
public Map<String, Object> postProcess(Map<String, Object> inputMap, Map<String, Object> outputMap) {
validateOutputs(outputMap);
public Map<String, Object> postProcess(AiInput aiInput, AiOutput aiOutput) {
validateOutputs(aiOutput.getOutputData());
Map<String, Object> combindedMap = new HashMap<>();
combindedMap.putAll(inputMap);
combindedMap.putAll(outputMap);
combindedMap.putAll(aiInput.getInputData());
combindedMap.putAll(aiOutput.getOutputData());
return combindedMap;
}

18
spring-ai-core/src/main/java/org/springframework/ai/chain/AiInput.java Normal file
View File

@@ -0,0 +1,18 @@
package org.springframework.ai.chain;
import java.util.List;
import java.util.Map;
public class AiInput {
private Map<String, Object> inputData;
public AiInput(Map<String, Object> inputData) {
this.inputData = inputData;
}
Map<String, Object> getInputData() {
return inputData;
}
}

18
spring-ai-core/src/main/java/org/springframework/ai/chain/AiOutput.java Normal file
View File

@@ -0,0 +1,18 @@
package org.springframework.ai.chain;
import java.util.List;
import java.util.Map;
public class AiOutput {
private final Map<String, Object> outputData;
public AiOutput(Map<String, Object> outputData) {
this.outputData = outputData;
}
Map<String, Object> getOutputData() {
return this.outputData;
}
}

6
spring-ai-core/src/main/java/org/springframework/ai/chain/Chain.java
View File

@@ -4,14 +4,14 @@ import java.util.List;
import java.util.Map;
import java.util.function.Function;
public interface Chain extends Function<Map<String, Object>, Map<String, Object>> {
public interface Chain extends Function<AiInput, AiOutput> {
List<String> getInputKeys();
List<String> getOutputKeys();
Map<String, Object> preProcess(Map<String, Object> inputMap);
AiInput preProcess(AiInput aiInput);
Map<String, Object> postProcess(Map<String, Object> inputMap, Map<String, Object> outputMap);
Map<String, Object> postProcess(AiInput aiInput, AiOutput aiOutput);
}

14
spring-ai-core/src/test/java/org/springframework/ai/chain/ChainTests.java
View File

@@ -14,15 +14,17 @@ class ChainTests {
@Test
void badInputs() {
Chain chain = new FakeChain();
assertThatThrownBy(() -> chain.apply(Map.of("foobar", "baz"))).isInstanceOf(IllegalArgumentException.class)
AiInput aiInput = new AiInput(Map.of("foobar", "baz"));
assertThatThrownBy(() -> chain.apply(aiInput)).isInstanceOf(IllegalArgumentException.class)
.hasMessageContaining("Missing some input keys");
}
@Test
void correctInputs() {
Chain chain = new FakeChain();
Map<String, Object> output = chain.apply(Map.of("foo", "bar"));
assertThat(output).containsEntry("foo", "bar").containsEntry("bar", "baz");
AiInput aiInput = new AiInput(Map.of("foo", "bar"));
AiOutput aiOutput = chain.apply(aiInput);
assertThat(aiOutput.getOutputData()).containsEntry("foo", "bar").containsEntry("bar", "baz");
}
class FakeChain extends AbstractChain {
@@ -53,12 +55,12 @@ class ChainTests {
}
@Override
protected Map<String, Object> doApply(Map<String, Object> inputMap) {
protected AiOutput doApply(AiInput aiInput) {
if (beCorrect) {
return Map.of("bar", "baz");
return new AiOutput(Map.of("bar", "baz"));
}
else {
return Map.of("baz", "bar");
return new AiOutput(Map.of("baz", "bar"));
}
}