|
| 1 | +package gospice |
| 2 | + |
| 3 | +import ( |
| 4 | + "bytes" |
| 5 | + "context" |
| 6 | + "encoding/json" |
| 7 | + "fmt" |
| 8 | + "io" |
| 9 | + "net/http" |
| 10 | +) |
| 11 | + |
| 12 | +// SearchRequest describes a search against the runtime's /v1/search endpoint. |
| 13 | +// |
| 14 | +// Only Text is required. Supplying Keywords adds a lexical pass, which the |
| 15 | +// runtime combines with the vector scores into a single hybrid ranking. |
| 16 | +type SearchRequest struct { |
| 17 | + // Text is the text to find similar documents for. Required. |
| 18 | + Text string `json:"text"` |
| 19 | + |
| 20 | + // Datasets restricts the search to the named datasets. When empty, the |
| 21 | + // runtime searches every searchable dataset. |
| 22 | + Datasets []string `json:"datasets,omitempty"` |
| 23 | + |
| 24 | + // Limit caps the number of matches returned per dataset. When nil, the |
| 25 | + // runtime applies its own default. |
| 26 | + Limit *int `json:"limit,omitempty"` |
| 27 | + |
| 28 | + // Where is a SQL predicate filtering candidate rows, without the leading |
| 29 | + // WHERE - for example "user_id = 42". |
| 30 | + Where *string `json:"where,omitempty"` |
| 31 | + |
| 32 | + // AdditionalColumns names extra columns to return with each match. A |
| 33 | + // primary key column is returned in SearchMatch.PrimaryKey, the rest in |
| 34 | + // SearchMatch.Data. |
| 35 | + AdditionalColumns []string `json:"additional_columns,omitempty"` |
| 36 | + |
| 37 | + // Keywords drives the lexical pass of a hybrid search. |
| 38 | + Keywords []string `json:"keywords,omitempty"` |
| 39 | +} |
| 40 | + |
| 41 | +// SearchMatch is a single document matched by Search. |
| 42 | +type SearchMatch struct { |
| 43 | + // Dataset is the dataset the match was found in. |
| 44 | + Dataset string `json:"dataset"` |
| 45 | + |
| 46 | + // Score is the match's similarity to the query. Higher is more similar. |
| 47 | + Score float64 `json:"_score"` |
| 48 | + |
| 49 | + // Matches holds the matched values keyed by the column they came from. |
| 50 | + // Each value is a slice because one column can contribute several chunks |
| 51 | + // to a single match. |
| 52 | + Matches map[string][]any `json:"matches"` |
| 53 | + |
| 54 | + // PrimaryKey identifies the matched row. Empty when the dataset declares |
| 55 | + // no primary key. |
| 56 | + PrimaryKey map[string]any `json:"primary_key"` |
| 57 | + |
| 58 | + // Data holds any AdditionalColumns that were requested. |
| 59 | + Data map[string]any `json:"data"` |
| 60 | + |
| 61 | + // Metadata holds extra per-match metadata the runtime attached. |
| 62 | + Metadata map[string]any `json:"metadata"` |
| 63 | +} |
| 64 | + |
| 65 | +// SearchResponse is the result of a single Search call. |
| 66 | +type SearchResponse struct { |
| 67 | + // Results are the matches, ordered by descending score. |
| 68 | + Results []SearchMatch `json:"results"` |
| 69 | + |
| 70 | + // DurationMs is how long the runtime reported the search took. |
| 71 | + DurationMs uint64 `json:"duration_ms"` |
| 72 | +} |
| 73 | + |
| 74 | +// Search finds documents similar to req.Text by calling the runtime's |
| 75 | +// /v1/search endpoint. |
| 76 | +// |
| 77 | +// It runs against datasets that have an embedding column and a loaded |
| 78 | +// embedding model. See https://docs.spice.ai/features/search-and-retrieval for |
| 79 | +// how to configure them. |
| 80 | +func (c *SpiceClient) Search(ctx context.Context, req *SearchRequest) (*SearchResponse, error) { |
| 81 | + if req == nil { |
| 82 | + return nil, fmt.Errorf("req is required") |
| 83 | + } |
| 84 | + if req.Text == "" { |
| 85 | + return nil, fmt.Errorf("req.Text is required and must be a non-empty search string") |
| 86 | + } |
| 87 | + if req.Limit != nil && *req.Limit < 1 { |
| 88 | + return nil, fmt.Errorf("req.Limit must be greater than 0, got %d", *req.Limit) |
| 89 | + } |
| 90 | + |
| 91 | + jsonData, err := json.Marshal(req) |
| 92 | + if err != nil { |
| 93 | + return nil, fmt.Errorf("error marshaling SearchRequest: %w", err) |
| 94 | + } |
| 95 | + |
| 96 | + url := fmt.Sprintf("%s/v1/search", c.baseHttpUrl) |
| 97 | + |
| 98 | + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewBuffer(jsonData)) |
| 99 | + if err != nil { |
| 100 | + return nil, fmt.Errorf("error creating request: %w", err) |
| 101 | + } |
| 102 | + |
| 103 | + httpReq = httpReq.WithContext(c.traceHttpRequest(ctx, "Search", httpReq)) |
| 104 | + |
| 105 | + httpReq.Header.Set("X-API-Key", c.apiKey) |
| 106 | + httpReq.Header.Set("Content-Type", "application/json") |
| 107 | + httpReq.Header.Set("user-agent", c.userAgent) |
| 108 | + |
| 109 | + resp, err := c.httpClient.Do(httpReq) |
| 110 | + if err != nil { |
| 111 | + return nil, fmt.Errorf("error executing request: %w", err) |
| 112 | + } |
| 113 | + defer func() { _ = resp.Body.Close() }() |
| 114 | + |
| 115 | + respBody, err := io.ReadAll(resp.Body) |
| 116 | + if err != nil { |
| 117 | + return nil, fmt.Errorf("error reading response from POST %s: %w", url, err) |
| 118 | + } |
| 119 | + |
| 120 | + if resp.StatusCode != http.StatusOK { |
| 121 | + // The runtime explains search failures in a plain-text body ("No data |
| 122 | + // sources provided"). Surface it rather than only the status code. |
| 123 | + return nil, fmt.Errorf("POST %s failed with status=%d: %s", url, resp.StatusCode, bytes.TrimSpace(respBody)) |
| 124 | + } |
| 125 | + |
| 126 | + var searchResp SearchResponse |
| 127 | + if err := json.Unmarshal(respBody, &searchResp); err != nil { |
| 128 | + return nil, fmt.Errorf("error decoding response from POST %s: %w", url, err) |
| 129 | + } |
| 130 | + |
| 131 | + return &searchResp, nil |
| 132 | +} |
0 commit comments