very basic version of few shot optimizer wrestling with types

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<string[]> {
+  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<number> {
+  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'));

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';

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<any, any>,
+  optimizer: OptimizerSpec<PromptTemplate<any, any>>,
+) => {
+  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<TInput = any, TOutput = any>(
     steps: Array<any>
   ) {
-    // Use the existing flow implementation from the flow module
     return createFlow<TInput, TOutput>(...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,
+
 };
 
 /**

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<I, O> {
+  trainset:   TrainExample<I, O>[];
+  metric:     MetricFn<O, O>;
+  maxDemos?:  number;     // default 4
+  trials?:    number;     // default 80
+  costCapUSD?: number;    // default Infinity
+}
+
+/** Factory exported to users. */
+export function fewShot<I, O>(
+  opts: FewShotOpts<I, O>
+): OptimizerSpec<PromptTemplate<O, I>> {
+
+  const {
+    trainset,
+    metric,
+    maxDemos = 4,
+    trials   = 80,
+    costCapUSD = Number.POSITIVE_INFINITY
+  } = opts;
+
+  return {
+    async run(base: PromptTemplate<O, I>): Promise<PromptTemplate<O, I>> {
+      let bestScore = -Infinity;
+      let bestClone: PromptTemplate<O, I> = base;
+      debug('optimizer:few-shot', 'Starting run with base prompt:', base);
+
+      const demoPool = trainset.slice(0, 12);            // cap brute pool
+      const demoCombos: TrainExample<I, O>[][] = [];
+
+      /* --- 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<I, O>(demos: TrainExample<I, O>[]): 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';
+}

src/lib/optimize/index.ts

@@ -0,0 +1,2 @@
+export * from './types';
+export { fewShot } from './few-shot';   // <— exact file name

src/lib/optimize/metric.ts

@@ -0,0 +1,19 @@
+import { MetricFn } from './types';
+
+/** Generic exact-match metric. */
+export function exactMatch(): MetricFn<any, any> {
+  return (pred, gold) => (JSON.stringify(pred) === JSON.stringify(gold) ? 1 : 0);
+}
+
+/** Simple F-1 for sets of strings. */
+export function f1(): MetricFn<string[] | string, string[] | string> {
+  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);
+  };
+}

src/lib/optimize/types.ts

@@ -0,0 +1,15 @@
+import { PromptTemplate } from '../prompts';
+
+/** Training example tuple. */
+export interface TrainExample<I = any, O = any> {
+  input: I;
+  output: O;
+}
+
+/** Metric: returns single scalar; higher is better. */
+export type MetricFn<P = any, G = any> = (prediction: P, gold: G) => number | Promise<number>;
+
+/** Every optimiser returns a tuned clone of the target. */
+export interface OptimizerSpec<T = PromptTemplate<any, any>> {
+  run(target: T): Promise<T>;
+}

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<any, PromptVariables>;
+
+  /**
+   * 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<string, string> }): void;
 }
+
+/** Helper namespace exported from ./optimize (e.g., fewShot) */
+type OptimizeHelpers = typeof import('./optimize');
+
+/** Callable optimise function merged with helpers */
+type OptimizeFn = (<TOut = any, TIn = any>(
+  target: import('./prompts/types').PromptTemplate<TOut, TIn>,
+  optimizer: import('./optimize/types').OptimizerSpec<import('./prompts/types').PromptTemplate<TOut, TIn>>
+) => Promise<import('./prompts/types').PromptTemplate<TOut, TIn>>) & OptimizeHelpers;

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<string, number> = {
+  '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;
+}

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);
+}