Repository navigation
fix: advertise Data API clause JSON in MCP schemas #2619
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
767afa4
d827066
acfbe5d
027c386
b2088a9
bdc563c
bbacbda
2f2720e
773477d
fdb5f1a
de4573f
c61a536
f665be8
f96ac86
ca50f43
c0944a2
2d4c204
e1c457e
41015d3
53ce63e
04bf945
ff0f145
d3e4ea7
15875b7
5fb06da
e4c6b92
0900d87
ab4d996
55d3d72
0b89e71
6e2cf87
d3c2429
dc4979e
9692939
94d14c7
14cabf2
941548a
a7534de
2de2a2f
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,11 +1,15 @@ | ||
| package io.stargate.sgv2.jsonapi.api.model.command.clause.sort; | ||
|
|
||
| import static io.stargate.sgv2.jsonapi.config.constants.DocumentConstants.Fields.*; | ||
|
|
||
| import com.fasterxml.jackson.annotation.JsonIgnore; | ||
| import com.fasterxml.jackson.databind.annotation.JsonDeserialize; | ||
| import io.stargate.sgv2.jsonapi.metrics.CommandFeatures; | ||
| import io.stargate.sgv2.jsonapi.util.recordable.Recordable; | ||
| import java.util.*; | ||
| import org.eclipse.microprofile.openapi.annotations.enums.SchemaType; | ||
| import org.eclipse.microprofile.openapi.annotations.media.Schema; | ||
| import org.eclipse.microprofile.openapi.annotations.media.SchemaProperty; | ||
|
|
||
| /** | ||
| * @param vectorizeSort If <code>null</code> then the $vectorize field was not present in the sort | ||
|
|
@@ -19,18 +23,84 @@ | |
| */ | ||
| @JsonDeserialize(using = FindAndRerankSortClauseDeserializer.class) | ||
| @Schema( | ||
| name = "FindAndRerankSort", | ||
| type = SchemaType.OBJECT, | ||
| implementation = Map.class, | ||
| examples = | ||
| """ | ||
| {"$sort" : {"$hybrid" : "Same query for vectorize and bm25 sorting"}}} | ||
| {"$sort" : {"$hybrid" : {"$vectorize" : "vectorize sort query" , "$lexical": "lexical sort" }}} | ||
| {"$sort" : {"$hybrid" : {"$vector" : [1,2,3] , "$lexical": "lexical sort" }}} | ||
| """) | ||
| nullable = true, | ||
| additionalProperties = Schema.False.class, | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. why do this ?
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The parser only accepts $hybrid and rejects any other key. This makes the schema reject the same payloads, so a client sees the problem before it calls the server. |
||
| properties = | ||
| @SchemaProperty( | ||
| name = HYBRID_FIELD, | ||
| anyOf = {String.class, FindAndRerankSort.HybridSortSchema.class}), | ||
| description = | ||
| "Use $hybrid with a query string or an object containing $vectorize or $vector, and optionally $lexical." | ||
| + " $vector and $vectorize cannot be used together.", | ||
| examples = { | ||
|
erichare marked this conversation as resolved.
Resolved
|
||
| """ | ||
| {"$hybrid": "Same query for vectorize and BM25 sorting"} | ||
| """, | ||
| """ | ||
| {"$hybrid": {"$vectorize": "vectorize sort query", "$lexical": "lexical sort"}} | ||
| """, | ||
| """ | ||
| {"$hybrid": {"$vector": [1, 2, 3], "$lexical": "lexical sort"}} | ||
| """, | ||
| """ | ||
| {"$hybrid": {"$vector": {"$binary": "P4AAAEAAAABAQAAA"}, "$lexical": "lexical query"}} | ||
| """ | ||
| }) | ||
| public record FindAndRerankSort( | ||
| String vectorizeSort, String lexicalSort, float[] vectorSort, CommandFeatures commandFeatures) | ||
| @JsonIgnore String vectorizeSort, | ||
| @JsonIgnore String lexicalSort, | ||
| @JsonIgnore float[] vectorSort, | ||
| @JsonIgnore CommandFeatures commandFeatures) | ||
| implements Recordable { | ||
|
|
||
| // Schema-only classes. @SchemaProperty can point at a nested schema only through a class, so | ||
| // these empty classes carry the @Schema for each nested part of the sort. Nothing creates them. | ||
|
|
||
| /** The object form of $hybrid, with $vector or $vectorize and an optional $lexical. */ | ||
| @Schema( | ||
| name = "HybridSort", | ||
| type = SchemaType.OBJECT, | ||
| implementation = Object.class, | ||
| additionalProperties = Schema.False.class, | ||
|
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. how does it look when we call with some {"foo": 1} with and without this additionalProperties? |
||
| properties = { | ||
| @SchemaProperty( | ||
| name = VECTOR_EMBEDDING_TEXT_FIELD, | ||
| type = SchemaType.STRING, | ||
| nullable = true), | ||
| @SchemaProperty(name = LEXICAL_CONTENT_FIELD, type = SchemaType.STRING, nullable = true), | ||
| @SchemaProperty( | ||
| name = VECTOR_EMBEDDING_FIELD, | ||
| nullable = true, | ||
| anyOf = {float[].class, BinaryVectorSchema.class}) | ||
| }, | ||
| not = ConflictingVectorSortSchema.class) | ||
| public static class HybridSortSchema {} | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. WHY does this class exists ?
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. It describes the object form of $hybrid for OpenAPI and MCP, since the record's own fields don't match the JSON clients send. I added a doc comment on it and on the two classes below it. |
||
|
|
||
| /** A binary encoded vector such as {"$binary": "P4AAAEAAAABAQAAA"}. */ | ||
| @Schema( | ||
| name = "BinaryVector", | ||
| type = SchemaType.OBJECT, | ||
| implementation = Object.class, | ||
| nullable = true, | ||
| additionalProperties = Schema.False.class, | ||
| requiredProperties = BINARY_VECTOR_TEXT_FIELD, | ||
| properties = @SchemaProperty(name = BINARY_VECTOR_TEXT_FIELD, type = SchemaType.STRING)) | ||
| public static class BinaryVectorSchema {} | ||
|
|
||
| /** | ||
| * Matches an object that has both $vector and $vectorize. HybridSort uses it as its {@code not} | ||
| * schema, so a sort that sets both fails validation. | ||
| */ | ||
| @Schema( | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. this really needs comment to explain what is happening
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Added a doc comment. HybridSort uses it as its not schema, so a sort with both $vector and $vectorize fails validation. |
||
| name = "ConflictingVectorSort", | ||
| type = SchemaType.OBJECT, | ||
| implementation = Object.class, | ||
| requiredProperties = {VECTOR_EMBEDDING_FIELD, VECTOR_EMBEDDING_TEXT_FIELD}) | ||
| public static class ConflictingVectorSortSchema {} | ||
|
|
||
| public static final FindAndRerankSort NO_ARG_SORT = | ||
| new FindAndRerankSort(null, null, null, CommandFeatures.EMPTY); | ||
|
|
||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,7 +1,13 @@ | ||
| package io.stargate.sgv2.jsonapi.api.model.command.impl; | ||
|
|
||
| import static io.stargate.sgv2.jsonapi.config.constants.DocumentConstants.Fields.*; | ||
| import static io.stargate.sgv2.jsonapi.config.constants.OpenApiConstants.Components.SCHEMA_REF_PREFIX; | ||
| import static io.stargate.sgv2.jsonapi.config.constants.ServiceDescConstants.AUTHENTICATION; | ||
| import static io.stargate.sgv2.jsonapi.config.constants.ServiceDescConstants.MODEL_NAME; | ||
| import static io.stargate.sgv2.jsonapi.config.constants.ServiceDescConstants.PARAMETERS; | ||
| import static io.stargate.sgv2.jsonapi.config.constants.ServiceDescConstants.PROVIDER; | ||
|
|
||
| import com.fasterxml.jackson.annotation.JsonIgnore; | ||
| import com.fasterxml.jackson.annotation.JsonProperty; | ||
| import com.fasterxml.jackson.annotation.JsonTypeName; | ||
| import com.fasterxml.jackson.core.JsonParser; | ||
|
|
@@ -28,10 +34,9 @@ | |
| import javax.annotation.Nullable; | ||
| import org.eclipse.microprofile.openapi.annotations.enums.SchemaType; | ||
| import org.eclipse.microprofile.openapi.annotations.media.Schema; | ||
| import org.eclipse.microprofile.openapi.annotations.media.SchemaProperty; | ||
|
|
||
| @Schema( | ||
| description = | ||
| "Finds documents using using vector and lexical sorting, then reranks the results.") | ||
| @Schema(description = "Finds documents using vector and lexical sorting, then reranks the results.") | ||
| @JsonTypeName(CommandName.Names.FIND_AND_RERANK) | ||
| public record FindAndRerankCommand( | ||
| @Valid @JsonProperty("filter") FilterDefinition filterDefinition, | ||
|
|
@@ -70,15 +75,7 @@ public record Options( | |
| implementation = Integer.class) | ||
| Integer limit, | ||
| /** ---- */ | ||
| @Schema( | ||
| description = | ||
| "The maximum number of documents to read for the vector and lexical queries that feed into the reranking. May be a number or an object with $vector and $lexical fields. The accepted range is determined by server configuration.", | ||
| examples = | ||
| """ | ||
| {"hybridLimits" : 100} | ||
| {"hybridLimits" : {"$vector" : 100, "$lexical" : 10}} | ||
| """) | ||
| HybridLimits hybridLimits, | ||
| @Schema(ref = HybridLimits.REF) HybridLimits hybridLimits, | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. nit: put the function on a new line
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I tried it, but google-java-format puts it back on one line and the CI format check would fail. |
||
| /** ---- */ | ||
| @Schema( | ||
| description = | ||
|
|
@@ -117,18 +114,67 @@ public record Options( | |
| "Optional override for the reranking service. Completely replaces the" | ||
| + " collection-level reranking configuration when provided. Both 'provider' and" | ||
| + " 'modelName' are required.", | ||
| implementation = CreateCollectionCommand.Options.RerankServiceDesc.class) | ||
| implementation = RerankOverrideSchema.class, | ||
| ref = RerankOverrideSchema.REF) | ||
| @JsonProperty("rerank") | ||
| CreateCollectionCommand.Options.RerankServiceDesc rerankServiceOverride) {} | ||
|
|
||
| // Matches the override validation in CollectionRerankDef. Update both together. | ||
| @Schema( | ||
| name = RerankOverrideSchema.NAME, | ||
| type = SchemaType.OBJECT, | ||
| implementation = Object.class, | ||
| nullable = true, | ||
| additionalProperties = Schema.False.class, | ||
| requiredProperties = {PROVIDER, MODEL_NAME}, | ||
| properties = { | ||
| @SchemaProperty(name = PROVIDER, type = SchemaType.STRING), | ||
| @SchemaProperty(name = MODEL_NAME, type = SchemaType.STRING), | ||
| @SchemaProperty( | ||
| name = AUTHENTICATION, | ||
| type = SchemaType.OBJECT, | ||
| nullable = true, | ||
| additionalProperties = String.class), | ||
| @SchemaProperty( | ||
| name = PARAMETERS, | ||
| type = SchemaType.OBJECT, | ||
| nullable = true, | ||
| additionalProperties = Schema.False.class) | ||
| }, | ||
| examples = | ||
| """ | ||
| {"provider": "nvidia", "modelName": "nvidia/llama-3.2-nv-rerankqa-1b-v2"} | ||
| """) | ||
| public static class RerankOverrideSchema { | ||
| public static final String NAME = "RerankServiceOverride"; | ||
| public static final String REF = SCHEMA_REF_PREFIX + NAME; | ||
| } | ||
|
|
||
| @JsonDeserialize(using = HybridLimitsDeserializer.class) | ||
| @Schema( | ||
| name = HybridLimits.NAME, | ||
| description = | ||
| "Candidate read limits as one integer for both legs, or an object containing both $vector and $lexical integer limits.", | ||
| // Void is the annotation default, but the OpenAPI scanner only sees values set in source. | ||
| // Setting it stops the scanner adding the record fields as an object type next to anyOf. | ||
| implementation = Void.class, | ||
| anyOf = {Integer.class, HybridLimits.ValuesSchema.class}, | ||
| examples = { | ||
| "50", | ||
| """ | ||
| {"$vector": 50, "$lexical": 10} | ||
| """ | ||
| }) | ||
| public record HybridLimits( | ||
| @JsonProperty(DocumentConstants.Fields.VECTOR_EMBEDDING_FIELD) int vectorLimit, | ||
| /** ---- */ | ||
| @JsonProperty(DocumentConstants.Fields.LEXICAL_CONTENT_FIELD) int lexicalLimit, | ||
| CommandFeatures commandFeatures) | ||
| @JsonIgnore CommandFeatures commandFeatures) | ||
| implements Recordable { | ||
|
|
||
| public static final String NAME = "HybridLimits"; | ||
| public static final String REF = SCHEMA_REF_PREFIX + NAME; | ||
|
|
||
| public static final HybridLimits DEFAULT = | ||
| new HybridLimits( | ||
| RerankingConstants.HybridSearchLimits.DEFAULT, | ||
|
|
@@ -142,6 +188,19 @@ public DataRecorder recordTo(DataRecorder dataRecorder) { | |
| .append("lexicalLimit", lexicalLimit) | ||
| .append("commandFeatures", commandFeatures); | ||
| } | ||
|
|
||
| /** Schema for the object form of the limits, the second branch of the anyOf above. */ | ||
| @Schema( | ||
| name = "HybridLimitValues", | ||
| type = SchemaType.OBJECT, | ||
| nullable = true, | ||
| additionalProperties = Schema.False.class, | ||
| requiredProperties = {VECTOR_EMBEDDING_FIELD, LEXICAL_CONTENT_FIELD}, | ||
| properties = { | ||
| @SchemaProperty(name = VECTOR_EMBEDDING_FIELD, type = SchemaType.INTEGER), | ||
| @SchemaProperty(name = LEXICAL_CONTENT_FIELD, type = SchemaType.INTEGER) | ||
| }) | ||
| public static class ValuesSchema {} | ||
| } | ||
|
|
||
| /** Deserializer for the `hybridLimits` option, the limits for the inner reads. */ | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -148,16 +148,23 @@ public Uni<ToolResponse> estimatedDocumentCount( | |
| return mcpResource.processCollectionCommand(keyspace, collection, command); | ||
| } | ||
|
|
||
| @Tool( | ||
| description = | ||
| "Finds documents using using vector and lexical sorting, then reranks the results.") | ||
| @Tool(description = "Finds documents using vector and lexical sorting, then reranks the results.") | ||
| public Uni<ToolResponse> findAndRerank( | ||
| @ToolArg(description = "Name of the keyspace") String keyspace, | ||
| @ToolArg(description = "Name of the collection/table") String collection, | ||
| @ToolArg(description = "filter", required = false) FilterDefinition filter, | ||
| @ToolArg(description = "projection", required = false) JsonNode projection, | ||
| @ToolArg(description = "sort", required = false) FindAndRerankSort sort, | ||
| @ToolArg(description = "options", required = false) FindAndRerankCommand.Options options) { | ||
| @ToolArg(description = "Name of the collection") String collection, | ||
| // filter and sort have no description here because it would replace the one on their | ||
| // classes, which Swagger also uses. | ||
| @ToolArg(required = false) FilterDefinition filter, | ||
| @ToolArg( | ||
| description = | ||
| "Fields to include or exclude in the returned documents, for example {\"title\": 1}", | ||
| required = false) | ||
| JsonNode projection, | ||
| @ToolArg(required = false) FindAndRerankSort sort, | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. bring all the parameters up to the same standard for use of description etc
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Will do. Do you mean every tool in this class, or only findAndRerank?
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. filter now uses the FilterDefinition description like sort does, and projection and options have real descriptions. The other tools stay as they were in this PR, and I can bring them to the same standard in a follow-up if you want. |
||
| @ToolArg( | ||
| description = "Options for result limits, reranking, and what the response includes", | ||
| required = false) | ||
| FindAndRerankCommand.Options options) { | ||
|
|
||
| var command = new FindAndRerankCommand(filter, projection, sort, options); | ||
| return mcpResource.processCollectionCommand(keyspace, collection, command); | ||
|
|
||
Uh oh!
There was an error while loading. Please reload this page.