chores: Fix checkstyle errors

This commit is contained in:
Ilayaperumal Gopinathan
2024-11-25 23:13:20 +00:00
parent e7f37a06b8
commit a782486b06
12 changed files with 307 additions and 303 deletions

View File

@@ -301,18 +301,18 @@ public class AnthropicApi {
public enum Role {
// @formatter:off
/**
* The user role.
/**
* The user role.
*/
@JsonProperty("user")
USER,
@JsonProperty("user")
USER,
/**
* The assistant role.
*/
@JsonProperty("assistant")
ASSISTANT
// @formatter:on
/**
* The assistant role.
*/
@JsonProperty("assistant")
ASSISTANT
// @formatter:on
}
@@ -441,18 +441,18 @@ public class AnthropicApi {
@JsonInclude(Include.NON_NULL)
public record ChatCompletionRequest(
// @formatter:off
@JsonProperty("model") String model,
@JsonProperty("messages") List<AnthropicMessage> messages,
@JsonProperty("system") String system,
@JsonProperty("max_tokens") Integer maxTokens,
@JsonProperty("metadata") Metadata metadata,
@JsonProperty("stop_sequences") List<String> stopSequences,
@JsonProperty("stream") Boolean stream,
@JsonProperty("temperature") Double temperature,
@JsonProperty("top_p") Double topP,
@JsonProperty("top_k") Integer topK,
@JsonProperty("tools") List<Tool> tools) {
// @formatter:on
@JsonProperty("model") String model,
@JsonProperty("messages") List<AnthropicMessage> messages,
@JsonProperty("system") String system,
@JsonProperty("max_tokens") Integer maxTokens,
@JsonProperty("metadata") Metadata metadata,
@JsonProperty("stop_sequences") List<String> stopSequences,
@JsonProperty("stream") Boolean stream,
@JsonProperty("temperature") Double temperature,
@JsonProperty("top_p") Double topP,
@JsonProperty("top_k") Integer topK,
@JsonProperty("tools") List<Tool> tools) {
// @formatter:on
public ChatCompletionRequest(String model, List<AnthropicMessage> messages, String system, Integer maxTokens,
Double temperature, Boolean stream) {
@@ -620,9 +620,9 @@ public class AnthropicApi {
@JsonInclude(Include.NON_NULL)
public record AnthropicMessage(
// @formatter:off
@JsonProperty("content") List<ContentBlock> content,
@JsonProperty("role") Role role) {
// @formatter:on
@JsonProperty("content") List<ContentBlock> content,
@JsonProperty("role") Role role) {
// @formatter:on
}
/**
@@ -644,23 +644,23 @@ public class AnthropicApi {
@JsonInclude(Include.NON_NULL)
public record ContentBlock(
// @formatter:off
@JsonProperty("type") Type type,
@JsonProperty("source") Source source,
@JsonProperty("text") String text,
@JsonProperty("type") Type type,
@JsonProperty("source") Source source,
@JsonProperty("text") String text,
// applicable only for streaming responses.
@JsonProperty("index") Integer index,
// applicable only for streaming responses.
@JsonProperty("index") Integer index,
// tool_use response only
@JsonProperty("id") String id,
@JsonProperty("name") String name,
@JsonProperty("input") Map<String, Object> input,
// tool_use response only
@JsonProperty("id") String id,
@JsonProperty("name") String name,
@JsonProperty("input") Map<String, Object> input,
// tool_result response only
@JsonProperty("tool_use_id") String toolUseId,
@JsonProperty("content") String content
) {
// @formatter:on
// tool_result response only
@JsonProperty("tool_use_id") String toolUseId,
@JsonProperty("content") String content
) {
// @formatter:on
/**
* Create content block
@@ -805,10 +805,10 @@ public class AnthropicApi {
@JsonInclude(Include.NON_NULL)
public record Source(
// @formatter:off
@JsonProperty("type") String type,
@JsonProperty("media_type") String mediaType,
@JsonProperty("data") String data) {
// @formatter:on
@JsonProperty("type") String type,
@JsonProperty("media_type") String mediaType,
@JsonProperty("data") String data) {
// @formatter:on
/**
* Create source
@@ -837,10 +837,10 @@ public class AnthropicApi {
@JsonInclude(Include.NON_NULL)
public record Tool(
// @formatter:off
@JsonProperty("name") String name,
@JsonProperty("description") String description,
@JsonProperty("input_schema") Map<String, Object> inputSchema) {
// @formatter:on
@JsonProperty("name") String name,
@JsonProperty("description") String description,
@JsonProperty("input_schema") Map<String, Object> inputSchema) {
// @formatter:on
}
// CB START EVENT
@@ -863,15 +863,15 @@ public class AnthropicApi {
@JsonInclude(Include.NON_NULL)
public record ChatCompletionResponse(
// @formatter:off
@JsonProperty("id") String id,
@JsonProperty("type") String type,
@JsonProperty("role") Role role,
@JsonProperty("content") List<ContentBlock> content,
@JsonProperty("model") String model,
@JsonProperty("stop_reason") String stopReason,
@JsonProperty("stop_sequence") String stopSequence,
@JsonProperty("usage") Usage usage) {
// @formatter:on
@JsonProperty("id") String id,
@JsonProperty("type") String type,
@JsonProperty("role") Role role,
@JsonProperty("content") List<ContentBlock> content,
@JsonProperty("model") String model,
@JsonProperty("stop_reason") String stopReason,
@JsonProperty("stop_sequence") String stopSequence,
@JsonProperty("usage") Usage usage) {
// @formatter:on
}
// CB DELTA EVENT
@@ -885,88 +885,88 @@ public class AnthropicApi {
@JsonInclude(Include.NON_NULL)
public record Usage(
// @formatter:off
@JsonProperty("input_tokens") Integer inputTokens,
@JsonProperty("output_tokens") Integer outputTokens) {
// @formatter:off
}
@JsonProperty("input_tokens") Integer inputTokens,
@JsonProperty("output_tokens") Integer outputTokens) {
// @formatter:off
}
/// ECB STOP
/**
* Special event used to aggregate multiple tool use events into a single event with
* list of aggregated ContentBlockToolUse.
*/
public static class ToolUseAggregationEvent implements StreamEvent {
/**
* Special event used to aggregate multiple tool use events into a single event with
* list of aggregated ContentBlockToolUse.
*/
public static class ToolUseAggregationEvent implements StreamEvent {
private Integer index;
private Integer index;
private String id;
private String id;
private String name;
private String name;
private String partialJson = "";
private String partialJson = "";
private List<ContentBlockStartEvent.ContentBlockToolUse> toolContentBlocks = new ArrayList<>();
private List<ContentBlockStartEvent.ContentBlockToolUse> toolContentBlocks = new ArrayList<>();
@Override
public EventType type() {
return EventType.TOOL_USE_AGGREGATE;
}
@Override
public EventType type() {
return EventType.TOOL_USE_AGGREGATE;
}
/**
/**
* Get tool content blocks.
* @return The tool content blocks.
*/
public List<ContentBlockStartEvent.ContentBlockToolUse> getToolContentBlocks() {
return this.toolContentBlocks;
}
*/
public List<ContentBlockStartEvent.ContentBlockToolUse> getToolContentBlocks() {
return this.toolContentBlocks;
}
/**
/**
* Check if the event is empty.
* @return True if the event is empty, false otherwise.
*/
public boolean isEmpty() {
return (this.index == null || this.id == null || this.name == null
|| !StringUtils.hasText(this.partialJson));
}
*/
public boolean isEmpty() {
return (this.index == null || this.id == null || this.name == null
|| !StringUtils.hasText(this.partialJson));
}
ToolUseAggregationEvent withIndex(Integer index) {
this.index = index;
return this;
}
ToolUseAggregationEvent withIndex(Integer index) {
this.index = index;
return this;
}
ToolUseAggregationEvent withId(String id) {
this.id = id;
return this;
}
ToolUseAggregationEvent withId(String id) {
this.id = id;
return this;
}
ToolUseAggregationEvent withName(String name) {
this.name = name;
return this;
}
ToolUseAggregationEvent withName(String name) {
this.name = name;
return this;
}
ToolUseAggregationEvent appendPartialJson(String partialJson) {
this.partialJson = this.partialJson + partialJson;
return this;
}
ToolUseAggregationEvent appendPartialJson(String partialJson) {
this.partialJson = this.partialJson + partialJson;
return this;
}
void squashIntoContentBlock() {
Map<String, Object> map = (StringUtils.hasText(this.partialJson))
? ModelOptionsUtils.jsonToMap(this.partialJson) : Map.of();
this.toolContentBlocks.add(new ContentBlockStartEvent.ContentBlockToolUse("tool_use", this.id, this.name, map));
this.index = null;
this.id = null;
this.name = null;
this.partialJson = "";
}
void squashIntoContentBlock() {
Map<String, Object> map = (StringUtils.hasText(this.partialJson))
? ModelOptionsUtils.jsonToMap(this.partialJson) : Map.of();
this.toolContentBlocks.add(new ContentBlockStartEvent.ContentBlockToolUse("tool_use", this.id, this.name, map));
this.index = null;
this.id = null;
this.name = null;
this.partialJson = "";
}
@Override
public String toString() {
return "EventToolUseBuilder [index=" + this.index + ", id=" + this.id + ", name=" + this.name + ", partialJson="
+ this.partialJson + ", toolUseMap=" + this.toolContentBlocks + "]";
}
@Override
public String toString() {
return "EventToolUseBuilder [index=" + this.index + ", id=" + this.id + ", name=" + this.name + ", partialJson="
+ this.partialJson + ", toolUseMap=" + this.toolContentBlocks + "]";
}
}
}
///////////////////////////////////////
/// MESSAGE EVENTS
@@ -974,54 +974,54 @@ public class AnthropicApi {
// MESSAGE START EVENT
/**
* Content block start event.
* @param type The event type.
* @param index The index of the content block.
* @param contentBlock The content block body.
*/
@JsonInclude(Include.NON_NULL)
public record ContentBlockStartEvent(
// @formatter:off
@JsonProperty("type") EventType type,
@JsonProperty("index") Integer index,
@JsonProperty("content_block") ContentBlockBody contentBlock) implements StreamEvent {
/**
* Content block start event.
* @param type The event type.
* @param index The index of the content block.
* @param contentBlock The content block body.
*/
@JsonInclude(Include.NON_NULL)
public record ContentBlockStartEvent(
// @formatter:off
@JsonProperty("type") EventType type,
@JsonProperty("index") Integer index,
@JsonProperty("content_block") ContentBlockBody contentBlock) implements StreamEvent {
@JsonTypeInfo(use = JsonTypeInfo.Id.NAME, include = JsonTypeInfo.As.EXISTING_PROPERTY, property = "type",
visible = true)
@JsonSubTypes({ @JsonSubTypes.Type(value = ContentBlockToolUse.class, name = "tool_use"),
@JsonSubTypes.Type(value = ContentBlockText.class, name = "text") })
public interface ContentBlockBody {
String type();
}
@JsonTypeInfo(use = JsonTypeInfo.Id.NAME, include = JsonTypeInfo.As.EXISTING_PROPERTY, property = "type",
visible = true)
@JsonSubTypes({ @JsonSubTypes.Type(value = ContentBlockToolUse.class, name = "tool_use"),
@JsonSubTypes.Type(value = ContentBlockText.class, name = "text") })
public interface ContentBlockBody {
String type();
}
/**
/**
* Tool use content block.
* @param type The content block type.
* @param id The tool use id.
* @param name The tool use name.
* @param input The tool use input.
*/
@JsonInclude(Include.NON_NULL)
public record ContentBlockToolUse(
@JsonProperty("type") String type,
@JsonProperty("id") String id,
@JsonProperty("name") String name,
@JsonProperty("input") Map<String, Object> input) implements ContentBlockBody {
}
*/
@JsonInclude(Include.NON_NULL)
public record ContentBlockToolUse(
@JsonProperty("type") String type,
@JsonProperty("id") String id,
@JsonProperty("name") String name,
@JsonProperty("input") Map<String, Object> input) implements ContentBlockBody {
}
/**
/**
* Text content block.
* @param type The content block type.
* @param text The text content.
*/
@JsonInclude(Include.NON_NULL)
public record ContentBlockText(
@JsonProperty("type") String type,
@JsonProperty("text") String text) implements ContentBlockBody {
}
}
// @formatter:on
*/
@JsonInclude(Include.NON_NULL)
public record ContentBlockText(
@JsonProperty("type") String type,
@JsonProperty("text") String text) implements ContentBlockBody {
}
}
// @formatter:on
// MESSAGE DELTA EVENT
@@ -1035,41 +1035,41 @@ public class AnthropicApi {
@JsonInclude(Include.NON_NULL)
public record ContentBlockDeltaEvent(
// @formatter:off
@JsonProperty("type") EventType type,
@JsonProperty("index") Integer index,
@JsonProperty("delta") ContentBlockDeltaBody delta) implements StreamEvent {
@JsonProperty("type") EventType type,
@JsonProperty("index") Integer index,
@JsonProperty("delta") ContentBlockDeltaBody delta) implements StreamEvent {
@JsonTypeInfo(use = JsonTypeInfo.Id.NAME, include = JsonTypeInfo.As.EXISTING_PROPERTY, property = "type",
visible = true)
@JsonSubTypes({ @JsonSubTypes.Type(value = ContentBlockDeltaText.class, name = "text_delta"),
@JsonSubTypes.Type(value = ContentBlockDeltaJson.class, name = "input_json_delta") })
public interface ContentBlockDeltaBody {
String type();
}
@JsonTypeInfo(use = JsonTypeInfo.Id.NAME, include = JsonTypeInfo.As.EXISTING_PROPERTY, property = "type",
visible = true)
@JsonSubTypes({ @JsonSubTypes.Type(value = ContentBlockDeltaText.class, name = "text_delta"),
@JsonSubTypes.Type(value = ContentBlockDeltaJson.class, name = "input_json_delta") })
public interface ContentBlockDeltaBody {
String type();
}
/**
* Text content block delta.
* @param type The content block type.
* @param text The text content.
*/
@JsonInclude(Include.NON_NULL)
public record ContentBlockDeltaText(
@JsonProperty("type") String type,
@JsonProperty("text") String text) implements ContentBlockDeltaBody {
}
/**
* Text content block delta.
* @param type The content block type.
* @param text The text content.
*/
@JsonInclude(Include.NON_NULL)
public record ContentBlockDeltaText(
@JsonProperty("type") String type,
@JsonProperty("text") String text) implements ContentBlockDeltaBody {
}
/**
/**
* JSON content block delta.
* @param type The content block type.
* @param partialJson The partial JSON content.
*/
@JsonInclude(Include.NON_NULL)
public record ContentBlockDeltaJson(
@JsonProperty("type") String type,
@JsonProperty("partial_json") String partialJson) implements ContentBlockDeltaBody {
}
}
// @formatter:on
@JsonInclude(Include.NON_NULL)
public record ContentBlockDeltaJson(
@JsonProperty("type") String type,
@JsonProperty("partial_json") String partialJson) implements ContentBlockDeltaBody {
}
}
// @formatter:on
// MESSAGE STOP EVENT
@@ -1082,10 +1082,10 @@ public class AnthropicApi {
@JsonInclude(Include.NON_NULL)
public record ContentBlockStopEvent(
// @formatter:off
@JsonProperty("type") EventType type,
@JsonProperty("index") Integer index) implements StreamEvent {
}
// @formatter:on
@JsonProperty("type") EventType type,
@JsonProperty("index") Integer index) implements StreamEvent {
}
// @formatter:on
/**
* Message start event.
@@ -1095,10 +1095,10 @@ public class AnthropicApi {
*/
@JsonInclude(Include.NON_NULL)
public record MessageStartEvent(// @formatter:off
@JsonProperty("type") EventType type,
@JsonProperty("message") ChatCompletionResponse message) implements StreamEvent {
}
// @formatter:on
@JsonProperty("type") EventType type,
@JsonProperty("message") ChatCompletionResponse message) implements StreamEvent {
}
// @formatter:on
/**
* Message delta event.
@@ -1110,31 +1110,30 @@ public class AnthropicApi {
@JsonInclude(Include.NON_NULL)
public record MessageDeltaEvent(
// @formatter:off
@JsonProperty("type") EventType type,
@JsonProperty("delta") MessageDelta delta,
@JsonProperty("usage") MessageDeltaUsage usage) implements StreamEvent {
@JsonProperty("type") EventType type,
@JsonProperty("delta") MessageDelta delta,
@JsonProperty("usage") MessageDeltaUsage usage) implements StreamEvent {
/**
* Message delta.
/**
* @param stopReason The stop reason.
* @param stopSequence The stop sequence.
*/
@JsonInclude(Include.NON_NULL)
public record MessageDelta(
@JsonProperty("stop_reason") String stopReason,
@JsonProperty("stop_sequence") String stopSequence) {
}
@JsonInclude(Include.NON_NULL)
public record MessageDelta(
@JsonProperty("stop_reason") String stopReason,
@JsonProperty("stop_sequence") String stopSequence) {
}
/**
* Message delta usage.
* @param outputTokens The output tokens.
*/
@JsonInclude(Include.NON_NULL)
public record MessageDeltaUsage(
@JsonProperty("output_tokens") Integer outputTokens) {
}
}
// @formatter:on
/**
* Message delta usage.
* @param outputTokens The output tokens.
*/
@JsonInclude(Include.NON_NULL)
public record MessageDeltaUsage(
@JsonProperty("output_tokens") Integer outputTokens) {
}
}
// @formatter:on
/**
* Message stop event.
@@ -1143,10 +1142,10 @@ public class AnthropicApi {
*/
@JsonInclude(Include.NON_NULL)
public record MessageStopEvent(
// @formatter:off
@JsonProperty("type") EventType type) implements StreamEvent {
}
// @formatter:on
//@formatter:off
@JsonProperty("type") EventType type) implements StreamEvent {
}
// @formatter:on
///////////////////////////////////////
/// ERROR EVENT
@@ -1160,21 +1159,21 @@ public class AnthropicApi {
@JsonInclude(Include.NON_NULL)
public record ErrorEvent(
// @formatter:off
@JsonProperty("type") EventType type,
@JsonProperty("error") Error error) implements StreamEvent {
@JsonProperty("type") EventType type,
@JsonProperty("error") Error error) implements StreamEvent {
/**
* Error body.
* @param type The error type.
* @param message The error message.
*/
@JsonInclude(Include.NON_NULL)
public record Error(
@JsonProperty("type") String type,
@JsonProperty("message") String message) {
}
}
// @formatter:on
/**
* Error body.
* @param type The error type.
* @param message The error message.
*/
@JsonInclude(Include.NON_NULL)
public record Error(
@JsonProperty("type") String type,
@JsonProperty("message") String message) {
}
}
// @formatter:on
///////////////////////////////////////
/// PING EVENT
@@ -1187,8 +1186,8 @@ public class AnthropicApi {
@JsonInclude(Include.NON_NULL)
public record PingEvent(
// @formatter:off
@JsonProperty("type") EventType type) implements StreamEvent {
}
// @formatter:on
@JsonProperty("type") EventType type) implements StreamEvent {
}
// @formatter:on
}
}

View File

@@ -18,7 +18,6 @@ package org.springframework.ai.azure.openai;
import com.azure.ai.openai.OpenAIClient;
import com.azure.ai.openai.OpenAIClientBuilder;
import com.azure.ai.openai.OpenAIServiceVersion;
import com.azure.core.credential.AzureKeyCredential;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;

View File

@@ -529,7 +529,6 @@ public class BedrockProxyChatModel extends AbstractToolCallSupport implements Ch
Flux<ChatResponse> chatResponseFlux = chatResponses.switchMap(chatResponse -> {
if (!this.isProxyToolCalls(prompt, this.defaultOptions) && chatResponse != null
&& this.isToolCall(chatResponse, Set.of(StopReason.TOOL_USE.toString()))) {
var toolCallConversation = this.handleToolCalls(prompt, chatResponse);
return this.internalStream(new Prompt(toolCallConversation, prompt.getOptions()), chatResponse);
}

View File

@@ -61,7 +61,7 @@ public class BedrockConverseUsageAggregationTests {
@BeforeEach
public void beforeEach() {
chatModel = new BedrockProxyChatModel(this.bedrockRuntimeClient, this.bedrockRuntimeAsyncClient,
this.chatModel = new BedrockProxyChatModel(this.bedrockRuntimeClient, this.bedrockRuntimeAsyncClient,
FunctionCallingOptions.builder().build(), null, List.of(), ObservationRegistry.NOOP);
}
@@ -90,9 +90,6 @@ public class BedrockConverseUsageAggregationTests {
assertThat(result.getMetadata().getUsage().getTotalTokens()).isEqualTo(30);
}
public record Request(String location, String unit) {
}
@Test
public void callWithToolUse() {
@@ -164,4 +161,7 @@ public class BedrockConverseUsageAggregationTests {
// TODO: Implement the test
}
public record Request(String location, String unit) {
}
}

View File

@@ -97,9 +97,9 @@ class BedrockCohereEmbeddingModelIT {
ArgumentCaptor<CohereEmbeddingBedrockApi.CohereEmbeddingRequest> requestCaptor = ArgumentCaptor
.forClass(CohereEmbeddingBedrockApi.CohereEmbeddingRequest.class);
EmbeddingResponse embeddingResponse = embeddingModel.embedForResponse(List.of(longText));
EmbeddingResponse embeddingResponse = this.embeddingModel.embedForResponse(List.of(longText));
verify(embeddingApi).embedding(requestCaptor.capture());
verify(this.embeddingApi).embedding(requestCaptor.capture());
CohereEmbeddingBedrockApi.CohereEmbeddingRequest capturedRequest = requestCaptor.getValue();
assertThat(capturedRequest.texts()).hasSize(1);
@@ -124,7 +124,7 @@ class BedrockCohereEmbeddingModelIT {
EmbeddingResponse embeddingResponse = this.embeddingModelStartTruncate.embedForResponse(List.of(longText));
// Verify truncation behavior
verify(embeddingApi).embedding(requestCaptor.capture());
verify(this.embeddingApi).embedding(requestCaptor.capture());
String truncatedText = requestCaptor.getValue().texts().get(0);
assertThat(truncatedText.length()).isLessThanOrEqualTo(2048);
assertThat(truncatedText).doesNotContain(startMarker);

View File

@@ -101,13 +101,15 @@ public class DefaultChatGenerationMetadata implements ChatGenerationMetadata {
@Override
public boolean equals(Object obj) {
if (this == obj)
if (this == obj) {
return true;
if (obj == null || getClass() != obj.getClass())
}
if (obj == null || getClass() != obj.getClass()) {
return false;
}
DefaultChatGenerationMetadata other = (DefaultChatGenerationMetadata) obj;
return Objects.equals(metadata, other.metadata) && Objects.equals(finishReason, other.finishReason)
&& Objects.equals(contentFilters, other.contentFilters);
return Objects.equals(this.metadata, other.metadata) && Objects.equals(this.finishReason, other.finishReason)
&& Objects.equals(this.contentFilters, other.contentFilters);
}
@Override

View File

@@ -137,11 +137,12 @@ public final class SimpleVectorStoreContent implements Content {
@Override
public boolean equals(Object o) {
if (this == o)
if (this == o) {
return true;
if (o == null || getClass() != o.getClass())
}
if (o == null || getClass() != o.getClass()) {
return false;
}
SimpleVectorStoreContent that = (SimpleVectorStoreContent) o;
return Objects.equals(this.id, that.id) && Objects.equals(this.content, that.content)
&& Objects.equals(this.metadata, that.metadata) && Arrays.equals(this.embedding, that.embedding);
@@ -159,7 +160,7 @@ public final class SimpleVectorStoreContent implements Content {
@Override
public String toString() {
return "SimpleVectorStoreContent{" + "id='" + this.id + '\'' + ", content='" + this.content + '\''
+ ", metadata=" + this.metadata + ", embedding=" + Arrays.toString(embedding) + '}';
+ ", metadata=" + this.metadata + ", embedding=" + Arrays.toString(this.embedding) + '}';
}
}

View File

@@ -16,8 +16,11 @@
package org.springframework.ai.chat.client.advisor;
import java.util.List;
import org.junit.jupiter.api.Test;
import org.mockito.ArgumentCaptor;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.model.ChatModel;
@@ -29,8 +32,6 @@ import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.preretrieval.query.transformation.QueryTransformer;
import org.springframework.ai.rag.retrieval.search.DocumentRetriever;
import java.util.List;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
import static org.mockito.BDDMockito.given;

View File

@@ -16,13 +16,14 @@
package org.springframework.ai.rag.orchestration.routing;
import org.junit.jupiter.api.Test;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.retrieval.search.DocumentRetriever;
import java.util.ArrayList;
import java.util.List;
import org.junit.jupiter.api.Test;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.retrieval.search.DocumentRetriever;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
import static org.mockito.Mockito.mock;

View File

@@ -16,14 +16,15 @@
package org.springframework.ai.rag.retrieval.join;
import org.junit.jupiter.api.Test;
import org.springframework.ai.document.Document;
import org.springframework.ai.rag.Query;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import org.junit.jupiter.api.Test;
import org.springframework.ai.document.Document;
import org.springframework.ai.rag.Query;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;

View File

@@ -16,14 +16,6 @@
package org.springframework.ai.vectorstore;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.io.CleanupMode;
import org.junit.jupiter.api.io.TempDir;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.core.io.Resource;
import java.io.File;
import java.io.IOException;
import java.nio.file.Path;
@@ -35,6 +27,15 @@ import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.io.CleanupMode;
import org.junit.jupiter.api.io.TempDir;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.core.io.Resource;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
import static org.mockito.ArgumentMatchers.any;
@@ -52,12 +53,11 @@ class SimpleVectorStoreTests {
@BeforeEach
void setUp() {
mockEmbeddingModel = mock(EmbeddingModel.class);
when(mockEmbeddingModel.dimensions()).thenReturn(3);
when(mockEmbeddingModel.embed(any(String.class))).thenReturn(new float[] { 0.1f, 0.2f, 0.3f });
when(mockEmbeddingModel.embed(any(Document.class))).thenReturn(new float[] { 0.1f, 0.2f, 0.3f });
vectorStore = new SimpleVectorStore(mockEmbeddingModel);
this.mockEmbeddingModel = mock(EmbeddingModel.class);
when(this.mockEmbeddingModel.dimensions()).thenReturn(3);
when(this.mockEmbeddingModel.embed(any(String.class))).thenReturn(new float[] { 0.1f, 0.2f, 0.3f });
when(this.mockEmbeddingModel.embed(any(Document.class))).thenReturn(new float[] { 0.1f, 0.2f, 0.3f });
this.vectorStore = new SimpleVectorStore(this.mockEmbeddingModel);
}
@Test
@@ -68,9 +68,9 @@ class SimpleVectorStoreTests {
.withMetadata(Map.of("key", "value"))
.build();
vectorStore.add(List.of(doc));
this.vectorStore.add(List.of(doc));
List<Document> results = vectorStore.similaritySearch("test content");
List<Document> results = this.vectorStore.similaritySearch("test content");
assertThat(results).hasSize(1).first().satisfies(result -> {
assertThat(result.getId()).isEqualTo("1");
assertThat(result.getContent()).isEqualTo("test content");
@@ -83,21 +83,22 @@ class SimpleVectorStoreTests {
List<Document> docs = Arrays.asList(Document.builder().withId("1").withContent("first").build(),
Document.builder().withId("2").withContent("second").build());
vectorStore.add(docs);
this.vectorStore.add(docs);
List<Document> results = vectorStore.similaritySearch("first");
List<Document> results = this.vectorStore.similaritySearch("first");
assertThat(results).hasSize(2).extracting(Document::getId).containsExactlyInAnyOrder("1", "2");
}
@Test
void shouldHandleEmptyDocumentList() {
assertThatThrownBy(() -> vectorStore.add(Collections.emptyList())).isInstanceOf(IllegalArgumentException.class)
assertThatThrownBy(() -> this.vectorStore.add(Collections.emptyList()))
.isInstanceOf(IllegalArgumentException.class)
.hasMessage("Documents list cannot be empty");
}
@Test
void shouldHandleNullDocumentList() {
assertThatThrownBy(() -> vectorStore.add(null)).isInstanceOf(NullPointerException.class)
assertThatThrownBy(() -> this.vectorStore.add(null)).isInstanceOf(NullPointerException.class)
.hasMessage("Documents list cannot be null");
}
@@ -105,32 +106,32 @@ class SimpleVectorStoreTests {
void shouldDeleteDocuments() {
Document doc = Document.builder().withId("1").withContent("test content").build();
vectorStore.add(List.of(doc));
assertThat(vectorStore.similaritySearch("test")).hasSize(1);
this.vectorStore.add(List.of(doc));
assertThat(this.vectorStore.similaritySearch("test")).hasSize(1);
vectorStore.delete(List.of("1"));
assertThat(vectorStore.similaritySearch("test")).isEmpty();
this.vectorStore.delete(List.of("1"));
assertThat(this.vectorStore.similaritySearch("test")).isEmpty();
}
@Test
void shouldHandleDeleteOfNonexistentDocument() {
vectorStore.delete(List.of("nonexistent-id"));
this.vectorStore.delete(List.of("nonexistent-id"));
// Should not throw exception and return true
assertThat(vectorStore.delete(List.of("nonexistent-id")).get()).isTrue();
assertThat(this.vectorStore.delete(List.of("nonexistent-id")).get()).isTrue();
}
@Test
void shouldPerformSimilaritySearchWithThreshold() {
// Configure mock to return different embeddings for different queries
when(mockEmbeddingModel.embed("query")).thenReturn(new float[] { 0.9f, 0.9f, 0.9f });
when(this.mockEmbeddingModel.embed("query")).thenReturn(new float[] { 0.9f, 0.9f, 0.9f });
Document doc = Document.builder().withId("1").withContent("test content").build();
vectorStore.add(List.of(doc));
this.vectorStore.add(List.of(doc));
SearchRequest request = SearchRequest.query("query").withSimilarityThreshold(0.99f).withTopK(5);
List<Document> results = vectorStore.similaritySearch(request);
List<Document> results = this.vectorStore.similaritySearch(request);
assertThat(results).isEmpty();
}
@@ -142,12 +143,12 @@ class SimpleVectorStoreTests {
.withMetadata(new HashMap<>(Map.of("key", "value")))
.build();
vectorStore.add(List.of(doc));
this.vectorStore.add(List.of(doc));
File saveFile = tempDir.resolve("vector-store.json").toFile();
vectorStore.save(saveFile);
File saveFile = this.tempDir.resolve("vector-store.json").toFile();
this.vectorStore.save(saveFile);
SimpleVectorStore loadedStore = new SimpleVectorStore(mockEmbeddingModel);
SimpleVectorStore loadedStore = new SimpleVectorStore(this.mockEmbeddingModel);
loadedStore.load(saveFile);
List<Document> results = loadedStore.similaritySearch("test content");
@@ -163,7 +164,7 @@ class SimpleVectorStoreTests {
Resource mockResource = mock(Resource.class);
when(mockResource.getInputStream()).thenThrow(new IOException("Resource not found"));
assertThatThrownBy(() -> vectorStore.load(mockResource)).isInstanceOf(RuntimeException.class)
assertThatThrownBy(() -> this.vectorStore.load(mockResource)).isInstanceOf(RuntimeException.class)
.hasCauseInstanceOf(IOException.class)
.hasMessageContaining("Resource not found");
}
@@ -172,7 +173,7 @@ class SimpleVectorStoreTests {
void shouldHandleSaveToInvalidLocation() {
File invalidFile = new File("/invalid/path/file.json");
assertThatThrownBy(() -> vectorStore.save(invalidFile)).isInstanceOf(RuntimeException.class)
assertThatThrownBy(() -> this.vectorStore.save(invalidFile)).isInstanceOf(RuntimeException.class)
.hasCauseInstanceOf(IOException.class);
}
@@ -185,7 +186,7 @@ class SimpleVectorStoreTests {
final String id = String.valueOf(i);
threads[i] = new Thread(() -> {
Document doc = Document.builder().withId(id).withContent("content " + id).build();
vectorStore.add(List.of(doc));
this.vectorStore.add(List.of(doc));
});
threads[i].start();
}
@@ -196,7 +197,7 @@ class SimpleVectorStoreTests {
SearchRequest request = SearchRequest.query("test").withTopK(numThreads);
List<Document> results = vectorStore.similaritySearch(request);
List<Document> results = this.vectorStore.similaritySearch(request);
assertThat(results).hasSize(numThreads);

View File

@@ -288,7 +288,7 @@ public class ChromaApi {
@JsonInclude(JsonInclude.Include.NON_NULL)
public record CreateCollectionRequest(// @formatter:off
@JsonProperty("name") String name,
@JsonProperty("metadata") Map<String, Object> metadata) {// @formatter:on
@JsonProperty("metadata") Map<String, Object> metadata) { // @formatter:on
public CreateCollectionRequest(String name) {
this(name, new HashMap<>(Map.of("hnsw:space", "cosine")));
@@ -314,7 +314,7 @@ public class ChromaApi {
@JsonProperty("ids") List<String> ids,
@JsonProperty("embeddings") List<float[]> embeddings,
@JsonProperty("metadatas") List<Map<String, Object>> metadata,
@JsonProperty("documents") List<String> documents) {// @formatter:on
@JsonProperty("documents") List<String> documents) { // @formatter:on
// Convenance for adding a single embedding.
public AddEmbeddingsRequest(String id, float[] embedding, Map<String, Object> metadata, String document) {
@@ -332,7 +332,7 @@ public class ChromaApi {
@JsonInclude(JsonInclude.Include.NON_NULL)
public record DeleteEmbeddingsRequest(// @formatter:off
@JsonProperty("ids") List<String> ids,
@JsonProperty("where") Map<String, Object> where) {// @formatter:on
@JsonProperty("where") Map<String, Object> where) { // @formatter:on
public DeleteEmbeddingsRequest(List<String> ids) {
this(ids, null);
@@ -356,7 +356,7 @@ public class ChromaApi {
@JsonProperty("where") Map<String, Object> where,
@JsonProperty("limit") Integer limit,
@JsonProperty("offset") Integer offset,
@JsonProperty("include") List<Include> include) {// @formatter:on
@JsonProperty("include") List<Include> include) { // @formatter:on
public GetEmbeddingsRequest(List<String> ids) {
this(ids, null, 10, 0, Include.all);
@@ -385,7 +385,7 @@ public class ChromaApi {
@JsonProperty("ids") List<String> ids,
@JsonProperty("embeddings") List<float[]> embeddings,
@JsonProperty("documents") List<String> documents,
@JsonProperty("metadatas") List<Map<String, String>> metadata) {// @formatter:on
@JsonProperty("metadatas") List<Map<String, String>> metadata) { // @formatter:on
}
/**
@@ -401,7 +401,7 @@ public class ChromaApi {
* [metadatas, documents, distances].
*/
@JsonInclude(JsonInclude.Include.NON_NULL)
public record QueryRequest( // @formatter:off
public record QueryRequest(// @formatter:off
@JsonProperty("query_embeddings") List<float[]> queryEmbeddings,
@JsonProperty("n_results") Integer nResults,
@JsonProperty("where") Map<String, Object> where,
@@ -454,7 +454,7 @@ public class ChromaApi {
@JsonProperty("embeddings") List<List<float[]>> embeddings,
@JsonProperty("documents") List<List<String>> documents,
@JsonProperty("metadatas") List<List<Map<String, Object>>> metadata,
@JsonProperty("distances") List<List<Double>> distances) {// @formatter:on
@JsonProperty("distances") List<List<Double>> distances) { // @formatter:on
}
/**
@@ -472,7 +472,7 @@ public class ChromaApi {
@JsonProperty("embedding") float[] embedding,
@JsonProperty("document") String document,
@JsonProperty("metadata") Map<String, Object> metadata,
@JsonProperty("distances") Double distances) {// @formatter:on
@JsonProperty("distances") Double distances) { // @formatter:on
}