lex-dataset 0.1.0 → 0.2.1

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 5e63278e11eb728b818a7519bc0a33fb713088d3d149100d4398414126e6ced8
-  data.tar.gz: a119bcde6cec1c5f253a46284ead75f5fe672f0cd6a75e7334a6d7b325e61255
+  metadata.gz: 736a22b87a807e24ded11c873e9a8d20dc754a7aabff14e54d33f6b889df2a1c
+  data.tar.gz: 702e7e82a1d51e996938034c04354bb467019bf99e116097b3a4da02e618e132
 SHA512:
-  metadata.gz: f97ed63b5cc90b64a8773186bdc7d2df964a8ba334c5827dff51653db83030c9e19f89fbd07b7537995fe23677cd22792241bcc53244bcd19bd0e7705c9c6caa
-  data.tar.gz: 71209de12c040ce5c4a09436cf8902fc517352b30f59072a9a0dcfd7db031d9279b2e108c2188d6a108d40c611d68306f67d1503d6aa38ee623d046b93a0d56b
+  metadata.gz: c74a390f975e215614623f06e9e2cb9a4af33137329e05d82719fc12a1fe71715079a448f18ef64ae9967492f8cc11a873fb1f670d320575b823568cb7902012
+  data.tar.gz: dbb2cedfd0146ab4c66230761f21ec01bcf9284a01fd699753184ccbf18d2e139620046e167f0792ec5ff3eaa171c31b53f7b0599a7161cb932ebc1f6fb4588d

data/README.md

@@ -0,0 +1,71 @@
+# lex-dataset
+
+Versioned dataset management for LegionIO. Provides immutable versioned dataset storage with CSV, JSON, and JSONL import/export and content-hash deduplication.
+
+## Overview
+
+`lex-dataset` stores named datasets with full version history. Each version is content-hashed — submitting the same rows twice results in no new version. Datasets consist of input/expected-output row pairs suitable for LLM evaluation workflows.
+
+## Installation
+
+```ruby
+gem 'lex-dataset'
+```
+
+## Usage
+
+```ruby
+require 'legion/extensions/dataset'
+
+client = Legion::Extensions::Dataset::Client.new
+
+# Create a dataset with inline rows
+client.create_dataset(
+  name: 'qa-pairs-v1',
+  description: 'Question-answer evaluation set',
+  rows: [
+    { input: 'What is BGP?', expected_output: 'Border Gateway Protocol' },
+    { input: 'What is OSPF?', expected_output: 'Open Shortest Path First' }
+  ]
+)
+# => { created: true, name: 'qa-pairs-v1', version: 1, row_count: 2 }
+
+# Import from file
+client.import_dataset(name: 'qa-from-file', path: '/data/qa.jsonl', format: 'jsonl')
+
+# Export a specific version
+client.export_dataset(name: 'qa-pairs-v1', path: '/tmp/export.json', format: 'json')
+
+# Retrieve rows
+client.get_dataset(name: 'qa-pairs-v1')
+# => { name: 'qa-pairs-v1', version: 1, row_count: 2, rows: [...] }
+
+# List all datasets
+client.list_datasets
+```
+
+## Supported Formats
+
+| Format | Description |
+|--------|-------------|
+| `json` | Array of row objects (default) |
+| `jsonl` | One JSON object per line |
+| `csv` | Header row + data rows |
+
+## Related Repos
+
+- `lex-eval` — uses datasets as input for LLM evaluation runs
+- `lex-prompt` — versioned prompt templates consumed alongside datasets in evaluation workflows
+- `legion-data` — underlying Sequel database connection (SQLite/PostgreSQL/MySQL)
+
+## Development
+
+```bash
+bundle install
+bundle exec rspec
+bundle exec rubocop
+```
+
+## License
+
+MIT

data/lib/legion/extensions/dataset/client.rb

@@ -6,6 +6,7 @@ module Legion
       class Client
         include Runners::Dataset
         include Runners::Experiment
+        include Runners::Sampling
 
         def initialize(db: nil, **opts)
           @db   = db

data/lib/legion/extensions/dataset/runners/dataset.rb

@@ -58,8 +58,94 @@ module Legion
               rows: rows.map { |r| { row_index: r[:row_index], input: r[:input], expected_output: r[:expected_output] } } }
           end
 
+          def generate_dataset(name:, description:, count: 10, schema: nil, model: nil, **)
+            return { error: 'legion-llm is not available' } unless llm_available?
+
+            rows = call_llm_for_rows(description: description, count: count, schema: schema, model: model)
+            return rows if rows.is_a?(Hash) && rows[:error]
+
+            result = create_dataset(name: name, description: description, rows: rows)
+            result.merge(generated: true)
+          end
+
           private
 
+          def llm_available?
+            defined?(Legion::LLM) && Legion::LLM.respond_to?(:started?) && Legion::LLM.started?
+          end
+
+          def call_llm_for_rows(description:, count:, schema:, model:)
+            prompt = build_generate_prompt(description: description, count: count, schema: schema)
+            llm_opts = model ? { model: model } : {}
+
+            response = invoke_llm(prompt: prompt, **llm_opts)
+            rows = parse_llm_rows(response)
+
+            if rows.nil?
+              retry_prompt = "#{prompt}\n\nIMPORTANT: Your previous response was not valid JSON. Return ONLY a valid JSON array."
+              response = invoke_llm(prompt: retry_prompt, **llm_opts)
+              rows = parse_llm_rows(response)
+            end
+
+            rows || { error: 'LLM did not return valid JSON after retry' }
+          end
+
+          def invoke_llm(prompt:, **llm_opts)
+            result = if Legion::LLM.respond_to?(:structured)
+                       Legion::LLM.structured(
+                         message: prompt,
+                         schema:  generate_schema,
+                         **llm_opts
+                       )
+                     else
+                       Legion::LLM.chat(message: prompt, **llm_opts)
+                     end
+            content = result.respond_to?(:content) ? result.content : result.to_s
+            content.strip.sub(/\A```(?:json)?\n?/, '').sub(/\n?```\z/, '')
+          end
+
+          def parse_llm_rows(content)
+            parsed = ::JSON.parse(content)
+            return nil unless parsed.is_a?(Array)
+
+            parsed.map do |item|
+              h = item.transform_keys(&:to_sym)
+              { input: h[:input].to_s, expected_output: h[:expected_output]&.to_s }
+            end
+          rescue ::JSON::ParserError
+            nil
+          end
+
+          def build_generate_prompt(description:, count:, schema:)
+            lines = []
+            lines << "You are a test case generator. Generate exactly #{count} test cases as a JSON array."
+            lines << 'Each test case must have "input" and "expected_output" fields.'
+            lines << ''
+            lines << "Description: #{description}"
+            if schema
+              lines << ''
+              lines << 'Schema guidance for inputs and outputs:'
+              lines << "```json\n#{::JSON.generate(schema)}\n```"
+            end
+            lines << ''
+            lines << 'Respond ONLY with a valid JSON array, no other text.'
+            lines.join("\n")
+          end
+
+          def generate_schema
+            {
+              type:  'array',
+              items: {
+                type:       'object',
+                properties: {
+                  input:           { type: 'string' },
+                  expected_output: { type: 'string' }
+                },
+                required:   %w[input expected_output]
+              }
+            }
+          end
+
           def create_version(dataset_id, rows)
             hash = OpenSSL::Digest.new('SHA256').hexdigest(rows.to_s)
             ver_num = (db[:dataset_versions].where(dataset_id: dataset_id).max(:version) || 0) + 1

data/lib/legion/extensions/dataset/runners/sampling.rb

@@ -0,0 +1,72 @@
+# frozen_string_literal: true
+
+module Legion
+  module Extensions
+    module Dataset
+      module Runners
+        module Sampling
+          def sample_from_traces(dataset_name:, source: :legion_data, filters: {},
+                                 sample_size: nil, strategy: :recent, **)
+            traces = fetch_traces(source, filters)
+            sampled = apply_strategy(traces, strategy, sample_size)
+            rows = sampled.map { |t| { input: t[:input], expected_output: nil, metadata: t[:span_kind] } }
+            create_dataset(name: dataset_name, rows: rows)
+          end
+
+          private
+
+          def fetch_traces(source, filters)
+            case source
+            when :legion_data then fetch_from_db(filters)
+            else raise ArgumentError, "unknown trace source: #{source}"
+            end
+          end
+
+          def fetch_from_db(filters)
+            query = db[:traces]
+            query = query.where(span_kind: filters[:span_kind]) if filters[:span_kind]
+            query = query.where(status: filters[:status]) if filters[:status]
+            if filters[:time_range]
+              cutoff = Time.now.utc - filters[:time_range]
+              query = query.where { created_at >= cutoff }
+            end
+            query.order(Sequel.desc(:created_at)).all
+          end
+
+          def apply_strategy(traces, strategy, sample_size)
+            case strategy.to_sym
+            when :random       then sample_random(traces, sample_size)
+            when :error_biased then sample_error_biased(traces, sample_size)
+            when :stratified   then sample_stratified(traces, sample_size)
+            else sample_recent(traces, sample_size)
+            end
+          end
+
+          def sample_recent(traces, size)
+            size ? traces.first(size) : traces
+          end
+
+          def sample_random(traces, size)
+            size ? traces.sample(size) : traces.shuffle
+          end
+
+          def sample_error_biased(traces, size)
+            errors, successes = traces.partition { |t| t[:status] == 'error' }
+            return traces unless size
+
+            half = size / 2
+            (errors.first(half) + successes.first(size - half)).first(size)
+          end
+
+          def sample_stratified(traces, size)
+            groups = traces.group_by { |t| t[:span_kind] }
+            return traces unless size
+
+            per_group = [size / [groups.size, 1].max, 1].max
+            groups.values.flat_map { |g| g.first(per_group) }.first(size)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/legion/extensions/dataset/version.rb

@@ -3,7 +3,7 @@
 module Legion
   module Extensions
     module Dataset
-      VERSION = '0.1.0'
+      VERSION = '0.2.1'
     end
   end
 end

data/lib/legion/extensions/dataset.rb

@@ -4,6 +4,7 @@ require_relative 'dataset/version'
 require_relative 'dataset/helpers/import_export'
 require_relative 'dataset/runners/dataset'
 require_relative 'dataset/runners/experiment'
+require_relative 'dataset/runners/sampling'
 require_relative 'dataset/client'
 
 module Legion

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: lex-dataset
 version: !ruby/object:Gem::Version
-  version: 0.1.0
+  version: 0.2.1
 platform: ruby
 authors:
 - Matthew Iverson
@@ -17,11 +17,13 @@ executables: []
 extensions: []
 extra_rdoc_files: []
 files:
+- README.md
 - lib/legion/extensions/dataset.rb
 - lib/legion/extensions/dataset/client.rb
 - lib/legion/extensions/dataset/helpers/import_export.rb
 - lib/legion/extensions/dataset/runners/dataset.rb
 - lib/legion/extensions/dataset/runners/experiment.rb
+- lib/legion/extensions/dataset/runners/sampling.rb
 - lib/legion/extensions/dataset/version.rb
 homepage: https://github.com/LegionIO/lex-dataset
 licenses: