vector_amp 0.1.0

data/lib/vector_amp/connections.rb

@@ -0,0 +1,49 @@
+# frozen_string_literal: true
+
+require_relative "utils"
+
+module VectorAmp
+  # Managed OAuth/credential connections used by ingestion sources.
+  #
+  # Connections live at the gateway root (`/connections`), not under
+  # `/ingestion`. A connection captures a provider authorization (e.g. Google
+  # Drive, Confluence) that ingestion sources can then reference via
+  # `connection_id` instead of embedding raw credentials.
+  class ConnectionsResource
+    # @param transport [#request] API transport.
+    # @return [ConnectionsResource]
+    def initialize(transport)
+      @transport = transport
+    end
+
+    # List connections, optionally filtered by provider.
+    # @param provider [String, nil] optional provider filter (e.g. `google_drive`, `confluence`).
+    # @return [Hash] response envelope with `connections`.
+    def list(provider: nil)
+      @transport.request(:get, "/connections", query: Utils.compact_hash(provider: provider))
+    end
+
+    # Create a connection and begin the provider authorization flow.
+    # @param provider [String] provider identifier (e.g. `google_drive`, `confluence`).
+    # @param source_type [String, nil] optional source type the connection will be used with.
+    # @return [Hash] created connection with `id`, `provider`, `status`, and `authorization_url`.
+    def create(provider, source_type: nil)
+      body = Utils.compact_hash(provider: provider, source_type: source_type)
+      @transport.request(:post, "/connections", body: body)
+    end
+
+    # Fetch a connection.
+    # @param connection_id [String] connection id.
+    # @return [Hash] connection resource.
+    def get(connection_id)
+      @transport.request(:get, "/connections/#{connection_id}")
+    end
+
+    # Delete a connection.
+    # @param connection_id [String] connection id.
+    # @return [Hash] delete response.
+    def delete(connection_id)
+      @transport.request(:delete, "/connections/#{connection_id}")
+    end
+  end
+end

data/lib/vector_amp/dataset.rb

@@ -0,0 +1,182 @@
+# frozen_string_literal: true
+
+module VectorAmp
+  # Rich dataset resource returned by DatasetsResource#create/get/list.
+  #
+  # It keeps the raw API payload while adding convenient instance methods that
+  # delegate to the existing service-style APIs.
+  class Dataset
+    # @return [DatasetsResource] backing dataset service.
+    # @return [Client, nil] client that created this object, required for convenience ingestion/ask helpers.
+    # @return [String] dataset id.
+    # @return [Hash] normalized raw API payload.
+    attr_reader :service, :client, :id, :data
+    alias raw_data data
+
+    # @param data [Hash] dataset API payload; `id` or `dataset_id` is required.
+    # @param service [DatasetsResource] backing dataset service.
+    # @param client [Client, nil] optional client for convenience helpers.
+    # @return [Dataset]
+    def initialize(data, service:, client: nil)
+      @data = normalize_data(data)
+      @service = service
+      @client = client
+      @id = extract_id(@data)
+      raise ArgumentError, "dataset id is required" if @id.nil? || @id.to_s.empty?
+    end
+
+    # Read a raw dataset field by string or symbol key.
+    # @param key [String, Symbol] field name.
+    # @return [Object, nil]
+    def [](key)
+      @data[key.to_s]
+    end
+
+    # Fetch a raw dataset field by string or symbol key.
+    # @param key [String, Symbol] field name.
+    # @return [Object]
+    def fetch(key, *args, &block)
+      @data.fetch(key.to_s, *args, &block)
+    end
+
+    # @param key [String, Symbol] field name.
+    # @return [Boolean] whether the raw payload has the field.
+    def key?(key)
+      @data.key?(key.to_s)
+    end
+    alias has_key? key?
+
+    # @return [Hash] shallow copy of the raw API payload.
+    def to_h
+      @data.dup
+    end
+    alias to_hash to_h
+
+    def inspect
+      "#<#{self.class} id=#{id.inspect} data=#{@data.inspect}>"
+    end
+
+    # Search this dataset.
+    # @param query_text [String, nil] optional text query; alternatively pass `query:` vector/options.
+    # @param options [Hash] forwarded to {DatasetsResource#search}.
+    # @return [Hash] search response.
+    def search(query_text = nil, **options)
+      service.search(id, query_text, **options)
+    end
+
+    # Insert vectors into this dataset.
+    # @param vectors [Array<Hash>] vector records with ids, values, and optional metadata.
+    # @return [Hash] insert response.
+    def insert(vectors:)
+      service.insert(id, vectors: vectors)
+    end
+
+    # Embed and insert texts into this dataset.
+    # @param texts_arg [Array<String>, nil] positional texts for convenience.
+    # @param texts [Array<String>, nil] keyword texts.
+    # @param ids [Array<String>, nil] optional ids; generated UUIDs when omitted.
+    # @param metadata [Hash, Array<Hash>, nil] metadata applied to all texts or per-text.
+    # @return [Hash] insert response.
+    def add_texts(texts_arg = nil, texts: nil, ids: nil, metadata: nil)
+      service.add_texts(id, texts_arg, texts: texts, ids: ids, metadata: metadata)
+    end
+
+    # Delete this dataset.
+    # @return [Hash] delete response.
+    def delete
+      service.delete(id)
+    end
+
+    # Fetch stats for this dataset.
+    # @return [Hash] dataset statistics.
+    def stats
+      service.stats(id)
+    end
+
+
+    # List retained source documents for this dataset using cursor pagination.
+    # @param limit [Integer, nil] maximum documents to return.
+    # @param cursor [String, nil] cursor from a previous response's `next_cursor`.
+    # @param status [String, nil] optional document status filter.
+    # @return [Hash] response envelope with `documents` and `next_cursor`.
+    def list_documents(limit: 50, cursor: nil, status: nil)
+      service.list_documents(id, limit: limit, cursor: cursor, status: status)
+    end
+
+    # Download retained original bytes for a dataset source document.
+    # @param document_id [String] document id returned by {#list_documents}.
+    # @return [String] raw document bytes.
+    def download_document(document_id)
+      service.download_document(id, document_id)
+    end
+
+    # Ask an intelligence question constrained to this dataset.
+    # @param query [String] natural-language question.
+    # @param options [Hash] forwarded to {Client#ask}; `dataset_id` is set to this dataset id.
+    # @return [Hash] intelligence response.
+    def ask(query, **options)
+      require_client!("ask")
+      client.ask(query, **options.merge(dataset_id: id))
+    end
+
+    # Upload local files by auto-creating a `file_upload` source, initializing presigned uploads, and completing the upload job.
+    # @param paths [String, Array<String>] local file paths to upload.
+    # @param source_name [String, nil] optional source name; defaults to timestamped Ruby SDK file-upload name.
+    # @param description [String, nil] optional source description.
+    # @param metadata [Hash] optional source metadata; dataset_id is added automatically.
+    # @return [Hash] upload completion/job response.
+    def ingest_files(paths:, source_name: nil, description: nil, metadata: {})
+      require_client!("ingest_files")
+      client.ingestion.ingest_files(
+        dataset_id: id,
+        paths: paths,
+        source_name: source_name,
+        description: description,
+        metadata: metadata
+      )
+    end
+
+    # Start an ingestion job from an existing source into this dataset.
+    # @param source_id [String, Source, Hash, nil] source id or object/hash containing an id.
+    # @param source [Source, Hash, nil] alternate source object/hash containing an id.
+    # @param pipeline_id [String, nil] optional pipeline id.
+    # @return [Hash] ingestion job response.
+    def ingest_source(source_id = nil, source: nil, pipeline_id: nil)
+      require_client!("ingest_source")
+      resolved_source_id = extract_source_id(source_id || source)
+      raise ArgumentError, "source_id is required" if resolved_source_id.nil? || resolved_source_id.to_s.empty?
+
+      client.ingestion.start_job(source_id: resolved_source_id, dataset_id: id, pipeline_id: pipeline_id)
+    end
+
+    private
+
+    def normalize_data(value)
+      case value
+      when Hash
+        value.each_with_object({}) { |(key, item), memo| memo[key.to_s] = item }
+      else
+        {}
+      end
+    end
+
+    def extract_id(hash)
+      return nil unless hash
+
+      hash["id"] || hash[:id] || hash["dataset_id"] || hash[:dataset_id]
+    end
+
+    def extract_source_id(value)
+      return value if value.is_a?(String) || value.is_a?(Symbol)
+      return value.id if value.respond_to?(:id)
+
+      extract_id(normalize_data(value))
+    end
+
+    def require_client!(method_name)
+      return if client
+
+      raise ArgumentError, "#{method_name} requires a Dataset created by VectorAmp::Client"
+    end
+  end
+end

data/lib/vector_amp/datasets.rb

@@ -0,0 +1,237 @@
+# frozen_string_literal: true
+
+require "securerandom"
+require_relative "dataset"
+require_relative "embedding"
+require_relative "utils"
+
+module VectorAmp
+  # Dataset API resource for create/list/get/search/insert operations.
+  class DatasetsResource
+    # @param transport [#request] API transport.
+    # @param client [Client, nil] optional client attached to returned Dataset objects.
+    # @return [DatasetsResource]
+    def initialize(transport, client: nil)
+      @transport = transport
+      @client = client
+    end
+
+    # List datasets.
+    # @param limit [Integer] page size; defaults to 50.
+    # @param offset [Integer] page offset; defaults to 0.
+    # @return [Hash] response envelope with `datasets` wrapped as Dataset objects when present.
+    def list(limit: 50, offset: 0)
+      wrap_list(@transport.request(:get, "/datasets", query: { limit: limit, offset: offset }))
+    end
+
+    # Fetch a dataset by id.
+    # @param dataset_id [String] dataset id.
+    # @return [Dataset]
+    def get(dataset_id)
+      wrap_dataset(@transport.request(:get, "/datasets/#{dataset_id}"))
+    end
+
+    # Delete a dataset by id.
+    # @param dataset_id [String] dataset id.
+    # @return [Hash] delete response.
+    def delete(dataset_id)
+      @transport.request(:delete, "/datasets/#{dataset_id}")
+    end
+
+    # Fetch dataset statistics.
+    # @param dataset_id [String] dataset id.
+    # @return [Hash] stats response.
+    def stats(dataset_id)
+      @transport.request(:get, "/datasets/#{dataset_id}/stats")
+    end
+
+    # Create a SABLE dataset. `index_type` is managed by the SDK and always sent as `sable`.
+    #
+    # Only `name` is required. The embedding defaults to the managed
+    # `vectoramp/VectorAmp-Embedding-4B` model and the dimension is inferred
+    # (2560 for the default model). Pass `embedding:` (a Hash or model String, or
+    # the result of {VectorAmp::Embedding.openai}) to use a different model; if
+    # the model's dimension is not known to the SDK you must also pass `dim:`.
+    #
+    # @param name [String] dataset name.
+    # @param dim [Integer, nil] embedding/vector dimension; inferred when omitted for known models.
+    # @param embedding [Hash, String, nil] embedding configuration; defaults to the managed VectorAmp model.
+    # @param metric [String] distance metric; defaults to `cosine`.
+    # @param hybrid [Boolean, nil] enable hybrid (dense + sparse) search; sends `hybrid: true`.
+    # @param filters [Hash, nil] optional filter schema/config.
+    # @param metadata_schema [Hash, nil] optional metadata schema.
+    # @param tuning [Hash, nil] optional SABLE tuning parameters.
+    # @param metadata [Hash, nil] optional dataset metadata.
+    # @return [Dataset] created dataset.
+    # @raise [ArgumentError] when `index_type` is supplied, dim cannot be inferred, or unknown options are passed.
+    def create(name:, dim: nil, embedding: nil, metric: "cosine", hybrid: nil, filters: nil, metadata_schema: nil, tuning: nil, metadata: nil, **unknown)
+      if unknown.key?(:index_type) || unknown.key?("index_type")
+        raise ArgumentError, "index_type is managed by VectorAmp Ruby SDK and is always 'sable'"
+      end
+      Utils.ensure_no_unknown!(unknown, "create")
+
+      resolved_embedding = Embedding.normalize(embedding)
+      resolved_dim = dim || Embedding.infer_dim(resolved_embedding)
+      if resolved_dim.nil?
+        model = resolved_embedding[:model] || resolved_embedding["model"]
+        raise ArgumentError, "dim is required for embedding model #{model.inspect}; pass dim: explicitly"
+      end
+
+      body = Utils.compact_hash(
+        name: name,
+        dim: resolved_dim,
+        metric: metric,
+        embedding: resolved_embedding,
+        index_type: "sable",
+        hybrid: hybrid,
+        filters: filters,
+        metadata_schema: metadata_schema,
+        tuning: tuning,
+        metadata: metadata
+      )
+      wrap_dataset(@transport.request(:post, "/datasets", body: body))
+    end
+
+
+    # List retained source documents for a dataset using cursor pagination.
+    # @param dataset_id [String] dataset id.
+    # @param limit [Integer, nil] maximum documents to return.
+    # @param cursor [String, nil] cursor from a previous response's `next_cursor`.
+    # @param status [String, nil] optional document status filter, e.g. `ready`.
+    # @return [Hash] response envelope with `documents` and `next_cursor`.
+    def list_documents(dataset_id, limit: 50, cursor: nil, status: nil)
+      @transport.request(:get, "/datasets/#{dataset_id}/documents", query: Utils.compact_hash(limit: limit, cursor: cursor, status: status))
+    end
+
+    # Download the retained original bytes for a dataset document.
+    # The HTTP transport follows redirects so this returns the final raw object bytes.
+    # @param dataset_id [String] dataset id.
+    # @param document_id [String] document id returned by {#list_documents}.
+    # @return [String] raw document bytes.
+    def download_document(dataset_id, document_id)
+      @transport.request(:get, "/datasets/#{dataset_id}/documents/#{document_id}/download", raw: true, headers: { Accept: "*/*" })
+    end
+
+    # Search a dataset by text or vector query.
+    # @param dataset_id [String] dataset id.
+    # @param query_text_or_options [String, Hash, nil] text query or legacy options hash.
+    # @param query [Array<Numeric>, nil] vector query.
+    # @param query_text [String, nil] explicit text query; overrides positional text when provided.
+    # @param search_text [String, nil] alias for query_text for one-field hybrid/BM25 search.
+    # @param top_k [Integer] number of results; defaults to 10.
+    # @param filters [Hash, nil] metadata filters.
+    # @param advanced_filters [Hash, nil] advanced filter expression.
+    # @param embedding_model [String, nil] optional embedding model override.
+    # @param embedding_provider [String, nil] optional embedding provider override.
+    # @param nprobe_override [Integer, nil] optional SABLE search probe override.
+    # @param rerank_depth_override [Integer, nil] optional rerank depth override.
+    # @param hybrid [Boolean, nil] enable hybrid search when supported.
+    # @param sparse_query [Hash, nil] sparse query for hybrid search.
+    # @param alpha [Numeric, nil] dense/sparse weighting for hybrid search.
+    # @param include_embeddings [Boolean, nil] include vector values in results.
+    # @param include_documents [Boolean, nil] include document fields in results.
+    # @param include_metadata [Boolean, nil] include metadata; API default is true.
+    # @param rerank [Boolean, Hash, nil] enable semantic reranking. `true` or `{ enabled: true }` uses vectoramp / VectorAmp-Rerank-v1.
+    # @return [Hash] search response.
+    def search(dataset_id, query_text_or_options = nil, query: nil, query_text: nil, search_text: nil, top_k: 10, filters: nil, advanced_filters: nil,
+               embedding_model: nil, embedding_provider: nil, nprobe_override: nil, rerank_depth_override: nil,
+               hybrid: nil, sparse_query: nil, alpha: nil, include_embeddings: nil, include_documents: nil,
+               include_metadata: nil, rerank: nil, **unknown)
+      if query_text_or_options.is_a?(Hash)
+        unknown = query_text_or_options.merge(unknown)
+        query_text_or_options = nil
+      end
+      Utils.ensure_no_unknown!(unknown, "search")
+      raise ArgumentError, "provide query_text or search_text, not both" if query_text && search_text
+      resolved_query_text = query_text || search_text || query_text_or_options
+      body = Utils.compact_hash(
+        query: query,
+        query_text: resolved_query_text,
+        top_k: top_k,
+        filters: filters,
+        advanced_filters: advanced_filters,
+        embedding_model: embedding_model,
+        embedding_provider: embedding_provider,
+        nprobe_override: nprobe_override,
+        rerank_depth_override: rerank_depth_override,
+        hybrid: hybrid,
+        sparse_query: sparse_query,
+        alpha: alpha,
+        include_embeddings: include_embeddings,
+        include_documents: include_documents,
+        include_metadata: include_metadata,
+        rerank: rerank
+      )
+      @transport.request(:post, "/datasets/#{dataset_id}/search", body: body)
+    end
+
+    # Insert vectors into a dataset.
+    #
+    # Vector record ids may be strings or integers. Integer ids are preserved as
+    # JSON numbers (not coerced to strings) so the API does not rewrite them.
+    # @param dataset_id [String] dataset id.
+    # @param vectors [Array<Hash>] vector records with ids, values, and optional metadata.
+    # @return [Hash] insert response.
+    def insert(dataset_id, vectors:)
+      @transport.request(:post, "/datasets/#{dataset_id}/insert", body: { vectors: Utils.normalize_vectors(vectors) })
+    end
+
+    # Generate embeddings using the dataset embedding configuration.
+    # @param dataset_id [String] dataset id.
+    # @param text [String, nil] single text to embed.
+    # @param texts [Array<String>, nil] multiple texts to embed.
+    # @return [Hash] embed response.
+    def embed(dataset_id, text: nil, texts: nil)
+      raise ArgumentError, "provide text or texts" if text.nil? && texts.nil?
+
+      @transport.request(:post, "/datasets/#{dataset_id}/embed", body: Utils.compact_hash(text: text, texts: texts))
+    end
+
+    # Embed and insert texts into a dataset.
+    #
+    # Accepts a single string or a list of strings. Ids are auto-generated when
+    # omitted; supplied ids may be strings or integers, and integer ids are
+    # preserved as JSON numbers. The source text is copied into `metadata.text`.
+    # @param dataset_id [String] dataset id.
+    # @param texts_arg [String, Array<String>, nil] positional text(s) for convenience.
+    # @param texts [String, Array<String>, nil] keyword text(s).
+    # @param ids [Array, nil] optional ids; generated UUIDs when omitted.
+    # @param metadata [Hash, Array<Hash>, nil] metadata applied to all texts or per-text.
+    # @return [Hash] insert response.
+    def add_texts(dataset_id, texts_arg = nil, texts: nil, ids: nil, metadata: nil)
+      texts = Array(texts || texts_arg)
+      raise ArgumentError, "texts must not be empty" if texts.empty?
+      raise ArgumentError, "ids length must match texts length" if ids && ids.length != texts.length
+      if metadata.is_a?(Array) && metadata.length != texts.length
+        raise ArgumentError, "metadata length must match texts length"
+      end
+
+      response = embed(dataset_id, texts: texts)
+      embeddings = response.fetch("embeddings")
+      vectors = texts.each_with_index.map do |text, index|
+        item_metadata = metadata.is_a?(Array) ? metadata[index] : metadata
+        {
+          id: ids ? Utils.coerce_vector_id(ids[index]) : SecureRandom.uuid,
+          values: embeddings[index],
+          metadata: Utils.compact_hash((item_metadata || {}).merge(text: text))
+        }
+      end
+      insert(dataset_id, vectors: vectors)
+    end
+
+    private
+
+    def wrap_dataset(data)
+      Dataset.new(data, service: self, client: @client)
+    end
+
+    def wrap_list(response)
+      return response unless response.is_a?(Hash)
+
+      datasets = response["datasets"] || response[:datasets]
+      return response unless datasets.is_a?(Array)
+
+      response.merge("datasets" => datasets.map { |dataset| wrap_dataset(dataset) })
+    end
+  end
+end

data/lib/vector_amp/embedding.rb

@@ -0,0 +1,67 @@
+# frozen_string_literal: true
+
+module VectorAmp
+  # Embedding configuration defaults and helpers.
+  #
+  # SDK callers rarely need to specify an embedding model. By default datasets use
+  # the managed VectorAmp embedding model and the SDK infers the vector dimension.
+  module Embedding
+    module_function
+
+    # Default embedding provider used when none is supplied.
+    DEFAULT_PROVIDER = "vectoramp"
+    # Default embedding model used when none is supplied.
+    DEFAULT_MODEL = "VectorAmp-Embedding-4B"
+
+    # Built-in dimension inference for known provider/model pairs.
+    DIM_TABLE = {
+      ["vectoramp", "VectorAmp-Embedding-4B"] => 2560,
+      ["openai", "text-embedding-3-small"] => 1536,
+      ["openai", "text-embedding-3-large"] => 3072,
+    }.freeze
+
+    # The fully managed default embedding configuration.
+    # @return [Hash] `{ provider:, model: }`
+    def default_embedding
+      { provider: DEFAULT_PROVIDER, model: DEFAULT_MODEL }
+    end
+
+    # Build an OpenAI embedding configuration.
+    # @param size [String, Symbol] `"small"` (text-embedding-3-small) or `"large"` (text-embedding-3-large).
+    # @return [Hash] `{ provider: "openai", model: ... }`
+    # @raise [ArgumentError] when size is not "small" or "large".
+    def openai(size = "small")
+      model = case size.to_s
+              when "small" then "text-embedding-3-small"
+              when "large" then "text-embedding-3-large"
+              else
+                raise ArgumentError, %(openai size must be "small" or "large", got #{size.inspect})
+              end
+      { provider: "openai", model: model }
+    end
+
+    # Normalize a user-supplied embedding value into a config hash.
+    # Accepts a Hash (`{ provider:, model:, secret_ref? }`) or a String model name
+    # (assumed to use the default provider).
+    # @param embedding [Hash, String, nil]
+    # @return [Hash] normalized embedding config with string keys for provider/model lookups.
+    def normalize(embedding)
+      case embedding
+      when nil then default_embedding
+      when String then { provider: DEFAULT_PROVIDER, model: embedding }
+      when Hash then embedding
+      else
+        raise ArgumentError, "embedding must be a Hash or String"
+      end
+    end
+
+    # Infer the vector dimension for an embedding config, or nil when unknown.
+    # @param embedding [Hash] normalized embedding config.
+    # @return [Integer, nil] inferred dimension or nil for custom/unknown models.
+    def infer_dim(embedding)
+      provider = embedding[:provider] || embedding["provider"]
+      model = embedding[:model] || embedding["model"]
+      DIM_TABLE[[provider.to_s, model.to_s]]
+    end
+  end
+end

data/lib/vector_amp/error.rb

@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+module VectorAmp
+  class Error < StandardError; end
+  class ConfigurationError < Error; end
+
+  class APIError < Error
+    attr_reader :status, :body, :headers
+
+    def initialize(message, status: nil, body: nil, headers: {})
+      super(message)
+      @status = status
+      @body = body
+      @headers = headers
+    end
+  end
+end