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

Abstract API for AI model clients

Browse Source 
* An abstract API for AI model clients
 * Providing portable client request options while still allowing vendor specific options when required.  Implemented only for StabilityAI/OpenAI ImageClient
 * Support for text->image for openai and stabilityai.

  Partial fix for #27 :  Text To Image and Fixes #266 and Fixes #261
This commit is contained in:
Mark Pollack
2024-01-15 00:06:24 -05:00
committed by Christian Tzolov
parent 08fa0e393c
commit 243cef976c
162 changed files with 3691 additions and 630 deletions
Download Patch File Download Diff File Expand all files Collapse all files

45
spring-ai-core/src/test/java/org/springframework/ai/chat/ChatClientTests.java
View File

@@ -21,20 +21,13 @@ import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.ArgumentMatchers.isA;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.doCallRealMethod;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.verifyNoMoreInteractions;
import java.util.Collections;
import static org.mockito.Mockito.*;
import org.junit.jupiter.api.Test;
import org.mockito.Mockito;
import org.springframework.ai.prompt.Prompt;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.prompt.Prompt;
/**
* Unit Tests for {@link ChatClient}.
@@ -51,10 +44,23 @@ class ChatClientTests {
String responseMessage = "All your bases are belong to us";
ChatClient mockClient = Mockito.mock(ChatClient.class);
Generation generation = spy(new Generation(responseMessage));
ChatResponse response = spy(new ChatResponse(Collections.singletonList(generation)));
doCallRealMethod().when(mockClient).generate(anyString());
AssistantMessage mockAssistantMessage = Mockito.mock(AssistantMessage.class);
when(mockAssistantMessage.getContent()).thenReturn(responseMessage);
// Create a mock Generation
Generation generation = Mockito.mock(Generation.class);
when(generation.getOutput()).thenReturn(mockAssistantMessage);
// Create a mock ChatResponse with the mock Generation
ChatResponse response = Mockito.mock(ChatResponse.class);
when(response.getResult()).thenReturn(generation);
// Generation generation = spy(new Generation(responseMessage));
// ChatResponse response = spy(new
// ChatResponse(Collections.singletonList(generation)));
doCallRealMethod().when(mockClient).call(anyString());
doAnswer(invocationOnMock -> {
@@ -65,14 +71,15 @@ class ChatClientTests {
return response;
}).when(mockClient).generate(any(Prompt.class));
}).when(mockClient).call(any(Prompt.class));
assertThat(mockClient.generate(userMessage)).isEqualTo(responseMessage);
assertThat(mockClient.call(userMessage)).isEqualTo(responseMessage);
verify(mockClient, times(1)).generate(eq(userMessage));
verify(mockClient, times(1)).generate(isA(Prompt.class));
verify(response, times(1)).getGeneration();
verify(generation, times(1)).getContent();
verify(mockClient, times(1)).call(eq(userMessage));
verify(mockClient, times(1)).call(isA(Prompt.class));
verify(response, times(1)).getResult();
verify(generation, times(1)).getOutput();
verify(mockAssistantMessage, times(1)).getContent();
verifyNoMoreInteractions(mockClient, generation, response);
}

1
spring-ai-core/src/test/java/org/springframework/ai/embedding/EmbeddingUtilTests.java
View File

@@ -18,6 +18,7 @@ package org.springframework.ai.embedding;
import java.util.List;
import org.junit.jupiter.api.Disabled;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.junit.jupiter.params.ParameterizedTest;

3
spring-ai-core/src/test/java/org/springframework/ai/metadata/PromptMetadataTests.java
View File

@@ -22,7 +22,8 @@ import static org.mockito.Mockito.mock;
import org.junit.jupiter.api.Test;
import org.springframework.ai.metadata.PromptMetadata.PromptFilterMetadata;
import org.springframework.ai.chat.metadata.PromptMetadata;
import org.springframework.ai.chat.metadata.PromptMetadata.PromptFilterMetadata;
/**
* Unit Tests for {@link PromptMetadata}.

1
spring-ai-core/src/test/java/org/springframework/ai/metadata/UsageTests.java
View File

@@ -25,6 +25,7 @@ import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.verifyNoMoreInteractions;
import org.junit.jupiter.api.Test;
import org.springframework.ai.chat.metadata.Usage;
/**
* Unit Tests for {@link Usage}.

20
spring-ai-core/src/test/java/org/springframework/ai/prompt/PromptTemplateTest.java
View File

@@ -2,6 +2,7 @@ package org.springframework.ai.prompt;
import org.junit.jupiter.api.Disabled;
import org.junit.jupiter.api.Test;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.core.io.InputStreamResource;
import org.springframework.core.io.Resource;
@@ -18,17 +19,18 @@ public class PromptTemplateTest {
@Test
public void testRender() {
// Create a map with string keys and object values to serve as a model for testing
// Create a map with string keys and object values to serve as a generative for
// testing
Map<String, Object> model = new HashMap<>();
model.put("key1", "value1");
model.put("key2", true);
model.put("key3", 100);
// Create a simple template with placeholders for keys in the model
// Create a simple template with placeholders for keys in the generative
String template = "This is a {key1}, it is {key2}, and it costs {key3}";
PromptTemplate promptTemplate = new PromptTemplate(template, model);
// The expected result after rendering the template with the model
// The expected result after rendering the template with the generative
String expected = "This is a value1, it is true, and it costs 100";
String result = promptTemplate.render();
@@ -44,7 +46,8 @@ public class PromptTemplateTest {
@Disabled("Need to improve PromptTemplate to better handle Resource toString and tracking with 'dynamicModel' for underlying StringTemplate")
@Test
public void testRenderResource() throws Exception {
// Create a map with string keys and object values to serve as a model for testing
// Create a map with string keys and object values to serve as a generative for
// testing
Map<String, Object> model = new HashMap<>();
model.put("key1", "value1");
model.put("key2", true);
@@ -55,11 +58,11 @@ public class PromptTemplateTest {
model.put("key3", resource);
// Create a simple template with placeholders for keys in the model
// Create a simple template with placeholders for keys in the generative
String template = "{key1}, {key2}, {key3}";
PromptTemplate promptTemplate = new PromptTemplate(template, model);
// The expected result after rendering the template with the model
// The expected result after rendering the template with the generative
String expected = "value1, true, it costs 100";
String result = promptTemplate.render();
@@ -69,11 +72,12 @@ public class PromptTemplateTest {
@Test
public void testRenderFailure() {
// Create a map with string keys and object values to serve as a model for testing
// Create a map with string keys and object values to serve as a generative for
// testing
Map<String, Object> model = new HashMap<>();
model.put("key1", "value1");
// Create a simple template that includes a key not present in the model
// Create a simple template that includes a key not present in the generative
String template = "This is a {key2}!";
PromptTemplate promptTemplate = new PromptTemplate(template, model);

7
spring-ai-core/src/test/java/org/springframework/ai/prompt/PromptTests.java
View File

@@ -18,6 +18,9 @@ package org.springframework.ai.prompt;
import org.assertj.core.api.Assertions;
import org.junit.jupiter.api.Test;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.chat.prompt.SystemPromptTemplate;
import java.util.HashMap;
import java.util.Map;
@@ -54,7 +57,7 @@ class PromptTests {
// to have access to Messages
Prompt prompt = pt.create(model);
assertThat(prompt.getContents()).isNotNull();
assertThat(prompt.getMessages()).isNotEmpty().hasSize(1);
assertThat(prompt.getInstructions()).isNotEmpty().hasSize(1);
System.out.println(prompt.getContents());
String systemTemplate = "You are a helpful assistant that translates {input_language} to {output_language}.";
@@ -86,7 +89,7 @@ class PromptTests {
// ChatPromptTemplate chatPromptTemplate = new ChatPromptTemplate(systemPrompt,
// humanPrompt);
// Prompt chatPrompt chatPromptTemplate.create(model);
// Prompt chatPrompt chatPromptTemplate.create(generative);
}