vectra-client 0.1.2

data/lib/vectra/providers/pgvector.rb

@@ -0,0 +1,297 @@
+# frozen_string_literal: true
+
+require_relative "pgvector/connection"
+require_relative "pgvector/sql_helpers"
+require_relative "pgvector/index_management"
+
+module Vectra
+  module Providers
+    # PostgreSQL with pgvector extension provider
+    #
+    # This provider uses PostgreSQL with the pgvector extension for vector
+    # similarity search. Each "index" maps to a PostgreSQL table.
+    #
+    # @example Table structure
+    #   CREATE EXTENSION IF NOT EXISTS vector;
+    #   CREATE TABLE my_index (
+    #     id TEXT PRIMARY KEY,
+    #     embedding vector(384),
+    #     metadata JSONB DEFAULT '{}',
+    #     namespace TEXT DEFAULT '',
+    #     created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+    #   );
+    #   CREATE INDEX ON my_index USING ivfflat (embedding vector_cosine_ops);
+    #
+    # @example Usage
+    #   client = Vectra.pgvector(
+    #     connection_url: "postgres://user:pass@localhost/mydb"
+    #   )
+    #   client.upsert(index: 'documents', vectors: [...])
+    #
+    class Pgvector < Base
+      include Connection
+      include SqlHelpers
+      include IndexManagement
+
+      DISTANCE_FUNCTIONS = {
+        "cosine" => "<=>",
+        "euclidean" => "<->",
+        "inner_product" => "<#>"
+      }.freeze
+
+      DEFAULT_METRIC = "cosine"
+
+      def initialize(config)
+        super
+        @connection = nil
+        @table_cache = {}
+      end
+
+      # @see Base#provider_name
+      def provider_name
+        :pgvector
+      end
+
+      # @see Base#upsert
+      def upsert(index:, vectors:, namespace: nil)
+        ensure_table_exists!(index)
+        normalized = normalize_vectors(vectors)
+        ns = namespace || ""
+
+        upserted = 0
+        normalized.each do |vec|
+          upsert_single_vector(index, vec, ns)
+          upserted += 1
+        end
+
+        log_debug("Upserted #{upserted} vectors to #{index}")
+        { upserted_count: upserted }
+      end
+
+      # @see Base#query
+      def query(index:, vector:, top_k: 10, namespace: nil, filter: nil,
+                include_values: false, include_metadata: true)
+        ensure_table_exists!(index)
+
+        distance_op = DISTANCE_FUNCTIONS[table_metric(index)]
+        vector_literal = format_vector(vector)
+
+        sql = build_query_sql(
+          index: index,
+          vector_literal: vector_literal,
+          distance_op: distance_op,
+          top_k: top_k,
+          namespace: namespace,
+          filter: filter,
+          include_values: include_values,
+          include_metadata: include_metadata
+        )
+
+        result = execute(sql)
+        matches = result.map { |row| build_match_from_row(row, include_values, include_metadata) }
+
+        log_debug("Query returned #{matches.size} results")
+        QueryResult.from_response(matches: matches, namespace: namespace)
+      end
+
+      # @see Base#fetch
+      def fetch(index:, ids:, namespace: nil)
+        ensure_table_exists!(index)
+
+        placeholders = ids.map.with_index { |_, i| "$#{i + 1}" }.join(", ")
+        sql = "SELECT id, embedding, metadata FROM #{quote_ident(index)} WHERE id IN (#{placeholders})"
+        sql += " AND namespace = $#{ids.size + 1}" if namespace
+
+        params = namespace ? ids + [namespace] : ids
+        result = execute(sql, params)
+
+        vectors = {}
+        result.each do |row|
+          vectors[row["id"]] = Vector.new(
+            id: row["id"],
+            values: parse_vector(row["embedding"]),
+            metadata: parse_json(row["metadata"])
+          )
+        end
+        vectors
+      end
+
+      # @see Base#update
+      def update(index:, id:, metadata: nil, values: nil, namespace: nil)
+        ensure_table_exists!(index)
+        updates, params, param_idx = build_update_params(metadata, values)
+
+        return { updated: false } if updates.empty?
+
+        sql = "UPDATE #{quote_ident(index)} SET #{updates.join(', ')} WHERE id = $#{param_idx}"
+        params << id
+
+        if namespace
+          sql += " AND namespace = $#{param_idx + 1}"
+          params << namespace
+        end
+
+        execute(sql, params)
+        log_debug("Updated vector #{id}")
+        { updated: true }
+      end
+
+      # @see Base#delete
+      def delete(index:, ids: nil, namespace: nil, filter: nil, delete_all: false)
+        ensure_table_exists!(index)
+
+        if delete_all || (namespace && ids.nil? && filter.nil?)
+          delete_all_vectors(index, namespace)
+        elsif ids
+          delete_by_ids(index, ids, namespace)
+        elsif filter
+          sql, params = build_filter_delete_sql(index, filter, namespace)
+          execute(sql, params)
+        end
+
+        log_debug("Deleted vectors from #{index}")
+        { deleted: true }
+      end
+
+      # @see Base#list_indexes
+      def list_indexes
+        sql = <<~SQL
+          SELECT table_name
+          FROM information_schema.columns
+          WHERE column_name = 'embedding'
+            AND data_type = 'USER-DEFINED'
+            AND table_schema = 'public'
+            AND udt_name = 'vector'
+        SQL
+
+        result = execute(sql)
+        result.map { |row| describe_index(index: row["table_name"]) }
+      end
+
+      # @see Base#describe_index
+      def describe_index(index:)
+        sql = <<~SQL
+          SELECT format_type(a.atttypid, a.atttypmod) as data_type
+          FROM pg_attribute a
+          JOIN pg_class c ON a.attrelid = c.oid
+          WHERE c.relname = $1 AND a.attname = 'embedding' AND a.attnum > 0
+        SQL
+
+        result = execute(sql, [index])
+        raise NotFoundError, "Index '#{index}' not found" if result.empty?
+
+        dimension = resolve_index_dimension(index, result)
+
+        { name: index, dimension: dimension, metric: table_metric(index), status: "ready" }
+      end
+
+      # Resolve vector dimension for an index from various sources
+      def resolve_index_dimension(index, pg_attribute_result)
+        type_info = pg_attribute_result.first["data_type"] || pg_attribute_result.first["udt_name"]
+        dim = extract_dimension_from_type(type_info) if type_info
+
+        return dim if dim
+
+        if @table_cache[index].is_a?(Hash)
+          return @table_cache[index][:dimension]
+        end
+
+        alt_sql = <<~SQL
+          SELECT udt_name, data_type FROM information_schema.columns
+          WHERE table_schema = 'public' AND table_name = $1 AND column_name = 'embedding'
+        SQL
+        alt_result = execute(alt_sql, [index])
+        udt = alt_result.first && (alt_result.first["udt_name"] || alt_result.first["data_type"])
+        extract_dimension_from_type(udt) if udt
+      end
+
+      # @see Base#stats
+      def stats(index:, namespace: nil)
+        ensure_table_exists!(index)
+
+        count_sql = "SELECT COUNT(*) as count FROM #{quote_ident(index)}"
+        count_sql += " WHERE namespace = $1" if namespace
+
+        count_result = execute(count_sql, namespace ? [namespace] : [])
+        total_count = count_result.first["count"].to_i
+
+        ns_sql = "SELECT namespace, COUNT(*) as count FROM #{quote_ident(index)} GROUP BY namespace"
+        ns_result = execute(ns_sql)
+        namespaces = ns_result.each_with_object({}) do |row, hash|
+          hash[row["namespace"] || ""] = { vector_count: row["count"].to_i }
+        end
+
+        info = describe_index(index: index)
+        { total_vector_count: total_count, dimension: info[:dimension], namespaces: namespaces }
+      end
+
+      private
+
+      # Build update parameters
+      def build_update_params(metadata, values)
+        updates = []
+        params = []
+        param_idx = 1
+
+        # Put embedding param first so tests expect embedding = $1::vector when provided
+        if values
+          updates << "embedding = $#{param_idx}::vector"
+          params << format_vector(values)
+          param_idx += 1
+        end
+
+        if metadata
+          updates << "metadata = metadata || $#{param_idx}::jsonb"
+          params << metadata.to_json
+          param_idx += 1
+        end
+
+        [updates, params, param_idx]
+      end
+
+      # Upsert a single vector
+      def upsert_single_vector(index, vec, namespace)
+        sql = <<~SQL
+          INSERT INTO #{quote_ident(index)} (id, embedding, metadata, namespace)
+          VALUES ($1, $2::vector, $3::jsonb, $4)
+          ON CONFLICT (id) DO UPDATE SET
+            embedding = EXCLUDED.embedding,
+            metadata = EXCLUDED.metadata,
+            namespace = EXCLUDED.namespace
+        SQL
+
+        params = [vec[:id], format_vector(vec[:values]), (vec[:metadata] || {}).to_json, namespace]
+        execute(sql, params)
+      end
+
+      # Delete all vectors from index
+      def delete_all_vectors(index, namespace)
+        sql = "DELETE FROM #{quote_ident(index)}"
+        sql += " WHERE namespace = $1" if namespace
+        execute(sql, namespace ? [namespace] : [])
+      end
+
+      # Delete vectors by IDs
+      def delete_by_ids(index, ids, namespace)
+        placeholders = ids.map.with_index { |_, i| "$#{i + 1}" }.join(", ")
+        sql = "DELETE FROM #{quote_ident(index)} WHERE id IN (#{placeholders})"
+        params = ids.dup
+
+        if namespace
+          sql += " AND namespace = $#{ids.size + 1}"
+          params << namespace
+        end
+
+        execute(sql, params)
+      end
+
+      # Override validate_config! for pgvector-specific validation
+      def validate_config!
+        raise ConfigurationError, "Provider must be configured" if config.provider.nil?
+        return if config.host
+
+        raise ConfigurationError, "Host (connection URL or hostname) must be configured for pgvector"
+      end
+    end
+  end
+end

data/lib/vectra/providers/pinecone.rb

@@ -0,0 +1,308 @@
+# frozen_string_literal: true
+
+module Vectra
+  module Providers
+    # Pinecone vector database provider
+    #
+    # @example
+    #   provider = Vectra::Providers::Pinecone.new(config)
+    #   provider.upsert(index: 'my-index', vectors: [...])
+    #
+    class Pinecone < Base
+      API_VERSION = "2024-07"
+
+      def initialize(config)
+        super
+        @control_plane_connection = nil
+        @data_plane_connections = {}
+      end
+
+      # @see Base#provider_name
+      def provider_name
+        :pinecone
+      end
+
+      # @see Base#upsert
+      def upsert(index:, vectors:, namespace: nil)
+        normalized = normalize_vectors(vectors)
+
+        body = { vectors: normalized }
+        body[:namespace] = namespace if namespace
+
+        response = data_connection(index).post("/vectors/upsert", body)
+
+        if response.success?
+          log_debug("Upserted #{normalized.size} vectors to #{index}")
+          {
+            upserted_count: response.body["upsertedCount"] || normalized.size
+          }
+        else
+          handle_error(response)
+        end
+      end
+
+      # @see Base#query
+      def query(index:, vector:, top_k: 10, namespace: nil, filter: nil,
+                include_values: false, include_metadata: true)
+        body = {
+          vector: vector.map(&:to_f),
+          topK: top_k,
+          includeValues: include_values,
+          includeMetadata: include_metadata
+        }
+        body[:namespace] = namespace if namespace
+        body[:filter] = transform_filter(filter) if filter
+
+        response = data_connection(index).post("/query", body)
+
+        if response.success?
+          log_debug("Query returned #{response.body['matches']&.size || 0} results")
+          QueryResult.from_response(
+            matches: transform_matches(response.body["matches"] || []),
+            namespace: response.body["namespace"],
+            usage: response.body["usage"]
+          )
+        else
+          handle_error(response)
+        end
+      end
+
+      # @see Base#fetch
+      def fetch(index:, ids:, namespace: nil)
+        params = { ids: ids }
+        params[:namespace] = namespace if namespace
+
+        response = data_connection(index).get("/vectors/fetch") do |req|
+          ids.each { |id| req.params.add("ids", id) }
+          req.params["namespace"] = namespace if namespace
+        end
+
+        if response.success?
+          vectors = {}
+          (response.body["vectors"] || {}).each do |id, data|
+            vectors[id] = Vector.new(
+              id: id,
+              values: data["values"],
+              metadata: data["metadata"],
+              sparse_values: data["sparseValues"]
+            )
+          end
+          vectors
+        else
+          handle_error(response)
+        end
+      end
+
+      # @see Base#update
+      def update(index:, id:, metadata: nil, values: nil, namespace: nil)
+        body = { id: id }
+        body[:setMetadata] = metadata if metadata
+        body[:values] = values.map(&:to_f) if values
+        body[:namespace] = namespace if namespace
+
+        response = data_connection(index).post("/vectors/update", body)
+
+        if response.success?
+          log_debug("Updated vector #{id}")
+          { updated: true }
+        else
+          handle_error(response)
+        end
+      end
+
+      # @see Base#delete
+      def delete(index:, ids: nil, namespace: nil, filter: nil, delete_all: false)
+        body = {}
+        body[:ids] = ids if ids
+        body[:namespace] = namespace if namespace
+        body[:filter] = transform_filter(filter) if filter
+        body[:deleteAll] = true if delete_all
+
+        response = data_connection(index).post("/vectors/delete", body)
+
+        if response.success?
+          log_debug("Deleted vectors from #{index}")
+          { deleted: true }
+        else
+          handle_error(response)
+        end
+      end
+
+      # @see Base#list_indexes
+      def list_indexes
+        response = control_connection.get("/indexes")
+
+        if response.success?
+          (response.body["indexes"] || []).map do |idx|
+            {
+              name: idx["name"],
+              dimension: idx["dimension"],
+              metric: idx["metric"],
+              host: idx["host"],
+              status: idx.dig("status", "ready") ? "ready" : "initializing"
+            }
+          end
+        else
+          handle_error(response)
+        end
+      end
+
+      # @see Base#describe_index
+      def describe_index(index:)
+        response = control_connection.get("/indexes/#{index}")
+
+        if response.success?
+          body = response.body
+          {
+            name: body["name"],
+            dimension: body["dimension"],
+            metric: body["metric"],
+            host: body["host"],
+            spec: body["spec"],
+            status: body["status"]
+          }
+        else
+          handle_error(response)
+        end
+      end
+
+      # @see Base#stats
+      def stats(index:, namespace: nil)
+        body = {}
+        body[:filter] = {} if namespace.nil?
+
+        response = data_connection(index).post("/describe_index_stats", body)
+
+        if response.success?
+          {
+            total_vector_count: response.body["totalVectorCount"],
+            dimension: response.body["dimension"],
+            index_fullness: response.body["indexFullness"],
+            namespaces: response.body["namespaces"]
+          }
+        else
+          handle_error(response)
+        end
+      end
+
+      # Create a new index
+      #
+      # @param name [String] index name
+      # @param dimension [Integer] vector dimension
+      # @param metric [String] similarity metric (cosine, euclidean, dotproduct)
+      # @param spec [Hash] index spec (serverless or pod configuration)
+      # @return [Hash] created index info
+      def create_index(name:, dimension:, metric: "cosine", spec: nil)
+        body = {
+          name: name,
+          dimension: dimension,
+          metric: metric
+        }
+
+        # Default to serverless spec if not provided
+        body[:spec] = spec || {
+          serverless: {
+            cloud: "aws",
+            region: config.environment || "us-east-1"
+          }
+        }
+
+        response = control_connection.post("/indexes", body)
+
+        if response.success?
+          log_debug("Created index #{name}")
+          describe_index(index: name)
+        else
+          handle_error(response)
+        end
+      end
+
+      # Delete an index
+      #
+      # @param name [String] index name
+      # @return [Hash] deletion result
+      def delete_index(name:)
+        response = control_connection.delete("/indexes/#{name}")
+
+        if response.success?
+          log_debug("Deleted index #{name}")
+          { deleted: true }
+        else
+          handle_error(response)
+        end
+      end
+
+      private
+
+      # Control plane connection (for index management)
+      def control_connection
+        @control_plane_connection ||= build_connection(
+          "https://api.pinecone.io",
+          {
+            "Api-Key" => config.api_key,
+            "X-Pinecone-API-Version" => API_VERSION
+          }
+        )
+      end
+
+      # Data plane connection (for vector operations)
+      # Each index has its own host
+      def data_connection(index)
+        @data_plane_connections[index] ||= begin
+          host = resolve_index_host(index)
+          build_connection(
+            "https://#{host}",
+            {
+              "Api-Key" => config.api_key,
+              "X-Pinecone-API-Version" => API_VERSION
+            }
+          )
+        end
+      end
+
+      # Resolve the host for an index
+      def resolve_index_host(index)
+        # If a direct host is configured, use that
+        return config.host if config.host
+
+        # Otherwise, fetch from the API
+        info = describe_index(index: index)
+        host = info[:host]
+
+        raise ConfigurationError, "Could not resolve host for index '#{index}'" unless host
+
+        host
+      end
+
+      # Transform metadata filter to Pinecone format
+      def transform_filter(filter)
+        return nil unless filter
+
+        # Simple key-value filters are wrapped in $eq
+        filter.transform_values do |value|
+          case value
+          when Hash
+            value # Already a filter operator
+          when Array
+            { "$in" => value }
+          else
+            { "$eq" => value }
+          end
+        end
+      end
+
+      # Transform matches from API response
+      def transform_matches(matches)
+        matches.map do |match|
+          {
+            id: match["id"],
+            score: match["score"],
+            values: match["values"],
+            metadata: match["metadata"],
+            sparse_values: match["sparseValues"]
+          }
+        end
+      end
+    end
+  end
+end

data/lib/vectra/providers/qdrant.rb

@@ -0,0 +1,48 @@
+# frozen_string_literal: true
+
+module Vectra
+  module Providers
+    # Qdrant vector database provider (planned for v0.2.0)
+    #
+    # @note This provider is not yet implemented
+    #
+    class Qdrant < Base
+      def provider_name
+        :qdrant
+      end
+
+      def upsert(index:, vectors:, namespace: nil)
+        raise NotImplementedError, "Qdrant provider is planned for v0.2.0"
+      end
+
+      def query(index:, vector:, top_k: 10, namespace: nil, filter: nil,
+                include_values: false, include_metadata: true)
+        raise NotImplementedError, "Qdrant provider is planned for v0.2.0"
+      end
+
+      def fetch(index:, ids:, namespace: nil)
+        raise NotImplementedError, "Qdrant provider is planned for v0.2.0"
+      end
+
+      def update(index:, id:, metadata:, namespace: nil)
+        raise NotImplementedError, "Qdrant provider is planned for v0.2.0"
+      end
+
+      def delete(index:, ids: nil, namespace: nil, filter: nil, delete_all: false)
+        raise NotImplementedError, "Qdrant provider is planned for v0.2.0"
+      end
+
+      def list_indexes
+        raise NotImplementedError, "Qdrant provider is planned for v0.2.0"
+      end
+
+      def describe_index(index:)
+        raise NotImplementedError, "Qdrant provider is planned for v0.2.0"
+      end
+
+      def stats(index:, namespace: nil)
+        raise NotImplementedError, "Qdrant provider is planned for v0.2.0"
+      end
+    end
+  end
+end

data/lib/vectra/providers/weaviate.rb

@@ -0,0 +1,48 @@
+# frozen_string_literal: true
+
+module Vectra
+  module Providers
+    # Weaviate vector database provider (planned for v0.3.0)
+    #
+    # @note This provider is not yet implemented
+    #
+    class Weaviate < Base
+      def provider_name
+        :weaviate
+      end
+
+      def upsert(index:, vectors:, namespace: nil)
+        raise NotImplementedError, "Weaviate provider is planned for v0.3.0"
+      end
+
+      def query(index:, vector:, top_k: 10, namespace: nil, filter: nil,
+                include_values: false, include_metadata: true)
+        raise NotImplementedError, "Weaviate provider is planned for v0.3.0"
+      end
+
+      def fetch(index:, ids:, namespace: nil)
+        raise NotImplementedError, "Weaviate provider is planned for v0.3.0"
+      end
+
+      def update(index:, id:, metadata:, namespace: nil)
+        raise NotImplementedError, "Weaviate provider is planned for v0.3.0"
+      end
+
+      def delete(index:, ids: nil, namespace: nil, filter: nil, delete_all: false)
+        raise NotImplementedError, "Weaviate provider is planned for v0.3.0"
+      end
+
+      def list_indexes
+        raise NotImplementedError, "Weaviate provider is planned for v0.3.0"
+      end
+
+      def describe_index(index:)
+        raise NotImplementedError, "Weaviate provider is planned for v0.3.0"
+      end
+
+      def stats(index:, namespace: nil)
+        raise NotImplementedError, "Weaviate provider is planned for v0.3.0"
+      end
+    end
+  end
+end