-
Notifications
You must be signed in to change notification settings - Fork 736
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add observability support to VectorStore
Implementation: - Introduce AbstractObservationVectorStore with instrumentation for add, delete, and similaritySearch methods - Create VectorStoreObservationContext to capture operation details - Implement DefaultVectorStoreObservationConvention for naming and tagging - Add VectorStoreObservationDocumentation for defining observation keys - Create VectorStoreObservationAutoConfiguration for auto-configuring observations - Add VectorStoreObservationProperties to control optional observation content filters - Update VectorStore interface with getName() method - Modify PgVectorStore and SimpleVectorStore to extend AbstractObservationVectorStore - Add vector_store Spring AI kind Filters: - Implement VectorStoreQueryResponseObservationFilter - Add VectorStoreDeleteRequestContentObservationFilter and VectorStoreAddRequestContentObservationFilter Enhancements: - Update PgVectorStoreAutoConfiguration to support observations - Add observation support to PgVectorStore's Builder - Add VectorStoreObservationContext.Operation enum with ADD, DELETE, and QUERY options Tests: - Add tests for VectorStore context, convention, and filters - Add VectorStoreObservationAutoConfiguration tests - Add PgVectorObservationIT Resolves #1205
- Loading branch information
1 parent
72369d5
commit bc3f9ac
Showing
23 changed files
with
1,620 additions
and
38 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
39 changes: 39 additions & 0 deletions
39
...ore/src/main/java/org/springframework/ai/observation/conventions/VectorStoreProvider.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,39 @@ | ||
/* | ||
* Copyright 2024 - 2024 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. | ||
* You may obtain a copy of the License at | ||
* | ||
* https://www.apache.org/licenses/LICENSE-2.0 | ||
* | ||
* Unless required by applicable law or agreed to in writing, software | ||
* distributed under the License is distributed on an "AS IS" BASIS, | ||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
* See the License for the specific language governing permissions and | ||
* limitations under the License. | ||
*/ | ||
package org.springframework.ai.observation.conventions; | ||
|
||
/** | ||
* @author Christian Tzolov | ||
* @since 1.0.0 | ||
*/ | ||
public enum VectorStoreProvider { | ||
|
||
// @formatter:off | ||
PG_VECTOR("pg_vector"), | ||
SIMPLE_VECTOR_STORE("simple_vector_store"); | ||
|
||
// @formatter:on | ||
private final String value; | ||
|
||
VectorStoreProvider(String value) { | ||
this.value = value; | ||
} | ||
|
||
public String value() { | ||
return this.value; | ||
} | ||
|
||
} |
54 changes: 54 additions & 0 deletions
54
...main/java/org/springframework/ai/observation/conventions/VectorStoreSimilarityMetric.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,54 @@ | ||
/* | ||
* Copyright 2024 - 2024 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. | ||
* You may obtain a copy of the License at | ||
* | ||
* https://www.apache.org/licenses/LICENSE-2.0 | ||
* | ||
* Unless required by applicable law or agreed to in writing, software | ||
* distributed under the License is distributed on an "AS IS" BASIS, | ||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
* See the License for the specific language governing permissions and | ||
* limitations under the License. | ||
*/ | ||
package org.springframework.ai.observation.conventions; | ||
|
||
/** | ||
* @author Christian Tzolov | ||
* @since 1.0.0 | ||
*/ | ||
public enum VectorStoreSimilarityMetric { | ||
|
||
// @formatter:off | ||
|
||
/** | ||
* The cosine metric. | ||
*/ | ||
COSINE("cosine"), | ||
/** | ||
* The euclidean distance metric. | ||
*/ | ||
EUCLIDEAN("euclidean"), | ||
/** | ||
* The manhattan distance metric. | ||
*/ | ||
MANHATTAN("manhattan"), | ||
/** | ||
* The dot product metric. | ||
*/ | ||
DOT("dot"); | ||
|
||
// @formatter:on | ||
private final String value; | ||
|
||
VectorStoreSimilarityMetric(String value) { | ||
this.value = value; | ||
} | ||
|
||
public String value() { | ||
return this.value; | ||
} | ||
|
||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
107 changes: 107 additions & 0 deletions
107
...n/java/org/springframework/ai/vectorstore/observation/AbstractObservationVectorStore.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,107 @@ | ||
/* | ||
* Copyright 2024 - 2024 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. | ||
* You may obtain a copy of the License at | ||
* | ||
* https://www.apache.org/licenses/LICENSE-2.0 | ||
* | ||
* Unless required by applicable law or agreed to in writing, software | ||
* distributed under the License is distributed on an "AS IS" BASIS, | ||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
* See the License for the specific language governing permissions and | ||
* limitations under the License. | ||
*/ | ||
package org.springframework.ai.vectorstore.observation; | ||
|
||
import java.util.List; | ||
import java.util.Optional; | ||
|
||
import org.springframework.ai.document.Document; | ||
import org.springframework.ai.vectorstore.SearchRequest; | ||
import org.springframework.ai.vectorstore.VectorStore; | ||
import org.springframework.lang.Nullable; | ||
|
||
import io.micrometer.observation.ObservationRegistry; | ||
|
||
/** | ||
* @author Christian Tzolov | ||
* @since 1.0.0 | ||
*/ | ||
public abstract class AbstractObservationVectorStore implements VectorStore { | ||
|
||
private static final VectorStoreObservationConvention DEFAULT_OBSERVATION_CONVENTION = new DefaultVectorStoreObservationConvention(); | ||
|
||
private final ObservationRegistry observationRegistry; | ||
|
||
@Nullable | ||
private final VectorStoreObservationConvention customObservationConvention; | ||
|
||
public AbstractObservationVectorStore() { | ||
this(ObservationRegistry.NOOP, null); | ||
} | ||
|
||
public AbstractObservationVectorStore(ObservationRegistry observationRegistry) { | ||
this(observationRegistry, null); | ||
} | ||
|
||
public AbstractObservationVectorStore(ObservationRegistry observationRegistry, | ||
VectorStoreObservationConvention customSearchObservationConvention) { | ||
this.observationRegistry = observationRegistry; | ||
this.customObservationConvention = customSearchObservationConvention; | ||
} | ||
|
||
@Override | ||
public void add(List<Document> documents) { | ||
|
||
VectorStoreObservationContext observationContext = this | ||
.createObservationContextBuilder(VectorStoreObservationContext.Operation.ADD.value()) | ||
.build(); | ||
|
||
VectorStoreObservationDocumentation.AI_VECTOR_STORE | ||
.observation(this.customObservationConvention, DEFAULT_OBSERVATION_CONVENTION, () -> observationContext, | ||
observationRegistry) | ||
.observe(() -> this.doAdd(documents)); | ||
} | ||
|
||
@Override | ||
public Optional<Boolean> delete(List<String> deleteDocIds) { | ||
|
||
VectorStoreObservationContext observationContext = this | ||
.createObservationContextBuilder(VectorStoreObservationContext.Operation.DELETE.value()) | ||
.build(); | ||
|
||
return VectorStoreObservationDocumentation.AI_VECTOR_STORE | ||
.observation(this.customObservationConvention, DEFAULT_OBSERVATION_CONVENTION, () -> observationContext, | ||
this.observationRegistry) | ||
.observe(() -> this.doDelete(deleteDocIds)); | ||
} | ||
|
||
@Override | ||
public List<Document> similaritySearch(SearchRequest request) { | ||
|
||
VectorStoreObservationContext searchObservationContext = this | ||
.createObservationContextBuilder(VectorStoreObservationContext.Operation.QUERY.value()) | ||
.withQueryRequest(request) | ||
.build(); | ||
|
||
return VectorStoreObservationDocumentation.AI_VECTOR_STORE | ||
.observation(this.customObservationConvention, DEFAULT_OBSERVATION_CONVENTION, | ||
() -> searchObservationContext, this.observationRegistry) | ||
.observe(() -> { | ||
var documents = this.doSimilaritySearch(request); | ||
searchObservationContext.setQueryResponse(documents); | ||
return documents; | ||
}); | ||
} | ||
|
||
public abstract void doAdd(List<Document> documents); | ||
|
||
public abstract Optional<Boolean> doDelete(List<String> idList); | ||
|
||
public abstract List<Document> doSimilaritySearch(SearchRequest request); | ||
|
||
public abstract VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName); | ||
|
||
} |
Oops, something went wrong.