From feca433a20690020c9f2735be088a49f5be300af Mon Sep 17 00:00:00 2001 From: viksit Date: Wed, 14 May 2025 22:46:25 -0400 Subject: [PATCH] very basic version of few shot optimizer wrestling with types --- examples/optimizer-fewshot.ts | 73 +++++++++++++++++++++ src/index.ts | 2 + src/lib/core.ts | 31 ++++++++- src/lib/optimize/few-shot.ts | 112 ++++++++++++++++++++++++++++++++ src/lib/optimize/index.ts | 2 + src/lib/optimize/metric.ts | 19 ++++++ src/lib/optimize/types.ts | 15 +++++ src/lib/types.ts | 16 +++++ src/lib/utils/costs.ts | 16 +++++ src/lib/utils/tokens.ts | 10 +++ tests/optimize/few-shot-test.ts | 88 +++++++++++++++++++++++++ 11 files changed, 383 insertions(+), 1 deletion(-) create mode 100644 examples/optimizer-fewshot.ts create mode 100644 src/lib/optimize/few-shot.ts create mode 100644 src/lib/optimize/index.ts create mode 100644 src/lib/optimize/metric.ts create mode 100644 src/lib/optimize/types.ts create mode 100644 src/lib/utils/costs.ts create mode 100644 src/lib/utils/tokens.ts create mode 100644 tests/optimize/few-shot-test.ts diff --git a/examples/optimizer-fewshot.ts b/examples/optimizer-fewshot.ts new file mode 100644 index 0000000..db54120 --- /dev/null +++ b/examples/optimizer-fewshot.ts @@ -0,0 +1,73 @@ +// examples/optimizer-fewshot.ts +import { selvedge as s } from '../src'; + +s.debug('*'); + +s.models({ gpt35: s.openai('gpt-3.5-turbo') }); + +const queryWriter = s.prompt` + QUESTION: ${ q => q } + + Give exactly three distinct web-search queries (one per line) that + would help answer the question. No URLs, only plain text queries. +` + .inputs({ q: s.schema.string() }) + .outputs({ queries: s.schema.array(s.schema.string()) }) + .using('gpt35'); + +/* ── 3. Stub retrieval: query → URLs (replace in prod) ─────── */ +async function search(query: string): Promise { + return [ + `https://example.com/${query.replace(/\s+/g, '_')}/1`, + `https://example.com/${query.replace(/\s+/g, '_')}/2` + ]; +} + +/* ── 4. Recall@k metric (k = all returned URLs) ─────────────── */ +async function recallMetric( + pred: { queries: string[] }, + goldUrls: string[] +): Promise { + const urls = (await Promise.all(pred.queries.map(search))).flat(); + if (!goldUrls.length) return 0; + const hits = urls.filter(u => goldUrls.includes(u)).length; + return hits / goldUrls.length; // 0 … 1 +} + +/* ── 5. Training data with *gold URLs* ─────────────────────── */ +const trainset = [ + { + input: { q: 'Why did Tesla stock drop in Jan 2023?' }, + goldUrls: [ + 'https://news.site/tesla-jan-2023.html', + 'https://finance.site/tesla-earnings-q4.html' + ] + }, + { + input: { q: 'When was Rust 1.0 released?' }, + goldUrls: [ + 'https://en.wikipedia.org/wiki/Rust_(programming_language)' + ] + } +] as any; + + +/* ── 6. Metric wrapper to match optimiser signature ──────────── */ +const metric = (pred, gold) => recallMetric(pred, gold.goldUrls); + +/* ── 7. Optimise (no few-shot demos, 1 trial) ─────────────────── */ + +const tuned = await s.optimize( + queryWriter, + s.optimize.fewShot({ + trainset, + metric, + maxDemos: 0, // <-- no confusing URL demos + trials: 1, + costCapUSD: 0.02 + }) +); + +/* ── 8. Run on a fresh question ───────────────────────────────── */ +const result = await tuned({ q: 'Impact of remote work on cybersecurity 2024' }); +console.log('Optimised queries:\n', result.queries.join('\n')); \ No newline at end of file diff --git a/src/index.ts b/src/index.ts index d16d9f9..34e4583 100644 --- a/src/index.ts +++ b/src/index.ts @@ -27,6 +27,8 @@ export * from './lib/core'; export * from './lib/types'; export * from './lib/prompts'; export * from './lib/schema'; +export * as optimize from './lib/optimize'; +export * as metric from './lib/optimize/metric'; export { ModelRegistry } from './lib/models'; export { Store, store } from './lib/storage'; export { SelvedgeManager, manager } from './lib/manager'; diff --git a/src/lib/core.ts b/src/lib/core.ts index bf944c6..d01726e 100644 --- a/src/lib/core.ts +++ b/src/lib/core.ts @@ -9,6 +9,31 @@ import { store } from './storage'; import { flow as createFlow } from './flow'; import { enableDebug, enableNamespace, parseDebugString, debug } from './utils/debug'; import schemaHelpers from './schema'; // Import the schema helpers +import * as optimizeModule from './optimize'; +import { OptimizerSpec } from './optimize/types'; + +// --- optimisation wrapper (placed *before* selvedge object) ------------- +const optimizeFn: any = async ( + target: PromptTemplate, + optimizer: OptimizerSpec>, +) => { + if (!optimizer || typeof optimizer.run !== 'function') { + throw new Error('Second argument to s.optimize must be an OptimizerSpec with a .run() method'); + } + return optimizer.run(target); +}; +// Attach helper constructors like fewShot onto the function object +Object.assign(optimizeFn, optimizeModule); +// Ensure the key helpers are enumerable (in case Object.assign skipped) +if (!Object.prototype.hasOwnProperty.call(optimizeFn, 'fewShot') && 'fewShot' in optimizeModule) { + Object.defineProperty(optimizeFn, 'fewShot', { + value: (optimizeModule as any).fewShot, + writable: false, + enumerable: true, + configurable: false, + }); +} +// ------------------------------------------------------------------------ /** * The main Selvedge instance that provides access to all library functionality @@ -80,7 +105,6 @@ export const selvedge: SelvedgeInstance = { flow( steps: Array ) { - // Use the existing flow implementation from the flow module return createFlow(...steps); }, @@ -436,6 +460,11 @@ export const selvedge: SelvedgeInstance = { ChainOfThought: (t: TemplateStringsArray, ...v: any[]) => selvedge.prompt(t, ...v).prefix('Think step-by-step before answering.\n'), + /** + * Optimisation entry-point + */ + optimize: optimizeFn, + }; /** diff --git a/src/lib/optimize/few-shot.ts b/src/lib/optimize/few-shot.ts new file mode 100644 index 0000000..b6d0af6 --- /dev/null +++ b/src/lib/optimize/few-shot.ts @@ -0,0 +1,112 @@ +import { z } from 'zod'; +import { PromptTemplate } from '../prompts'; +import { TrainExample, MetricFn, OptimizerSpec } from './types'; +import { estimateTokens } from '../utils/tokens'; +import { openaiCostUSD } from '../utils/costs'; +import { debug } from '../utils/debug'; + +/** Options for the few-shot optimiser. */ +export interface FewShotOpts { + trainset: TrainExample[]; + metric: MetricFn; + maxDemos?: number; // default 4 + trials?: number; // default 80 + costCapUSD?: number; // default Infinity +} + +/** Factory exported to users. */ +export function fewShot( + opts: FewShotOpts +): OptimizerSpec> { + + const { + trainset, + metric, + maxDemos = 4, + trials = 80, + costCapUSD = Number.POSITIVE_INFINITY + } = opts; + + return { + async run(base: PromptTemplate): Promise> { + let bestScore = -Infinity; + let bestClone: PromptTemplate = base; + debug('optimizer:few-shot', 'Starting run with base prompt:', base); + + const demoPool = trainset.slice(0, 12); // cap brute pool + const demoCombos: TrainExample[][] = []; + + /* --- enumerate / sample demo sets --------------------------------- */ + const brute = demoPool.length <= maxDemos && demoPool.length <= 6; + if (brute) { + // brute force every subset up to maxDemos + debug('optimizer:few-shot', 'Using brute force for demo sets.'); + const subsets = (arr: any[]): any[][] => + arr.length === 0 + ? [[]] + : subsets(arr.slice(1)).flatMap(s => [s, [arr[0], ...s]]); + subsets(demoPool).forEach(set => { + if (set.length && set.length <= maxDemos) demoCombos.push(set); + }); + } else { + // random sample + debug('optimizer:few-shot', 'Using random sampling for demo sets.'); + for (let i = 0; i < trials; i++) { + const shuffled = [...demoPool].sort(() => Math.random() - 0.5); + demoCombos.push(shuffled.slice(0, maxDemos)); + } + } + debug('optimizer:few-shot', `Generated ${demoCombos.length} demo combinations.`); + + /* --- evaluate each candidate -------------------------------------- */ + let runningCost = 0; + + for (const demos of demoCombos) { + debug('optimizer:few-shot', 'Evaluating demo set:', demos); + if (demos.length === 0) { + debug('optimizer:few-shot', 'Skipping empty demo set.'); + continue; + } + const promptCost = + (estimateTokens(base.render(demos[0].input)) / 1000) * + openaiCostUSD['gpt-4o']; + if (runningCost + promptCost > costCapUSD) { + debug('optimizer:few-shot', `Cost cap exceeded: ${runningCost + promptCost} > ${costCapUSD}. Stopping evaluation.`); + break; + } + + // Build a few‑shot prefix manually (safer than relying on .train()) + const fewShotPrefix = buildFewShotPrefix(demos); + + const variant = base.clone().prefix(fewShotPrefix); + + let scoreSum = 0; + for (const ex of trainset) { + const pred = await variant(ex.input); + scoreSum += await metric(pred, ex.output); + } + const avg = scoreSum / trainset.length; + + if (avg > bestScore) { + bestScore = avg; + bestClone = variant; + debug('optimizer:few-shot', `New best score: ${bestScore}. Variant:`, bestClone); + } + runningCost += promptCost; + } + + debug('optimizer:few-shot', `Finished run. Best score: ${bestScore}. Returning prompt:`, bestClone); + return bestClone; + } + }; +} + +/* Helper: convert demo tuples to readable prompt text */ +function buildFewShotPrefix(demos: TrainExample[]): string { + const blocks = demos.map(({ input, output }) => { + const inStr = typeof input === 'string' ? input : JSON.stringify(input); + const outStr = typeof output === 'string' ? output : JSON.stringify(output); + return `### Example\nInput: ${inStr}\nOutput: ${outStr}\n`; + }); + return blocks.join('\n') + '\n\n'; +} \ No newline at end of file diff --git a/src/lib/optimize/index.ts b/src/lib/optimize/index.ts new file mode 100644 index 0000000..9d23470 --- /dev/null +++ b/src/lib/optimize/index.ts @@ -0,0 +1,2 @@ +export * from './types'; +export { fewShot } from './few-shot'; // <— exact file name \ No newline at end of file diff --git a/src/lib/optimize/metric.ts b/src/lib/optimize/metric.ts new file mode 100644 index 0000000..e752f7e --- /dev/null +++ b/src/lib/optimize/metric.ts @@ -0,0 +1,19 @@ +import { MetricFn } from './types'; + +/** Generic exact-match metric. */ +export function exactMatch(): MetricFn { + return (pred, gold) => (JSON.stringify(pred) === JSON.stringify(gold) ? 1 : 0); +} + +/** Simple F-1 for sets of strings. */ +export function f1(): MetricFn { + return (p, g) => { + const pred = Array.isArray(p) ? new Set(p) : new Set([p]); + const gold = Array.isArray(g) ? new Set(g) : new Set([g]); + const hits = [...pred].filter(x => gold.has(x)).length; + if (!hits) return 0; + const precision = hits / pred.size; + const recall = hits / gold.size; + return 2 * precision * recall / (precision + recall); + }; +} \ No newline at end of file diff --git a/src/lib/optimize/types.ts b/src/lib/optimize/types.ts new file mode 100644 index 0000000..b51cad5 --- /dev/null +++ b/src/lib/optimize/types.ts @@ -0,0 +1,15 @@ +import { PromptTemplate } from '../prompts'; + +/** Training example tuple. */ +export interface TrainExample { + input: I; + output: O; +} + +/** Metric: returns single scalar; higher is better. */ +export type MetricFn

= (prediction: P, gold: G) => number | Promise; + +/** Every optimiser returns a tuned clone of the target. */ +export interface OptimizerSpec> { + run(target: T): Promise; +} \ No newline at end of file diff --git a/src/lib/types.ts b/src/lib/types.ts index cc59a5f..00c6705 100644 --- a/src/lib/types.ts +++ b/src/lib/types.ts @@ -33,6 +33,7 @@ export interface ModelDefinition { * The core Selvedge instance interface */ export interface SelvedgeInstance { + [x: string]: any; /** * Register models with simple alias names */ @@ -138,8 +139,14 @@ export interface SelvedgeInstance { * Create a Chain of Thought prompt */ ChainOfThought: (strings: TemplateStringsArray, ...values: any[]) => PromptTemplate; + + /** + * Optimise a prompt template using the given optimiser spec + */ + optimize: OptimizeFn; } + /** * Type for the schema helper functions provided by Selvedge. */ @@ -188,3 +195,12 @@ export interface ModelAdapter { /** Optional method to set mock responses for testing */ setResponses?(responses: { completion?: string; chat?: string | ((messages: any[]) => string); promptMap?: Record }): void; } + +/** Helper namespace exported from ./optimize (e.g., fewShot) */ +type OptimizeHelpers = typeof import('./optimize'); + +/** Callable optimise function merged with helpers */ +type OptimizeFn = (( + target: import('./prompts/types').PromptTemplate, + optimizer: import('./optimize/types').OptimizerSpec> +) => Promise>) & OptimizeHelpers; diff --git a/src/lib/utils/costs.ts b/src/lib/utils/costs.ts new file mode 100644 index 0000000..acb07eb --- /dev/null +++ b/src/lib/utils/costs.ts @@ -0,0 +1,16 @@ +/** + * Minimal per-1K-tokens price map (USD). + * Extend as your model registry grows; numbers are illustrative. + */ +export const openaiCostUSD: Record = { + 'gpt-3.5-turbo': 0.0015, + 'gpt-4': 0.03, // input price per 1K tokens + 'gpt-4o': 0.015 +}; + +/** + * Get cost for model id or fallback to 0.02 USD / 1K tokens. + */ +export function pricePer1K(model: string): number { + return openaiCostUSD[model] ?? 0.02; +} \ No newline at end of file diff --git a/src/lib/utils/tokens.ts b/src/lib/utils/tokens.ts new file mode 100644 index 0000000..a0fe5c9 --- /dev/null +++ b/src/lib/utils/tokens.ts @@ -0,0 +1,10 @@ + +/** + * Ultra-cheap token estimator. + * Good enough for cost-guardrails; replace with tiktoken if you need accuracy. + */ +export function estimateTokens(text: string): number { + if (!text) return 0; + // rough heuristic: 1 token ≈ 4 characters for English prose + return Math.ceil(text.length / 4); +} \ No newline at end of file diff --git a/tests/optimize/few-shot-test.ts b/tests/optimize/few-shot-test.ts new file mode 100644 index 0000000..2baa22e --- /dev/null +++ b/tests/optimize/few-shot-test.ts @@ -0,0 +1,88 @@ +/** + * Optimizer (few-shot) unit tests + * + * Run with: bun test + */ + +// @ts-ignore – Bun test global typings +import { expect, describe, it, beforeEach } from 'bun:test'; +import { selvedge } from '../../src/lib/core'; +import { ModelRegistry } from '../../src/lib/models'; +import { ModelProvider } from '../../src/lib/types'; +import { MockModelAdapter } from '../../src/lib/providers/mock/mock'; +import * as z from 'zod'; + +/* ── Dummy search helper (no network) ─────────────────────────── */ +async function searchStub(query: string): Promise { + return [`https://stub/${query.replace(/\s+/g, '_')}`]; +} + +/* ── Recall metric for unit test (k = all URLs) ───────────────── */ +async function recallMetric( + pred: { queries: string[] }, + goldUrls: string[] +): Promise { + const returned = (await Promise.all(pred.queries.map(searchStub))).flat(); + const hits = returned.filter(u => goldUrls.includes(u)).length; + return goldUrls.length ? hits / goldUrls.length : 0; +} + +describe('Few-shot optimiser', () => { + beforeEach(() => { + ModelRegistry.clear(); + selvedge.models({ + testModel: selvedge.mock('test-model') + }); + }); + + it('optimises a prompt and preserves callable behaviour', async () => { + /* 1. Set up mock model response (always 3 queries) */ + const mockAdapter = ModelRegistry.getAdapter({ + provider: ModelProvider.MOCK, + model: 'test-model' + }) as MockModelAdapter; + + mockAdapter.setResponses({ + chat: JSON.stringify({ + queries: ['alpha query', 'beta query', 'gamma query'] + }) + }); + + /* 2. Base prompt */ + const writeQueries = selvedge.prompt` + QUESTION: ${q => q} + Return JSON { "queries": [string, string, string] } + ` + .inputs({ q: selvedge.schema.string() }) + .outputs({ queries: selvedge.schema.array(selvedge.schema.string()) }) + .using('testModel'); + + /* 3. Tiny trainset (question + gold URLs) */ + const trainset = [ + { + input: { q: 'Why did Tesla stock drop?' }, + goldUrls: ['https://stub/alpha_query'] + } + ] as any; // cast to satisfy FewShot typings + + const metric = (pred, gold) => recallMetric(pred, gold.goldUrls); + + /* 4. Optimise with zero demos (just metric pass-through) */ + const tuned = await selvedge.optimize( + writeQueries, + selvedge.optimize.fewShot({ + trainset, + metric, + maxDemos: 0, + trials: 1 + }) + ); + + /* 5. Call tuned prompt and check structure */ + const res = await tuned({ q: 'Any question' }); + + expect(Array.isArray(res.queries)).toBe(true); + expect(res.queries.length).toBe(3); + expect(res.queries[0]).toBe('alpha query'); + }); +}); \ No newline at end of file