|
| 1 | +/* |
| 2 | +Copyright 2024 The Spice.ai OSS Authors |
| 3 | +
|
| 4 | +Permission is hereby granted, free of charge, to any person obtaining a copy |
| 5 | +of this software and associated documentation files (the "Software"), to deal |
| 6 | +in the Software without restriction, including without limitation the rights |
| 7 | +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell |
| 8 | +copies of the Software, and to permit persons to whom the Software is |
| 9 | +furnished to do so, subject to the following conditions: |
| 10 | +
|
| 11 | +The above copyright notice and this permission notice shall be included in all |
| 12 | +copies or substantial portions of the Software. |
| 13 | +
|
| 14 | +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR |
| 15 | +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, |
| 16 | +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE |
| 17 | +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER |
| 18 | +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, |
| 19 | +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE |
| 20 | +SOFTWARE. |
| 21 | +*/ |
| 22 | + |
| 23 | +using System.Text.Json; |
| 24 | +using System.Text.Json.Serialization; |
| 25 | + |
| 26 | +namespace Spice.Nsql; |
| 27 | + |
| 28 | +/// <summary> |
| 29 | +/// A natural-language query against the runtime's <c>/v1/nsql</c> endpoint. |
| 30 | +/// </summary> |
| 31 | +/// <remarks> |
| 32 | +/// Only <see cref="Query"/> is required. The runtime needs an LLM model configured in the |
| 33 | +/// Spicepod to translate it; when exactly one is configured, <see cref="Model"/> may be left |
| 34 | +/// unset and the runtime selects it. |
| 35 | +/// </remarks> |
| 36 | +public class NsqlRequest |
| 37 | +{ |
| 38 | + /// <summary> |
| 39 | + /// Creates an NSQL request. |
| 40 | + /// </summary> |
| 41 | + /// <param name="query">The question to answer, in natural language</param> |
| 42 | + public NsqlRequest(string query) |
| 43 | + { |
| 44 | + Query = query; |
| 45 | + } |
| 46 | + |
| 47 | + /// <summary> |
| 48 | + /// The question to answer, in natural language. Required. |
| 49 | + /// </summary> |
| 50 | + [JsonPropertyName("query")] |
| 51 | + public string Query { get; set; } |
| 52 | + |
| 53 | + /// <summary> |
| 54 | + /// Names the LLM used to generate SQL. When unset, the runtime uses the only compatible |
| 55 | + /// model configured in the Spicepod, and reports an error if there is not exactly one. |
| 56 | + /// </summary> |
| 57 | + [JsonPropertyName("model")] |
| 58 | + [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)] |
| 59 | + public string? Model { get; set; } |
| 60 | + |
| 61 | + /// <summary> |
| 62 | + /// Hints which datasets to sample when building model context. This is a sampling hint |
| 63 | + /// only - it does not restrict which tables the generated query may reference. When |
| 64 | + /// unset, all datasets are used. |
| 65 | + /// </summary> |
| 66 | + [JsonPropertyName("datasets")] |
| 67 | + [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)] |
| 68 | + public IReadOnlyList<string>? Datasets { get; set; } |
| 69 | + |
| 70 | + /// <summary> |
| 71 | + /// Includes sample rows in the context given to the model. It improves generation on |
| 72 | + /// ambiguous schemas at the cost of sending data values to the model. |
| 73 | + /// </summary> |
| 74 | + [JsonPropertyName("sample_data_enabled")] |
| 75 | + [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingDefault)] |
| 76 | + public bool SampleDataEnabled { get; set; } |
| 77 | + |
| 78 | + /// <summary> |
| 79 | + /// A stable key forwarded to the model provider for prompt caching. Reuse it across |
| 80 | + /// related requests to benefit from it. |
| 81 | + /// </summary> |
| 82 | + [JsonPropertyName("prompt_cache_key")] |
| 83 | + [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)] |
| 84 | + public string? PromptCacheKey { get; set; } |
| 85 | +} |
| 86 | + |
| 87 | +/// <summary> |
| 88 | +/// Describes one column of an <see cref="NsqlResponse"/>. |
| 89 | +/// </summary> |
| 90 | +public class NsqlField |
| 91 | +{ |
| 92 | + /// <summary> |
| 93 | + /// The column name. |
| 94 | + /// </summary> |
| 95 | + [JsonPropertyName("name")] |
| 96 | + public string Name { get; set; } = string.Empty; |
| 97 | + |
| 98 | + /// <summary> |
| 99 | + /// The column's Arrow type in its JSON encoding. Simple types encode as a quoted string |
| 100 | + /// (<c>"Utf8"</c>, <c>"Int64"</c>); parameterized ones as an object (for example |
| 101 | + /// <c>{"Timestamp":["Nanosecond",null]}</c>), which is why this is captured as raw JSON |
| 102 | + /// rather than a fixed shape. |
| 103 | + /// </summary> |
| 104 | + [JsonPropertyName("data_type")] |
| 105 | + public JsonElement DataType { get; set; } |
| 106 | + |
| 107 | + /// <summary> |
| 108 | + /// Whether the column admits nulls. |
| 109 | + /// </summary> |
| 110 | + [JsonPropertyName("nullable")] |
| 111 | + public bool Nullable { get; set; } |
| 112 | +} |
| 113 | + |
| 114 | +/// <summary> |
| 115 | +/// The schema of the rows an NSQL call returned. |
| 116 | +/// </summary> |
| 117 | +/// <remarks> |
| 118 | +/// <see cref="Fields"/> is empty when the generated query returned no rows - the runtime |
| 119 | +/// omits the schema body in that case. |
| 120 | +/// </remarks> |
| 121 | +public class NsqlSchema |
| 122 | +{ |
| 123 | + /// <summary> |
| 124 | + /// The columns of the result, in order. |
| 125 | + /// </summary> |
| 126 | + [JsonPropertyName("fields")] |
| 127 | + public IReadOnlyList<NsqlField> Fields { get; set; } = new List<NsqlField>(); |
| 128 | +} |
| 129 | + |
| 130 | +/// <summary> |
| 131 | +/// The result of running a natural-language query via <see cref="SpiceClient.NsqlAsync"/>. |
| 132 | +/// </summary> |
| 133 | +public class NsqlResponse |
| 134 | +{ |
| 135 | + /// <summary> |
| 136 | + /// The query the model generated. It is worth logging: a surprising result is usually a |
| 137 | + /// surprising query. |
| 138 | + /// </summary> |
| 139 | + [JsonPropertyName("sql")] |
| 140 | + public string SQL { get; set; } = string.Empty; |
| 141 | + |
| 142 | + /// <summary> |
| 143 | + /// The number of rows returned. |
| 144 | + /// </summary> |
| 145 | + [JsonPropertyName("row_count")] |
| 146 | + public int RowCount { get; set; } |
| 147 | + |
| 148 | + /// <summary> |
| 149 | + /// Describes the columns in <see cref="Data"/>. |
| 150 | + /// </summary> |
| 151 | + [JsonPropertyName("schema")] |
| 152 | + public NsqlSchema Schema { get; set; } = new(); |
| 153 | + |
| 154 | + /// <summary> |
| 155 | + /// The rows, each keyed by column name. Values are decoded from JSON, so they carry |
| 156 | + /// JSON's types rather than the Arrow types named in <see cref="Schema"/> - numbers |
| 157 | + /// arrive as <see cref="JsonElement"/> holding a number. Use |
| 158 | + /// <see cref="SpiceClient.NsqlGenerateSqlAsync"/> with <see cref="SpiceClient.Query"/> |
| 159 | + /// when Arrow-typed results matter. |
| 160 | + /// </summary> |
| 161 | + [JsonPropertyName("data")] |
| 162 | + public IReadOnlyList<IReadOnlyDictionary<string, JsonElement>> Data { get; set; } |
| 163 | + = new List<IReadOnlyDictionary<string, JsonElement>>(); |
| 164 | +} |
0 commit comments