|
| 1 | +package gospice |
| 2 | + |
| 3 | +import ( |
| 4 | + "bytes" |
| 5 | + "context" |
| 6 | + "encoding/json" |
| 7 | + "fmt" |
| 8 | + "io" |
| 9 | + "net/http" |
| 10 | +) |
| 11 | + |
| 12 | +// nsqlJSONMediaType asks the runtime for the envelope carrying the generated |
| 13 | +// SQL alongside the results. Without it /v1/nsql returns a bare array of rows |
| 14 | +// and the generated SQL is lost. |
| 15 | +const nsqlJSONMediaType = "application/vnd.spiceai.nsql.v1+json" |
| 16 | + |
| 17 | +// nsqlSQLMediaType asks the runtime to generate SQL without executing it. |
| 18 | +const nsqlSQLMediaType = "application/sql" |
| 19 | + |
| 20 | +// NsqlRequest describes a natural-language query against the runtime's |
| 21 | +// /v1/nsql endpoint. |
| 22 | +// |
| 23 | +// Only Query is required. The runtime needs an LLM model configured in the |
| 24 | +// Spicepod to translate it; when exactly one is configured, Model may be left |
| 25 | +// empty and the runtime selects it. |
| 26 | +type NsqlRequest struct { |
| 27 | + // Query is the question to answer, in natural language. Required. |
| 28 | + Query string `json:"query"` |
| 29 | + |
| 30 | + // Model names the LLM used to generate SQL. When empty, the runtime uses |
| 31 | + // the only compatible model configured in the Spicepod, and reports an |
| 32 | + // error if there is not exactly one. |
| 33 | + Model string `json:"model,omitempty"` |
| 34 | + |
| 35 | + // Datasets hints which datasets to sample when building model context. |
| 36 | + // This is a sampling hint only - it does not restrict which tables the |
| 37 | + // generated query may reference. When empty, all datasets are used. |
| 38 | + Datasets []string `json:"datasets,omitempty"` |
| 39 | + |
| 40 | + // SampleDataEnabled includes sample rows in the context given to the |
| 41 | + // model. It improves generation on ambiguous schemas at the cost of |
| 42 | + // sending data values to the model. |
| 43 | + SampleDataEnabled bool `json:"sample_data_enabled,omitempty"` |
| 44 | + |
| 45 | + // PromptCacheKey is a stable key forwarded to the model provider for |
| 46 | + // prompt caching. Reuse it across related requests to benefit from it. |
| 47 | + PromptCacheKey string `json:"prompt_cache_key,omitempty"` |
| 48 | +} |
| 49 | + |
| 50 | +// NsqlField describes one column of an NsqlResponse. |
| 51 | +type NsqlField struct { |
| 52 | + // Name is the column name. |
| 53 | + Name string `json:"name"` |
| 54 | + |
| 55 | + // DataType is the column's Arrow type in its JSON encoding. Simple types |
| 56 | + // encode as a quoted string ("Utf8", "Int64"); parameterized ones as an |
| 57 | + // object (for example {"Timestamp":["Nanosecond",null]}). |
| 58 | + DataType json.RawMessage `json:"data_type"` |
| 59 | + |
| 60 | + // Nullable reports whether the column admits nulls. |
| 61 | + Nullable bool `json:"nullable"` |
| 62 | +} |
| 63 | + |
| 64 | +// NsqlSchema is the schema of the rows an Nsql call returned. |
| 65 | +// |
| 66 | +// Fields is empty when the generated query returned no rows - the runtime |
| 67 | +// omits the schema body in that case. |
| 68 | +type NsqlSchema struct { |
| 69 | + Fields []NsqlField `json:"fields"` |
| 70 | +} |
| 71 | + |
| 72 | +// NsqlResponse is the result of running a natural-language query. |
| 73 | +type NsqlResponse struct { |
| 74 | + // SQL is the query the model generated. It is worth logging: a surprising |
| 75 | + // result is usually a surprising query. |
| 76 | + SQL string `json:"sql"` |
| 77 | + |
| 78 | + // RowCount is the number of rows returned. |
| 79 | + RowCount int `json:"row_count"` |
| 80 | + |
| 81 | + // Schema describes the columns in Data. |
| 82 | + Schema NsqlSchema `json:"schema"` |
| 83 | + |
| 84 | + // Data holds the rows, each keyed by column name. Values are decoded from |
| 85 | + // JSON, so they carry JSON's types rather than the Arrow types named in |
| 86 | + // Schema - numbers arrive as float64. Use NsqlGenerateSQL with Query when |
| 87 | + // Arrow-typed results matter. |
| 88 | + Data []map[string]any `json:"data"` |
| 89 | +} |
| 90 | + |
| 91 | +// Nsql answers req.Query by having the runtime's configured LLM generate SQL, |
| 92 | +// then running it. |
| 93 | +// |
| 94 | +// The generated SQL is returned in NsqlResponse.SQL. The runtime executes it |
| 95 | +// read-only and retries generation when the query fails to run, so a returned |
| 96 | +// error means generation or execution failed repeatedly. |
| 97 | +// |
| 98 | +// Nsql requires an LLM model in the Spicepod. See |
| 99 | +// https://docs.spice.ai/features/text-to-sql for how to configure one. |
| 100 | +func (c *SpiceClient) Nsql(ctx context.Context, req *NsqlRequest) (*NsqlResponse, error) { |
| 101 | + respBody, err := c.doNsqlRequest(ctx, req, "Nsql", nsqlJSONMediaType) |
| 102 | + if err != nil { |
| 103 | + return nil, err |
| 104 | + } |
| 105 | + |
| 106 | + var nsqlResp NsqlResponse |
| 107 | + if err := json.Unmarshal(respBody, &nsqlResp); err != nil { |
| 108 | + return nil, fmt.Errorf("error decoding response from POST %s/v1/nsql: %w", c.baseHttpUrl, err) |
| 109 | + } |
| 110 | + |
| 111 | + return &nsqlResp, nil |
| 112 | +} |
| 113 | + |
| 114 | +// NsqlGenerateSQL translates req.Query into SQL without running it. |
| 115 | +// |
| 116 | +// Use it to inspect or edit the query before running it, or to run it through |
| 117 | +// Query or Sql so the results arrive as Arrow rather than decoded JSON. |
| 118 | +func (c *SpiceClient) NsqlGenerateSQL(ctx context.Context, req *NsqlRequest) (string, error) { |
| 119 | + respBody, err := c.doNsqlRequest(ctx, req, "NsqlGenerateSQL", nsqlSQLMediaType) |
| 120 | + if err != nil { |
| 121 | + return "", err |
| 122 | + } |
| 123 | + |
| 124 | + return string(bytes.TrimSpace(respBody)), nil |
| 125 | +} |
| 126 | + |
| 127 | +// doNsqlRequest posts req to /v1/nsql asking for accept, and returns the |
| 128 | +// response body when the runtime answered 200. |
| 129 | +func (c *SpiceClient) doNsqlRequest(ctx context.Context, req *NsqlRequest, operation string, accept string) ([]byte, error) { |
| 130 | + if req == nil { |
| 131 | + return nil, fmt.Errorf("req is required") |
| 132 | + } |
| 133 | + if req.Query == "" { |
| 134 | + return nil, fmt.Errorf("req.Query is required and must be a non-empty natural language query") |
| 135 | + } |
| 136 | + |
| 137 | + jsonData, err := json.Marshal(req) |
| 138 | + if err != nil { |
| 139 | + return nil, fmt.Errorf("error marshaling NsqlRequest: %w", err) |
| 140 | + } |
| 141 | + |
| 142 | + url := fmt.Sprintf("%s/v1/nsql", c.baseHttpUrl) |
| 143 | + |
| 144 | + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewBuffer(jsonData)) |
| 145 | + if err != nil { |
| 146 | + return nil, fmt.Errorf("error creating request: %w", err) |
| 147 | + } |
| 148 | + |
| 149 | + httpReq = httpReq.WithContext(c.traceHttpRequest(ctx, operation, httpReq)) |
| 150 | + |
| 151 | + httpReq.Header.Set("Content-Type", "application/json") |
| 152 | + httpReq.Header.Set("Accept", accept) |
| 153 | + httpReq.Header.Set("user-agent", c.userAgent) |
| 154 | + // Only send the key when there is one — an empty X-API-Key reads as a |
| 155 | + // supplied-but-invalid credential to auth middleware, which is different |
| 156 | + // from omitting the header. Matches IsSpiceReady in client.go. |
| 157 | + if c.apiKey != "" { |
| 158 | + httpReq.Header.Set("X-API-Key", c.apiKey) |
| 159 | + } |
| 160 | + |
| 161 | + resp, err := c.httpClient.Do(httpReq) |
| 162 | + if err != nil { |
| 163 | + return nil, fmt.Errorf("error executing request: %w", err) |
| 164 | + } |
| 165 | + defer func() { _ = resp.Body.Close() }() |
| 166 | + |
| 167 | + respBody, err := io.ReadAll(resp.Body) |
| 168 | + if err != nil { |
| 169 | + return nil, fmt.Errorf("error reading response from POST %s: %w", url, err) |
| 170 | + } |
| 171 | + |
| 172 | + if resp.StatusCode != http.StatusOK { |
| 173 | + // The runtime explains NSQL failures in a plain-text body - a missing |
| 174 | + // or ambiguous model, or SQL that would not run. Surface it rather |
| 175 | + // than only the status code. |
| 176 | + return nil, fmt.Errorf("POST %s failed with status=%d: %s", url, resp.StatusCode, bytes.TrimSpace(respBody)) |
| 177 | + } |
| 178 | + |
| 179 | + return respBody, nil |
| 180 | +} |
0 commit comments