diff --git a/.goreleaser.yaml b/.goreleaser.yaml index df44a0b4..64d6400a 100644 --- a/.goreleaser.yaml +++ b/.goreleaser.yaml @@ -31,7 +31,7 @@ universal_binaries: - replace: true archives: - - format: tar.gz + - formats: ['tar.gz'] # this name template makes the OS and Arch compatible with the results of `uname`. name_template: >- {{ .ProjectName }}_ @@ -45,7 +45,7 @@ archives: # use zip for windows archives format_overrides: - goos: windows - format: zip + formats: ['zip'] report_sizes: true @@ -69,7 +69,7 @@ release: owner: pinecone-io name: cli -brews: +homebrew_casks: - name: pinecone homepage: "https://www.pinecone.io" description: "Pinecone CLI" @@ -81,21 +81,12 @@ brews: commit_author: name: goreleaserbot email: bot@goreleaser.com - commit_msg_template: "Brew formula update for pinecone version {{ .Tag }}" + commit_msg_template: "Brew cask update for pinecone version {{ .Tag }}" skip_upload: auto - directory: Formula - license: "Apache-2.0" - test: | - system "#{bin}/pc --help" - install: | - bin.install "pc" - bin.install_symlink "pc" => "pinecone" - - # Install man pages - man1.install Dir["man/man1/*.1"] - - # Add aliases: pc*.1 -> pinecone*.1, etc - Dir[man1/"pc*.1"].each do |src| - dest = File.basename(src).sub(/\Apc\b/, "pinecone") # Replace leading "pc" with "pinecone" - man1.install_symlink src => dest - end + directory: Casks + binaries: + - pc + manpages: + - "man/man1/*.1" + custom_block: | + binary "pc", target: "pinecone" diff --git a/cmd/gen-manpages/main.go b/cmd/gen-manpages/main.go index 25ebcf3a..0ba998b0 100644 --- a/cmd/gen-manpages/main.go +++ b/cmd/gen-manpages/main.go @@ -50,6 +50,23 @@ func main() { os.Exit(1) } + // Create pinecone*.1 symlinks for each pc*.1 man page so that + // `man pinecone` and `man pinecone-index` etc. work in addition to `man pc`. + pcPages, err := filepath.Glob(filepath.Join(*output, "pc*.1")) + if err != nil { + fmt.Fprintf(os.Stderr, "Warning: Failed to list man page files: %s\n", err) + } else { + for _, src := range pcPages { + base := filepath.Base(src) + dest := "pinecone" + base[len("pc"):] + destPath := filepath.Join(*output, dest) + os.Remove(destPath) + if err := os.Symlink(base, destPath); err != nil { + fmt.Fprintf(os.Stderr, "Warning: Failed to create symlink %s -> %s: %s\n", dest, base, err) + } + } + } + // List generated files for verbose if *verbose { files, err := filepath.Glob(filepath.Join(*output, "*.1")) diff --git a/go.mod b/go.mod index dd685817..8f7e0a23 100644 --- a/go.mod +++ b/go.mod @@ -10,7 +10,7 @@ require ( github.com/charmbracelet/lipgloss v0.10.0 github.com/fatih/color v1.16.0 github.com/golang-jwt/jwt/v5 v5.2.2 - github.com/pinecone-io/go-pinecone/v5 v5.2.0 + github.com/pinecone-io/go-pinecone/v5 v5.4.1 github.com/rs/zerolog v1.32.0 github.com/spf13/cobra v1.8.0 github.com/spf13/viper v1.18.2 diff --git a/go.sum b/go.sum index 59d64829..30e7d993 100644 --- a/go.sum +++ b/go.sum @@ -84,8 +84,8 @@ github.com/oapi-codegen/runtime v1.1.2 h1:P2+CubHq8fO4Q6fV1tqDBZHCwpVpvPg7oKiYzQ github.com/oapi-codegen/runtime v1.1.2/go.mod h1:SK9X900oXmPWilYR5/WKPzt3Kqxn/uS/+lbpREv+eCg= github.com/pelletier/go-toml/v2 v2.1.0 h1:FnwAJ4oYMvbT/34k9zzHuZNrhlz48GB3/s6at6/MHO4= github.com/pelletier/go-toml/v2 v2.1.0/go.mod h1:tJU2Z3ZkXwnxa4DPO899bsyIoywizdUvyaeZurnPPDc= -github.com/pinecone-io/go-pinecone/v5 v5.2.0 h1:QlkWUdjctCq+846m7bVhGNor+6PS2fubhjf55Uz4glM= -github.com/pinecone-io/go-pinecone/v5 v5.2.0/go.mod h1:6Fg85fcyvMUQFf9KW7zniN81kelSYvsjF+KPLdc1MGA= +github.com/pinecone-io/go-pinecone/v5 v5.4.1 h1:JJJ4VIu5NpFc3BIRcjc93n/XxYtACwRjRI/e6eHoOIU= +github.com/pinecone-io/go-pinecone/v5 v5.4.1/go.mod h1:6Fg85fcyvMUQFf9KW7zniN81kelSYvsjF+KPLdc1MGA= github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U= diff --git a/internal/pkg/cli/command/index/cmd.go b/internal/pkg/cli/command/index/cmd.go index b605b7be..d6818783 100644 --- a/internal/pkg/cli/command/index/cmd.go +++ b/internal/pkg/cli/command/index/cmd.go @@ -2,6 +2,7 @@ package index import ( "github.com/pinecone-io/cli/internal/pkg/cli/command/index/namespace" + "github.com/pinecone-io/cli/internal/pkg/cli/command/index/record" "github.com/pinecone-io/cli/internal/pkg/cli/command/index/vector" "github.com/pinecone-io/cli/internal/pkg/utils/help" "github.com/spf13/cobra" @@ -42,6 +43,7 @@ func NewIndexCmd() *cobra.Command { cmd.AddCommand(NewDescribeIndexStatsCmd()) cmd.AddGroup(help.GROUP_INDEX_DATA) + cmd.AddCommand(record.NewRecordCmd()) cmd.AddCommand(vector.NewVectorCmd()) cmd.AddGroup(help.GROUP_INDEX_NAMESPACE) diff --git a/internal/pkg/cli/command/index/record/cmd.go b/internal/pkg/cli/command/index/record/cmd.go new file mode 100644 index 00000000..7a840d66 --- /dev/null +++ b/internal/pkg/cli/command/index/record/cmd.go @@ -0,0 +1,31 @@ +package record + +import ( + "github.com/pinecone-io/cli/internal/pkg/utils/help" + "github.com/spf13/cobra" +) + +func NewRecordCmd() *cobra.Command { + cmd := &cobra.Command{ + Use: "record", + Short: "Work with records in an index", + Long: help.Long(` + Work with records in an integrated Pinecone index. + + Use these commands to upsert raw records and run semantic search against + them. All commands require --index-name and may optionally target a + --namespace. + `), + Example: help.Examples(` + pc index record upsert --index-name my-index --namespace my-namespace --body ./records.jsonl + pc index record search --index-name my-index --namespace my-namespace --inputs '{"text":"search query"}' + pc index record search --index-name my-index --namespace my-namespace --body ./search.json + `), + GroupID: help.GROUP_INDEX_DATA.ID, + } + + cmd.AddCommand(NewUpsertCmd()) + cmd.AddCommand(NewSearchCmd()) + + return cmd +} diff --git a/internal/pkg/cli/command/index/record/search.go b/internal/pkg/cli/command/index/record/search.go new file mode 100644 index 00000000..8d321ad6 --- /dev/null +++ b/internal/pkg/cli/command/index/record/search.go @@ -0,0 +1,189 @@ +package record + +import ( + "encoding/json" + + "github.com/pinecone-io/cli/internal/pkg/utils/argio" + "github.com/pinecone-io/cli/internal/pkg/utils/exit" + "github.com/pinecone-io/cli/internal/pkg/utils/flags" + "github.com/pinecone-io/cli/internal/pkg/utils/help" + "github.com/pinecone-io/cli/internal/pkg/utils/msg" + "github.com/pinecone-io/cli/internal/pkg/utils/pcio" + "github.com/pinecone-io/cli/internal/pkg/utils/presenters" + "github.com/pinecone-io/cli/internal/pkg/utils/sdk" + "github.com/pinecone-io/cli/internal/pkg/utils/style" + "github.com/pinecone-io/cli/internal/pkg/utils/text" + "github.com/spf13/cobra" + + "github.com/pinecone-io/go-pinecone/v5/pinecone" +) + +const defaultSearchTopK = 10 + +type searchCmdOptions struct { + indexName string + namespace string + topK int + inputs flags.JSONObject + filter flags.JSONObject + rerank flags.JSONObject + id string + fields flags.StringList + body string + json bool +} + +func NewSearchCmd() *cobra.Command { + options := searchCmdOptions{ + topK: defaultSearchTopK, + } + + cmd := &cobra.Command{ + Use: "search", + Short: "Search records in an integrated index", + Long: help.Long(` + Run semantic search against records in an integrated index. + + Provide query text via --inputs (inline JSON, ./path.json, or '-' for stdin). + Narrow results with --filter (metadata filter as a JSON object). + Rerank results with --rerank (JSON object with required fields: model, rank_fields). + Use --body to supply a full request body for advanced parameters like + vector overrides or match_terms. + + When a flag and --body both specify the same field, the flag takes precedence. + `), + Example: help.Examples(` + pc index record search --index-name my-index --namespace my-namespace --inputs '{"text":"find similar"}' + pc index record search --index-name my-index --inputs '{"text":"disease prevention"}' --filter '{"category":"health"}' + pc index record search --index-name my-index --inputs '{"text":"find similar"}' --rerank '{"model":"bge-reranker-v2-m3","rank_fields":["chunk_text"]}' + echo '{"text":"disease prevention"}' | pc index record search --index-name my-index --inputs - + pc index record search --index-name my-index --namespace my-namespace --body ./search.json + `), + Run: func(cmd *cobra.Command, args []string) { + runSearchCmd(cmd, options) + }, + } + + cmd.Flags().StringVarP(&options.indexName, "index-name", "n", "", "name of the index to search") + cmd.Flags().StringVar(&options.namespace, "namespace", "__default__", "namespace to search") + cmd.Flags().IntVarP(&options.topK, "top-k", "k", defaultSearchTopK, "number of results to return") + cmd.Flags().Var(&options.inputs, "inputs", "query inputs for search (inline JSON, ./path.json, or '-' for stdin)") + cmd.Flags().Var(&options.filter, "filter", "metadata filter (inline JSON, ./path.json, or '-' for stdin)") + cmd.Flags().Var(&options.rerank, "rerank", "rerank results (inline JSON, ./path.json, or '-' for stdin); required fields: model (string), rank_fields (string array)") + cmd.Flags().StringVar(&options.id, "id", "", "use an existing record's vector by ID for the query") + cmd.Flags().Var(&options.fields, "fields", "fields to return in results (JSON string array, ./path.json, or '-' for stdin)") + cmd.Flags().StringVar(&options.body, "body", "", "request body JSON (inline, ./path.json, or '-' for stdin; only one argument may use stdin)") + cmd.Flags().BoolVarP(&options.json, "json", "j", false, "output as JSON") + + _ = cmd.MarkFlagRequired("index-name") + + return cmd +} + +func runSearchCmd(cmd *cobra.Command, options searchCmdOptions) { + ctx := cmd.Context() + pc := sdk.NewPineconeClient(ctx) + + // Build req from flags. + req := pinecone.SearchRecordsRequest{ + Query: pinecone.SearchRecordsQuery{ + TopK: int32(options.topK), + }, + } + + if options.id != "" { + req.Query.Id = &options.id + } + if options.inputs != nil { + inputs := map[string]interface{}(options.inputs) + req.Query.Inputs = &inputs + } + if options.filter != nil { + filter := map[string]interface{}(options.filter) + req.Query.Filter = &filter + } + if len(options.fields) > 0 { + fieldsCopy := make([]string, len(options.fields)) + copy(fieldsCopy, options.fields) + req.Fields = &fieldsCopy + } + if options.rerank != nil { + b, err := json.Marshal(options.rerank) + if err != nil { + msg.FailMsg("Failed to encode --rerank value: %s", err) + exit.Error(err, "Failed to encode --rerank value") + } + var rerank pinecone.SearchRecordsRerank + if err := json.Unmarshal(b, &rerank); err != nil { + msg.FailMsg("Failed to parse --rerank value: %s", err) + exit.Error(err, "Failed to parse --rerank value") + } + req.Rerank = &rerank + } + + // Merge --body into req. For fields that have a dedicated flag, the flag + // takes precedence; body only fills in values that weren't explicitly set. + // Fields with no dedicated flag (Vector, MatchTerms, Rerank) always come + // from body. + if options.body != "" { + b, src, err := argio.DecodeJSONArg[pinecone.SearchRecordsRequest](options.body) + if err != nil { + msg.FailMsg("Failed to parse search body (%s): %s", style.Emphasis(src.Label), err) + exit.Errorf(err, "Failed to parse search body (%s)", src.Label) + } + if b != nil { + if !cmd.Flags().Changed("top-k") && b.Query.TopK > 0 { + req.Query.TopK = b.Query.TopK + } + if !cmd.Flags().Changed("id") && b.Query.Id != nil { + req.Query.Id = b.Query.Id + } + if !cmd.Flags().Changed("inputs") && b.Query.Inputs != nil { + req.Query.Inputs = b.Query.Inputs + } + if !cmd.Flags().Changed("filter") && b.Query.Filter != nil { + req.Query.Filter = b.Query.Filter + } + if !cmd.Flags().Changed("fields") && b.Fields != nil { + req.Fields = b.Fields + } + if b.Query.Vector != nil { + req.Query.Vector = b.Query.Vector + } + if b.Query.MatchTerms != nil { + req.Query.MatchTerms = b.Query.MatchTerms + } + if !cmd.Flags().Changed("rerank") && b.Rerank != nil { + req.Rerank = b.Rerank + } + } + } + + if req.Query.TopK <= 0 { + msg.FailMsg("Top-k must be greater than 0") + exit.ErrorMsg("Invalid top-k value") + } + + if req.Query.Id == nil && req.Query.Inputs == nil && req.Query.Vector == nil && req.Query.MatchTerms == nil { + msg.FailMsg("Provide --inputs, --id, or a body with vector/match_terms") + exit.ErrorMsg("Missing query inputs for search") + } + + ic, err := sdk.NewIndexConnection(ctx, pc, options.indexName, options.namespace) + if err != nil { + msg.FailMsg("Failed to create index connection: %s", err) + exit.Error(err, "Failed to create index connection") + } + + resp, err := ic.SearchRecords(ctx, &req) + if err != nil { + exit.Error(err, "Failed to search records") + } + + if options.json { + json := text.IndentJSON(resp) + pcio.Println(json) + } else { + presenters.PrintSearchRecordsTable(resp) + } +} diff --git a/internal/pkg/cli/command/index/record/upsert.go b/internal/pkg/cli/command/index/record/upsert.go new file mode 100644 index 00000000..a4de70d1 --- /dev/null +++ b/internal/pkg/cli/command/index/record/upsert.go @@ -0,0 +1,181 @@ +package record + +import ( + "bytes" + "context" + "encoding/json" + "io" + + "github.com/pinecone-io/cli/internal/pkg/utils/argio" + "github.com/pinecone-io/cli/internal/pkg/utils/exit" + "github.com/pinecone-io/cli/internal/pkg/utils/help" + "github.com/pinecone-io/cli/internal/pkg/utils/msg" + "github.com/pinecone-io/cli/internal/pkg/utils/pcio" + "github.com/pinecone-io/cli/internal/pkg/utils/sdk" + "github.com/pinecone-io/cli/internal/pkg/utils/style" + "github.com/pinecone-io/cli/internal/pkg/utils/text" + "github.com/spf13/cobra" + + "github.com/pinecone-io/go-pinecone/v5/pinecone" +) + +// UpsertRecordsBody is the JSON payload for --body/--file. +// It accepts either {"records": [...]} where each element is an IntegratedRecord, +// a JSON array of IntegratedRecord objects, or a JSONL stream of IntegratedRecord objects. +type UpsertRecordsBody struct { + Records []pinecone.IntegratedRecord `json:"records"` +} + +type upsertCmdOptions struct { + file string + indexName string + namespace string + batchSize int + json bool +} + +func NewUpsertCmd() *cobra.Command { + options := upsertCmdOptions{} + + cmd := &cobra.Command{ + Use: "upsert", + Short: "Upsert records into an index from a JSON/JSONL payload", + Long: help.Long(` + Upsert records into an index namespace from a JSON or JSONL payload. + + The request --body/--file may be a JSON object containing "records": [...] + (a list of IntegratedRecord objects), a raw JSON array of records, or a + JSONL stream of IntegratedRecord objects. Bodies can be inline JSON, + loaded from ./file.json[l], or read from stdin with '-'. + + Body schema: UpsertRecordsBody (records shaped like pinecone.IntegratedRecord: + https://pkg.go.dev/github.com/pinecone-io/go-pinecone/v5/pinecone#IntegratedRecord) + `), + Example: help.Examples(` + pc index record upsert --index-name my-index --namespace my-namespace --body ./records.json + pc index record upsert --index-name my-index --namespace my-namespace --body ./records.jsonl + cat records.jsonl | pc index record upsert --index-name my-index --namespace my-namespace --body - + `), + Run: func(cmd *cobra.Command, args []string) { + runUpsertCmd(cmd.Context(), options) + }, + } + + cmd.Flags().StringVarP(&options.indexName, "index-name", "n", "", "name of index to upsert into") + cmd.Flags().StringVar(&options.namespace, "namespace", "__default__", "namespace to upsert into") + cmd.Flags().StringVar(&options.file, "file", "", "request body JSON or JSONL (inline, ./path.json[l], or '-' for stdin; only one argument may use stdin)") + cmd.Flags().StringVar(&options.file, "body", "", "alias for --file") + cmd.Flags().IntVarP(&options.batchSize, "batch-size", "b", 96, "records per batch (max 96)") + cmd.Flags().BoolVarP(&options.json, "json", "j", false, "output as JSON") + + _ = cmd.MarkFlagRequired("index-name") + + return cmd +} + +func runUpsertCmd(ctx context.Context, options upsertCmdOptions) { + if options.file == "" { + msg.FailMsg("Either --file or --body must be provided") + exit.ErrorMsg("Either --file or --body must be provided") + } + + b, src, err := argio.ReadAll(options.file) + if err != nil { + msg.FailMsg("Failed to read upsert body (%s): %s", style.Emphasis(src.Label), err) + exit.Error(err, "Failed to read upsert body") + } + + payload, err := parseUpsertRecordsBody(b) + if err != nil { + msg.FailMsg("Failed to parse upsert body (%s): %s", style.Emphasis(src.Label), err) + exit.Error(err, "Failed to parse upsert body") + } + + pc := sdk.NewPineconeClient(ctx) + ic, err := sdk.NewIndexConnection(ctx, pc, options.indexName, options.namespace) + if err != nil { + msg.FailMsg("Failed to create index connection: %s", err) + exit.Error(err, "Failed to create index connection") + } + + if len(payload.Records) == 0 { + msg.FailMsg("No records found in %s", style.Emphasis(src.Label)) + exit.ErrorMsg("No records provided for upsert") + } + + records := make([]*pinecone.IntegratedRecord, 0, len(payload.Records)) + for i := range payload.Records { + records = append(records, &payload.Records[i]) + } + + if options.batchSize <= 0 { + options.batchSize = len(records) + } + + batches := make([][]*pinecone.IntegratedRecord, 0, (len(records)+options.batchSize-1)/options.batchSize) + for i := 0; i < len(records); i += options.batchSize { + end := i + options.batchSize + if end > len(records) { + end = len(records) + } + batches = append(batches, records[i:end]) + } + + for i, batch := range batches { + err := ic.UpsertRecords(ctx, batch) + if err != nil { + msg.FailMsg("Failed to upsert %d records in batch %d: %s", len(batch), i+1, err) + exit.Errorf(err, "Failed to upsert %d records in batch %d", len(batch), i+1) + } else if options.json { + summary := map[string]any{ + "batch": i + 1, + "batches": len(batches), + "records": len(batch), + "namespace": options.namespace, + } + pcio.Println(text.IndentJSON(summary)) + } else { + msg.SuccessMsg("Upserted %d records into namespace %s (batch %d of %d)", len(batch), options.namespace, i+1, len(batches)) + } + } +} + +func parseUpsertRecordsBody(b []byte) (*UpsertRecordsBody, error) { + // First try JSON object: {"records":[...]} + { + var payload UpsertRecordsBody + dec := json.NewDecoder(bytes.NewReader(b)) + dec.DisallowUnknownFields() + if err := dec.Decode(&payload); err == nil && len(payload.Records) > 0 { + return &payload, nil + } + } + + // Second try: raw JSON array of IntegratedRecord objects + { + var records []pinecone.IntegratedRecord + dec := json.NewDecoder(bytes.NewReader(b)) + dec.DisallowUnknownFields() + if err := dec.Decode(&records); err == nil && len(records) > 0 { + return &UpsertRecordsBody{Records: records}, nil + } + } + + // Fallback: JSONL/stream of pinecone.IntegratedRecord values + var records []pinecone.IntegratedRecord + dec := json.NewDecoder(bytes.NewReader(b)) + dec.DisallowUnknownFields() + for { + var rec pinecone.IntegratedRecord + if err := dec.Decode(&rec); err == io.EOF { + break + } else if err != nil { + return nil, err + } + records = append(records, rec) + } + if len(records) == 0 { + return nil, io.EOF + } + return &UpsertRecordsBody{Records: records}, nil +} diff --git a/internal/pkg/cli/command/index/vector/upsert.go b/internal/pkg/cli/command/index/vector/upsert.go index 03dd947c..53c520b6 100644 --- a/internal/pkg/cli/command/index/vector/upsert.go +++ b/internal/pkg/cli/command/index/vector/upsert.go @@ -22,7 +22,7 @@ import ( // UpsertBody is the JSON payload for --file / --body. // It accepts either {"vectors": [...]} with elements shaped like pinecone.Vector // (see https://pkg.go.dev/github.com/pinecone-io/go-pinecone/v5/pinecone#Vector), -// or a JSONL stream of pinecone.Vector objects. +// a raw JSON array of pinecone.Vector objects, or a JSONL stream of pinecone.Vector objects. type UpsertBody struct { Vectors []pinecone.Vector `json:"vectors"` } @@ -44,7 +44,7 @@ func NewUpsertCmd() *cobra.Command { Long: help.Long(` Upsert vectors into an index namespace from a JSON or JSONL payload. - The request --file may be a JSON object containing "vectors": [...] or a JSONL stream of Vector objects. + The request --file may be a JSON object containing "vectors": [...], a raw JSON array of Vector objects, or a JSONL stream of Vector objects. Control batch size with --batch-size. Bodies can be inline JSON, loaded from ./file.json[l], or read from stdin with '-'. Body schema: UpsertBody (vectors shaped like pinecone.Vector: https://pkg.go.dev/github.com/pinecone-io/go-pinecone/v5/pinecone#Vector) @@ -150,9 +150,9 @@ func runUpsertCmd(ctx context.Context, options upsertCmdOptions) { } func parseUpsertBody(b []byte) (*UpsertBody, error) { - var payload UpsertBody // First try and decode as JSON: {"vectors":[...]} { + var payload UpsertBody dec := json.NewDecoder(bytes.NewReader(b)) dec.DisallowUnknownFields() if err := dec.Decode(&payload); err == nil && len(payload.Vectors) > 0 { @@ -160,6 +160,16 @@ func parseUpsertBody(b []byte) (*UpsertBody, error) { } } + // Second try: raw JSON array of Vector objects + { + var vectors []pinecone.Vector + dec := json.NewDecoder(bytes.NewReader(b)) + dec.DisallowUnknownFields() + if err := dec.Decode(&vectors); err == nil && len(vectors) > 0 { + return &UpsertBody{Vectors: vectors}, nil + } + } + // Fallback: JSONL/stream of pinecone.Vector values var vectors []pinecone.Vector dec := json.NewDecoder(bytes.NewReader(b)) diff --git a/internal/pkg/utils/presenters/search_records.go b/internal/pkg/utils/presenters/search_records.go new file mode 100644 index 00000000..af5693e3 --- /dev/null +++ b/internal/pkg/utils/presenters/search_records.go @@ -0,0 +1,68 @@ +package presenters + +import ( + "sort" + "strings" + + "github.com/pinecone-io/cli/internal/pkg/utils/pcio" + "github.com/pinecone-io/cli/internal/pkg/utils/text" + "github.com/pinecone-io/go-pinecone/v5/pinecone" +) + +func PrintSearchRecordsTable(resp *pinecone.SearchRecordsResponse) { + writer := NewTabWriter() + if resp == nil { + PrintEmptyState(writer, "search results") + return + } + + if resp.Usage.ReadUnits > 0 { + pcio.Fprintf(writer, "Usage: %d (read units)\n", resp.Usage.ReadUnits) + } + if resp.Usage.EmbedTotalTokens != nil { + pcio.Fprintf(writer, "Embed tokens: %d\n", *resp.Usage.EmbedTotalTokens) + } + if resp.Usage.RerankUnits != nil { + pcio.Fprintf(writer, "Rerank units: %d\n", *resp.Usage.RerankUnits) + } + + pcio.Fprintln(writer, "ID\tSCORE\tFIELDS") + + for _, hit := range resp.Result.Hits { + fields := previewFields(hit.Fields, 3) + pcio.Fprintf(writer, "%s\t%f\t%s\n", hit.Id, hit.Score, fields) + } + + writer.Flush() +} + +func previewFields(fields map[string]any, limit int) string { + if len(fields) == 0 { + return "" + } + + keys := make([]string, 0, len(fields)) + for k := range fields { + keys = append(keys, k) + } + sort.Strings(keys) + + show := keys + truncated := false + if len(show) > limit { + show = show[:limit] + truncated = true + } + + limited := make(map[string]any, len(show)) + for _, k := range show { + limited[k] = fields[k] + } + + out := text.InlineJSON(limited) + if truncated && strings.HasSuffix(out, "}") { + out = strings.TrimRight(out, "}") + ", ...}" + } + + return out +}