prescient 0.2.0 → 0.4.0

data/lib/prescient/pgvector.rb

@@ -0,0 +1,194 @@
+# frozen_string_literal: true
+
+require 'json'
+
+# PostgreSQL pgvector integration.
+#
+# The integration accepts a PG-compatible connection object, so applications
+# choose and manage their own PostgreSQL driver and connection lifecycle.
+# Stores provider embeddings and performs nearest-neighbor searches.
+class Prescient::Pgvector::Store
+  # @return [Hash<Symbol, String>] Supported pgvector distance operators
+  METRICS = {
+    cosine:        '<=>',
+    euclidean:     '<->',
+    inner_product: '<#>',
+  }.freeze
+
+  # @return [Integer] Required vector dimensions
+  attr_reader :dimensions
+
+  # @return [String] Embeddings table name
+  attr_reader :table
+
+  # @param connection [Object] PG-compatible object responding to `exec`
+  #   and `exec_params`
+  # @param dimensions [Integer] Required dimensions for every embedding
+  # @param table [String, Symbol] Safe PostgreSQL table identifier
+  def initialize(connection:, dimensions:, table: 'prescient_embeddings')
+    @connection = connection
+    @dimensions = validate_dimensions(dimensions)
+    @table = validate_table(table)
+  end
+
+  # Create the pgvector extension and the embeddings table.
+  #
+  # @return [void]
+  def install!
+    @connection.exec('CREATE EXTENSION IF NOT EXISTS vector')
+    @connection.exec(<<~SQL)
+      CREATE TABLE IF NOT EXISTS #{table} (
+        id text PRIMARY KEY,
+        provider text NOT NULL,
+        model text NOT NULL,
+        dimensions integer NOT NULL CHECK (dimensions = #{dimensions}),
+        embedding vector(#{dimensions}) NOT NULL,
+        content text,
+        metadata jsonb NOT NULL DEFAULT '{}'::jsonb,
+        created_at timestamptz NOT NULL DEFAULT CURRENT_TIMESTAMP,
+        updated_at timestamptz NOT NULL DEFAULT CURRENT_TIMESTAMP
+      )
+    SQL
+  end
+
+  # Create an HNSW index for the selected distance metric.
+  #
+  # @param metric [Symbol] `:cosine`, `:euclidean`, or `:inner_product`
+  # @return [void]
+  def create_index!(metric: :cosine)
+    metric_name = validate_metric(metric)
+    @connection.exec(<<~SQL)
+      CREATE INDEX IF NOT EXISTS #{table}_#{metric_name}_embedding_idx
+      ON #{table} USING hnsw (embedding #{metric_operator_class(metric_name)})
+    SQL
+  end
+
+  # Insert or replace an embedding record.
+  #
+  # @return [Hash] Stored record metadata
+  def upsert(id:, embedding:, provider:, model:, content: nil, metadata: {})
+    vector = serialize_vector(embedding)
+    parameters = [id.to_s, provider.to_s, model.to_s, dimensions, vector, content, JSON.generate(metadata)]
+    result = @connection.exec_params(<<~SQL, parameters)
+      INSERT INTO #{table} (id, provider, model, dimensions, embedding, content, metadata)
+      VALUES ($1, $2, $3, $4, $5::vector, $6, $7::jsonb)
+      ON CONFLICT (id) DO UPDATE SET
+        provider = EXCLUDED.provider,
+        model = EXCLUDED.model,
+        dimensions = EXCLUDED.dimensions,
+        embedding = EXCLUDED.embedding,
+        content = EXCLUDED.content,
+        metadata = EXCLUDED.metadata,
+        updated_at = CURRENT_TIMESTAMP
+      RETURNING id, provider, model, dimensions, content, metadata
+    SQL
+
+    record_from(result.first)
+  end
+
+  # Find the nearest stored embeddings.
+  #
+  # @param embedding [Array<Numeric>] Query vector
+  # @param limit [Integer] Maximum result count
+  # @param metric [Symbol] Distance metric
+  # @param provider [String, Symbol, nil] Optional provider filter
+  # @param model [String, nil] Optional model filter
+  # @return [Array<Hash>] Records ordered by ascending distance
+  def search(embedding:, limit: 10, metric: :cosine, provider: nil, model: nil)
+    vector = serialize_vector(embedding)
+    limit = validate_limit(limit)
+    metric = validate_metric(metric)
+    filters, parameters = search_filters(provider, model)
+    result = @connection.exec_params(search_query(metric, filters), [vector, limit, *parameters])
+
+    result.map { |row| record_from(row) }
+  end
+
+  private
+
+  def validate_dimensions(value)
+    return value if value.is_a?(Integer) && value.positive?
+
+    raise ArgumentError, 'dimensions must be a positive integer'
+  end
+
+  def validate_table(value)
+    table = value.to_s
+    return table if /\A[a-z_][a-z0-9_]*\z/.match?(table)
+
+    raise ArgumentError, 'table must be a lowercase PostgreSQL identifier'
+  end
+
+  def serialize_vector(embedding)
+    unless embedding.is_a?(Array) && embedding.length == dimensions
+      raise Prescient::InvalidVectorError, "embedding must contain exactly #{dimensions} values"
+    end
+
+    values = embedding.map { |value| Float(value) }
+    raise Prescient::InvalidVectorError, 'embedding values must be finite' unless values.all?(&:finite?)
+
+    "[#{values.join(',')}]"
+  rescue ArgumentError, TypeError
+    raise Prescient::InvalidVectorError, 'embedding values must be numeric'
+  end
+
+  def validate_limit(value)
+    return value if value.is_a?(Integer) && value.positive?
+
+    raise ArgumentError, 'limit must be a positive integer'
+  end
+
+  def validate_metric(value)
+    return value if METRICS.key?(value)
+
+    raise ArgumentError, "unsupported distance metric: #{value}"
+  end
+
+  def metric_operator_class(metric)
+    {
+      cosine:        'vector_cosine_ops',
+      euclidean:     'vector_l2_ops',
+      inner_product: 'vector_ip_ops',
+    }.fetch(metric)
+  end
+
+  def search_filters(provider, model)
+    filters = []
+    parameters = []
+    if provider
+      filters << "provider = $#{parameters.length + 3}"
+      parameters << provider.to_s
+    end
+    if model
+      filters << "model = $#{parameters.length + 3}"
+      parameters << model
+    end
+
+    [filters, parameters]
+  end
+
+  def search_query(metric, filters)
+    where = filters.empty? ? '' : "WHERE #{filters.join(' AND ')}"
+
+    <<~SQL
+      SELECT id, provider, model, dimensions, content, metadata,
+             embedding #{METRICS.fetch(metric)} $1::vector AS distance
+      FROM #{table}
+      #{where}
+      ORDER BY embedding #{METRICS.fetch(metric)} $1::vector
+      LIMIT $2
+    SQL
+  end
+
+  def record_from(row)
+    {
+      id:         row.fetch('id'),
+      provider:   row.fetch('provider'),
+      model:      row.fetch('model'),
+      dimensions: Integer(row.fetch('dimensions')),
+      content:    row.fetch('content'),
+      metadata:   JSON.parse(row.fetch('metadata')),
+      distance:   row['distance']&.to_f,
+    }
+  end
+end

data/lib/prescient/provider/anthropic.rb

@@ -2,6 +2,7 @@
 
 require 'httparty'
 
+# Anthropic Messages API provider adapter.
 class Prescient::Provider::Anthropic < Prescient::Base
   include HTTParty
 
@@ -12,12 +13,19 @@ class Prescient::Provider::Anthropic < Prescient::Base
     self.class.default_timeout(@options[:timeout] || 60)
   end
 
+  # Anthropic does not provide embeddings through this adapter.
+  #
+  # @raise [Prescient::Error] Always, because Anthropic embeddings are not supported
   def generate_embedding(_text, **_options)
     # Anthropic doesn't provide embedding API, raise error
     raise Prescient::Error,
           'Anthropic provider does not support embeddings. Use OpenAI or HuggingFace for embeddings.'
   end
 
+  # Generate a response using Anthropic's Messages API.
+  # @param prompt [String] Prompt to send
+  # @param context_items [Array<Hash, String>] Optional context items
+  # @return [Hash] Normalized response data
   def generate_response(prompt, context_items = [], **options)
     handle_errors do
       formatted_prompt = build_prompt(prompt, context_items)
@@ -29,7 +37,7 @@ class Prescient::Provider::Anthropic < Prescient::Base
                                    'anthropic-version' => '2023-06-01',
                                  },
                                  body:    {
-                                   model:       @options[:model],
+                                   model:       options[:model] || @options[:model],
                                    max_tokens:  options[:max_tokens] || 2000,
                                    temperature: options[:temperature] || 0.7,
                                    messages:    [
@@ -47,7 +55,7 @@ class Prescient::Provider::Anthropic < Prescient::Base
 
       {
         response:        content.strip,
-        model:           @options[:model],
+        model:           options[:model] || @options[:model],
         provider:        'anthropic',
         processing_time: nil,
         metadata:        {
@@ -57,58 +65,63 @@ class Prescient::Provider::Anthropic < Prescient::Base
     end
   end
 
+  # Check Anthropic API availability using the non-generating `/v1/models` endpoint.
+  #
+  # `reachable` indicates the API answered successfully. `ready` indicates that
+  # the configured model appears in the returned model list.
+  #
+  # @return [Hash] Provider health information
   def health_check
     handle_errors do
-      # Test with a simple message
-      response = self.class.post('/v1/messages',
-                                 headers: {
-                                   'Content-Type'      => 'application/json',
-                                   'x-api-key'         => @options[:api_key],
-                                   'anthropic-version' => '2023-06-01',
-                                 },
-                                 body:    {
-                                   model:      @options[:model],
-                                   max_tokens: 10,
-                                   messages:   [
-                                     {
-                                       role:    'user',
-                                       content: 'Test',
-                                     },
-                                   ],
-                                 }.to_json)
+      response = self.class.get('/v1/models', headers: api_headers)
 
       if response.success?
+        models = response.parsed_response['data'] || []
+        model_available = models.any? { |model| model['id'] == @options[:model] }
         {
-          status:   'healthy',
-          provider: 'anthropic',
-          model:    @options[:model],
-          ready:    true,
+          status:           'healthy',
+          provider:         'anthropic',
+          reachable:        true,
+          models_available: models.map { |model| model['id'] },
+          model:            { name: @options[:model], available: model_available },
+          ready:            model_available,
         }
       else
         {
-          status:   'unhealthy',
-          provider: 'anthropic',
-          error:    "HTTP #{response.code}",
-          message:  response.message,
+          status: 'unhealthy', provider: 'anthropic', reachable: true,
+          error: "HTTP #{response.code}", message: response.message, ready: false
         }
       end
     end
-  rescue Prescient::ConnectionError => e
+  rescue Prescient::Error => e
     {
-      status:   'unavailable',
-      provider: 'anthropic',
-      error:    e.class.name,
-      message:  e.message,
+      status:    'unavailable',
+      provider:  'anthropic',
+      reachable: false,
+      error:     e.class.name,
+      message:   e.message,
+      ready:     false,
     }
   end
 
+  # Return models available to the configured Anthropic account.
+  # @return [Array<Hash>] Model descriptors
   def list_models
-    # Anthropic doesn't provide a models list API
-    [
-      { name: 'claude-3-haiku-20240307', type: 'text' },
-      { name: 'claude-3-sonnet-20240229', type: 'text' },
-      { name: 'claude-3-opus-20240229', type: 'text' },
-    ]
+    handle_errors do
+      response = self.class.get('/v1/models', headers: api_headers)
+      validate_response!(response, 'model listing')
+
+      (response.parsed_response['data'] || []).map do |model|
+        {
+          name:             model['id'],
+          type:             'text',
+          display_name:     model['display_name'],
+          created_at:       model['created_at'],
+          max_input_tokens: model['max_input_tokens'],
+          max_tokens:       model['max_tokens'],
+        }.compact
+      end
+    end
   end
 
   protected
@@ -124,23 +137,11 @@ class Prescient::Provider::Anthropic < Prescient::Base
 
   private
 
-  def validate_response!(response, operation)
-    return if response.success?
-
-    case response.code
-    when 400
-      raise Prescient::Error, "Bad request for #{operation}: #{response.body}"
-    when 401
-      raise Prescient::AuthenticationError, "Authentication failed for #{operation}"
-    when 403
-      raise Prescient::AuthenticationError, "Forbidden access for #{operation}"
-    when 429
-      raise Prescient::RateLimitError, "Rate limit exceeded for #{operation}"
-    when 500..599
-      raise Prescient::Error, "Anthropic server error during #{operation}: #{response.body}"
-    else
-      raise Prescient::Error,
-            "Anthropic request failed for #{operation}: HTTP #{response.code} - #{response.message}"
-    end
+  def api_headers
+    {
+      'Content-Type'      => 'application/json',
+      'x-api-key'         => @options[:api_key],
+      'anthropic-version' => '2023-06-01',
+    }
   end
 end

data/lib/prescient/provider/huggingface.rb

@@ -2,11 +2,25 @@
 
 require 'httparty'
 
+# Hugging Face router-backed Inference Providers API adapter.
 class Prescient::Provider::HuggingFace < Prescient::Base
   include HTTParty
 
-  base_uri 'https://api-inference.huggingface.co'
+  base_uri 'https://router.huggingface.co'
 
+  # Router path for the Hugging Face feature-extraction provider.
+  # @return [String] Feature-extraction endpoint template
+  FEATURE_EXTRACTION_PATH = '/hf-inference/models/%<model>s/pipeline/feature-extraction'
+
+  # OpenAI-compatible router path for Hugging Face chat completions.
+  # @return [String] Chat-completions endpoint path
+  CHAT_COMPLETIONS_PATH = '/v1/chat/completions'
+
+  # OpenAI-compatible router path for listing available chat models.
+  # @return [String] Model-list endpoint path
+  MODEL_LIST_PATH = '/v1/models'
+
+  # Known embedding dimensions for commonly used models.
   EMBEDDING_DIMENSIONS = {
     'sentence-transformers/all-MiniLM-L6-v2'     => 384,
     'sentence-transformers/all-mpnet-base-v2'    => 768,
@@ -18,21 +32,20 @@ class Prescient::Provider::HuggingFace < Prescient::Base
     self.class.default_timeout(@options[:timeout] || 60)
   end
 
-  def generate_embedding(text, **_options)
+  # Generate an embedding through Hugging Face feature extraction.
+  # @param text [String] Text to embed
+  # @return [Array<Float>] Embedding vector
+  def generate_embedding(text, **options)
     handle_errors do
       clean_text_input = clean_text(text)
 
-      response = self.class.post("/pipeline/feature-extraction/#{@options[:embedding_model]}",
+      embedding_model = options[:model] || @options[:embedding_model]
+      response = self.class.post(FEATURE_EXTRACTION_PATH % { model: embedding_model },
                                  headers: {
                                    'Content-Type'  => 'application/json',
                                    'Authorization' => "Bearer #{@options[:api_key]}",
                                  },
-                                 body:    {
-                                   inputs:  clean_text_input,
-                                   options: {
-                                     wait_for_model: true,
-                                   },
-                                 }.to_json)
+                                 body:    { inputs: clean_text_input }.to_json)
 
       validate_response!(response, 'embedding generation')
 
@@ -42,82 +55,79 @@ class Prescient::Provider::HuggingFace < Prescient::Base
 
       raise Prescient::InvalidResponseError, 'No embedding returned' unless embedding_data.is_a?(Array)
 
-      expected_dimensions = EMBEDDING_DIMENSIONS[@options[:embedding_model]] || 384
-      normalize_embedding(embedding_data, expected_dimensions)
+      expected_dimensions = EMBEDDING_DIMENSIONS[embedding_model] || @options[:embedding_dimensions]
+      unless expected_dimensions
+        raise Prescient::Error,
+              "Embedding dimensions are required for model #{embedding_model}"
+      end
+
+      validate_embedding_dimensions(embedding_data, expected_dimensions)
     end
   end
 
+  # Generate text through a Hugging Face text-generation model.
+  # @param prompt [String] Prompt to send
+  # @param context_items [Array<Hash, String>] Optional context items
+  # @return [Hash] Normalized response data
   def generate_response(prompt, context_items = [], **options)
     handle_errors do
       formatted_prompt = build_prompt(prompt, context_items)
 
-      response = self.class.post("/models/#{@options[:chat_model]}",
+      response = self.class.post(CHAT_COMPLETIONS_PATH,
                                  headers: {
                                    'Content-Type'  => 'application/json',
                                    'Authorization' => "Bearer #{@options[:api_key]}",
                                  },
                                  body:    {
-                                   inputs:     formatted_prompt,
-                                   parameters: {
-                                     max_new_tokens:   options[:max_tokens] || 2000,
-                                     temperature:      options[:temperature] || 0.7,
-                                     top_p:            options[:top_p] || 0.9,
-                                     return_full_text: false,
-                                   },
-                                   options:    {
-                                     wait_for_model: true,
-                                   },
+                                   model:       options[:model] || @options[:chat_model],
+                                   messages:    [{ role: 'user', content: formatted_prompt }],
+                                   max_tokens:  options[:max_tokens] || 2000,
+                                   temperature: options[:temperature] || 0.7,
+                                   top_p:       options[:top_p] || 0.9,
                                  }.to_json)
 
       validate_response!(response, 'text generation')
 
-      # HuggingFace returns different formats depending on the model
-      generated_text = nil
       parsed_response = response.parsed_response
-
-      if parsed_response.is_a?(Array) && parsed_response.first.is_a?(Hash)
-        generated_text = parsed_response.first['generated_text']
-      elsif parsed_response.is_a?(Hash)
-        generated_text = parsed_response['generated_text'] || parsed_response['text']
-      end
-
+      generated_text = parsed_response.dig('choices', 0, 'message', 'content') if parsed_response.is_a?(Hash)
       raise Prescient::InvalidResponseError, 'No response generated' unless generated_text
 
       {
         response:        generated_text.strip,
-        model:           @options[:chat_model],
+        model:           options[:model] || @options[:chat_model],
         provider:        'huggingface',
         processing_time: nil,
-        metadata:        {},
+        metadata:        {
+          usage:         parsed_response['usage'],
+          finish_reason: parsed_response.dig('choices', 0, 'finish_reason'),
+        },
       }
     end
   end
 
+  # Check availability of the configured embedding and text models.
+  #
+  # The embedding model is checked against the model metadata API on
+  # `huggingface.co`, while the chat model is checked against the router's
+  # OpenAI-compatible `/v1/models` listing. `ready` requires both checks to
+  # succeed.
+  #
+  # @return [Hash] Provider health information
   def health_check
     handle_errors do
-      # Test embedding model
-      embedding_response = self.class.post("/pipeline/feature-extraction/#{@options[:embedding_model]}",
-                                           headers: {
-                                             'Authorization' => "Bearer #{@options[:api_key]}",
-                                           },
-                                           body:    { inputs: 'test' }.to_json)
-
-      # Test chat model
-      chat_response = self.class.post("/models/#{@options[:chat_model]}",
-                                      headers: {
-                                        'Authorization' => "Bearer #{@options[:api_key]}",
-                                      },
-                                      body:    {
-                                        inputs:     'test',
-                                        parameters: { max_new_tokens: 5 },
-                                      }.to_json)
+      embedding_response = self.class.get("https://huggingface.co/api/models/#{@options[:embedding_model]}",
+                                          headers: { 'Authorization' => "Bearer #{@options[:api_key]}" })
+      chat_response = self.class.get(MODEL_LIST_PATH,
+                                     headers: { 'Authorization' => "Bearer #{@options[:api_key]}" })
 
       embedding_healthy = embedding_response.success?
-      chat_healthy = chat_response.success?
+      chat_models = chat_response.parsed_response['data'] || []
+      chat_healthy = chat_response.success? && chat_models.any? { |model| model['id'] == @options[:chat_model] }
 
       {
         status:          embedding_healthy && chat_healthy ? 'healthy' : 'partial',
         provider:        'huggingface',
+        reachable:       true,
         embedding_model: {
           name:      @options[:embedding_model],
           available: embedding_healthy,
@@ -129,15 +139,23 @@ class Prescient::Provider::HuggingFace < Prescient::Base
         ready:           embedding_healthy && chat_healthy,
       }
     end
-  rescue Prescient::ConnectionError => e
+  rescue Prescient::Error => e
     {
-      status:   'unavailable',
-      provider: 'huggingface',
-      error:    e.class.name,
-      message:  e.message,
+      status:    'unavailable',
+      provider:  'huggingface',
+      reachable: false,
+      error:     e.class.name,
+      message:   e.message,
+      ready:     false,
     }
   end
 
+  # Return the configured Hugging Face models.
+  #
+  # This method does not query the Hugging Face APIs. It reflects the current
+  # adapter configuration only.
+  #
+  # @return [Array<Hash>] Model descriptors
   def list_models
     # HuggingFace doesn't provide a simple API to list all models
     # Return the configured models
@@ -165,36 +183,19 @@ class Prescient::Provider::HuggingFace < Prescient::Base
 
   private
 
-  def validate_response!(response, operation)
-    return if response.success?
-
-    case response.code
-    when 400
-      raise Prescient::Error, "Bad request for #{operation}: #{response.body}"
-    when 401
-      raise Prescient::AuthenticationError, "Authentication failed for #{operation}"
-    when 403
-      raise Prescient::AuthenticationError, "Forbidden access for #{operation}"
-    when 429
-      raise Prescient::RateLimitError, "Rate limit exceeded for #{operation}"
-    when 503
+  def provider_error(message, response, operation:, provider: nil, error_class: Prescient::ProviderError)
+    if response.code == 503
       # HuggingFace model loading
       error_body = begin
         response.parsed_response
       rescue StandardError
-        response.body
+        nil
       end
       if error_body.is_a?(Hash) && error_body['error']&.include?('loading')
-        raise Prescient::Error, 'Model is loading, please try again later'
+        message = 'Model is loading, please try again later'
       end
-
-      raise Prescient::Error, "HuggingFace service unavailable for #{operation}"
-
-    when 500..599
-      raise Prescient::Error, "HuggingFace server error during #{operation}: #{response.body}"
-    else
-      raise Prescient::Error,
-            "HuggingFace request failed for #{operation}: HTTP #{response.code} - #{response.message}"
     end
+
+    super
   end
 end