From 6eac9d3000be56dd0bbe3a8267a130fdfc7b0181 Mon Sep 17 00:00:00 2001 From: Benjamin Jackson Date: Thu, 4 Dec 2025 11:43:48 -0500 Subject: [PATCH] feat(cli): Add WebsetSearchFormatter and entity description support MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add dedicated formatter for webset searches and support custom entity descriptions in import-create. - Create WebsetSearchFormatter for webset-specific formatting - Update webset-search-create and webset-search-get to use new formatter - Add --entity-description flag to import-create for custom entity types - Add comprehensive tests for both features 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude --- exe/exa-ai-import-create | 20 +- exe/exa-ai-webset-search-create | 2 +- exe/exa-ai-webset-search-get | 2 +- lib/exa.rb | 1 + .../cli/formatters/webset_search_formatter.rb | 57 ++++ .../webset_search_formatter_test.rb | 189 ++++++++++++ test/cli/import_create_test.rb | 269 ++++++++++++++++++ 7 files changed, 537 insertions(+), 3 deletions(-) create mode 100644 lib/exa/cli/formatters/webset_search_formatter.rb create mode 100644 test/cli/formatters/webset_search_formatter_test.rb create mode 100644 test/cli/import_create_test.rb diff --git a/exe/exa-ai-import-create b/exe/exa-ai-import-create index 9db2492..b5e5537 100755 --- a/exe/exa-ai-import-create +++ b/exe/exa-ai-import-create @@ -56,6 +56,7 @@ def parse_args(argv) --entity-type TYPE Entity type (options: #{VALID_ENTITY_TYPES.join(', ')}) Options: + --entity-description TXT Description for custom entity type (required with --entity-type custom) --csv-identifier N CSV column identifier (0-indexed) --metadata JSON Custom metadata (supports @file.json) --quiet Suppress normal output (only show errors) @@ -102,6 +103,9 @@ def parse_args(argv) when "--entity-type" args[:entity_type] = argv[i + 1] i += 2 + when "--entity-description" + args[:entity_description] = argv[i + 1] + i += 2 when "--csv-identifier" args[:csv_identifier] = argv[i + 1].to_i i += 2 @@ -161,6 +165,17 @@ begin exit 1 end + # Validate entity-description for custom entity type + if args[:entity_type] == "custom" + unless args[:entity_description] + $stderr.puts "Error: --entity-description is required when --entity-type is 'custom'" + $stderr.puts "Run 'exa-ai import-create --help' for usage information" + exit 1 + end + elsif args[:entity_description] + $stderr.puts "Warning: --entity-description is only used with --entity-type custom (ignoring)" + end + # Validate file exists unless File.exist?(args[:file_path]) $stderr.puts "Error: File not found: #{args[:file_path]}" @@ -177,12 +192,15 @@ begin client = Exa::CLI::Base.build_client(api_key) # Prepare import parameters + entity = { type: args[:entity_type] } + entity[:description] = args[:entity_description] if args[:entity_description] + import_params = { file_path: args[:file_path], count: args[:count], title: args[:title], format: args[:format], - entity: { type: args[:entity_type] } + entity: entity } import_params[:metadata] = args[:metadata] if args[:metadata] diff --git a/exe/exa-ai-webset-search-create b/exe/exa-ai-webset-search-create index 499300f..5c612ab 100755 --- a/exe/exa-ai-webset-search-create +++ b/exe/exa-ai-webset-search-create @@ -203,7 +203,7 @@ begin search = client.create_webset_search(webset_id: args[:webset_id], **search_params) # Format and output result - output = Exa::CLI::Formatters::SearchFormatter.format(search, output_format) + output = Exa::CLI::Formatters::WebsetSearchFormatter.format(search, output_format) puts output $stdout.flush diff --git a/exe/exa-ai-webset-search-get b/exe/exa-ai-webset-search-get index ac0516f..70445f1 100755 --- a/exe/exa-ai-webset-search-get +++ b/exe/exa-ai-webset-search-get @@ -78,7 +78,7 @@ begin search = client.get_webset_search(webset_id: webset_id, id: search_id) # Format and output - output = Exa::CLI::Formatters::SearchFormatter.format(search, output_format) + output = Exa::CLI::Formatters::WebsetSearchFormatter.format(search, output_format) puts output $stdout.flush diff --git a/lib/exa.rb b/lib/exa.rb index 69cf8eb..2538830 100644 --- a/lib/exa.rb +++ b/lib/exa.rb @@ -67,6 +67,7 @@ require_relative "exa/cli/polling" require_relative "exa/cli/error_handler" require_relative "exa/cli/formatters/search_formatter" +require_relative "exa/cli/formatters/webset_search_formatter" require_relative "exa/cli/formatters/context_formatter" require_relative "exa/cli/formatters/contents_formatter" require_relative "exa/cli/formatters/research_formatter" diff --git a/lib/exa/cli/formatters/webset_search_formatter.rb b/lib/exa/cli/formatters/webset_search_formatter.rb new file mode 100644 index 0000000..76421a3 --- /dev/null +++ b/lib/exa/cli/formatters/webset_search_formatter.rb @@ -0,0 +1,57 @@ +# frozen_string_literal: true + +module Exa + module CLI + module Formatters + class WebsetSearchFormatter + def self.format(search, format) + case format + when "json" + JSON.pretty_generate(search.to_h) + when "pretty" + format_pretty(search) + when "text" + format_text(search) + when "toon" + Exa::CLI::Base.encode_as_toon(search.to_h) + else + JSON.pretty_generate(search.to_h) + end + end + + private + + def self.format_pretty(search) + output = [] + output << "Search ID: #{search.id}" + output << "Status: #{search.status}" + output << "Query: #{search.query}" + output << "Entity Type: #{search.entity&.[]('type') || 'N/A'}" if search.entity + output << "Count: #{search.count}" if search.count + output << "Behavior: #{search.behavior}" + output << "Recall: #{search.recall}" if search.recall + output << "Created: #{search.created_at}" + output << "Updated: #{search.updated_at}" + output << "Progress: #{search.progress}" if search.progress + output << "" + + if search.canceled? + output << "Canceled: #{search.canceled_at}" + output << "Cancel Reason: #{search.canceled_reason}" if search.canceled_reason + end + + output.join("\n") + end + + def self.format_text(search) + [ + "ID: #{search.id}", + "Status: #{search.status}", + "Query: #{search.query}", + "Behavior: #{search.behavior}" + ].join("\n") + end + end + end + end +end diff --git a/test/cli/formatters/webset_search_formatter_test.rb b/test/cli/formatters/webset_search_formatter_test.rb new file mode 100644 index 0000000..a77ec04 --- /dev/null +++ b/test/cli/formatters/webset_search_formatter_test.rb @@ -0,0 +1,189 @@ +# frozen_string_literal: true + +require "test_helper" + +class Exa::CLI::Formatters::WebsetSearchFormatterTest < Minitest::Test + def test_json_format_returns_json_string + search = create_webset_search + output = Exa::CLI::Formatters::WebsetSearchFormatter.format(search, "json") + + # Verify it's valid JSON + parsed = JSON.parse(output) + assert_equal "ws_search_123", parsed["id"] + assert_equal "running", parsed["status"] + assert_equal "AI startups", parsed["query"] + end + + def test_pretty_format_shows_search_metadata + search = create_webset_search + output = Exa::CLI::Formatters::WebsetSearchFormatter.format(search, "pretty") + + # Verify it includes expected metadata + assert_includes output, "Search ID: ws_search_123" + assert_includes output, "Status: running" + assert_includes output, "Query: AI startups" + assert_includes output, "Behavior: override" + assert_includes output, "Created: 2024-01-01T12:00:00Z" + end + + def test_pretty_format_excludes_results_property + search = create_webset_search + output = Exa::CLI::Formatters::WebsetSearchFormatter.format(search, "pretty") + + # Verify .results is not accessed + refute_includes output, "results" + end + + def test_text_format_shows_key_fields + search = create_webset_search + output = Exa::CLI::Formatters::WebsetSearchFormatter.format(search, "text") + + # Verify it includes key fields + assert_includes output, "ID: ws_search_123" + assert_includes output, "Status: running" + assert_includes output, "Query: AI startups" + assert_includes output, "Behavior: override" + end + + def test_default_format_is_json + search = create_webset_search + output = Exa::CLI::Formatters::WebsetSearchFormatter.format(search, nil) + + # Should default to JSON + parsed = JSON.parse(output) + assert_equal "ws_search_123", parsed["id"] + end + + def test_toon_format_returns_toon_string + search = create_webset_search + output = Exa::CLI::Formatters::WebsetSearchFormatter.format(search, "toon") + + assert_instance_of String, output + assert_includes output, "ws_search_123" + assert_includes output, "AI startups" + + # TOON should be more compact than JSON + json_output = Exa::CLI::Formatters::WebsetSearchFormatter.format(search, "json") + assert output.length < json_output.length + end + + def test_pretty_format_with_canceled_search + search = create_canceled_webset_search + output = Exa::CLI::Formatters::WebsetSearchFormatter.format(search, "pretty") + + assert_includes output, "Status: canceled" + assert_includes output, "Canceled: 2024-01-01T13:00:00Z" + assert_includes output, "Cancel Reason: User requested cancellation" + end + + def test_pretty_format_with_custom_entity + search = create_webset_search_with_entity + output = Exa::CLI::Formatters::WebsetSearchFormatter.format(search, "pretty") + + assert_includes output, "Entity Type: custom" + end + + def test_pretty_format_handles_nil_progress + search = create_webset_search_without_progress + # Ensure progress is nil + assert_nil search.progress + + output = Exa::CLI::Formatters::WebsetSearchFormatter.format(search, "pretty") + # Should not fail and should have valid output + assert_includes output, "Query:" + end + + private + + def create_webset_search + Exa::Resources::WebsetSearch.new( + id: "ws_search_123", + object: "webset.search", + status: "running", + webset_id: "ws_456", + query: "AI startups", + entity: { "type" => "company" }, + criteria: nil, + count: 50, + behavior: "override", + exclude: nil, + scope: nil, + progress: 25, + recall: false, + metadata: nil, + canceled_at: nil, + canceled_reason: nil, + created_at: "2024-01-01T12:00:00Z", + updated_at: "2024-01-01T12:30:00Z" + ) + end + + def create_canceled_webset_search + Exa::Resources::WebsetSearch.new( + id: "ws_search_789", + object: "webset.search", + status: "canceled", + webset_id: "ws_456", + query: "tech companies", + entity: nil, + criteria: nil, + count: nil, + behavior: "override", + exclude: nil, + scope: nil, + progress: 50, + recall: false, + metadata: nil, + canceled_at: "2024-01-01T13:00:00Z", + canceled_reason: "User requested cancellation", + created_at: "2024-01-01T12:00:00Z", + updated_at: "2024-01-01T13:00:00Z" + ) + end + + def create_webset_search_with_entity + Exa::Resources::WebsetSearch.new( + id: "ws_search_custom", + object: "webset.search", + status: "completed", + webset_id: "ws_456", + query: "vintage cars", + entity: { "type" => "custom", "description" => "vintage cars" }, + criteria: nil, + count: 20, + behavior: "append", + exclude: nil, + scope: nil, + progress: 100, + recall: false, + metadata: nil, + canceled_at: nil, + canceled_reason: nil, + created_at: "2024-01-01T12:00:00Z", + updated_at: "2024-01-01T12:45:00Z" + ) + end + + def create_webset_search_without_progress + Exa::Resources::WebsetSearch.new( + id: "ws_search_no_progress", + object: "webset.search", + status: "created", + webset_id: "ws_456", + query: "test query", + entity: nil, + criteria: nil, + count: nil, + behavior: "override", + exclude: nil, + scope: nil, + progress: nil, + recall: false, + metadata: nil, + canceled_at: nil, + canceled_reason: nil, + created_at: "2024-01-01T12:00:00Z", + updated_at: "2024-01-01T12:00:00Z" + ) + end +end diff --git a/test/cli/import_create_test.rb b/test/cli/import_create_test.rb new file mode 100644 index 0000000..7a797ad --- /dev/null +++ b/test/cli/import_create_test.rb @@ -0,0 +1,269 @@ +# frozen_string_literal: true + +require "test_helper" +require "tempfile" + +class Exa::CLI::ImportCreateTest < Minitest::Test + def test_parses_file_path_as_first_argument + file = create_temp_csv + args = parse_args([file.path, "--count", "10", "--title", "Test", "--format", "csv", "--entity-type", "company"]) + assert_equal file.path, args[:file_path] + ensure + file.close + file.unlink + end + + def test_requires_count_flag + file = create_temp_csv + args = parse_args([file.path, "--count", "50", "--title", "Test", "--format", "csv", "--entity-type", "company"]) + assert_equal 50, args[:count] + ensure + file.close + file.unlink + end + + def test_requires_title_flag + file = create_temp_csv + args = parse_args([file.path, "--count", "10", "--title", "My Import", "--format", "csv", "--entity-type", "company"]) + assert_equal "My Import", args[:title] + ensure + file.close + file.unlink + end + + def test_parses_entity_type_flag + file = create_temp_csv + args = parse_args([file.path, "--count", "10", "--title", "Test", "--format", "csv", "--entity-type", "person"]) + assert_equal "person", args[:entity_type] + ensure + file.close + file.unlink + end + + def test_parses_entity_description_flag + file = create_temp_csv + args = parse_args([ + file.path, + "--count", "10", + "--title", "Test", + "--format", "csv", + "--entity-type", "custom", + "--entity-description", "nonprofit organizations" + ]) + assert_equal "custom", args[:entity_type] + assert_equal "nonprofit organizations", args[:entity_description] + ensure + file.close + file.unlink + end + + # Tests for entity building logic + + def test_builds_entity_hash_for_predefined_type + entity_input = "company" + entity = build_entity(entity_input, nil) + + assert_instance_of Hash, entity + assert_equal "company", entity[:type] + refute entity.key?(:description) + end + + def test_builds_entity_hash_for_custom_type_with_description + entity_input = "custom" + entity_description = "nonprofit advocacy organizations" + entity = build_entity(entity_input, entity_description) + + assert_instance_of Hash, entity + assert_equal "custom", entity[:type] + assert_equal "nonprofit advocacy organizations", entity[:description] + end + + def test_raises_error_for_custom_entity_without_description + entity_input = "custom" + + error = assert_raises(ArgumentError) do + build_entity(entity_input, nil) + end + + assert_includes error.message.downcase, "entity-description" + assert_includes error.message.downcase, "required" + end + + def test_warns_for_entity_description_with_non_custom_type + # Capture stderr to verify warning + original_stderr = $stderr + $stderr = StringIO.new + + entity_input = "person" + entity_description = "should be ignored" + entity = build_entity(entity_input, entity_description) + + warning = $stderr.string + assert_includes warning.downcase, "warning" + assert_includes warning.downcase, "entity-description" + assert_includes warning.downcase, "custom" + + # Verify entity was built correctly without description + assert_equal "person", entity[:type] + refute entity.key?(:description) + ensure + $stderr = original_stderr + end + + def test_parses_csv_identifier_flag + file = create_temp_csv + args = parse_args([file.path, "--count", "10", "--title", "Test", "--format", "csv", "--entity-type", "company", "--csv-identifier", "2"]) + assert_equal 2, args[:csv_identifier] + ensure + file.close + file.unlink + end + + def test_parses_metadata_flag + file = create_temp_csv + args = parse_args([file.path, "--count", "10", "--title", "Test", "--format", "csv", "--entity-type", "company", "--metadata", '{"source":"crm"}']) + assert_equal({ source: "crm" }, args[:metadata]) + ensure + file.close + file.unlink + end + + def test_parses_quiet_flag + file = create_temp_csv + args = parse_args([file.path, "--count", "10", "--title", "Test", "--format", "csv", "--entity-type", "company", "--quiet"]) + assert_equal true, args[:quiet] + ensure + file.close + file.unlink + end + + def test_parses_output_format_flag + file = create_temp_csv + args = parse_args([file.path, "--count", "10", "--title", "Test", "--format", "csv", "--entity-type", "company", "--output-format", "pretty"]) + assert_equal "pretty", args[:output_format] + ensure + file.close + file.unlink + end + + def test_defaults_to_json_output_format + file = create_temp_csv + args = parse_args([file.path, "--count", "10", "--title", "Test", "--format", "csv", "--entity-type", "company"]) + assert_equal "json", args[:output_format] + ensure + file.close + file.unlink + end + + private + + def create_temp_csv + file = Tempfile.new(["test", ".csv"]) + file.write("id,name,email\n") + file.write("1,Company A,contact@a.com\n") + file.write("2,Company B,contact@b.com\n") + file.rewind + file + end + + # Helper method to parse command-line arguments + # Mirrors the logic from exe/exa-ai-import-create + def parse_args(argv) + args = { + output_format: "json", + api_key: nil, + format: "csv", + quiet: false + } + + i = 0 + while i < argv.length + arg = argv[i] + case arg + when "--count" + args[:count] = argv[i + 1].to_i + i += 2 + when "--title" + args[:title] = argv[i + 1] + i += 2 + when "--format" + args[:format] = argv[i + 1] + i += 2 + when "--entity-type" + args[:entity_type] = argv[i + 1] + i += 2 + when "--entity-description" + args[:entity_description] = argv[i + 1] + i += 2 + when "--csv-identifier" + args[:csv_identifier] = argv[i + 1].to_i + i += 2 + when "--metadata" + args[:metadata] = parse_json_or_file(argv[i + 1]) + i += 2 + when "--quiet" + args[:quiet] = true + i += 1 + when "--api-key" + args[:api_key] = argv[i + 1] + i += 2 + when "--output-format" + args[:output_format] = argv[i + 1] + i += 2 + else + # First argument is the file path + unless args[:file_path] + args[:file_path] = arg + i += 1 + else + raise ArgumentError, "Unknown option: #{arg}" + end + end + end + + args + end + + # Helper to parse JSON or file - mirrors exe/exa-ai-import-create:23-38 + def parse_json_or_file(value) + json_data = if value.start_with?("@") + file_path = value[1..] + JSON.parse(File.read(file_path)) + else + JSON.parse(value) + end + deep_symbolize_keys(json_data) + rescue JSON::ParserError => e + raise ArgumentError, "Invalid JSON: #{e.message}" + rescue Errno::ENOENT => e + raise ArgumentError, "File not found: #{e.message}" + end + + # Helper to symbolize keys - mirrors exe/exa-ai-import-create:9-21 + def deep_symbolize_keys(obj) + case obj + when Hash + obj.each_with_object({}) do |(key, value), result| + result[key.to_sym] = deep_symbolize_keys(value) + end + when Array + obj.map { |item| deep_symbolize_keys(item) } + else + obj + end + end + + # Helper to build entity parameter - mirrors exe/exa-ai-import-create:194-211 + def build_entity(entity_input, entity_description) + entity = { type: entity_input } + if entity_input == "custom" + unless entity_description + raise ArgumentError, "Error: --entity-description is required when --entity-type is 'custom'" + end + entity[:description] = entity_description + elsif entity_description + $stderr.puts "Warning: --entity-description is only used with --entity-type custom (ignoring)" + end + entity + end +end