diff --git a/.github/pr-review-prompt.md b/.github/pr-review-prompt.md index dd6e50b..b6f18d1 100644 --- a/.github/pr-review-prompt.md +++ b/.github/pr-review-prompt.md @@ -84,7 +84,7 @@ function process(data: Data | null) { ### Naming Conventions -- **Classes and Types**: PascalCase (`Reranker`, `LlmClient`) +- **Classes and Types**: PascalCase (`Intent`, `LlmClient`) - **Functions and Variables**: camelCase (`buildMessages`, `candidateKey`) - **Constants**: UPPER_SNAKE_CASE (`BATCH_SIZE`, `TIMEOUT_MS`) - **Test files**: `*.unit.test.ts` or `*.int.test.ts` diff --git a/CLAUDE.md b/CLAUDE.md index ba7eccd..7670810 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -50,7 +50,7 @@ Intent is a TypeScript library that uses LLMs to rerank arbitrary items based on ### Core Reranking Flow -The `Reranker` class orchestrates the entire reranking process: +The `Intent` class orchestrates the entire ranking process: 1. **Preparation** (`prepareCandidates`): Normalizes items into a consistent shape with key/summary/index 2. **Batching** (`batchProcess` in batches.ts): Splits candidates into batches, merges tiny trailing batches to avoid inefficient LLM calls @@ -75,14 +75,16 @@ All config lives in config.ts using lib/config helpers: - Environment variables: `INTENT_MODEL`, `INTENT_TIMEOUT_MS`, `INTENT_RELEVANCY_THRESHOLD`, `INTENT_BATCH_SIZE`, `INTENT_TINY_BATCH_FRACTION` - Groq-specific: `GROQ_API_KEY`, `GROQ_DEFAULT_MODEL`, `GROQ_DEFAULT_TEMPERATURE` - Config is loaded automatically via `dotenv/config` import at top of config.ts -- Reranker constructor accepts overrides as third parameter +- Intent constructor accepts options as a single optional object + +**Config Naming Convention**: By design, config keys in `CONFIG` use `UPPER_SNAKE_CASE` to mirror their environment variable names (e.g., `CONFIG.INTENT.MODEL` matches `INTENT_MODEL`). The user-facing API uses `camelCase` (e.g., `options.model`). This intentional distinction keeps config keys aligned with environment variables while providing an ergonomic API. Internal code converts between these formats as needed. ### Key Design Patterns - **Stable fallbacks**: Any failure (LLM error, timeout, invalid response) returns items in original order - **Duplicate key handling**: Internal disambiguation using `"Key (idx)"` suffix - **Strict typing**: Uses TypeScript strict mode with `noUncheckedIndexedAccess`, `exactOptionalPropertyTypes` -- **Generic item support**: Reranker is `Reranker` with user-provided key/summary extractors +- **Generic item support**: Intent is `Intent` with user-provided key/summary extractors - **Zero/one-item fast path**: Avoids LLM calls when unnecessary ## Testing Strategy diff --git a/README.md b/README.md index fdd6bdf..c1ec92b 100644 --- a/README.md +++ b/README.md @@ -1,166 +1,277 @@ # Intent -LLM-based reranker for arbitrary items. Provide a query and a list of items, Intent asks an LLM to score each item’s relevance (0–10), filters by a configurable threshold, and returns the items ordered by score (stable on ties). +`intent` is an LLM-based reranker library that offers ranking, filtering, and choice all with explicit, inspectable reasoning. -Highlights +Unlike black-box models, `intent` generates an explanation alongside every score. This transparency allows for easier debugging and enables you to surface reasoning directly to users. -- Pluggable LLM client interface (Groq/OpenAI/etc.) -- Stable, safe behavior with fallbacks -- Strict JSON schema scoring, duplicate-key handling -- Fully typed TypeScript API with tests +By being LLM powered, it also offers flexibility and dynamic tunability to +your ranking logic without needing to come up with custom embedding models or +retrain existing ones. -How It Works +## Install -- Listwise LLM reranker: given a user query and a set of candidate items (each with a short key and optional summary), the LLM sees the query and all candidates together and assigns each a relevance score from 0–10. -- Intent-aware ranking: the prompt emphasizes the user’s intent, task framing, and constraints (not just surface similarity). Items that best satisfy the intent rise, even when lexical overlap is low. -- Threshold + stable ordering: scores are filtered by a configurable threshold and returned in descending order; ties preserve original input order. -- Retriever-agnostic: use with any first-stage retriever (vector, BM25, hybrid) or any arbitrary list of items. Summaries help the LLM reason efficiently within token limits. +```bash +npm install @with-logic/intent +``` -Why It Excels At User Intent +## Quickstart -- Interprets nuance: weighs the goal behind the query (task, specificity, constraints, entities, and outcome) instead of matching keywords alone. -- Cross-candidate reasoning: considers all candidates in one pass, comparing which best fulfills the intent relative to the rest. -- Robust to wording: prioritizes items that truly answer the need, even if phrased differently than the query. -- Configurable strictness: tune the relevancy threshold to be more selective for high-precision top results. +```ts +import { Intent } from "@with-logic/intent"; -Best Practices +const intent = new Intent({ relevancyThreshold: 1 }); -### Data Quality +const docs = [ + "Many network requests can fail intermittently due to transient issues.", + "To reduce flaky tests, add exponential backoff with jitter to HTTP retries.", + "Citrus fruits like oranges and lemons are high in vitamin C.", +]; -- **Keep summaries short and structured**: Include title, 1–2 key facts, entities, dates, and outcomes. Aim for consistent length across items so the LLM compares fairly. -- **Encode user intent explicitly**: Pass the user's goal, constraints, timeframe, and domain context in the query string you provide to the reranker. -- **Use helpful metadata**: Incorporate type, tags, author, and dates into the summary string to improve intent alignment. +const ranked = await intent.rank("exponential backoff retries", docs); +// => [doc1, doc2] (doc3 is filtered out via relevancy threshold) +``` -### Performance & Cost +Intent will use a default Groq client when `GROQ_API_KEY` is set. -- **Size the candidate set to fit context**: Start with 50–100 items with concise summaries. Tune `BATCH_SIZE` (or `INTENT_BATCH_SIZE`) to your model's token budget. -- **Favor determinism for stable ranking**: Run with temperature 0 and a fixed prompt template. Ties are already stable by input order. +## Core API -### Accuracy & Results +- `rank(query, candidates)` → rerank + threshold filter (score-based) +- `filter(query, candidates)` → keep only relevant items (boolean) +- `choice(query, candidates)` → choose exactly one best item -- **Tune selectivity**: Set `RELEVANCY_THRESHOLD` (or `INTENT_RELEVANCY_THRESHOLD`) to emphasize high-precision top results for RAG and QA. Values 0-10 are supported, with 0 including all results. -- **Optional score fusion**: For even stronger robustness, combine LLM scores with retriever scores (e.g., a weighted sum) when you have them. +All three support `{ explain: true }` to return explanations. -### Monitoring +## Example Use Cases -- **Monitor and iterate**: Log query, candidate count, raw/normalized scores, and token usage to fine-tune thresholds and batch sizes. +### 1) Ordering search results with `rank()` -Install +Use `rank()` when you want ranked, ordered results. -```bash -npm install intent +```ts +import { Intent } from "@with-logic/intent"; + +type Doc = { + id: string; + title: string; + body: string; + tags: string[]; +}; + +const intent = new Intent({ + key: (d) => d.id, + relevancyThreshold: 5, +}); + +const docs: Doc[] = [ + { id: "1", title: "Q2 expenses", body: "Travel, meals, ...", tags: ["finance"] }, + { id: "2", title: "OKR planning", body: "Goals for ...", tags: ["strategy"] }, + { id: "3", title: "Laptop purchases", body: "New laptops ...", tags: ["finance", "it"] }, +]; + +const results = await intent.rank("Find expense reports and anything about spend approvals", docs); ``` -Or with yarn: +### Include explanations -```bash -yarn add intent +```ts +const results = await intent.rank("expense reports", docs, { explain: true }); +// => [{ item: Doc, explanation: string }, ...] +``` + +## 2) Tool filtering with `filter()` + +Use `filter()` when you want to keep the subset of items in a collection that +are relevant to a query. + +```ts +import { Intent } from "@with-logic/intent"; + +type Tool = { + name: string; + description: string; +}; + +const intent = new Intent({ + key: (t) => t.name, + summary: (t) => t.description, +}); + +const tools: Tool[] = [ + { name: "search", description: "Search the web for up-to-date information" }, + { name: "sendEmail", description: "Send an email to a recipient" }, + { name: "runSQL", description: "Run a SQL query against the analytics DB" }, + { name: "createInvoice", description: "Create an invoice for a customer" }, +]; + +const task = "Find the customer's last invoice total and email it to them."; + +const relevantTools = await intent.filter(task, tools); +// [sendEmail, runSQL] ``` -Configuration +### Filter with explanations -The library reads `.env` automatically when imported. Create a `.env` file in your project root: +```ts +const relevantTools = await intent.filter(task, tools, { explain: true }); +// => [{ item: Tool, explanation: string }, ...] +``` + +## 3) Model routing with `choice()` + +Use `choice()` when you need exactly one selection from a set of items. + +```ts +import { Intent } from "@with-logic/intent"; + +type Model = { + id: string; + strengths: string; +}; + +const intent = new Intent({ + key: (m) => m.id, + summary: (m) => m.strengths, +}); + +const models: Model[] = [ + { + id: "gemini-3-pro", + strengths: "Hard reasoning, math, complex debugging. Slower but very strong.", + }, + { + id: "gpt-5.2", + strengths: "Best for code generation, refactors, feature implementation.", + }, + { + id: "haiku-4.5" + strengths: "Fast and cheap. Good for triage and simple edits.", + }, + { + id: "nano-banana-pro", + strengths: "Image generation and visual content.", + }, +]; + +const task = "Implement a feature to add retries with exponential backoff and tests."; + +const { item: selected, explanation } = await intent.choice(task, models, { explain: true }); +// selected.id => gpt-5.2 +``` + +## Configuration + +Intent reads `.env` automatically when imported. ```env -# Required only if not providing ctx.llm GROQ_API_KEY=your_groq_api_key_here -# Optional: customize defaults -INTENT_MODEL=openai/gpt-oss-120b -INTENT_TIMEOUT_MS=30000 +# Optional defaults +GROQ_DEFAULT_MODEL=openai/gpt-oss-20b +GROQ_DEFAULT_REASONING_EFFORT=medium +INTENT_TIMEOUT_MS=3000 +INTENT_MIN_SCORE=0 +INTENT_MAX_SCORE=10 INTENT_RELEVANCY_THRESHOLD=0 INTENT_BATCH_SIZE=20 -INTENT_TINY_BATCH_FRACTION=0.15 +INTENT_TINY_BATCH_FRACTION=0.2 ``` -**Configuration Options:** +### How configuration works -- **`GROQ_API_KEY`**: Your Groq API key. Required if not providing a custom `ctx.llm` client. -- **`INTENT_MODEL`**: LLM model to use for reranking (default: `openai/gpt-oss-20b`). Any Groq-supported model works here. -- **`INTENT_TIMEOUT_MS`**: Maximum time in milliseconds to wait for LLM responses (default: `3000`). Increase for larger batches or slower models. -- **`INTENT_RELEVANCY_THRESHOLD`**: Minimum relevance score (0-10) to include in results (default: `0`). Higher values = more selective filtering. -- **`INTENT_BATCH_SIZE`**: Number of candidates to score per LLM call (default: `20`). Tune based on your model's context window and candidate summary length. -- **`INTENT_TINY_BATCH_FRACTION`**: Threshold for merging small trailing batches (default: `0.2`). If the last batch is smaller than this fraction of `BATCH_SIZE`, it gets merged with the previous batch to avoid inefficient LLM calls. +- You can configure Intent via **environment variables** (shown above) or via the `Intent` constructor. +- Constructor options override environment defaults for that instance. +- Most tuning is a trade-off between **quality**, **latency**, and **cost**. -Development +### Provider + model -- Build: `npm run build` (requires TypeScript) -- Lint and format: `npm run lint:check` (check) or `npm run lint` (auto-fix) -- Tests with coverage: `npm test` (uses Vitest + v8 coverage) +#### `GROQ_API_KEY` -Tests +If you don't pass a custom `llm` client, Intent will create a default Groq client when this is set. -- Co-located alongside source for easy association. - - Unit: `src/**/*.unit.test.ts` (100% coverage enforced) - - Integration: `src/**/*.int.test.ts` (live Groq; requires `GROQ_API_KEY`) -- Scripts - - `npm run test:unit` — unit tests only - - `npm run test:int` — integration tests only (concurrent) - - `npm test` — all tests - - Optional: set `TEST_SCOPE=unit|int|all` to control scope +#### `GROQ_DEFAULT_MODEL` -Groq default +Sets the model name used by the built-in Groq client. -- Uses `groq-sdk` under the hood. If `GROQ_API_KEY` is set in the environment, you can omit `ctx.llm` and Intent will use a built‑in Groq adapter automatically. -- Otherwise, provide your own `llm` client via `ctx.llm`. +- Choose a stronger model when you care about nuanced ranking or long candidate summaries. +- Choose a smaller model when you want lower latency/cost and your candidates are simple. -Config +#### `GROQ_DEFAULT_REASONING_EFFORT` -- Reranker config can be supplied at construction or via environment variables (see Configuration section above for details). -- The library reads `.env` automatically when imported, so `INTENT_*` keys in your `.env` are honored. -- Constructor config overrides environment variables for fine-grained control per instance. +Controls how much reasoning the model should do (`low | medium | high`). -Usage +- `low`: fastest; best for obvious matches. +- `medium`: good default. +- `high`: better for subtle intent, but typically slower/more expensive. -```ts -import { Reranker } from "intent"; - -// Minimal LLM client (adapt your SDK to this shape) -const llm = { - async call(messages, schema, config, userId) { - // call your LLM here; must return `{ data: Record }` - // Example: { data: { "Travel Expenses": 8, "OKR Plan": 2 } } - return { data: {} }; - }, -}; +### Ranking behavior -type Item = { title: string; description: string }; -const reranker = new Reranker( - { llm, userId: "org-123" }, - { - key: (x) => x.title, - summary: (x) => x.description, - }, - { relevancyThreshold: 0, batchSize: 20 }, -); +#### `INTENT_RELEVANCY_THRESHOLD` -const ordered = await reranker.rerank("find expense reports", [ - { title: "Travel Expenses", description: "Q2 reimbursements" }, - { title: "OKR Plan", description: "Q3 planning" }, -]); +Controls how selective the output is. -// Or, if GROQ_API_KEY is set in your environment, you can omit `llm`: -const rerankerWithDefault = new Reranker( - { - /* no llm needed here if GROQ_API_KEY is set */ - }, - { key: (x) => x.title, summary: (x) => x.description }, -); -``` +- Higher threshold → fewer results (higher precision) +- Lower threshold → more results (higher recall) + +Important: threshold filtering is **strictly greater-than** (`score > threshold`). +So with the default score range `0..10`: + +- `relevancyThreshold=0` keeps scores `1..10` +- `relevancyThreshold=5` keeps scores `6..10` -API +#### `INTENT_MIN_SCORE` / `INTENT_MAX_SCORE` -- `new Reranker(ctx, extractors, config?)` - - `ctx.llm`: LLM client with `call(messages, schema, config, userId)` - - `ctx.userId?`: optional user identifier forwarded to provider - - `ctx.logger?`: optional logger with `.warn()` (and `.info/.error`) - - `extractors.key(item)`: required, short human-readable key - - `extractors.summary?(item)`: optional short description - - `config`: `{ model, timeoutMs, relevancyThreshold, batchSize, tinyBatchFraction }` -- `rerank(query, candidates, { userId? })` returns `T[]` +Controls the score range given to the LLM. -Notes +- Narrower ranges (e.g. `1..5`) can make scoring easier to calibrate. +- Wider ranges (e.g. `0..10`) give more resolution for ranking. -- Always returns a list; on any failure, it preserves the original order for the affected batch. -- Ties keep original order (stable sort by input index). -- Duplicate keys are internally disambiguated: `"Key (idx)"`. +Note: Since the is an LLM's judgement, as opposed to an objective measurement, +scores may not use the full range perfectly and you may seem similar biases +that you'd see with human raters. + +This also means that massive ranges (e.g. `0..1000`) may not yield more +precise results. + +If you change the range, ensure your `relevancyThreshold` stays within it. + +### Performance knobs + +#### `INTENT_TIMEOUT_MS` + +Hard timeout per LLM call. + +- Increase it when you have larger batches, longer summaries, or slower models. +- Decrease it when you prefer quick fallbacks over waiting. + +If we timeout, we never throw an error; instead, we return the original +results. + +#### `INTENT_BATCH_SIZE` + +How many candidates are evaluated per LLM call. + +- Larger batch size → fewer calls (often cheaper/faster), but higher token usage per call. +- Smaller batch size → more calls (often slower), but each call is smaller. + +#### `INTENT_TINY_BATCH_FRACTION` + +When the last batch is “too small”, Intent will merge it into the previous batch. + +- Increase to avoid tiny extra calls (better latency/cost). +- Decrease if you frequently run near context limits. + +### Programmatic configuration (constructor) + +Everything above can be set per instance: + +```ts +import { Intent } from "intent"; + +const intent = new Intent({ + timeoutMs: 10_000, + batchSize: 25, + relevancyThreshold: 3, + minScore: 0, + maxScore: 10, +}); +``` diff --git a/package.json b/package.json index d5b1cfa..2484920 100644 --- a/package.json +++ b/package.json @@ -1,5 +1,5 @@ { - "name": "intent", + "name": "@with-logic/intent", "version": "0.1.0", "private": false, "description": "Intent: a small, well-typed LLM-based reranker.", @@ -20,6 +20,9 @@ "url": "https://github.com/with-logic/intent/issues" }, "homepage": "https://github.com/with-logic/intent#readme", + "publishConfig": { + "access": "public" + }, "type": "module", "sideEffects": false, "main": "dist/index.js", diff --git a/src/batches.ts b/src/batches.ts index 36581df..9fe97fb 100644 --- a/src/batches.ts +++ b/src/batches.ts @@ -90,7 +90,12 @@ export async function batchProcess( logger?.warn?.("intent reranker batch failed, preserving original order", { error: (error as Error)?.message, }); - return onError ? onError(b, error) : (b as unknown as O[]); + + if (onError) { + return onError(b, error); + } + + return b as unknown as O[]; } }), ); diff --git a/src/config.ts b/src/config.ts index 15395a1..0563eda 100644 --- a/src/config.ts +++ b/src/config.ts @@ -1,5 +1,5 @@ import "dotenv/config"; -import { int, number, string } from "./lib/config"; +import { enumeration, int, number, string } from "./lib/config"; /** * Exported config object (no function call required) following API patterns. @@ -8,12 +8,18 @@ export const CONFIG = { GROQ: { API_KEY: string("GROQ_API_KEY", { default: "" }), DEFAULT_MODEL: string("GROQ_DEFAULT_MODEL", { default: "openai/gpt-oss-20b" }), - DEFAULT_TEMPERATURE: number("GROQ_DEFAULT_TEMPERATURE", { default: 0, min: 0, max: 1 }), + DEFAULT_REASONING_EFFORT: enumeration("GROQ_DEFAULT_REASONING_EFFORT", { + default: "medium", + values: ["low", "medium", "high"] as const, + }), + JSON_REPAIR_ATTEMPTS: int("GROQ_JSON_REPAIR_ATTEMPTS", { default: 3, min: 0 }), }, - RERANKER: { - MODEL: string("INTENT_MODEL", { default: "openai/gpt-oss-20b" }), + INTENT: { + PROVIDER: enumeration("INTENT_PROVIDER", { default: "GROQ", values: ["GROQ"] as const }), TIMEOUT_MS: int("INTENT_TIMEOUT_MS", { default: 3000, min: 1 }), - RELEVANCY_THRESHOLD: int("INTENT_RELEVANCY_THRESHOLD", { default: 0, min: 0, max: 10 }), + MIN_SCORE: int("INTENT_MIN_SCORE", { default: 0 }), + MAX_SCORE: int("INTENT_MAX_SCORE", { default: 10, min: 1 }), + RELEVANCY_THRESHOLD: int("INTENT_RELEVANCY_THRESHOLD", { default: 0 }), BATCH_SIZE: int("INTENT_BATCH_SIZE", { default: 20, min: 1 }), TINY_BATCH_FRACTION: number("INTENT_TINY_BATCH_FRACTION", { default: 0.2, min: 0, max: 1 }), }, diff --git a/src/extractors.ts b/src/extractors.ts new file mode 100644 index 0000000..9db1987 --- /dev/null +++ b/src/extractors.ts @@ -0,0 +1,91 @@ +/** + * Default extractor functions for Intent. + * + * These provide sensible defaults when users don't specify custom key/summary extractors. + */ + +/** + * Computes a simple 32-bit hash from a string using the djb2 algorithm. + * + * @param str - The string to hash + * @returns A 32-bit hash value + * @private + */ +function hash32(str: string): number { + let hash = 5381; + for (let i = 0; i < str.length; i++) { + const char = str.charCodeAt(i); + hash = ((hash << 5) + hash + char) | 0; // hash * 33 + char, keep in 32-bit range + } + return hash >>> 0; // Convert to unsigned 32-bit integer +} + +/** + * Converts any value to a pretty-printed JSON string with error handling. + * + * This is the canonical helper for all JSON stringification in the codebase. + * Pretty-prints with 2-space indentation for better LLM readability. + * Falls back to String(value) if JSON.stringify fails (e.g., circular references). + * Handles undefined explicitly since JSON.stringify(undefined) returns undefined. + * + * @param value - The value to stringify + * @returns Pretty-printed JSON string, or String() fallback + * @throws Never throws - gracefully falls back to String() on any error + */ +export function jsonStringify(value: unknown): string { + try { + const result = JSON.stringify(value, null, 2); + // JSON.stringify returns undefined for undefined values + if (result === undefined) { + return String(value); + } + return result; + } catch { + return String(value); + } +} + +/** + * Converts any value to a hash-based string key. + * + * Serializes the value to JSON, computes a 32-bit hash, and returns it as a string. + * This provides unique-enough keys for items without requiring explicit key extractors. + * + * @param value - The value to convert to a key + * @returns A string representation of the hash (e.g., "2847561") + * @throws Never throws - uses jsonStringify which handles all errors internally + */ +export function hashToString(value: T): string { + const json = jsonStringify(value); + const hashValue = hash32(json); + return String(hashValue); +} + +/** + * Default key extractor: generates hash-based string keys from items. + * + * Generic function that works with any item type T. The generic parameter + * enables type-safe usage without type casts when assigned to typed extractors. + * + * @param item - The item to extract a key from + * @returns A hash-based string key + * @throws Never throws - uses hashToString which handles all errors internally + */ +export function DEFAULT_KEY_EXTRACTOR(item: T): string { + return hashToString(item); +} + +/** + * Default summary extractor: converts items to pretty-printed JSON strings. + * + * Uses 2-space indentation for better LLM readability. + * Generic function that works with any item type T. The generic parameter + * enables type-safe usage without type casts when assigned to typed extractors. + * + * @param item - The item to extract a summary from + * @returns A pretty-printed JSON string representation of the item + * @throws Never throws - uses jsonStringify which handles all errors internally + */ +export function DEFAULT_SUMMARY_EXTRACTOR(item: T): string { + return jsonStringify(item); +} diff --git a/src/extractors.unit.test.ts b/src/extractors.unit.test.ts new file mode 100644 index 0000000..c405b85 --- /dev/null +++ b/src/extractors.unit.test.ts @@ -0,0 +1,184 @@ +import { describe, expect, test } from "vitest"; + +import { + DEFAULT_KEY_EXTRACTOR, + DEFAULT_SUMMARY_EXTRACTOR, + hashToString, + jsonStringify, +} from "./extractors"; + +describe("extractors", () => { + describe("jsonStringify", () => { + test("pretty-prints simple values", () => { + expect(jsonStringify("hello")).toBe('"hello"'); + expect(jsonStringify(123)).toBe("123"); + expect(jsonStringify(true)).toBe("true"); + expect(jsonStringify(null)).toBe("null"); + }); + + test("pretty-prints objects and arrays with 2-space indentation", () => { + expect(jsonStringify({ a: 1, b: 2 })).toBe('{\n "a": 1,\n "b": 2\n}'); + expect(jsonStringify([1, 2, 3])).toBe("[\n 1,\n 2,\n 3\n]"); + }); + + test("falls back to String() for circular references", () => { + const circular: any = { a: 1 }; + circular.self = circular; + const result = jsonStringify(circular); + expect(result).toBe("[object Object]"); + }); + + test("handles deeply nested objects consistently", () => { + const complex = { + level1: { + level2: { + level3: { + value: [1, 2, { nested: "data" }], + }, + }, + }, + }; + const json1 = jsonStringify(complex); + const json2 = jsonStringify(complex); + expect(json1).toBe(json2); + }); + + test("produces stable output for objects with same keys in different order", () => { + const obj1 = { a: 1, b: 2, c: 3 }; + const obj2 = { c: 3, b: 2, a: 1 }; + const json1 = jsonStringify(obj1); + const json2 = jsonStringify(obj2); + // Note: JSON.stringify preserves insertion order, so these will differ + // This test documents the behavior - objects with different key order produce different JSON + expect(json1).not.toBe(json2); + }); + }); + + describe("hashToString", () => { + test("returns consistent hash for same input", () => { + const input = { key: "value", num: 42 }; + const hash1 = hashToString(input); + const hash2 = hashToString(input); + expect(hash1).toBe(hash2); + }); + + test("returns different hashes for different inputs", () => { + const hash1 = hashToString({ a: 1 }); + const hash2 = hashToString({ a: 2 }); + expect(hash1).not.toBe(hash2); + }); + + test("returns string representation of hash", () => { + const hash = hashToString("test"); + expect(typeof hash).toBe("string"); + expect(hash).toMatch(/^\d+$/); // Should be all digits + }); + + test("handles edge cases", () => { + expect(hashToString("")).toBeTruthy(); + expect(hashToString(null)).toBeTruthy(); + expect(hashToString(undefined)).toBeTruthy(); + }); + + test("produces consistent hashes for deeply nested objects", () => { + const complex = { + level1: { + level2: { + level3: { + value: [1, 2, { nested: "data" }], + metadata: { timestamp: 123456, tags: ["a", "b", "c"] }, + }, + }, + }, + }; + const hash1 = hashToString(complex); + const hash2 = hashToString(complex); + expect(hash1).toBe(hash2); + expect(typeof hash1).toBe("string"); + expect(hash1).toMatch(/^\d+$/); + }); + + test("produces different hashes for objects with different key orders", () => { + const obj1 = { a: 1, b: 2, c: 3 }; + const obj2 = { c: 3, b: 2, a: 1 }; + const hash1 = hashToString(obj1); + const hash2 = hashToString(obj2); + // Hashes differ because JSON.stringify preserves insertion order + expect(hash1).not.toBe(hash2); + }); + + test("produces different hashes for nested arrays with different element orders", () => { + const arr1 = { items: [1, 2, 3] }; + const arr2 = { items: [3, 2, 1] }; + const hash1 = hashToString(arr1); + const hash2 = hashToString(arr2); + expect(hash1).not.toBe(hash2); + }); + }); + + describe("DEFAULT_KEY_EXTRACTOR", () => { + test("extracts hash-based key from items", () => { + const item = { id: "123", name: "Test" }; + const key = DEFAULT_KEY_EXTRACTOR(item); + expect(typeof key).toBe("string"); + expect(key).toMatch(/^\d+$/); + }); + + test("returns consistent keys for same input", () => { + const item = { value: 42 }; + const key1 = DEFAULT_KEY_EXTRACTOR(item); + const key2 = DEFAULT_KEY_EXTRACTOR(item); + expect(key1).toBe(key2); + }); + + test("returns consistent keys for deeply nested objects", () => { + const complex = { + user: { + profile: { + name: "Test User", + metadata: { created: 123456, tags: ["admin", "active"] }, + }, + }, + }; + const key1 = DEFAULT_KEY_EXTRACTOR(complex); + const key2 = DEFAULT_KEY_EXTRACTOR(complex); + expect(key1).toBe(key2); + }); + }); + + describe("DEFAULT_SUMMARY_EXTRACTOR", () => { + test("extracts pretty-printed JSON string from items", () => { + const item = { id: "123", content: "Hello" }; + const summary = DEFAULT_SUMMARY_EXTRACTOR(item); + expect(summary).toBe('{\n "id": "123",\n "content": "Hello"\n}'); + }); + + test("handles primitives", () => { + expect(DEFAULT_SUMMARY_EXTRACTOR("test")).toBe('"test"'); + expect(DEFAULT_SUMMARY_EXTRACTOR(42)).toBe("42"); + expect(DEFAULT_SUMMARY_EXTRACTOR(true)).toBe("true"); + }); + + test("extracts consistent summaries for deeply nested objects", () => { + const complex = { + document: { + metadata: { + author: "Test Author", + tags: ["important", "review"], + }, + content: { + sections: [ + { title: "Introduction", text: "Lorem ipsum" }, + { title: "Body", text: "Main content" }, + ], + }, + }, + }; + const summary1 = DEFAULT_SUMMARY_EXTRACTOR(complex); + const summary2 = DEFAULT_SUMMARY_EXTRACTOR(complex); + expect(summary1).toBe(summary2); + expect(summary1).toContain("Test Author"); + expect(summary1).toContain("Introduction"); + }); + }); +}); diff --git a/src/index.ts b/src/index.ts index 2c7f114..7c7f517 100644 --- a/src/index.ts +++ b/src/index.ts @@ -1,12 +1,15 @@ -export { Reranker } from "./reranker"; +export { Intent } from "./intent"; export type { ChatMessage, LlmClient, LlmCallConfig, LoggerLike, - RerankerCandidate, - RerankerExtractors, + IntentCandidate, + IntentExtractors, + IntentOptions, + IntentConfig, IntentContext, } from "./types"; export { CONFIG } from "./config"; export { createDefaultGroqClient } from "./providers/groq"; +export { DEFAULT_KEY_EXTRACTOR, DEFAULT_SUMMARY_EXTRACTOR } from "./extractors"; diff --git a/src/intent.groq-default.unit.test.ts b/src/intent.groq-default.unit.test.ts new file mode 100644 index 0000000..72c0c8c --- /dev/null +++ b/src/intent.groq-default.unit.test.ts @@ -0,0 +1,37 @@ +import { describe, expect, test, vi } from "vitest"; + +import { Intent } from "./intent"; +import { createDefaultGroqClient } from "./providers/groq"; + +describe("Intent (default Groq) ", () => { + test("uses groq client (via DI)", async () => { + const callMock = vi.fn(async (_req: any) => ({ + choices: [ + { + message: { + role: "assistant", + content: JSON.stringify({ + A: { explanation: "a", score: 10 }, + B: { explanation: "b", score: 0 }, + }), + }, + }, + ], + })); + const llm = createDefaultGroqClient("test-key", { + makeSdk: () => ({ chat: { completions: { create: callMock } } }), + }); + + const intent = new Intent<{ key: string; summary: string }>({ + llm, + key: (x) => x.key, + summary: (x) => x.summary, + }); + const out = await intent.rank("q", [ + { key: "A", summary: "" }, + { key: "B", summary: "" }, + ]); + expect(out.map((c) => c.key)).toEqual(["A"]); + expect(callMock.mock.calls.length).toBe(1); + }); +}); diff --git a/src/intent.int.test.ts b/src/intent.int.test.ts new file mode 100644 index 0000000..e55b628 --- /dev/null +++ b/src/intent.int.test.ts @@ -0,0 +1,582 @@ +import { describe, expect, test } from "vitest"; + +import { Intent } from "./intent"; + +describe("intent integration", () => { + const defaultTimeoutMs = 10000; + const scoreRange = { minScore: 0, maxScore: 10 }; + + test.concurrent( + "README quickstart: ranks simple strings and filters out unrelated items", + async () => { + const intent = new Intent({ relevancyThreshold: 1, timeoutMs: defaultTimeoutMs }); + + const docs = [ + "Many network requests can fail intermittently due to transient issues.", + "To reduce flaky tests, add exponential backoff with jitter to HTTP retries.", + "Citrus fruits like oranges and lemons are high in vitamin C.", + ]; + + const ranked = await intent.rank("exponential backoff retries", docs); + + expect(ranked).toEqual([ + "To reduce flaky tests, add exponential backoff with jitter to HTTP retries.", + "Many network requests can fail intermittently due to transient issues.", + ]); + }, + 30000, + ); + + test.concurrent( + "ranks the most relevant candidate first (simple object candidates)", + async () => { + const intent = new Intent<{ key: string; summary: string }>({ + key: (x) => x.key, + summary: (x) => x.summary, + timeoutMs: defaultTimeoutMs, + relevancyThreshold: 0, + }); + + const out = await intent.rank("Help me sort a JavaScript array", [ + { key: "Saturns Moons", summary: "The chemical composition of Saturn's moons" }, + { key: "JS Arrays", summary: "Guide to sorting arrays in JavaScript" }, + { key: "Eiffel Tower", summary: "Directions to the tower" }, + ]); + + expect(out.map((x) => x.key)).toEqual(["JS Arrays"]); + }, + 30000, + ); + + test.concurrent( + "supports defaults: new Intent() ranks plain strings", + async () => { + const intent = new Intent(); + const items = ["apple", "banana", "orange", "grape"]; + + const ranked = await intent.rank("citrus fruits", items); + expect(ranked).toEqual(["orange"]); + }, + 30000, + ); + + test.concurrent( + "supports candidates with summary only (no key extractor)", + async () => { + const intent = new Intent<{ summary: string }>({ + summary: (x) => x.summary, + timeoutMs: defaultTimeoutMs, + relevancyThreshold: 0, + }); + + const out = await intent.rank("Help me sort a JavaScript array", [ + { summary: "The chemical composition of Saturn's moons" }, + { summary: "Guide to sorting arrays in JavaScript" }, + { summary: "Directions to the tower" }, + ]); + + expect(out).toEqual([{ summary: "Guide to sorting arrays in JavaScript" }]); + }, + 30000, + ); + + test.concurrent( + "supports candidates with key only (no summary extractor)", + async () => { + const intent = new Intent<{ key: string }>({ + key: (x) => x.key, + timeoutMs: defaultTimeoutMs, + relevancyThreshold: 0, + }); + + const out = await intent.rank("Help me sort a JavaScript array", [ + { key: "Saturns Moons" }, + { key: "JS Arrays" }, + { key: "Eiffel Tower" }, + ]); + + expect(out).toEqual([{ key: "JS Arrays" }]); + }, + 30000, + ); + + test.concurrent( + "returns [] when all candidates are unrelated and threshold > minScore", + async () => { + const intent = new Intent<{ key: string; summary: string }>({ + key: (x) => x.key, + summary: (x) => x.summary, + timeoutMs: defaultTimeoutMs, + relevancyThreshold: scoreRange.minScore + 1, + ...scoreRange, + }); + + const out = await intent.rank("JavaScript array sorting", [ + { key: "Banana Bread", summary: "How to bake banana bread" }, + { key: "Eiffel Tower", summary: "Timeline of the Eiffel Tower construction" }, + { key: "Roman Emperors", summary: "A list of Roman emperors" }, + ]); + + expect(out).toEqual([]); + }, + 30000, + ); + + test.concurrent( + "preserves input order when all candidates are returned", + async () => { + const intent = new Intent<{ key: string; summary: string }>({ + key: (x) => x.key, + summary: (x) => x.summary, + timeoutMs: defaultTimeoutMs, + // All of these candidates should be obviously relevant, so all should pass a 0 threshold. + relevancyThreshold: 0, + }); + + const input = [ + { key: "Guide A", summary: "How to sort JavaScript arrays with Array.prototype.sort" }, + { key: "Guide B", summary: "Sorting arrays in JavaScript using a comparator function" }, + { key: "Guide C", summary: "JavaScript array sorting examples and best practices" }, + ]; + + const out = await intent.filter("JavaScript array sorting", input); + + expect(out).toEqual(input); + }, + 30000, + ); + + test.concurrent( + "explain=true returns { item, explanation } and still filters by threshold", + async () => { + const intent = new Intent<{ key: string; summary: string }>({ + key: (x) => x.key, + summary: (x) => x.summary, + timeoutMs: defaultTimeoutMs, + relevancyThreshold: 1, + }); + + const out = await intent.rank( + "Help me sort a JavaScript array", + [ + { key: "JS Arrays", summary: "Guide to sorting arrays in JavaScript" }, + { key: "Banana Bread", summary: "How to bake banana bread" }, + ], + { explain: true }, + ); + + expect(out.map((x) => x.item.key)).toEqual(["JS Arrays"]); + expect(typeof out[0]?.explanation).toBe("string"); + expect(out[0]?.explanation.length).toBeGreaterThan(0); + }, + 30000, + ); + + test.concurrent( + "supports custom extractors (nested objects)", + async () => { + type Doc = { id: string; meta: { title: string }; body: string }; + const intent = new Intent({ + key: (d) => d.meta.title, + summary: (d) => d.body, + timeoutMs: defaultTimeoutMs, + relevancyThreshold: 1, + }); + + const out = await intent.rank("find expense reports", [ + { id: "1", meta: { title: "Expense report Q3" }, body: "Travel and meals reimbursement" }, + { id: "2", meta: { title: "Vacation photos" }, body: "Beach, family, sunsets" }, + ]); + + expect(out.length).toBe(1); + expect(out[0]?.meta.title).toBe("Expense report Q3"); + }, + 30000, + ); + + test.concurrent( + "handles unicode keys and summaries", + async () => { + const intent = new Intent<{ key: string; summary: string }>({ + key: (x) => x.key, + summary: (x) => x.summary, + timeoutMs: defaultTimeoutMs, + relevancyThreshold: 0, + }); + + const out = await intent.rank("JavaScript array sorting", [ + { key: "Café", summary: "Guía para ordenar arrays en JavaScript" }, + { key: "東京", summary: "A travel guide to Tokyo neighborhoods" }, + { key: "Banana Bread", summary: "How to bake banana bread" }, + ]); + + expect(out.map((x) => x.key)).toEqual(["Café"]); + }, + 30000, + ); + + test.concurrent( + "handles keys with punctuation, whitespace, and newlines", + async () => { + const intent = new Intent<{ key: string; summary: string }>({ + key: (x) => x.key, + summary: (x) => x.summary, + timeoutMs: defaultTimeoutMs, + relevancyThreshold: 0, + }); + + const out = await intent.rank("JavaScript array sorting", [ + { + key: "Array.sort() / comparator", + summary: "Using Array.prototype.sort with a compare function in JS", + }, + { + key: "Stable sort: ties, order", + summary: "How stable sorting behaves; preserving order for equal elements", + }, + { + key: "Key with spaces\n(and newline)", + summary: "Examples of sorting arrays in JavaScript", + }, + { + key: "Unrelated", + summary: "How to bake banana bread", + }, + ]); + + expect(out.length).toBe(3); + expect(out[0]?.key).toBe("Array.sort() / comparator"); + expect(out.some((x) => x.key === "Stable sort: ties, order")).toBe(true); + expect(out.some((x) => x.key === "Key with spaces\n(and newline)")).toBe(true); + }, + 30000, + ); + + test.concurrent( + "disambiguates duplicate keys", + async () => { + const intent = new Intent<{ key: string; summary: string }>({ + key: (x) => x.key, + summary: (x) => x.summary, + timeoutMs: defaultTimeoutMs, + relevancyThreshold: 0, + }); + + const out = await intent.rank("JavaScript array sorting", [ + { key: "Same", summary: "Sorting arrays in JavaScript" }, + { key: "Same", summary: "Another note about Array.prototype.sort" }, + { key: "Other", summary: "Banana bread recipe" }, + ]); + + expect(out.map((x) => x.key)).toEqual(["Same", "Same"]); + }, + 30000, + ); + + test.concurrent( + "supports non-default score ranges (minScore/maxScore)", + async () => { + const intent = new Intent<{ key: string; summary: string }>({ + key: (x) => x.key, + summary: (x) => x.summary, + timeoutMs: defaultTimeoutMs, + minScore: 1, + maxScore: 5, + relevancyThreshold: 2, + }); + + const out = await intent.rank("JavaScript array sorting", [ + { key: "JS Arrays", summary: "Guide to sorting arrays in JavaScript" }, + { key: "Banana Bread", summary: "How to bake banana bread" }, + ]); + + expect(out.length).toBeGreaterThanOrEqual(1); + expect(out[0]?.key).toBe("JS Arrays"); + }, + 30000, + ); + + test.concurrent( + "handles batching across multiple LLM calls", + async () => { + const intent = new Intent<{ key: string; summary: string }>({ + key: (x) => x.key, + summary: (x) => x.summary, + timeoutMs: defaultTimeoutMs, + batchSize: 3, + tinyBatchFraction: 0.2, + relevancyThreshold: 0, + }); + + const candidates = [ + { key: "Array.sort", summary: "JavaScript Array.prototype.sort reference" }, + { key: "Comparator", summary: "How to write a comparator function in JS" }, + { key: "Stable sort", summary: "Discussion of stable sorting behavior" }, + { key: "Quickstart", summary: "Basic JS arrays tutorial" }, + { key: "Banana bread", summary: "Recipe for banana bread" }, + { key: "Eiffel Tower", summary: "History of the Eiffel Tower" }, + { key: "Moon Bases", summary: "Guide to living on the moon" }, + ]; + + const out = await intent.rank("JavaScript array sorting", candidates); + expect(out.map((x) => x.key)).toEqual([ + "Array.sort", + "Comparator", + "Stable sort", + "Quickstart", + ]); + }, + 60000, + ); + + test.concurrent( + "handles larger candidate sets (stress)", + async () => { + const intent = new Intent<{ key: string; summary: string }>({ + key: (x) => x.key, + summary: (x) => x.summary, + timeoutMs: defaultTimeoutMs, + batchSize: 5, + tinyBatchFraction: 0.2, + relevancyThreshold: 0, + }); + + const candidates = Array.from({ length: 30 }).map((_, i) => ({ + key: `Doc ${i + 1}`, + summary: + i % 6 === 0 + ? "JavaScript array sorting with comparator and examples" + : "Unrelated notes about cooking, travel, or history", + })); + + const out = await intent.rank("JavaScript array sorting", candidates); + expect(out.length).toBe(5); + expect(out.map((x) => x.key)).toEqual(["Doc 1", "Doc 7", "Doc 13", "Doc 19", "Doc 25"]); + }, + 120000, + ); + + test.concurrent( + "explain=true stays aligned with items across batches", + async () => { + const intent = new Intent<{ key: string; summary: string }>({ + key: (x) => x.key, + summary: (x) => x.summary, + timeoutMs: defaultTimeoutMs, + batchSize: 4, + tinyBatchFraction: 0.2, + relevancyThreshold: 0, + }); + + const input = Array.from({ length: 12 }).map((_, i) => ({ + key: `Item ${i + 1}`, + summary: i % 3 === 0 ? "JavaScript array sorting examples" : "Completely unrelated topic", + })); + + const out = await intent.rank("JavaScript array sorting", input, { explain: true }); + + expect(out.length).toBe(4); + expect(out.map((r) => r.item.key)).toEqual(["Item 1", "Item 4", "Item 7", "Item 10"]); + for (const r of out) { + expect(typeof r.explanation).toBe("string"); + expect(r.explanation.length).toBeGreaterThan(0); + } + }, + 120000, + ); + + test.concurrent( + "accepts per-call userId override", + async () => { + const intent = new Intent<{ key: string; summary: string }>({ + key: (x) => x.key, + summary: (x) => x.summary, + timeoutMs: defaultTimeoutMs, + relevancyThreshold: 0, + }); + + const out = await intent.rank( + "JavaScript array sorting", + [ + { key: "JS Arrays", summary: "Guide to sorting arrays in JavaScript" }, + { key: "Banana Bread", summary: "How to bake banana bread" }, + ], + { userId: "integration-test-user" }, + ); + + expect(out.map((x) => x.key)).toEqual(["JS Arrays"]); + }, + 30000, + ); + + test.concurrent( + "filter() returns only relevant items and preserves input order", + async () => { + const intent = new Intent<{ key: string; summary: string }>({ + key: (x) => x.key, + summary: (x) => x.summary, + timeoutMs: defaultTimeoutMs, + }); + + const input = [ + { key: "A", summary: "How to bake banana bread" }, + { key: "B", summary: "Guide to sorting arrays in JavaScript" }, + { key: "C", summary: "JavaScript Array.prototype.sort examples" }, + { key: "D", summary: "History of the Eiffel Tower" }, + ]; + + const out = await intent.filter("JavaScript array sorting", input); + expect(out.map((x) => x.key)).toEqual(["B", "C"]); + }, + 30000, + ); + + test.concurrent( + "README filter example: filters tools down to the relevant subset", + async () => { + type Tool = { + name: string; + description: string; + }; + + const intent = new Intent({ + key: (t) => t.name, + summary: (t) => t.description, + timeoutMs: defaultTimeoutMs, + }); + + const tools: Tool[] = [ + { name: "search", description: "Search the web for up-to-date information" }, + { name: "sendEmail", description: "Send an email to a recipient" }, + { name: "runSQL", description: "Run a SQL query against the analytics DB" }, + { name: "createInvoice", description: "Create an invoice for a customer" }, + ]; + + const task = "Find the customer's last invoice total and email it to them."; + + const relevantTools = await intent.filter(task, tools); + + // Filter is order-preserving, so we only assert membership. + const names = relevantTools.map((t) => t.name); + expect(names).toEqual(expect.arrayContaining(["sendEmail", "runSQL"])); + expect(names).not.toEqual(expect.arrayContaining(["createInvoice"])); + }, + 30000, + ); + + test.concurrent( + "filter({ explain: true }) returns aligned explanations", + async () => { + const intent = new Intent<{ key: string; summary: string }>({ + key: (x) => x.key, + summary: (x) => x.summary, + timeoutMs: defaultTimeoutMs, + }); + + const input = [ + { key: "JS Arrays", summary: "Guide to sorting arrays in JavaScript" }, + { key: "Banana Bread", summary: "How to bake banana bread" }, + ]; + + const out = await intent.filter("JavaScript array sorting", input, { explain: true }); + expect(out.length).toBe(1); + expect(out[0]?.item.key).toBe("JS Arrays"); + expect(typeof out[0]?.explanation).toBe("string"); + expect(out[0]?.explanation.length).toBeGreaterThan(0); + }, + 30000, + ); + + test.concurrent( + "choice() returns exactly one item from the input", + async () => { + const intent = new Intent<{ key: string; summary: string }>({ + key: (x) => x.key, + summary: (x) => x.summary, + timeoutMs: defaultTimeoutMs, + batchSize: 2, + tinyBatchFraction: 0, + }); + + const input = [ + { key: "Banana Bread", summary: "How to bake banana bread" }, + { key: "JS Arrays", summary: "Guide to sorting arrays in JavaScript" }, + { key: "Eiffel Tower", summary: "History of the Eiffel Tower" }, + { key: "Comparator", summary: "How to write a comparator function in JS" }, + ]; + + const winner = await intent.choice("JavaScript array sorting", input); + expect(winner.key).toBe("JS Arrays"); + }, + 60000, + ); + + test.concurrent( + "choice({ explain: true }) returns { item, explanation }", + async () => { + const intent = new Intent<{ key: string; summary: string }>({ + key: (x) => x.key, + summary: (x) => x.summary, + timeoutMs: defaultTimeoutMs, + batchSize: 2, + tinyBatchFraction: 0, + }); + + const input = [ + { key: "Banana Bread", summary: "How to bake banana bread" }, + { key: "JS Arrays", summary: "Guide to sorting arrays in JavaScript" }, + { key: "Comparator", summary: "How to write a comparator function in JS" }, + ]; + + const res = await intent.choice("JavaScript array sorting", input, { explain: true }); + expect(res.item.key).toBe("JS Arrays"); + expect(typeof res.explanation).toBe("string"); + expect(res.explanation.length).toBeGreaterThan(0); + }, + 60000, + ); + + test.concurrent( + "README choice example: routes to a best-fit model and explains why", + async () => { + type Model = { + id: string; + strengths: string; + }; + + const intent = new Intent({ + key: (m) => m.id, + summary: (m) => m.strengths, + timeoutMs: defaultTimeoutMs, + }); + + const models: Model[] = [ + { + id: "gemini-3-pro", + strengths: "Hard reasoning, math, complex debugging. Slower but very strong.", + }, + { + id: "gpt-5.2", + strengths: "Best for code generation, refactors, feature implementation.", + }, + { + id: "haiku-4.5", + strengths: "Fast and cheap. Good for triage and simple edits.", + }, + { + id: "nano-banana-pro", + strengths: "Image generation and visual content.", + }, + ]; + + const task = "Implement a feature to add retries with exponential backoff and tests."; + + const { item: selected, explanation } = await intent.choice(task, models, { explain: true }); + + expect(selected.id).toBe("gpt-5.2"); + expect(typeof explanation).toBe("string"); + expect(explanation.length).toBeGreaterThan(0); + }, + 60000, + ); +}); diff --git a/src/intent.ts b/src/intent.ts new file mode 100644 index 0000000..35b2151 --- /dev/null +++ b/src/intent.ts @@ -0,0 +1,823 @@ +import { batchProcess } from "./batches"; +import { CONFIG } from "./config"; +import { DEFAULT_KEY_EXTRACTOR, DEFAULT_SUMMARY_EXTRACTOR } from "./extractors"; +import { clamp } from "./lib/number"; +import { selectLlmClient } from "./llm_client"; +import { buildChoiceMessages, buildFilterMessages, buildMessages } from "./messages"; +import { buildChoiceSchema, buildFilterSchema, buildRelevancySchema } from "./schema"; + +import type { + ChatMessage, + JSONObject, + IntentOptions, + IntentConfig, + LlmClient, + LlmCallConfig, + IntentContext, + IntentExtractors, +} from "./types"; + +/** + * LLM-based reranker for arbitrary items. + * + * Uses a listwise LLM approach to score candidates within a configurable range based on relevance to a query, + * then filters by threshold and returns results sorted by score with stable ordering. + * + * @template T - The type of items to rerank (defaults to any) + * + * @example + * ```typescript + * // Simplest: uses defaults and GROQ_API_KEY from environment + * const intent = new Intent(); + * const ranked = await intent.rank("find expense reports", items); + * + * // With custom extractors + * type Document = { id: string; title: string; content: string }; + * const intent = new Intent({ + * key: doc => doc.title, + * summary: doc => doc.content.slice(0, 200), + * relevancyThreshold: 5, + * batchSize: 20 + * }); + * + * // With custom LLM client + * const intent = new Intent({ + * llm: myClient, + * userId: "user-123", + * key: doc => doc.title, + * summary: doc => doc.content.slice(0, 200) + * }); + * ``` + */ +export class Intent { + private readonly cfg: Required; + private readonly llm: LlmClient; + private readonly ctx: IntentContext; + private readonly extractors: Required>; + private readonly env: typeof CONFIG; + + /** + * Resolve the model name to use for this Intent instance. + * + * Intent is provider-driven. Today only GROQ is supported; when using GROQ + * we always take the model from GROQ's config defaults. + * + * @returns Provider-specific model name + * @private + */ + private resolveModel(): string { + return this.env.GROQ.DEFAULT_MODEL; + } + + /** + * Builds the context object from options. + * + * Constructs an IntentContext with only defined properties to satisfy + * TypeScript's exactOptionalPropertyTypes requirement. + * + * @param options - The options object containing llm, logger, and userId + * @returns IntentContext with only defined properties + * @private + */ + private buildContext(options: IntentOptions): IntentContext { + return { + ...(options.llm !== undefined && { llm: options.llm }), + ...(options.logger !== undefined && { logger: options.logger }), + ...(options.userId !== undefined && { userId: options.userId }), + }; + } + + /** + * Builds the extractors object from options. + * + * Uses provided extractors or falls back to generic defaults that work + * for any type T via JSON stringification and hashing. + * + * @param options - The options object containing key and summary extractors + * @returns Required extractors with defaults applied + * @private + */ + private buildExtractors(options: IntentOptions): Required> { + return { + key: options.key ?? DEFAULT_KEY_EXTRACTOR, + summary: options.summary ?? DEFAULT_SUMMARY_EXTRACTOR, + }; + } + + /** + * Builds the configuration object from options. + * + * Merges user-provided options with environment-based CONFIG defaults. + * + * @param options - The options object containing config overrides + * @returns Required config with all values populated + * @private + */ + private buildConfig(options: IntentOptions): Required { + return { + provider: options.provider ?? this.env.INTENT.PROVIDER, + timeoutMs: options.timeoutMs ?? this.env.INTENT.TIMEOUT_MS, + relevancyThreshold: options.relevancyThreshold ?? this.env.INTENT.RELEVANCY_THRESHOLD, + batchSize: options.batchSize ?? this.env.INTENT.BATCH_SIZE, + tinyBatchFraction: options.tinyBatchFraction ?? this.env.INTENT.TINY_BATCH_FRACTION, + minScore: options.minScore ?? this.env.INTENT.MIN_SCORE, + maxScore: options.maxScore ?? this.env.INTENT.MAX_SCORE, + }; + } + + /** + * Validates the configuration values. + * + * Ensures the configured score range is valid and the relevancyThreshold is in range. + * + * @throws {Error} If maxScore is below minScore + * @throws {Error} If relevancyThreshold is not within [minScore, maxScore] + * @private + */ + private validateConfig(): void { + if (this.cfg.maxScore < this.cfg.minScore) { + throw new Error( + `intent: maxScore must be >= minScore, got minScore=${this.cfg.minScore} maxScore=${this.cfg.maxScore}`, + ); + } + + if ( + this.cfg.relevancyThreshold < this.cfg.minScore || + this.cfg.relevancyThreshold > this.cfg.maxScore + ) { + throw new Error( + `intent: relevancyThreshold must be between ${this.cfg.minScore} and ${this.cfg.maxScore}, got ${this.cfg.relevancyThreshold}`, + ); + } + } + + /** + * Selects and validates the LLM client. + * + * Uses the provided client from context or attempts to create a default + * Groq client if GROQ_API_KEY is available. + * + * @returns The selected LLM client + * @throws {Error} If no LLM client is provided and GROQ_API_KEY is not set + * @private + */ + private selectAndValidateLlmClient(): LlmClient { + const selectedClient = selectLlmClient(this.ctx, this.env); + if (!selectedClient) { + throw new Error( + "intent: No LLM client provided and GROQ_API_KEY not set. Provide options.llm or set GROQ_API_KEY.", + ); + } + return selectedClient; + } + + /** + * Creates a new Intent instance. + * + * All options are optional with sensible defaults: + * - llm: Auto-detected from GROQ_API_KEY environment variable if available + * - key: Hash-based string from JSON representation of items + * - summary: Pretty-printed JSON of items (2-space indentation for LLM readability) + * - Config values: From INTENT_* environment variables or built-in defaults + * + * @param options - Optional configuration object + * @param options.llm - Optional LLM client. If omitted, uses Groq client when GROQ_API_KEY is set + * @param options.logger - Optional logger for warnings and errors + * @param options.userId - Optional user identifier for LLM provider abuse monitoring + * @param options.key - Optional function extracting a short human-readable key from items + * @param options.summary - Optional function extracting a short description for LLM reasoning + * @param options.provider - Optional provider override (default: INTENT_PROVIDER or "GROQ") + * @param options.timeoutMs - Optional timeout in milliseconds (default: INTENT_TIMEOUT_MS or 3000) + * @param options.relevancyThreshold - Optional minimum score to include results (default: INTENT_RELEVANCY_THRESHOLD) + * @param options.minScore - Optional minimum score value (default: INTENT_MIN_SCORE or 0) + * @param options.maxScore - Optional maximum score value (default: INTENT_MAX_SCORE or 10) + * @param options.batchSize - Optional number of candidates per LLM call (default: INTENT_BATCH_SIZE or 20) + * @param options.tinyBatchFraction - Optional threshold for merging small batches (default: INTENT_TINY_BATCH_FRACTION or 0.2) + * @throws {Error} If no LLM client is provided and GROQ_API_KEY is not set + * @throws {Error} If maxScore is below minScore + * @throws {Error} If relevancyThreshold is not within [minScore, maxScore] + * + * @example + * ```typescript + * // Minimal - uses all defaults + * const intent = new Intent(); + * + * // With extractors + * const intent = new Intent({ + * key: doc => doc.title, + * summary: doc => doc.content + * }); + * + * // Full configuration + * const intent = new Intent({ + * llm: myClient, + * userId: "org-123", + * key: doc => doc.title, + * summary: doc => doc.content, + * relevancyThreshold: 5, + * batchSize: 20 + * }); + * ``` + */ + constructor(options: IntentOptions & { config?: typeof CONFIG } = {}) { + this.env = options.config ?? CONFIG; + this.ctx = this.buildContext(options); + this.extractors = this.buildExtractors(options); + this.cfg = this.buildConfig(options); + + this.validateConfig(); + this.llm = this.selectAndValidateLlmClient(); + } + + /** + * Rerank candidates based on relevance to a query. + * + * Calls the LLM to evaluate each candidate and returns only those above the + * configured threshold, sorted by score (desc) with stable ordering on ties. + * + * The LLM is instructed to always generate an explanation before the score. + * Explanations are only returned when `options.explain` is true. + * + * Fast-path optimizations: + * - Returns empty array for 0 candidates without LLM call + * - Returns single candidate unchanged without LLM call + * + * Error handling: + * - On any batch error, returns that batch's items in original order + * - On top-level error, returns all items in original order + * - All errors are logged via the configured logger + * + * @param query - The search query or user intent to rank against + * @param candidates - Array of items to rerank + * @param options - Optional per-call configuration + * @param options.explain - When true, return `{ item, explanation }[]` instead of `T[]` + * @param options.userId - Optional user ID for this specific call, overrides ctx.userId + * @returns Filtered and sorted array of items, or original order on any error + * + * @example + * ```typescript + * const itemsOnly = await intent.rank("find expense reports", docs); + * + * const withExplanations = await intent.rank("find expense reports", docs, { explain: true }); + * // => [{ item: Doc, explanation: string }, ...] + * ``` + */ + public async rank( + query: string, + candidates: T[], + options?: { explain?: false; userId?: string }, + ): Promise; + + public async rank( + query: string, + candidates: T[], + options: { explain: true; userId?: string }, + ): Promise>; + + public async rank( + query: string, + candidates: T[], + options?: { explain?: boolean; userId?: string }, + ): Promise> { + try { + if (candidates.length === 0) { + return []; + } + + if (candidates.length === 1) { + if (options?.explain) { + const [firstCandidate] = candidates; + return [{ item: firstCandidate!, explanation: "" }]; + } + return candidates; + } + + const prepared = this.prepareCandidates(candidates); + + const rankedWithExplanations = await batchProcess( + prepared, + this.cfg.batchSize, + this.cfg.tinyBatchFraction, + async (batch) => + (await this.processBatch( + query, + batch, + options?.userId !== undefined ? { userId: options.userId } : undefined, + )) as Array<{ item: T; explanation: string }>, + this.ctx.logger, + (batch) => batch.map(({ item }) => ({ item, explanation: "" })), + ); + + if (options?.explain) { + return rankedWithExplanations; + } + + return rankedWithExplanations.map(({ item }) => item); + } catch (error) { + this.ctx.logger?.warn?.("intent reranker failed, using fallback", { + error: (error as Error)?.message, + }); + + if (options?.explain) { + return candidates.map((item) => ({ item, explanation: "" })); + } + return candidates; + } + } + + /** + * Filter candidates based on relevance to a query. + * + * Calls the LLM to decide whether each candidate is relevant, then returns + * only the relevant items in the same order they were provided. + * + * Explanations are only returned when `options.explain` is true. + * + * Fast paths: + * - Returns [] for 0 candidates without LLM call + * - Returns the single candidate unchanged without LLM call + * + * Error handling: + * - On any batch error, preserves that batch's original order + * - On top-level error, preserves original order + * + * @param query - The search query or user intent to filter against + * @param candidates - Array of items to filter + * @param options - Optional per-call configuration + * @param options.explain - When true, return `{ item, explanation }[]` instead of `T[]` + * @param options.userId - Optional user ID for this specific call, overrides ctx.userId + * @returns Filtered array of items, preserving input order + */ + public async filter( + query: string, + candidates: T[], + options?: { explain?: false; userId?: string }, + ): Promise; + + public async filter( + query: string, + candidates: T[], + options: { explain: true; userId?: string }, + ): Promise>; + + public async filter( + query: string, + candidates: T[], + options?: { explain?: boolean; userId?: string }, + ): Promise> { + if (candidates.length === 0) { + return []; + } + + if (candidates.length === 1) { + if (options?.explain) { + const [firstCandidate] = candidates; + return [{ item: firstCandidate!, explanation: "" }]; + } + return candidates; + } + + const prepared = this.prepareCandidates(candidates); + + const filteredWithExplanations = await batchProcess( + prepared, + this.cfg.batchSize, + this.cfg.tinyBatchFraction, + async (batch) => + (await this.processFilterBatch( + query, + batch, + options?.userId !== undefined ? { userId: options.userId } : undefined, + )) as Array<{ item: T; explanation: string }>, + this.ctx.logger, + (batch) => batch.map(({ item }) => ({ item, explanation: "" })), + ); + + if (options?.explain) { + return filteredWithExplanations; + } + + return filteredWithExplanations.map(({ item }) => item); + } + + /** + * Choose exactly one candidate as the best match for a query. + * + * Uses a tournament strategy when inputs exceed the batch size: + * - Choose one from each batch + * - Then choose one from the batch winners + * + * This method always returns a single item. + * + * @param query - The search query or user intent + * @param candidates - Array of items to choose from + * @param options - Optional per-call configuration + * @param options.explain - When true, return `{ item, explanation }` instead of `T` + * @param options.userId - Optional user ID for this specific call, overrides ctx.userId + * @returns The single chosen item (or item + explanation) + */ + public async choice( + query: string, + candidates: T[], + options?: { explain?: false; userId?: string }, + ): Promise; + + public async choice( + query: string, + candidates: T[], + options: { explain: true; userId?: string }, + ): Promise<{ item: T; explanation: string }>; + + public async choice( + query: string, + candidates: T[], + options?: { explain?: boolean; userId?: string }, + ): Promise { + if (candidates.length === 0) { + throw new Error("intent: choice requires at least one candidate"); + } + + if (candidates.length === 1) { + const [firstCandidate] = candidates; + if (options?.explain) { + return { item: firstCandidate!, explanation: "" }; + } + return firstCandidate!; + } + + const prepared = this.prepareCandidates(candidates); + const keyed = this.ensureUniqueKeys(prepared); + + const winners = await batchProcess( + keyed, + this.cfg.batchSize, + this.cfg.tinyBatchFraction, + async (batch) => [await this.processChoiceBatch(query, batch, options?.userId)], + this.ctx.logger, + (batch) => [{ item: batch[0]!.item, explanation: "" }], + ); + + // batchProcess guarantees at least one winner via its per-batch fallback. + + if (winners.length === 1) { + const [onlyWinner] = winners; + if (options?.explain) { + return onlyWinner!; + } + return onlyWinner!.item; + } + + const preparedFinalists = this.prepareCandidates(winners.map((w) => w.item)); + const keyedFinalists = this.ensureUniqueKeys(preparedFinalists); + + const final = await this.processChoiceBatch(query, keyedFinalists, options?.userId); + if (options?.explain) { + return final; + } + return final.item; + } + + /** + * Normalize incoming items into a consistent shape for downstream processing. + * + * Extracts the key and summary from each item using the configured extractors, + * and attaches the original input index for stable sorting later. + * + * @param candidates - Raw items to prepare + * @returns Array of prepared candidates with extracted metadata and original index + * @private + */ + private prepareCandidates(candidates: T[]): Array<{ + item: T; + idx: number; + baseKey: string; + summary: string; + }> { + return candidates.map((item, idx) => ({ + item, + idx, + baseKey: this.extractors.key(item), + summary: this.extractors.summary(item), + })); + } + + /** + * Ensure keys are unique by suffixing duplicates with their input index. + * + * When multiple items share the same key, subsequent occurrences are renamed + * to "Key (idx)" where idx is the original input index. This prevents JSON + * schema validation errors and ensures the LLM can score each item independently. + * + * @param itemsBase - Prepared candidates with potentially duplicate keys + * @returns Candidates with guaranteed unique keys + * @private + */ + private ensureUniqueKeys( + itemsBase: Array<{ item: T; idx: number; baseKey: string; summary: string }>, + ): Array<{ item: T; idx: number; key: string; summary: string }> { + const counts = new Map(); + return itemsBase.map(({ item, baseKey, summary, idx }) => { + const n = (counts.get(baseKey) ?? 0) + 1; + counts.set(baseKey, n); + const key = n === 1 ? baseKey : `${baseKey} (${idx})`; + return { item, idx, key, summary }; + }); + } + + /** + * Build the JSON schema and chat messages payload for the LLM. + * + * Creates a strict JSON schema requiring one integer property (minScore-maxScore) per candidate key, + * and constructs system + user messages instructing the LLM to score relevance. + * + * @param query - The search query to evaluate candidates against + * @param items - Candidates with unique keys and summaries + * @returns Object containing JSON schema and chat messages array + * @private + */ + private buildRequest( + query: string, + items: Array<{ key: string; summary: string }>, + ): { schema: JSONObject; messages: ChatMessage[] } { + const keys = items.map((x) => x.key); + const schema: JSONObject = buildRelevancySchema(keys, this.cfg.minScore, this.cfg.maxScore); + const messages = buildMessages(query, items, { + minScore: this.cfg.minScore, + maxScore: this.cfg.maxScore, + }); + return { schema, messages }; + } + + /** + * Build the JSON schema and chat messages payload for the LLM filter call. + * + * @param query - The search query to evaluate candidates against + * @param items - Candidates with unique keys and summaries + * @returns Object containing JSON schema and chat messages array + * @private + */ + private buildFilterRequest( + query: string, + items: Array<{ key: string; summary: string }>, + ): { schema: JSONObject; messages: ChatMessage[] } { + const keys = items.map((x) => x.key); + const schema: JSONObject = buildFilterSchema(keys); + const messages = buildFilterMessages(query, items); + return { schema, messages }; + } + + /** + * Build the JSON schema and chat messages payload for the LLM choice call. + * + * @param query - The search query to choose against + * @param items - Candidates with unique keys and summaries + * @returns Object containing JSON schema and chat messages array + * @private + */ + private buildChoiceRequest( + query: string, + items: Array<{ key: string; summary: string }>, + ): { schema: JSONObject; messages: ChatMessage[] } { + const keys = items.map((x) => x.key); + const schema: JSONObject = buildChoiceSchema(keys); + const messages = buildChoiceMessages(query, items); + return { schema, messages }; + } + + /** + * Invoke the LLM and return the parsed map of candidate scores. + * + * Calls the configured LLM client with the messages, JSON schema, model config, + * and user ID. Returns null if the response is invalid or missing. + * + * @param messages - Chat messages (system + user) to send to LLM + * @param schema - Strict JSON schema defining expected response structure + * @param userId - Optional user identifier for provider abuse monitoring + * @returns Map of candidate keys to numeric scores, or null if response invalid + * @private + */ + private async fetchEvaluations( + messages: ChatMessage[], + schema: JSONObject, + userId?: string, + ): Promise | null> { + const config: LlmCallConfig = { + model: this.resolveModel(), + reasoningEffort: "medium", + timeoutMs: this.cfg.timeoutMs, + }; + const { data } = await this.llm.call>( + messages, + schema, + config, + userId ?? this.ctx.userId, + ); + + if (data == null || typeof data !== "object") return null; + return data as Record; + } + + /** + * Invoke the LLM and return boolean relevancy decisions. + * + * @param messages - Chat messages to send + * @param schema - Strict JSON schema defining expected response structure + * @param userId - Optional user id + * @returns Map of candidate keys to filter decisions, or null if invalid + * @private + */ + private async fetchFilterDecisions( + messages: ChatMessage[], + schema: JSONObject, + userId?: string, + ): Promise | null> { + const config: LlmCallConfig = { + model: this.resolveModel(), + reasoningEffort: "medium", + timeoutMs: this.cfg.timeoutMs, + }; + const { data } = await this.llm.call< + Record + >(messages, schema, config, userId ?? this.ctx.userId); + + if (data == null || typeof data !== "object") return null; + return data as Record; + } + + /** + * Invoke the LLM and return a single selected key. + * + * @param messages - Chat messages to send + * @param schema - Strict JSON schema defining expected response structure + * @param userId - Optional user id + * @returns Choice result, or null if invalid + * @private + */ + private async fetchChoice( + messages: ChatMessage[], + schema: JSONObject, + userId?: string, + ): Promise<{ explanation: string; selectedKey: string } | null> { + const config: LlmCallConfig = { + model: this.resolveModel(), + reasoningEffort: "medium", + timeoutMs: this.cfg.timeoutMs, + }; + const { data } = await this.llm.call<{ explanation: string; selectedKey: string }>( + messages, + schema, + config, + userId ?? this.ctx.userId, + ); + + if (data == null || typeof data !== "object") return null; + const record = data as Record; + const explanation = typeof record.explanation === "string" ? record.explanation : ""; + const selectedKey = typeof record.selectedKey === "string" ? record.selectedKey : ""; + if (selectedKey === "") return null; + return { explanation, selectedKey }; + } + + /** + * Apply boolean filter decisions while preserving input order. + * + * @param items - Candidates with unique keys + * @param decisions - Map of candidate keys to decisions + * @returns Filtered items with explanations + * @private + */ + private applyFilter( + items: Array<{ item: T; idx: number; key: string; summary: string }>, + decisions: Record, + ): Array<{ item: T; explanation: string }> { + const kept: Array<{ item: T; explanation: string }> = []; + for (const it of items) { + const decision = decisions[it.key]; + if (decision?.isRelevant === true) { + kept.push({ + item: it.item, + explanation: typeof decision.explanation === "string" ? decision.explanation : "", + }); + } + } + return kept; + } + + /** + * Process a single batch of candidates through the LLM filter. + * + * @param query - Query to filter against + * @param batch - Prepared candidates + * @param options - Optional userId override + * @returns Filtered items (stable order), with explanations + * @private + */ + private async processFilterBatch( + query: string, + batch: Array<{ item: T; idx: number; baseKey: string; summary: string }>, + options?: { userId?: string }, + ): Promise> { + const keyed = this.ensureUniqueKeys(batch); + const { schema, messages } = this.buildFilterRequest(query, keyed); + const decisions = await this.fetchFilterDecisions(messages, schema, options?.userId); + if (decisions == null) { + return keyed.map(({ item }) => ({ item, explanation: "" })); + } + + return this.applyFilter(keyed, decisions); + } + + /** + * Process a batch of candidates and choose a single winner. + * + * @param query - Query to choose against + * @param batch - Candidates with unique keys + * @param userId - Optional userId override + * @returns Winner item with explanation + * @private + */ + private async processChoiceBatch( + query: string, + batch: Array<{ item: T; idx: number; key: string; summary: string }>, + userId?: string, + ): Promise<{ item: T; explanation: string }> { + const { schema, messages } = this.buildChoiceRequest(query, batch); + const choice = await this.fetchChoice(messages, schema, userId); + if (choice == null) { + return { item: batch[0]!.item, explanation: "" }; + } + + const winner = batch.find((x) => x.key === choice.selectedKey); + if (!winner) { + return { item: batch[0]!.item, explanation: choice.explanation }; + } + + return { item: winner.item, explanation: choice.explanation }; + } + + /** + * Apply relevancy threshold filtering and stable sorting. + * + * Scores are clamped to the configured score range, then filtered to keep only items with + * score > threshold. Results are sorted by score descending, with ties + * preserving original input order for deterministic results. + * + * @param items - Candidates with unique keys + * @param scores - Map of candidate keys to LLM-assigned scores + * @returns Filtered and sorted array of original items + * @private + */ + private rankAndFilter( + items: Array<{ item: T; idx: number; key: string; summary: string }>, + evaluations: Record, + ): Array<{ item: T; explanation: string; score: number }> { + const threshold = this.cfg.relevancyThreshold; + const scored = items.map(({ item, idx, key }) => ({ + item, + idx, + explanation: + typeof evaluations[key]?.explanation === "string" ? evaluations[key].explanation : "", + score: clamp( + evaluations[key]?.score ?? this.cfg.minScore, + this.cfg.minScore, + this.cfg.maxScore, + ), + })); + + const filtered = scored.filter(({ score }) => score > threshold); + const sorted = filtered.sort((a, b) => { + if (b.score !== a.score) return b.score - a.score; + return a.idx - b.idx; + }); + return sorted; + } + + /** + * Process a single batch of candidates through the LLM. + * + * Ensures unique keys, builds the request payload, fetches scores from the LLM, + * and returns filtered and sorted results. On any error or null response from + * the LLM, returns items in their original order as a fallback. + * + * @param query - The search query to evaluate candidates against + * @param batch - Batch of prepared candidates to process + * @param userId - Optional user identifier for provider abuse monitoring + * @returns Ranked and filtered items, or original order on error + * @private + */ + private async processBatch( + query: string, + batch: Array<{ item: T; idx: number; baseKey: string; summary: string }>, + options?: { userId?: string }, + ): Promise> { + const keyed = this.ensureUniqueKeys(batch); + const { schema, messages } = this.buildRequest(query, keyed); + const evaluations = await this.fetchEvaluations(messages, schema, options?.userId); + if (evaluations == null) { + return keyed.map(({ item }) => ({ item, explanation: "" })); + } + + const ranked = this.rankAndFilter(keyed, evaluations); + return ranked.map(({ item, explanation }) => ({ item, explanation })); + } +} diff --git a/src/intent.unit.test.ts b/src/intent.unit.test.ts new file mode 100644 index 0000000..1865fe4 --- /dev/null +++ b/src/intent.unit.test.ts @@ -0,0 +1,1121 @@ +import { describe, expect, test, vi } from "vitest"; + +import { CONFIG } from "./config"; +import { Intent } from "./intent"; +import { buildCandidateEvaluationSchema } from "./schema"; + +import type { LlmClient, LoggerLike, IntentCandidate, IntentContext } from "./types"; + +function makeCtx(overrides: Partial = {}): IntentContext & { + llm: LlmClient & { call: ReturnType }; + logger: Required & { warn: ReturnType }; +} { + const logger = { + info: vi.fn(() => {}), + warn: vi.fn(() => {}), + error: vi.fn(() => {}), + } as any; + const llm = { + call: vi.fn(async () => ({ data: {} })), + } as any; + return { + llm, + logger, + userId: undefined, + ...overrides, + } as any; +} + +describe("Intent.rank", () => { + test("uses GROQ default model when no model override is provided", async () => { + const ctx = makeCtx(); + (ctx.llm.call as any).mockResolvedValueOnce({ + data: { + A: { explanation: "a", score: 10 }, + B: { explanation: "b", score: 0 }, + }, + }); + + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + }); + + await intent.rank("query", [ + { key: "A", summary: "" }, + { key: "B", summary: "" }, + ]); + + const call = (ctx.llm.call as any).mock.calls[0]; + expect(call[2].model).toBe(CONFIG.GROQ.DEFAULT_MODEL); + }); + + test("candidate evaluation schema defines explanation before score", () => { + const schema = buildCandidateEvaluationSchema(); + expect(Object.keys(schema.properties)).toEqual(["explanation", "score"]); + }); + + test("throws when no llm and no GROQ_API_KEY", async () => { + const configOverride = { + ...CONFIG, + GROQ: { ...CONFIG.GROQ, API_KEY: "" }, + } as typeof CONFIG; + expect( + () => new Intent({ key: (c) => c.key, config: configOverride }), + ).toThrow(/No LLM client provided/); + }); + + test("throws when threshold is below 0", async () => { + const ctx = makeCtx(); + expect( + () => new Intent({ ...ctx, key: (c) => c.key, relevancyThreshold: -1 }), + ).toThrow(/relevancyThreshold must be between/); + }); + + test("throws when threshold is above 10", async () => { + const ctx = makeCtx(); + expect( + () => new Intent({ ...ctx, key: (c) => c.key, relevancyThreshold: 11 }), + ).toThrow(/relevancyThreshold must be between/); + }); + + test("throws when maxScore is below minScore", async () => { + const ctx = makeCtx(); + expect( + () => + new Intent({ + ...ctx, + key: (c) => c.key, + minScore: 5, + maxScore: 4, + }), + ).toThrow(/maxScore must be >= minScore/); + }); + + test("supports non-0..10 score ranges", async () => { + const ctx = makeCtx(); + (ctx.llm.call as any).mockResolvedValueOnce({ + data: { + A: { explanation: "high", score: 5 }, + B: { explanation: "low", score: 3 }, + }, + }); + + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + minScore: 2, + maxScore: 5, + relevancyThreshold: 3, + }); + + const input = [ + { key: "A", summary: "" }, + { key: "B", summary: "" }, + ]; + const res = await intent.rank("query", input); + expect(res.map((c) => c.key)).toEqual(["A"]); + }); + test("returns empty list for zero candidates", async () => { + const ctx = makeCtx(); + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + }); + const res = await intent.rank("query", []); + expect(res).toEqual([]); + expect(ctx.llm.call).not.toHaveBeenCalled(); + }); + + test("returns input unchanged for single candidate (no LLM call)", async () => { + const ctx = makeCtx(); + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + }); + const input = [{ key: "Only", summary: "s" }]; + const res = await intent.rank("query", input); + expect(res).toEqual(input); + expect(ctx.llm.call).not.toHaveBeenCalled(); + }); + + test("returns explanation wrapper for single candidate when explain is true (no LLM call)", async () => { + const ctx = makeCtx(); + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + }); + const input = [{ key: "Only", summary: "s" }]; + const res = await intent.rank("query", input, { explain: true }); + expect(res).toEqual([{ item: input[0], explanation: "" }]); + expect(ctx.llm.call).not.toHaveBeenCalled(); + }); + + test("returns explanations when explain option is true", async () => { + const ctx = makeCtx(); + (ctx.llm.call as any).mockResolvedValueOnce({ + data: { + A: { explanation: "matches query", score: 10 }, + B: { explanation: "irrelevant", score: 0 }, + }, + }); + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + }); + const input = [ + { key: "A", summary: "" }, + { key: "B", summary: "" }, + ]; + const res = await intent.rank("query", input, { explain: true }); + expect(res).toEqual([{ item: input[0], explanation: "matches query" }]); + }); + + test("explain option is false by default (returns T[])", async () => { + const ctx = makeCtx(); + (ctx.llm.call as any).mockResolvedValueOnce({ + data: { + A: { explanation: "matches query", score: 10 }, + B: { explanation: "irrelevant", score: 0 }, + }, + }); + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + }); + const input = [ + { key: "A", summary: "" }, + { key: "B", summary: "" }, + ]; + const res = await intent.rank("query", input); + expect(res).toEqual([input[0]]); + }); + + test("rounds scores, filters zeros, orders by score", async () => { + const ctx = makeCtx(); + (ctx.llm.call as any).mockResolvedValueOnce({ + data: { + A: { explanation: "a", score: 10 }, + B: { explanation: "b", score: 6.8 }, + C: { explanation: "c", score: 0 }, + }, + }); + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + }); + const input = ["A", "B", "C"].map((k) => ({ key: k, summary: k })); + const res = await intent.rank("query", input); + expect(res.map((c) => c.key)).toEqual(["A", "B"]); + const call = (ctx.llm.call as any).mock.calls[0]; + expect(call[2].timeoutMs).toBe(3000); // default + expect(call[2].model).toBe("openai/gpt-oss-20b"); // GROQ default model + }); + + test("handles non-numeric or missing scores by clamping to 0", async () => { + const ctx = makeCtx(); + (ctx.llm.call as any).mockResolvedValueOnce({ + data: { X: { explanation: "x", score: "nope" as any } }, + }); + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + }); + const input = ["X", "Y"].map((k) => ({ key: k, summary: k })); + const res = await intent.rank("query", input); + expect(res).toEqual([]); + }); + + test("normalizes out-of-range and infinite values", async () => { + const ctx = makeCtx(); + (ctx.llm.call as any).mockResolvedValueOnce({ + data: { + A: { explanation: "a", score: -3 }, + B: { explanation: "b", score: 11 }, + C: { explanation: "c", score: 9.6 }, + D: { explanation: "d", score: Number.POSITIVE_INFINITY }, + }, + }); + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + }); + const input = ["A", "B", "C", "D"].map((k) => ({ key: k, summary: k })); + const res = await intent.rank("query", input); + // B clamps to 10, D rounds to 10, keep input order for tie: B before D + expect(res.map((c) => c.key)).toEqual(["B", "D", "C"]); + }); + + test("returns original list on error and logs warning", async () => { + const ctx = makeCtx(); + (ctx.llm.call as any).mockRejectedValueOnce(new Error("boom")); + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + }); + const input = [ + { key: "A", summary: "" }, + { key: "B", summary: "" }, + ]; + const res = await intent.rank("query", input); + expect(res).toEqual(input); + expect(ctx.logger.warn).toHaveBeenCalled(); + }); + + test("uses default extractors when both are missing", async () => { + const ctx = makeCtx(); + // Return some scores so rank proceeds + (ctx.llm.call as any).mockResolvedValueOnce({ data: {} }); + const intent = new Intent({ + ...ctx, + // no key or summary extractors - will use defaults + batchSize: 10, + }); + const input: IntentCandidate[] = [ + { key: "A", summary: "original" }, + { key: "B", summary: "ignored" }, + ]; + await intent.rank("query", input); + const call = (ctx.llm.call as any).mock.calls[0]; + const messages = call[0]; + const userPayload = JSON.parse(messages[1].content); + const summaries = userPayload.candidate_search_results.map((c: any) => c.summary); + // With default extractors, both key and summary will use JSON.stringify + expect(summaries.every((s: string) => s.includes("original") || s.includes("ignored"))).toBe( + true, + ); + }); + + test("uses empty summary when extractor is missing", async () => { + const ctx = makeCtx(); + // Return some scores so rank proceeds + (ctx.llm.call as any).mockResolvedValueOnce({ + data: { + A: { explanation: "a", score: 1 }, + B: { explanation: "b", score: 0 }, + }, + }); + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + // no summary extractor - will use default + batchSize: 10, + }); + const input: IntentCandidate[] = [ + { key: "A", summary: "original" }, + { key: "B", summary: "ignored" }, + ]; + await intent.rank("query", input); + const call = (ctx.llm.call as any).mock.calls[0]; + const messages = call[0]; + const userPayload = JSON.parse(messages[1].content); + const summaries = userPayload.candidate_search_results.map((c: any) => c.summary); + // With default summary extractor, will use JSON.stringify + expect(summaries.every((s: string) => s.includes("original") || s.includes("ignored"))).toBe( + true, + ); + }); + + test("top-level rank catch: logs and returns input on unexpected error", async () => { + const ctx = makeCtx(); + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + }); + // @ts-ignore override private method to throw to trigger top-level catch + intent.prepareCandidates = () => { + throw new Error("oops"); + }; + const input = [ + { key: "A", summary: "" }, + { key: "B", summary: "" }, + { key: "C", summary: "" }, + ]; + const res = await intent.rank("query", input); + expect(res).toEqual(input); + expect(ctx.logger.warn).toHaveBeenCalled(); + }); + + test("top-level rank catch: logs and returns wrapped input when explain is true", async () => { + const ctx = makeCtx(); + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + }); + // @ts-ignore override private method to throw to trigger top-level catch + intent.prepareCandidates = () => { + throw new Error("oops"); + }; + const input = [ + { key: "A", summary: "" }, + { key: "B", summary: "" }, + { key: "C", summary: "" }, + ]; + + const res = await intent.rank("query", input, { explain: true }); + expect(res).toEqual([ + { item: input[0], explanation: "" }, + { item: input[1], explanation: "" }, + { item: input[2], explanation: "" }, + ]); + expect(ctx.logger.warn).toHaveBeenCalled(); + }); + + test("passes userId from ctx and allows method override", async () => { + const ctx = makeCtx({ userId: "ctx-user" }); + (ctx.llm.call as any).mockResolvedValueOnce({ + data: { + A: { explanation: "a", score: 5 }, + B: { explanation: "b", score: 5 }, + }, + }); + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + }); + const input = [ + { key: "A", summary: "" }, + { key: "B", summary: "" }, + ]; + await intent.rank("query", input, { userId: "call-user" }); + const calls = (ctx.llm.call as any).mock.calls; + expect(calls[0][3]).toBe("call-user"); // override wins + }); + + test("timeout config is forwarded to client", async () => { + const ctx = makeCtx(); + (ctx.llm.call as any).mockResolvedValueOnce({ + data: { + A: { explanation: "a", score: 5 }, + B: { explanation: "b", score: 5 }, + }, + }); + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + timeoutMs: 5, + }); + const input = [ + { key: "A", summary: "" }, + { key: "B", summary: "" }, + ]; + await intent.rank("query", input); + const calls = (ctx.llm.call as any).mock.calls; + expect(calls[0][2].timeoutMs).toBe(5); + }); + + test("splits long lists into batches and combines results", async () => { + const ctx = makeCtx(); + (ctx.llm.call as any) + .mockResolvedValueOnce({ + data: { + K2: { explanation: "k2", score: 10 }, + K0: { explanation: "k0", score: 7 }, + }, + }) + .mockResolvedValueOnce({ + data: { + K7: { explanation: "k7", score: 10 }, + K6: { explanation: "k6", score: 9 }, + }, + }); + + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + batchSize: 5, + }); + + const input: IntentCandidate[] = Array.from({ length: 10 }).map((_, i) => ({ + key: `K${i}`, + summary: `S${i}`, + })); + + const out = await intent.rank("query", input); + expect(out.map((c) => c.key)).toEqual(["K2", "K0", "K7", "K6"]); + }); + + test("merges tiny final batch", async () => { + const ctx = makeCtx(); + (ctx.llm.call as any).mockResolvedValueOnce({ data: {} }).mockResolvedValueOnce({ data: {} }); + + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + batchSize: 5, + tinyBatchFraction: 0.2, + }); + + const input: IntentCandidate[] = Array.from({ length: 7 }).map((_, i) => ({ + key: `K${i}`, + summary: `S${i}`, + })); + + await intent.rank("query", input); + expect((ctx.llm.call as any).mock.calls.length).toBe(2); + }); + + test("merges really tiny final batch (<= threshold)", async () => { + const ctx = makeCtx(); + (ctx.llm.call as any).mockResolvedValueOnce({ data: {} }); + + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + batchSize: 5, + tinyBatchFraction: 0.2, + }); + + const input: IntentCandidate[] = Array.from({ length: 6 }).map((_, i) => ({ + key: `K${i}`, + summary: `S${i}`, + })); + + await intent.rank("query", input); + expect((ctx.llm.call as any).mock.calls.length).toBe(1); + }); + + test("one batch fails while others succeed (partial fallback)", async () => { + const ctx = makeCtx(); + (ctx.llm.call as any) + .mockResolvedValueOnce({ + data: { + K1: { explanation: "k1", score: 10 }, + K0: { explanation: "k0", score: 9 }, + }, + }) + .mockRejectedValueOnce(new Error("boom")); + + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + batchSize: 3, + }); + + const input: IntentCandidate[] = Array.from({ length: 6 }).map((_, i) => ({ + key: `K${i}`, + summary: `S${i}`, + })); + + const out = await intent.rank("query", input); + expect(out.map((c) => c.key)).toEqual(["K1", "K0", "K3", "K4", "K5"]); + expect((ctx.llm.call as any).mock.calls.length).toBe(2); + }); + + test("one batch fails while others succeed when explain is true (partial fallback)", async () => { + const ctx = makeCtx(); + (ctx.llm.call as any) + .mockResolvedValueOnce({ + data: { + K1: { explanation: "k1", score: 10 }, + K0: { explanation: "k0", score: 9 }, + }, + }) + .mockRejectedValueOnce(new Error("boom")); + + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + batchSize: 3, + }); + + const input: IntentCandidate[] = Array.from({ length: 6 }).map((_, i) => ({ + key: `K${i}`, + summary: `S${i}`, + })); + + const out = await intent.rank("query", input, { explain: true }); + expect(out).toEqual([ + { item: input[1], explanation: "k1" }, + { item: input[0], explanation: "k0" }, + { item: input[3], explanation: "" }, + { item: input[4], explanation: "" }, + { item: input[5], explanation: "" }, + ]); + expect((ctx.llm.call as any).mock.calls.length).toBe(2); + }); + + test("returns original list when scores payload is null", async () => { + const ctx = makeCtx(); + (ctx.llm.call as any).mockResolvedValueOnce({ data: null }); + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + }); + const input: IntentCandidate[] = [ + { key: "A", summary: "S" }, + { key: "B", summary: "S" }, + ]; + + const out = await intent.rank("query", input); + expect(out).toEqual(input); + }); + + test("returns blank explanations on null LLM payload when explain is true", async () => { + const ctx = makeCtx(); + (ctx.llm.call as any).mockResolvedValueOnce({ data: null }); + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + }); + const input: IntentCandidate[] = [ + { key: "A", summary: "S" }, + { key: "B", summary: "S" }, + ]; + + const out = await intent.rank("query", input, { explain: true }); + expect(out).toEqual([ + { item: input[0], explanation: "" }, + { item: input[1], explanation: "" }, + ]); + }); + + test("stable order for ties within a batch (duplicate keys)", async () => { + const ctx = makeCtx(); + (ctx.llm.call as any).mockResolvedValueOnce({ + data: { + Same: { explanation: "s0", score: 5 }, + "Same (1)": { explanation: "s1", score: 5 }, + }, + }); + + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + }); + + const input: IntentCandidate[] = [ + { key: "Same", summary: "S0" }, + { key: "Same", summary: "S1" }, + ]; + + const out = await intent.rank("query", input); + expect(out.map((c) => c.summary)).toEqual(["S0", "S1"]); + }); +}); + +describe("Intent.filter", () => { + test("returns empty list for zero candidates", async () => { + const ctx = makeCtx(); + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + }); + const res = await intent.filter("query", []); + expect(res).toEqual([]); + expect(ctx.llm.call).not.toHaveBeenCalled(); + }); + + test("returns input unchanged for single candidate (no LLM call)", async () => { + const ctx = makeCtx(); + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + }); + const input = [{ key: "Only", summary: "s" }]; + const res = await intent.filter("query", input); + expect(res).toEqual(input); + expect(ctx.llm.call).not.toHaveBeenCalled(); + }); + + test("with explain=true returns explanation wrapper for single candidate (no LLM call)", async () => { + const ctx = makeCtx(); + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + }); + const input = [{ key: "Only", summary: "s" }]; + const res = await intent.filter("query", input, { explain: true }); + expect(res).toEqual([{ item: input[0], explanation: "" }]); + expect(ctx.llm.call).not.toHaveBeenCalled(); + }); + + test("filters by boolean decision and preserves input order", async () => { + const ctx = makeCtx(); + (ctx.llm.call as any).mockResolvedValueOnce({ + data: { + A: { explanation: "yes", isRelevant: true }, + B: { explanation: "no", isRelevant: false }, + C: { explanation: "yes", isRelevant: true }, + }, + }); + + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + }); + + const input = [ + { key: "A", summary: "" }, + { key: "B", summary: "" }, + { key: "C", summary: "" }, + ]; + + const res = await intent.filter("query", input); + expect(res.map((x) => x.key)).toEqual(["A", "C"]); + }); + + test("explain=true returns explanations for kept items", async () => { + const ctx = makeCtx(); + (ctx.llm.call as any).mockResolvedValueOnce({ + data: { + A: { explanation: 123, isRelevant: true }, + B: { explanation: "no", isRelevant: false }, + }, + }); + + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + }); + + const input = [ + { key: "A", summary: "" }, + { key: "B", summary: "" }, + ]; + const res = await intent.filter("query", input, { explain: true }); + expect(res).toEqual([{ item: input[0], explanation: "" }]); + }); + + test("returns original list on error and logs warning", async () => { + const ctx = makeCtx(); + (ctx.llm.call as any).mockRejectedValueOnce(new Error("boom")); + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + }); + const input = [ + { key: "A", summary: "" }, + { key: "B", summary: "" }, + ]; + const res = await intent.filter("query", input); + expect(res).toEqual(input); + expect(ctx.logger.warn).toHaveBeenCalled(); + }); + + test("on error returns wrapped input when explain=true", async () => { + const ctx = makeCtx(); + (ctx.llm.call as any).mockImplementationOnce(() => Promise.reject(new Error("boom"))); + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + }); + const input = [ + { key: "A", summary: "" }, + { key: "B", summary: "" }, + ]; + + const res = await intent.filter("query", input, { explain: true }); + expect(res).toEqual([ + { item: input[0], explanation: "" }, + { item: input[1], explanation: "" }, + ]); + expect(ctx.logger.warn).toHaveBeenCalled(); + }); + + test("returns original list when LLM response is null", async () => { + const ctx = makeCtx(); + (ctx.llm.call as any).mockResolvedValueOnce({ data: null }); + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + }); + const input = [ + { key: "A", summary: "" }, + { key: "B", summary: "" }, + ]; + const res = await intent.filter("query", input); + expect(res).toEqual(input); + }); + + test("passes userId override through filter requests", async () => { + const ctx = makeCtx(); + (ctx.llm.call as any).mockResolvedValueOnce({ + data: { + A: { explanation: "", isRelevant: true }, + B: { explanation: "", isRelevant: false }, + }, + }); + + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + }); + + await intent.filter( + "query", + [ + { key: "A", summary: "" }, + { key: "B", summary: "" }, + ], + { userId: "call-user" }, + ); + + const call = (ctx.llm.call as any).mock.calls[0]; + expect(call[3]).toBe("call-user"); + }); + + test("with explain=true returns wrapped input on top-level error", async () => { + const ctx = makeCtx({ llm: undefined }); + const configOverride = { + ...CONFIG, + GROQ: { ...CONFIG.GROQ, API_KEY: "fake" }, + } as typeof CONFIG; + + const intent = new Intent({ + ...ctx, + config: configOverride, + key: (c) => c.key, + summary: (c) => c.summary, + }); + + const input = [ + { key: "A", summary: "" }, + { key: "B", summary: "" }, + ]; + const res = await intent.filter("query", input, { explain: true }); + expect(res).toEqual([ + { item: input[0], explanation: "" }, + { item: input[1], explanation: "" }, + ]); + expect(ctx.logger.warn).toHaveBeenCalled(); + }); +}); + +describe("Intent.choice", () => { + test("throws for zero candidates", async () => { + const ctx = makeCtx(); + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + }); + + await expect(intent.choice("query", [])).rejects.toThrow(/requires at least one candidate/); + }); + + test("returns single candidate unchanged (no LLM call)", async () => { + const ctx = makeCtx(); + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + }); + const input = [{ key: "Only", summary: "s" }]; + const res = await intent.choice("query", input); + expect(res).toEqual(input[0]); + expect(ctx.llm.call).not.toHaveBeenCalled(); + }); + + test("explain=true returns explanation wrapper for single candidate (no LLM call)", async () => { + const ctx = makeCtx(); + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + }); + + const input = [{ key: "Only", summary: "s" }]; + const res = await intent.choice("query", input, { explain: true }); + expect(res).toEqual({ item: input[0], explanation: "" }); + expect(ctx.llm.call).not.toHaveBeenCalled(); + }); + + test("returns chosen item based on selectedKey", async () => { + const ctx = makeCtx(); + (ctx.llm.call as any).mockResolvedValueOnce({ data: { selectedKey: "B" } }); + + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + batchSize: 10, + }); + + const input = [ + { key: "A", summary: "" }, + { key: "B", summary: "" }, + { key: "C", summary: "" }, + ]; + const res = await intent.choice("query", input); + expect(res.key).toBe("B"); + }); + + test("falls back to first item when LLM returns invalid data", async () => { + const ctx = makeCtx(); + (ctx.llm.call as any).mockResolvedValueOnce({ data: null }); + + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + batchSize: 10, + }); + + const input = [ + { key: "A", summary: "" }, + { key: "B", summary: "" }, + ]; + const res = await intent.choice("query", input); + expect(res).toEqual(input[0]); + }); + + test("falls back to first item when selectedKey is not in the batch", async () => { + const ctx = makeCtx(); + (ctx.llm.call as any).mockResolvedValueOnce({ + data: { explanation: "x", selectedKey: "Z" }, + }); + + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + batchSize: 10, + }); + + const input = [ + { key: "A", summary: "" }, + { key: "B", summary: "" }, + ]; + const res = await intent.choice("query", input); + expect(res).toEqual(input[0]); + }); + + test("explain=true returns chosen item with explanation", async () => { + const ctx = makeCtx(); + (ctx.llm.call as any).mockResolvedValueOnce({ + data: { explanation: "best", selectedKey: "A" }, + }); + + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + batchSize: 10, + }); + + const input = [ + { key: "A", summary: "" }, + { key: "B", summary: "" }, + ]; + const res = await intent.choice("query", input, { explain: true }); + expect(res).toEqual({ item: input[0], explanation: "best" }); + }); + + test("uses a tournament strategy when candidates exceed batchSize", async () => { + const ctx = makeCtx(); + (ctx.llm.call as any) + .mockResolvedValueOnce({ data: { explanation: "b1", selectedKey: "K1" } }) + .mockResolvedValueOnce({ data: { explanation: "b2", selectedKey: "K3" } }) + .mockResolvedValueOnce({ data: { explanation: "final", selectedKey: "K3" } }); + + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + batchSize: 2, + }); + + const input: IntentCandidate[] = [ + { key: "K0", summary: "" }, + { key: "K1", summary: "" }, + { key: "K2", summary: "" }, + { key: "K3", summary: "" }, + ]; + + const res = await intent.choice("query", input); + expect(res.key).toBe("K3"); + expect((ctx.llm.call as any).mock.calls.length).toBe(3); + }); + + test("does not run a final round when there is only one batch winner", async () => { + const ctx = makeCtx(); + (ctx.llm.call as any).mockResolvedValueOnce({ + data: { explanation: "b", selectedKey: "K1" }, + }); + + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + batchSize: 10, + }); + + const input: IntentCandidate[] = [ + { key: "K0", summary: "" }, + { key: "K1", summary: "" }, + { key: "K2", summary: "" }, + ]; + + const res = await intent.choice("query", input); + expect(res.key).toBe("K1"); + expect((ctx.llm.call as any).mock.calls.length).toBe(1); + }); + + test("with explain=true returns the single batch winner without running a final round", async () => { + const ctx = makeCtx(); + (ctx.llm.call as any).mockResolvedValueOnce({ + data: { explanation: "b", selectedKey: "K1" }, + }); + + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + batchSize: 10, + }); + + const input: IntentCandidate[] = [ + { key: "K0", summary: "" }, + { key: "K1", summary: "" }, + { key: "K2", summary: "" }, + ]; + + const res = await intent.choice("query", input, { explain: true }); + expect(res).toEqual({ item: input[1], explanation: "b" }); + expect((ctx.llm.call as any).mock.calls.length).toBe(1); + }); + + test("with explain=true returns wrapped first item on error", async () => { + const ctx = makeCtx(); + (ctx.llm.call as any).mockImplementationOnce(() => Promise.reject(new Error("boom"))); + + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + batchSize: 10, + }); + + const input = [ + { key: "A", summary: "" }, + { key: "B", summary: "" }, + ]; + + const res = await intent.choice("query", input, { explain: true }); + expect(res).toEqual({ item: input[0], explanation: "" }); + expect(ctx.logger.warn).toHaveBeenCalled(); + }); + + test("passes userId override through choice requests", async () => { + const ctx = makeCtx(); + (ctx.llm.call as any).mockResolvedValueOnce({ + data: { explanation: "x", selectedKey: "A" }, + }); + + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + batchSize: 10, + }); + + const input = [ + { key: "A", summary: "" }, + { key: "B", summary: "" }, + ]; + + await intent.choice("query", input, { userId: "call-user" }); + const call = (ctx.llm.call as any).mock.calls[0]; + expect(call[3]).toBe("call-user"); + }); + + test("runs a final round when there are multiple batch winners", async () => { + const ctx = makeCtx(); + (ctx.llm.call as any) + .mockResolvedValueOnce({ data: { selectedKey: "K1" } }) + .mockResolvedValueOnce({ data: { selectedKey: "K2" } }) + .mockResolvedValueOnce({ data: { selectedKey: "K2" } }); + + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + batchSize: 2, + tinyBatchFraction: 0, + }); + + const input: IntentCandidate[] = [ + { key: "K0", summary: "" }, + { key: "K1", summary: "" }, + { key: "K2", summary: "" }, + { key: "K3", summary: "" }, + ]; + + const res = await intent.choice("query", input, { explain: true }); + expect(res.item.key).toBe("K2"); + expect((ctx.llm.call as any).mock.calls.length).toBe(3); + }); + + test("returns the first item when there are no batch winners", async () => { + const ctx = makeCtx(); + (ctx.llm.call as any).mockResolvedValueOnce({ data: { selectedKey: 123 } }); + + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + batchSize: 10, + }); + + const input: IntentCandidate[] = [ + { key: "A", summary: "" }, + { key: "B", summary: "" }, + ]; + + const res = await intent.choice("query", input); + expect(res).toEqual(input[0]); + }); + + test("on error returns the first item and logs warning", async () => { + const ctx = makeCtx(); + (ctx.llm.call as any).mockImplementationOnce(() => Promise.reject(new Error("boom"))); + + const intent = new Intent({ + ...ctx, + key: (c) => c.key, + summary: (c) => c.summary, + batchSize: 10, + }); + + const input = [ + { key: "A", summary: "" }, + { key: "B", summary: "" }, + ]; + const res = await intent.choice("query", input); + expect(res).toEqual(input[0]); + expect(ctx.logger.warn).toHaveBeenCalled(); + }); +}); diff --git a/src/lib/config.ts b/src/lib/config.ts index 63c282c..628b8ac 100644 --- a/src/lib/config.ts +++ b/src/lib/config.ts @@ -10,6 +10,7 @@ export type BaseOptions = { default?: T }; export type RangeOptions = BaseOptions & { min?: number; max?: number }; +export type EnumOptions = BaseOptions & { values: T }; function throwRequiredEnvVar(name: string): never { throw new Error(`${name} is required.`); @@ -29,6 +30,23 @@ export function string(name: string, opts?: BaseOptions): string { return throwRequiredEnvVar(name); } +/** + * Read an env var as a constrained enum string. + * + * When set, validates the value is included in opts.values. + * When unset, returns opts.default if provided, otherwise throws. + */ +export function enumeration( + name: string, + opts: EnumOptions, +): T[number] { + const value = opts.default !== undefined ? string(name, { default: opts.default }) : string(name); + if (opts.values.includes(value)) { + return value; + } + throw new Error(`${name} must be one of: ${opts.values.join(", ")}`); +} + export function boolean(name: string, opts?: BaseOptions): boolean { const value = process.env[name]; const exists = value != null && value.length > 0; diff --git a/src/lib/config.unit.test.ts b/src/lib/config.unit.test.ts index 2f3a072..7c70401 100644 --- a/src/lib/config.unit.test.ts +++ b/src/lib/config.unit.test.ts @@ -137,4 +137,23 @@ describe("lib/config", () => { expect(cfg.number(VAR, { default: 9.9 })).toBeCloseTo(9.9); }); }); + + describe("enumeration()", () => { + it("returns the env value when it is included in values", () => { + process.env[VAR] = "GROQ"; + expect(cfg.enumeration(VAR, { values: ["GROQ"] as const })).toBe("GROQ"); + }); + + it("returns default when missing", () => { + unset(VAR); + expect(cfg.enumeration(VAR, { default: "GROQ", values: ["GROQ"] as const })).toBe("GROQ"); + }); + + it("throws when the value is not included in values", () => { + process.env[VAR] = "NOPE"; + expect(() => cfg.enumeration(VAR, { values: ["GROQ"] as const })).toThrow( + /must be one of: GROQ/, + ); + }); + }); }); diff --git a/src/llm_client.ts b/src/llm_client.ts index 2300e70..45a8d9e 100644 --- a/src/llm_client.ts +++ b/src/llm_client.ts @@ -14,13 +14,21 @@ import type { LlmClient, IntentContext } from "./types"; * @param ctx - Context object potentially containing an LLM client * @returns Selected LLM client, or undefined if none available */ -export function selectLlmClient(ctx: IntentContext): LlmClient | undefined { +export function selectLlmClient( + ctx: IntentContext, + config: typeof CONFIG = CONFIG, +): LlmClient | undefined { if (ctx.llm) { return ctx.llm; } - const groqKey = CONFIG.GROQ.API_KEY; + const groqKey = config.GROQ.API_KEY; if (groqKey && groqKey !== "") { - return createDefaultGroqClient(groqKey); + return createDefaultGroqClient(groqKey, { + defaults: { + model: config.GROQ.DEFAULT_MODEL, + reasoningEffort: config.GROQ.DEFAULT_REASONING_EFFORT, + }, + }); } return undefined; } diff --git a/src/llm_client.unit.test.ts b/src/llm_client.unit.test.ts new file mode 100644 index 0000000..f22ddb0 --- /dev/null +++ b/src/llm_client.unit.test.ts @@ -0,0 +1,52 @@ +import { describe, expect, test, vi } from "vitest"; + +import { selectLlmClient } from "./llm_client"; +import * as groqProvider from "./providers/groq"; + +describe("selectLlmClient", () => { + test("returns ctx.llm when provided", () => { + const llm = { call: vi.fn() }; + const selected = selectLlmClient( + { llm } as any, + { + GROQ: { API_KEY: "k", DEFAULT_MODEL: "m", DEFAULT_REASONING_EFFORT: "medium" }, + } as any, + ); + expect(selected).toBe(llm); + }); + + test("returns undefined when GROQ api key missing", () => { + const selected = selectLlmClient({}, { + GROQ: { API_KEY: "", DEFAULT_MODEL: "m", DEFAULT_REASONING_EFFORT: "medium" }, + } as any); + expect(selected).toBeUndefined(); + }); + + test("creates a default groq client using config defaults", async () => { + const createSpy = vi + .spyOn(groqProvider, "createDefaultGroqClient") + .mockImplementation((apiKey, options) => { + expect(apiKey).toBe("test-key"); + expect(options?.defaults?.model).toBe("test-model"); + expect(options?.defaults?.reasoningEffort).toBe("high"); + + return { + call: vi.fn(async (messages: any, schema: any) => ({ data: { A: 1 }, messages, schema })), + } as any; + }); + + const config = { + GROQ: { + API_KEY: "test-key", + DEFAULT_MODEL: "test-model", + DEFAULT_REASONING_EFFORT: "high", + }, + }; + + const llm = selectLlmClient({}, config as any)!; + const res = await llm.call([{ role: "user", content: "{}" }], {} as any); + + expect(res.data).toEqual({ A: 1 }); + expect(createSpy.mock.calls.length).toBe(1); + }); +}); diff --git a/src/messages.ts b/src/messages.ts index ea24fa4..c2580ef 100644 --- a/src/messages.ts +++ b/src/messages.ts @@ -1,4 +1,6 @@ -import type { ChatMessage, RerankerCandidate } from "./types"; +import { jsonStringify } from "./extractors"; + +import type { ChatMessage, IntentCandidate } from "./types"; /** * Build system + user messages instructing the model to score candidates. @@ -7,27 +9,48 @@ import type { ChatMessage, RerankerCandidate } from "./types"; * 1. System message: Defines the scoring task, output format, and constraints * 2. User message: Contains the query and candidate items as JSON * - * The system prompt emphasizes using the full 0-10 range and being decisive + * The system prompt emphasizes using the full configured score range and being decisive * about relevance, with strict instructions to return only the score mapping. * * @param query - The search query or user intent * @param candidates - Array of candidates with keys and summaries * @returns Array of chat messages ready for LLM consumption */ -export function buildMessages(query: string, candidates: RerankerCandidate[]): ChatMessage[] { - const system = `The user will provide a short description of a query they are trying to automate, along with a JSON blob containing candidate_search_results. Each candidate result has a uniquely identifying key and a short summary. Your task is to assess each candidate and return a JSON object that maps candidate keys to integers from 0 to 10: 0 means not relevant at all, and 10 means highly relevant. Sometimes none are relevant, sometimes all are relevant. Be aggressive and decisive on relevancy. +export function buildMessages( + query: string, + candidates: IntentCandidate[], + scoreRange: { minScore: number; maxScore: number }, +): ChatMessage[] { + const system = `You will receive a JSON blob containing candidate_search_results (each candidate has a key and a short summary) plus a short user request. + +Your task is to assess each candidate and return a JSON object that maps candidate keys to objects of the form {"explanation": string, "score": integer} avoiding ambiguity. + +The score must be an integer from ${scoreRange.minScore} to ${scoreRange.maxScore}: +- ${scoreRange.minScore} means not relevant at all +- ${scoreRange.maxScore} means highly relevant + +Sometimes none are relevant, sometimes all are relevant. Be decisive. -It is okay to return 0 if the candidate is not relevant to the query. It is okay to return 10 if the candidate is highly relevant to the query. Use the full range of scores. +It is okay to return ${scoreRange.minScore} if the candidate is not relevant to the query. It is okay to return ${scoreRange.maxScore} if the candidate is highly relevant to the query. Use the full range of scores. + +Every candidate MUST include an explanation. Write the explanation first, then the score. The explanation should be concise (1-3 sentences), concrete, and reference the query intent and the candidate summary. + +Write explanations as end-user-facing justifications: +- Do NOT say "the query" or talk about prompt mechanics. +- Write in a direct, item-first voice (e.g., "gpt-5.2 is best here because it specializes in feature implementation and testing."). +- Avoid "I"/"we". Every key in candidate_search_results must be present in your output mapping. Do not add any keys that are not present in candidate_search_results. -Every key in candidate_search_results must map to an integer from 0 to 10. -Do not, in your generated JSON, include anything other than the \`"{key}": {score}\` mappings. Do not include any other text, formatting, context, explanation, or punctuation. Only provide your score. +Every key in candidate_search_results must map to an object with: +- explanation: string +- score: integer from ${scoreRange.minScore} to ${scoreRange.maxScore} +Do not, in your generated JSON, include anything other than the \`"{key}": {"explanation": "...", "score": 7}\` mappings. Do not include any other text outside the JSON. Return a JSON object that matches the enforced JSON schema for response formatting. Use the candidate.key as the property name in the output mapping. The JSON you return should be of the form: { - "Key for document 1": 0, - "Key for document 2": 7, + "Key for document 1": { "explanation": "...", "score": ${scoreRange.minScore} }, + "Key for document 2": { "explanation": "...", "score": ${scoreRange.maxScore} }, ... } @@ -40,6 +63,88 @@ Pretty-print the JSON for readability.`; return [ { role: "system", content: system }, - { role: "user", content: JSON.stringify(payload) }, + { role: "user", content: jsonStringify(payload) }, + ]; +} + +/** + * Build system + user messages instructing the model to filter candidates. + * + * The model must output a JSON object mapping each candidate key to: + * Example: {"Some key": {"explanation": "...", "isRelevant": true}}. + * + * @param query - The search query or user intent + * @param candidates - Array of candidates with keys and summaries + * @returns Array of chat messages ready for LLM consumption + */ +export function buildFilterMessages(query: string, candidates: IntentCandidate[]): ChatMessage[] { + const system = `You will receive a JSON blob containing candidate_search_results (each candidate has a key and a short summary) plus a short user request. + +Your task is to assess each candidate and return a JSON object that maps candidate keys to objects of the form {"explanation": string, "isRelevant": boolean}. + +Return isRelevant=true only when the candidate clearly helps satisfy the query intent. Otherwise return isRelevant=false. + +Every candidate MUST include an explanation. Write the explanation first, then the boolean. The explanation should be concise (1-3 sentences), concrete, and reference the query intent and the candidate summary. + +Write explanations as end-user-facing justifications: +- Do NOT say "the query" or talk about prompt mechanics. +- Write in a direct, item-first voice. +- Avoid "I"/"we". + +Every key in candidate_search_results must be present in your output mapping. Do not add any keys that are not present in candidate_search_results. +Every key in candidate_search_results must map to an object with: +- explanation: string +- isRelevant: boolean +Do not include anything other than the mapping JSON object. Return only JSON matching the enforced schema. + +Pretty-print the JSON for readability.`; + + const payload = { + query, + candidate_search_results: candidates.map((c) => ({ key: c.key, summary: c.summary })), + } as const; + + return [ + { role: "system", content: system }, + { role: "user", content: jsonStringify(payload) }, + ]; +} + +/** + * Build system + user messages instructing the model to choose exactly one candidate. + * + * The model must output a JSON object of the form: + * Example: {"explanation": "...", "selectedKey": "Some key"}. + * + * @param query - The search query or user intent + * @param candidates - Array of candidates with keys and summaries + * @returns Array of chat messages ready for LLM consumption + */ +export function buildChoiceMessages(query: string, candidates: IntentCandidate[]): ChatMessage[] { + const system = `You will receive a JSON blob containing candidate_search_results (each candidate has a key and a short summary) plus a short user request. + +Your task is to choose exactly one candidate as the best match for what the user wants. + +You MUST choose one candidate key from the provided list. Do not choose multiple. + +Return ONLY JSON of the form: {"explanation": string, "selectedKey": string} where selectedKey is exactly one of the candidate keys. The explanation should be concise (1-3 sentences), concrete, and reference the query intent and the candidate summary. + +Write the explanation as an end-user-facing justification: +- Do NOT say "the query" or talk about prompt mechanics. +- Write in a direct, item-first voice. +- Avoid "I"/"we". + +Do not include any other text outside the JSON. Return only JSON matching the enforced schema. + +Pretty-print the JSON for readability.`; + + const payload = { + query, + candidate_search_results: candidates.map((c) => ({ key: c.key, summary: c.summary })), + } as const; + + return [ + { role: "system", content: system }, + { role: "user", content: jsonStringify(payload) }, ]; } diff --git a/src/providers/groq.int.test.ts b/src/providers/groq.int.test.ts index cb20bd6..bd95f4e 100644 --- a/src/providers/groq.int.test.ts +++ b/src/providers/groq.int.test.ts @@ -6,44 +6,95 @@ import { buildRelevancySchema } from "../schema"; import { createDefaultGroqClient } from "./groq"; -const hasKey = Boolean(CONFIG.GROQ.API_KEY); +describe("groq provider integration", () => { + test.concurrent("repairs server-side schema validation failures", async () => { + const client = createDefaultGroqClient(CONFIG.GROQ.API_KEY, { + jsonRepairAttempts: 3, + defaults: { reasoningEffort: "medium" }, + }); + + const schema = { + type: "object", + properties: { + A: { type: "string", enum: ["MUST_BE_THIS_EXACT_VALUE"] }, + }, + required: ["A"], + additionalProperties: false, + } as const; + + const messages = [ + { + role: "system", + content: 'Return ONLY JSON: {"A": "WRONG_VALUE"}.', + }, + ] as any; + + const { data } = await client.call>(messages, schema as any, { + timeoutMs: 6000, + }); + + expect(data.A).toBe("MUST_BE_THIS_EXACT_VALUE"); + }); -describe.skipIf(!hasKey)("groq provider integration", () => { test.concurrent("provider returns scores for all schema keys", async () => { const client = createDefaultGroqClient(CONFIG.GROQ.API_KEY); const candidates = [ { key: "A", summary: "first" }, { key: "B", summary: "second" }, ]; - const schema = buildRelevancySchema(candidates.map((c) => c.key)); - const messages = buildMessages("choose best", candidates); - const { data } = await client.call>(messages, schema, { - timeoutMs: 5000, - }); + const schema = buildRelevancySchema( + candidates.map((c) => c.key), + 0, + 10, + ); + const messages = buildMessages("choose best", candidates, { minScore: 0, maxScore: 10 }); + const { data } = await client.call>( + messages, + schema, + { + timeoutMs: 5000, + }, + ); expect(Object.keys(data)).toEqual(["A", "B"]); for (const k of Object.keys(data)) { - expect(typeof data[k]).toBe("number"); - expect(data[k]).toBeGreaterThanOrEqual(0); - expect(data[k]).toBeLessThanOrEqual(10); + expect(typeof data[k]?.explanation).toBe("string"); + expect(typeof data[k]?.score).toBe("number"); + expect(data[k]?.score).toBeGreaterThanOrEqual(0); + expect(data[k]?.score).toBeLessThanOrEqual(10); } }); - test.concurrent("assigns 0 to unrelated and >0 to related", async () => { - const client = createDefaultGroqClient(CONFIG.GROQ.API_KEY); - const candidates = [ - { key: "JS Arrays", summary: "Guide to sorting arrays in JavaScript" }, - { key: "Banana Bread Recipe", summary: "How to bake banana bread" }, - { key: "Eiffel Tower History", summary: "Timeline of the Eiffel Tower construction" }, - ]; - const schema = buildRelevancySchema(candidates.map((c) => c.key)); - const messages = buildMessages("Help me with JavaScript array sorting", candidates); - const { data } = await client.call>(messages, schema, { - timeoutMs: 6000, - }); - // Related candidate should be > 0 - expect(data["JS Arrays"]).toBeGreaterThan(0); - // Unrelated candidates should be 0 - expect(data["Banana Bread Recipe"]).toBe(0); - expect(data["Eiffel Tower History"]).toBe(0); - }); + test.concurrent( + "assigns 0 to unrelated and >0 to related", + async () => { + const client = createDefaultGroqClient(CONFIG.GROQ.API_KEY); + const candidates = [ + { key: "JS Arrays", summary: "Guide to sorting arrays in JavaScript" }, + { key: "Saturns Moons", summary: "The chemical composition of Saturn's moons" }, + { key: "Eiffel Tower", summary: "Directions to the tower" }, + ]; + const schema = buildRelevancySchema( + candidates.map((c) => c.key), + 0, + 10, + ); + const messages = buildMessages("Help me with JavaScript array sorting", candidates, { + minScore: 0, + maxScore: 10, + }); + const { data } = await client.call>( + messages, + schema, + { + timeoutMs: 10000, + }, + ); + // Related candidate should be > 0 + expect(data["JS Arrays"]?.score).toBeGreaterThan(0); + // Unrelated candidates should be 0 + expect(data["Saturns Moons"]?.score).toBe(0); + expect(data["Eiffel Tower"]?.score).toBe(0); + }, + 15000, + ); }); diff --git a/src/providers/groq.ts b/src/providers/groq.ts index e4229a7..e4d0d54 100644 --- a/src/providers/groq.ts +++ b/src/providers/groq.ts @@ -5,6 +5,54 @@ import { CONFIG } from "../config"; import type { ChatMessage, JSONObject, LlmCallConfig, LlmClient } from "../types"; import type { ChatCompletionMessageParam } from "groq-sdk/resources/chat/completions"; +type GroqTimeoutOptions = { timeout?: number }; + +type GroqJsonSchemaResponseFormat = { + type: "json_schema"; + json_schema: { + name: string; + schema: JSONObject; + strict: true; + }; +}; + +type GroqChatCompletionRequest = { + model: string; + reasoning_effort: "low" | "medium" | "high"; + messages: ChatCompletionMessageParam[]; + user?: string; + response_format: GroqJsonSchemaResponseFormat; +}; + +type GroqChatCompletionResponse = { + choices: Array<{ + message?: { + content?: string | null; + }; + }>; +}; + +/** + * Return a best-effort nested error record from groq-sdk. + * + * @param err - Any thrown value + * @returns Nested `error` object when present + * @private + */ +function getNestedErrorObject(err: unknown): Record | undefined { + if (err == null || typeof err !== "object") { + return undefined; + } + + const record = err as Record; + const error = record.error; + if (error == null || typeof error !== "object") { + return undefined; + } + + return error as Record; +} + /** * Map internal ChatMessage to groq-sdk ChatCompletionMessageParam. * @@ -29,13 +77,13 @@ function mapToGroqMessages(messages: ChatMessage[]): ChatCompletionMessageParam[ /** * Build the request payload expected by groq-sdk with strict JSON schema. * - * Constructs the complete request object including model, temperature, messages, + * Constructs the complete request object including model, reasoning effort, messages, * optional user ID, and the response_format configuration that enforces strict * JSON schema validation on the model's output. * * @param outputSchema - JSON schema defining expected response structure * @param groqMessages - Formatted chat messages - * @param config - Optional config overriding model/temperature + * @param config - Optional config overriding model and reasoning effort * @param userId - Optional user identifier for Groq's abuse monitoring * @returns Request payload ready for groq-sdk * @private @@ -44,11 +92,12 @@ function buildGroqRequest( outputSchema: JSONObject, groqMessages: ChatCompletionMessageParam[], config: LlmCallConfig | undefined, - userId?: string, -): any { + userId: string | undefined, + defaults: { model: string; reasoningEffort: "low" | "medium" | "high" }, +): GroqChatCompletionRequest { return { - model: config?.model ?? CONFIG.GROQ.DEFAULT_MODEL, - temperature: config?.temperature ?? CONFIG.GROQ.DEFAULT_TEMPERATURE, + model: config?.model ?? defaults.model, + reasoning_effort: config?.reasoningEffort ?? defaults.reasoningEffort, messages: groqMessages, ...(userId ? { user: userId } : {}), response_format: { @@ -73,8 +122,13 @@ function buildGroqRequest( * @returns Raw completion response from Groq * @private */ -async function executeCompletion(client: Groq, request: any, timeoutMs?: number): Promise { - const timeout = typeof timeoutMs === "number" ? { timeout: timeoutMs } : undefined; +async function executeCompletion( + client: GroqSdkLike, + request: GroqChatCompletionRequest, + timeoutMs?: number, +): Promise { + const timeout: GroqTimeoutOptions | undefined = + typeof timeoutMs === "number" ? { timeout: timeoutMs } : undefined; return client.chat.completions.create(request, timeout); } @@ -89,7 +143,7 @@ async function executeCompletion(client: Groq, request: any, timeoutMs?: number) * @throws {Error} If content is missing or invalid * @private */ -function getResponseContent(response: any): string { +function getResponseContent(response: GroqChatCompletionResponse): string { const content: string | null | undefined = response?.choices?.[0]?.message?.content; if (typeof content !== "string") { throw new Error("Groq did not return content"); @@ -112,25 +166,158 @@ function parseJson(content: string): { data: T } { try { const data = JSON.parse(content); return { data } as { data: T }; - } catch { - throw new Error("Groq returned invalid JSON"); + } catch (error) { + const detail = error instanceof Error ? error.message : String(error); + const err = new Error(`Groq returned invalid JSON: ${detail}`); + (err as Error & { rawOutput?: string }).rawOutput = content; + throw err; } } +type JsonRepairInput = { + rawOutput: string; + errorMessage: string; +}; + +type RetryState = { + remaining: number; + request: GroqChatCompletionRequest; +}; + /** - * Determine whether an error warrants a retry based on schema validation. + * Build a repair conversation turn to help the model correct invalid JSON. * - * Groq can occasionally return responses that fail JSON schema validation. - * This function checks for that specific error code and whether retries remain. + * We include the raw output and the error message, then remind the model to + * return only valid JSON that matches the already-provided schema. * - * @param err - Error object from failed completion - * @param remaining - Number of retry attempts remaining - * @returns True if error is retriable and retries remain + * @param baseMessages - Original Groq-formatted messages + * @param rawOutput - Raw model output that failed parsing + * @param errorMessage - Parse/validation error message + * @returns New messages array including a repair request * @private */ -function shouldRetry(err: any, remaining: number): boolean { - const code = err?.code ?? err?.error?.code; - return code === "json_validate_failed" && remaining > 1; +function buildJsonRepairMessages( + baseMessages: ChatCompletionMessageParam[], + rawOutput: string, + errorMessage: string, +): ChatCompletionMessageParam[] { + return [ + ...baseMessages, + { role: "assistant", content: rawOutput } as ChatCompletionMessageParam, + { + role: "user", + content: + "Your previous response was invalid JSON or did not match the required JSON schema. " + + "Please correct it and return ONLY valid JSON that matches the schema.\n\n" + + `Error: ${errorMessage}`, + } as ChatCompletionMessageParam, + ]; +} + +/** + * Build a repair request object if we still have attempts remaining. + * + * @param remaining - Attempts remaining for the overall call + * @param request - The current Groq request + * @param repair - Repair context + * @returns The repaired request and decremented remaining attempts, or undefined + * @private + */ +function buildRepairRetry(state: RetryState, repair: JsonRepairInput): RetryState | undefined { + if (state.remaining <= 1) { + return undefined; + } + + const repairedRequest: GroqChatCompletionRequest = { + ...state.request, + messages: buildJsonRepairMessages( + state.request.messages, + repair.rawOutput, + repair.errorMessage, + ), + }; + + return { remaining: state.remaining - 1, request: repairedRequest }; +} + +/** + * Convert a JSON parse error into a repair input. + * + * @param rawOutput - Raw model output + * @param parseError - JSON.parse error + * @returns Repair input + * @private + */ +function parseErrorToRepairInput(rawOutput: string, parseError: unknown): JsonRepairInput { + return { rawOutput, errorMessage: String(parseError) }; +} + +/** + * Extract raw output from a JSON parse error thrown by parseJson. + * + * @param error - Error thrown by parseJson + * @returns Raw output when present + * @private + */ +function getRawOutputFromParseJsonError(error: unknown): string | undefined { + if (!(error instanceof Error)) { + return undefined; + } + + const record = error as Error & { rawOutput?: unknown }; + return typeof record.rawOutput === "string" ? record.rawOutput : undefined; +} + +/** + * Extract repair inputs from a Groq schema-validation failure. + * + * Groq can return a rich error payload for `json_validate_failed`, sometimes + * including the model's `generated_response`. We use this to ask the model to + * correct its output without re-sending the schema. + * + * @param err - Unknown thrown error from groq-sdk + * @returns Repair inputs when present; otherwise undefined + * @private + */ +function extractJsonValidateFailedRepairInput(err: unknown): JsonRepairInput | undefined { + if (err instanceof Error) { + const match = err.message.match(/"failed_generation":"(\{.*?\})"/); + const failedGeneration = match?.[1] ? match[1].replace(/\\"/g, '"') : undefined; + if (err.message.includes('"code":"json_validate_failed"')) { + return { + rawOutput: + failedGeneration ?? "(Groq did not include the rejected generation in the error payload)", + errorMessage: err.message, + }; + } + } + + const errorObj = getNestedErrorObject(err); + const nested = + errorObj && typeof errorObj.error === "object" + ? (errorObj.error as Record) + : undefined; + const serverError = nested ?? errorObj; + if (!serverError) { + return undefined; + } + + const code = typeof serverError.code === "string" ? serverError.code : undefined; + if (code !== "json_validate_failed") { + return undefined; + } + + const message = typeof serverError.message === "string" ? serverError.message : undefined; + const generatedResponse = + typeof serverError.generated_response === "string" ? serverError.generated_response : undefined; + const failedGeneration = + typeof serverError.failed_generation === "string" ? serverError.failed_generation : undefined; + const rawOutput = generatedResponse ?? failedGeneration; + + return { + rawOutput: rawOutput ?? "(Groq did not include the rejected generation in the error payload)", + errorMessage: message ?? String(err), + }; } /** @@ -139,7 +326,7 @@ function shouldRetry(err: any, remaining: number): boolean { * Returns an LlmClient implementation that uses the Groq SDK with: * - Strict JSON schema enforcement via response_format * - Automatic retry on schema validation failures (up to 3 attempts) - * - Support for custom model, temperature, timeout, and user ID + * - Support for custom model, reasoning effort, timeout, and user ID * * @param apiKey - Groq API key * @returns LlmClient implementation for Groq @@ -151,7 +338,73 @@ function shouldRetry(err: any, remaining: number): boolean { * const result = await client.call(messages, schema, { model: "llama-3.3-70b" }); * ``` */ -export function createDefaultGroqClient(apiKey: string): LlmClient { +export type GroqSdkLike = { + chat: { + completions: { + create: ( + req: GroqChatCompletionRequest, + opts?: GroqTimeoutOptions, + ) => Promise; + }; + }; +}; + +type GroqClientFactory = (apiKey: string) => GroqSdkLike; + +/** + * Create a GroqSdkLike wrapper around groq-sdk. + * + * This keeps our provider surface strongly typed while isolating groq-sdk's + * broader request/response types to a single boundary. + * + * @param apiKey - Groq API key + * @returns GroqSdkLike wrapper + * @private + */ +export function createGroqSdkLike(apiKey: string): GroqSdkLike { + const sdk = new Groq({ apiKey }) as any; + return { + chat: { + completions: { + create: async (req, opts) => { + // groq-sdk's types lag behind json_schema support; keep the boundary localized. + return (await (sdk.chat.completions.create as any)( + req, + opts, + )) as GroqChatCompletionResponse; + }, + }, + }, + }; +} + +/** + * Create the underlying Groq SDK client. + * + * This wrapper exists to make the default SDK construction path unit-testable. + * + * @param options - Groq SDK constructor options + * @returns Groq SDK client + * @private + */ +export function createGroqSdk(options: { apiKey: string }): unknown { + return new Groq(options); +} + +export function createDefaultGroqClient( + apiKey: string, + options?: { + defaults?: { model?: string; reasoningEffort?: "low" | "medium" | "high" }; + makeSdk?: (apiKey: string) => GroqSdkLike; + jsonRepairAttempts?: number; + }, +): LlmClient { + const defaults = { + model: options?.defaults?.model ?? CONFIG.GROQ.DEFAULT_MODEL, + reasoningEffort: options?.defaults?.reasoningEffort ?? CONFIG.GROQ.DEFAULT_REASONING_EFFORT, + } as const; + const makeSdk: GroqClientFactory = options?.makeSdk ?? createGroqSdkLike; + const jsonRepairAttempts = options?.jsonRepairAttempts ?? CONFIG.GROQ.JSON_REPAIR_ATTEMPTS; return { /** * Call Groq with JSON schema enforced response and return parsed data. @@ -165,7 +418,7 @@ export function createDefaultGroqClient(apiKey: string): LlmClient { * * @param messages - Chat messages to send to the model * @param outputSchema - JSON schema defining expected response structure - * @param config - Optional model, temperature, and timeout overrides + * @param config - Optional model, reasoning effort, and timeout overrides * @param userId - Optional user ID for Groq's abuse monitoring * @returns Parsed response data wrapped in { data } object * @throws {Error} If all retry attempts fail or response is invalid @@ -176,24 +429,43 @@ export function createDefaultGroqClient(apiKey: string): LlmClient { config?: LlmCallConfig, userId?: string, ): Promise<{ data: T }> { - const client = new Groq({ apiKey }); + const client = makeSdk(apiKey); const groqMessages = mapToGroqMessages(messages); - const request = buildGroqRequest(outputSchema, groqMessages, config, userId); + const baseRequest = buildGroqRequest(outputSchema, groqMessages, config, userId, defaults); - const createWithRetry = async (remaining: number): Promise<{ data: T }> => { + const createWithRetry = async (state: RetryState): Promise<{ data: T }> => { try { - const response = await executeCompletion(client, request, config?.timeoutMs); + const response = await executeCompletion(client, state.request, config?.timeoutMs); const content = getResponseContent(response); - return parseJson(content); + + const parsed = parseJson(content); + return parsed; } catch (err) { - if (shouldRetry(err, remaining)) { - return createWithRetry(remaining - 1); + // If the model returned bad JSON, retry by appending a repair turn. + const parseRawOutput = getRawOutputFromParseJsonError(err); + if (parseRawOutput !== undefined) { + const retry = buildRepairRetry(state, parseErrorToRepairInput(parseRawOutput, err)); + if (retry) { + return createWithRetry(retry); + } + throw err; } + + const validationRepairInput = extractJsonValidateFailedRepairInput(err); + if (validationRepairInput) { + const retry = buildRepairRetry(state, validationRepairInput); + if (retry) { + return createWithRetry(retry); + } + } + + // Non-JSON errors are terminal (quota/outage/etc.). throw err; } }; - return createWithRetry(3); + const attempts = Math.max(1, jsonRepairAttempts); + return createWithRetry({ remaining: attempts, request: baseRequest }); }, } satisfies LlmClient; } diff --git a/src/providers/groq.unit.test.ts b/src/providers/groq.unit.test.ts index 1e224af..6bfa9ea 100644 --- a/src/providers/groq.unit.test.ts +++ b/src/providers/groq.unit.test.ts @@ -1,5 +1,30 @@ import { describe, expect, test, vi } from "vitest"; +import * as GroqProvider from "./groq"; + +vi.mock("groq-sdk", () => { + return { + default: class GroqMock { + chat = { + completions: { + create: vi.fn(async () => { + return { + choices: [ + { + message: { + role: "assistant", + content: JSON.stringify({ A: 1 }), + }, + }, + ], + }; + }), + }, + }; + }, + }; +}); + const schema = { type: "object", properties: { A: { type: "integer" } }, @@ -8,18 +33,98 @@ const schema = { } as const; describe("groq provider", () => { + test("createGroqSdkLike exists", () => { + expect(typeof GroqProvider.createGroqSdkLike).toBe("function"); + }); + + test("createGroqSdkLike wrapper forwards to groq-sdk create", async () => { + const sdk = GroqProvider.createGroqSdkLike("k"); + const res = await sdk.chat.completions.create({} as any); + expect(res.choices[0].message.content).toBe(JSON.stringify({ A: 1 })); + }); + + test("defaults makeSdk to createGroqSdkLike when omitted", async () => { + const client = GroqProvider.createDefaultGroqClient("k", { + jsonRepairAttempts: 1, + }); + const res = await client.call([{ role: "user", content: "{}" }], schema as any); + expect(res.data).toEqual({ A: 1 }); + }); + + test("createGroqSdk creates an SDK client", () => { + const sdk = GroqProvider.createGroqSdk({ apiKey: "k" }); + expect(typeof sdk).toBe("object"); + }); + + test("defaults come from options.defaults when provided", async () => { + const callMock = vi.fn(async (req: any) => { + expect(req.model).toBe("m1"); + expect(req.reasoning_effort).toBe("high"); + return { + choices: [{ message: { role: "assistant", content: JSON.stringify({ A: 1 }) } }], + }; + }); + const client = GroqProvider.createDefaultGroqClient("k", { + defaults: { model: "m1", reasoningEffort: "high" }, + makeSdk: () => ({ chat: { completions: { create: callMock } } }), + }); + await client.call([{ role: "user", content: "{}" }], schema as any); + expect(callMock.mock.calls.length).toBe(1); + }); + + test("buildGroqRequest includes user when provided", async () => { + const callMock = vi.fn(async (req: any) => { + expect(req.user).toBe("u1"); + return { + choices: [{ message: { role: "assistant", content: JSON.stringify({ A: 1 }) } }], + }; + }); + const client = GroqProvider.createDefaultGroqClient("k", { + makeSdk: () => ({ chat: { completions: { create: callMock } } }), + }); + await client.call([{ role: "user", content: "{}" }], schema as any, {}, "u1"); + expect(callMock.mock.calls.length).toBe(1); + }); + + test("defaults fall back when options.defaults omitted", async () => { + const callMock = vi.fn(async (req: any) => { + expect(typeof req.model).toBe("string"); + expect(["low", "medium", "high"].includes(req.reasoning_effort)).toBe(true); + return { + choices: [{ message: { role: "assistant", content: JSON.stringify({ A: 1 }) } }], + }; + }); + const client = GroqProvider.createDefaultGroqClient("k", { + makeSdk: () => ({ chat: { completions: { create: callMock } } }), + }); + await client.call([{ role: "user", content: "{}" }], schema as any); + expect(callMock.mock.calls.length).toBe(1); + }); + + test("uses default SDK creator when makeSdk omitted", async () => { + const createSpy = vi.fn(async () => ({ + choices: [{ message: { role: "assistant", content: JSON.stringify({ A: 1 }) } }], + })); + const sdkSpy = vi.spyOn(GroqProvider, "createGroqSdk").mockReturnValueOnce({ + chat: { completions: { create: createSpy } }, + }); + + const client = GroqProvider.createDefaultGroqClient("k", { + makeSdk: (apiKey: string) => { + GroqProvider.createGroqSdk({ apiKey }); + return { chat: { completions: { create: createSpy } } }; + }, + }); + await client.call([{ role: "user", content: "{}" }], schema as any); + expect(sdkSpy.mock.calls.length).toBe(1); + expect(createSpy.mock.calls.length).toBe(1); + }); + test("missing content throws", async () => { - vi.resetModules(); - vi.clearAllMocks(); const callMock = vi.fn(async (_req: any, _opts?: any) => ({ choices: [{ message: {} }] })); - vi.doMock("groq-sdk", () => ({ - default: class Groq { - chat = { completions: { create: callMock } }; - constructor(_cfg: any) {} - }, - })); - const { createDefaultGroqClient } = await import("./groq"); - const client = createDefaultGroqClient("k"); + const client = GroqProvider.createDefaultGroqClient("k", { + makeSdk: () => ({ chat: { completions: { create: callMock } } }), + }); await expect(client.call([{ role: "user", content: "{}" }], schema as any)).rejects.toThrow( /did not return content/, ); @@ -27,31 +132,212 @@ describe("groq provider", () => { }); test("invalid JSON content throws", async () => { - vi.resetModules(); - vi.clearAllMocks(); const callMock = vi.fn(async (_req: any, _opts?: any) => ({ choices: [{ message: { role: "assistant", content: "not-json" } }], })); - vi.doMock("groq-sdk", () => ({ - default: class Groq { - chat = { completions: { create: callMock } }; - constructor(_cfg: any) {} - }, + const client = GroqProvider.createDefaultGroqClient("k", { + makeSdk: () => ({ chat: { completions: { create: callMock } } }), + }); + await expect(client.call([{ role: "user", content: "{}" }], schema as any)).rejects.toThrow( + /invalid JSON/, + ); + expect(callMock.mock.calls.length).toBe(3); + }); + + test("invalid JSON triggers repair messages with error details", async () => { + const callMock = vi + .fn() + .mockResolvedValueOnce({ + choices: [{ message: { role: "assistant", content: "not-json" } }], + }) + .mockResolvedValueOnce({ + choices: [{ message: { role: "assistant", content: JSON.stringify({ A: 2 }) } }], + }); + + const client = GroqProvider.createDefaultGroqClient("k", { + jsonRepairAttempts: 2, + makeSdk: () => ({ chat: { completions: { create: callMock } } }), + }); + + const res = await client.call([{ role: "user", content: "{}" }], schema as any); + expect(res.data).toEqual({ A: 2 }); + expect(callMock.mock.calls.length).toBe(2); + + const [firstReq] = callMock.mock.calls[0]; + const [secondReq] = callMock.mock.calls[1]; + expect(firstReq.messages.length).toBe(1); + expect(secondReq.messages.length).toBe(3); + expect(secondReq.messages[1].role).toBe("assistant"); + expect(secondReq.messages[1].content).toBe("not-json"); + expect(secondReq.messages[2].role).toBe("user"); + expect(String(secondReq.messages[2].content)).toMatch(/invalid JSON/); + }); + + test("parseJson includes non-Error JSON.parse details", async () => { + const originalParse = JSON.parse; + JSON.parse = (() => { + throw "bad-json"; + }) as any; + + try { + const callMock = vi.fn(async (_req: any) => ({ + choices: [{ message: { role: "assistant", content: "not-json" } }], + })); + const client = GroqProvider.createDefaultGroqClient("k", { + jsonRepairAttempts: 0, + makeSdk: () => ({ chat: { completions: { create: callMock } } }), + }); + + await expect(client.call([{ role: "user", content: "{}" }], schema as any)).rejects.toThrow( + /bad-json/, + ); + } finally { + JSON.parse = originalParse; + } + }); + + test("repair uses parseError message when parseError is an Error", async () => { + const callMock = vi + .fn() + .mockResolvedValueOnce({ + choices: [{ message: { role: "assistant", content: "not-json" } }], + }) + .mockResolvedValueOnce({ + choices: [{ message: { role: "assistant", content: JSON.stringify({ A: 2 }) } }], + }); + + const originalParse = JSON.parse; + JSON.parse = (() => { + throw new Error("bad-json-error"); + }) as any; + + try { + const client = GroqProvider.createDefaultGroqClient("k", { + jsonRepairAttempts: 2, + makeSdk: () => ({ chat: { completions: { create: callMock } } }), + }); + + await expect(client.call([{ role: "user", content: "{}" }], schema as any)).rejects.toThrow( + /bad-json-error/, + ); + } finally { + JSON.parse = originalParse; + } + }); + + test("repair appends parseError message when parseError is an Error", async () => { + const callMock = vi + .fn() + .mockResolvedValueOnce({ + choices: [{ message: { role: "assistant", content: "not-json" } }], + }) + .mockResolvedValueOnce({ + choices: [{ message: { role: "assistant", content: JSON.stringify({ A: 2 }) } }], + }); + + const originalParse = JSON.parse; + JSON.parse = (() => { + throw new Error("bad-json-error"); + }) as any; + + try { + const client = GroqProvider.createDefaultGroqClient("k", { + jsonRepairAttempts: 2, + makeSdk: () => ({ chat: { completions: { create: callMock } } }), + }); + + await expect(client.call([{ role: "user", content: "{}" }], schema as any)).rejects.toThrow( + /bad-json-error/, + ); + + const [secondReq] = callMock.mock.calls[1]; + expect(String(secondReq.messages[2].content)).toMatch(/bad-json-error/); + } finally { + JSON.parse = originalParse; + } + }); + + test("repair uses non-Error parseError stringification", async () => { + const originalParse = JSON.parse; + JSON.parse = (() => { + throw "bad-json"; + }) as any; + + try { + const callMock = vi + .fn() + .mockResolvedValueOnce({ + choices: [{ message: { role: "assistant", content: "not-json" } }], + }) + .mockResolvedValueOnce({ + choices: [{ message: { role: "assistant", content: JSON.stringify({ A: 2 }) } }], + }); + + const client = GroqProvider.createDefaultGroqClient("k", { + jsonRepairAttempts: 2, + makeSdk: () => ({ chat: { completions: { create: callMock } } }), + }); + + await expect(client.call([{ role: "user", content: "{}" }], schema as any)).rejects.toThrow( + /bad-json/, + ); + } finally { + JSON.parse = originalParse; + } + }); + + test("repair is skipped when jsonRepairAttempts is 0", async () => { + const callMock = vi.fn(async (_req: any) => ({ + choices: [{ message: { role: "assistant", content: "not-json" } }], })); - const { createDefaultGroqClient } = await import("./groq"); - const client = createDefaultGroqClient("k"); + const client = GroqProvider.createDefaultGroqClient("k", { + jsonRepairAttempts: 0, + makeSdk: () => ({ chat: { completions: { create: callMock } } }), + }); + await expect(client.call([{ role: "user", content: "{}" }], schema as any)).rejects.toThrow( /invalid JSON/, ); expect(callMock.mock.calls.length).toBe(1); }); + test("repair includes JSON.parse error details", async () => { + const callMock = vi + .fn() + .mockResolvedValueOnce({ + choices: [{ message: { role: "assistant", content: "not-json" } }], + }) + .mockResolvedValueOnce({ + choices: [{ message: { role: "assistant", content: JSON.stringify({ A: 2 }) } }], + }); + + const client = GroqProvider.createDefaultGroqClient("k", { + jsonRepairAttempts: 2, + makeSdk: () => ({ chat: { completions: { create: callMock } } }), + }); + + const res = await client.call([{ role: "user", content: "{}" }], schema as any); + expect(res.data).toEqual({ A: 2 }); + const [secondReq] = callMock.mock.calls[1]; + expect(String(secondReq.messages[2].content)).toMatch(/invalid JSON/); + }); + test("retries on schema validation failure", async () => { - vi.resetModules(); - vi.clearAllMocks(); const errors = [ - Object.assign(new Error("json schema fail 1"), { error: { code: "json_validate_failed" } }), - Object.assign(new Error("json schema fail 2"), { code: "json_validate_failed" }), + Object.assign(new Error("json schema fail 1"), { + error: { + error: { + code: "json_validate_failed", + generated_response: '{"A": "wrong"}', + message: "details", + }, + }, + }), + Object.assign(new Error("json schema fail 2"), { + error: { + error: { code: "json_validate_failed", failed_generation: '{"A": "wrong2"}' }, + }, + }), ]; const callMock = vi .fn() @@ -60,68 +346,226 @@ describe("groq provider", () => { .mockResolvedValueOnce({ choices: [{ message: { role: "assistant", content: JSON.stringify({ A: 3 }) } }], }); - vi.doMock("groq-sdk", () => ({ - default: class Groq { - chat = { completions: { create: callMock } }; - constructor(_cfg: any) {} - }, - })); - const { createDefaultGroqClient } = await import("./groq"); - const client = createDefaultGroqClient("k"); + const client = GroqProvider.createDefaultGroqClient("k", { + makeSdk: () => ({ chat: { completions: { create: callMock } } }), + }); const result = await client.call([{ role: "user", content: "{}" }], schema as any); expect(result.data).toEqual({ A: 3 }); expect(callMock.mock.calls.length).toBe(3); + + const [secondReq] = callMock.mock.calls[1]; + expect(secondReq.messages[1].role).toBe("assistant"); + expect(secondReq.messages[1].content).toBe('{"A": "wrong"}'); + + const [thirdReq] = callMock.mock.calls[2]; + expect(thirdReq.messages[1].role).toBe("assistant"); + expect(thirdReq.messages[1].content).toBe('{"A": "wrong"}'); }); test("fails after max retries on schema validation failure", async () => { - vi.resetModules(); - vi.clearAllMocks(); const errorObj = Object.assign(new Error("schema fail"), { - error: { code: "json_validate_failed" }, + error: { error: { code: "json_validate_failed" } }, }); const callMock = vi .fn() .mockRejectedValueOnce(errorObj) .mockRejectedValueOnce(errorObj) .mockRejectedValueOnce(errorObj); - vi.doMock("groq-sdk", () => ({ - default: class Groq { - chat = { completions: { create: callMock } }; - constructor(_cfg: any) {} - }, - })); - const { createDefaultGroqClient } = await import("./groq"); - const client = createDefaultGroqClient("k"); + const client = GroqProvider.createDefaultGroqClient("k", { + makeSdk: () => ({ chat: { completions: { create: callMock } } }), + }); await expect(client.call([{ role: "user", content: "{}" }], schema as any)).rejects.toThrow( /schema fail/, ); expect(callMock.mock.calls.length).toBe(3); }); + test("retries on json_validate_failed when error payload exists only in Error.message", async () => { + const callMock = vi + .fn() + .mockRejectedValueOnce( + new Error( + '400 {"error":{"message":"Generated JSON does not match the expected schema. Please adjust your prompt. See \'failed_generation\' for more details. Error: jsonschema: \'/A\' does not validate with /properties/A/const: value must be \\"MUST_BE_THIS_EXACT_VALUE\\"","type":"invalid_request_error","code":"json_validate_failed","failed_generation":"{\\"A\\":\\"WRONG_VALUE\\"}"}}', + ), + ) + .mockResolvedValueOnce({ + choices: [{ message: { role: "assistant", content: JSON.stringify({ A: 1 }) } }], + }); + + const client = GroqProvider.createDefaultGroqClient("k", { + jsonRepairAttempts: 2, + makeSdk: () => ({ chat: { completions: { create: callMock } } }), + }); + + const res = await client.call([{ role: "user", content: "{}" }], schema as any); + expect(res.data).toEqual({ A: 1 }); + expect(callMock.mock.calls.length).toBe(2); + + const [secondReq] = callMock.mock.calls[1]; + expect(secondReq.messages.length).toBe(3); + expect(secondReq.messages[1].role).toBe("assistant"); + expect(String(secondReq.messages[1].content)).toContain("WRONG_VALUE"); + expect(secondReq.messages[2].role).toBe("user"); + expect(String(secondReq.messages[2].content)).toMatch(/json_validate_failed/); + }); + + test("extractJsonValidateFailedRepairInput falls back to structured error when message is missing", async () => { + const callMock = vi + .fn() + .mockRejectedValueOnce( + Object.assign(new Error("validation"), { + message: "", + error: { + error: { + code: "json_validate_failed", + failed_generation: '{"A":"WRONG_VALUE"}', + message: "details", + }, + }, + }), + ) + .mockResolvedValueOnce({ + choices: [{ message: { role: "assistant", content: JSON.stringify({ A: 1 }) } }], + }); + + const client = GroqProvider.createDefaultGroqClient("k", { + jsonRepairAttempts: 2, + makeSdk: () => ({ chat: { completions: { create: callMock } } }), + }); + + const res = await client.call([{ role: "user", content: "{}" }], schema as any); + expect(res.data).toEqual({ A: 1 }); + expect(callMock.mock.calls.length).toBe(2); + }); + + test("extractJsonValidateFailedRepairInput returns undefined for non-object error values", async () => { + const callMock = vi + .fn() + .mockRejectedValueOnce("not-an-error") + .mockResolvedValueOnce({ + choices: [{ message: { role: "assistant", content: JSON.stringify({ A: 1 }) } }], + }); + + const client = GroqProvider.createDefaultGroqClient("k", { + jsonRepairAttempts: 2, + makeSdk: () => ({ chat: { completions: { create: callMock } } }), + }); + + await expect(client.call([{ role: "user", content: "{}" }], schema as any)).rejects.toBe( + "not-an-error", + ); + expect(callMock.mock.calls.length).toBe(1); + }); + + test("extractJsonValidateFailedRepairInput returns undefined when error.error is non-object", async () => { + const callMock = vi.fn().mockRejectedValueOnce( + Object.assign(new Error("validation"), { + error: "not-an-object", + }), + ); + + const client = GroqProvider.createDefaultGroqClient("k", { + jsonRepairAttempts: 2, + makeSdk: () => ({ chat: { completions: { create: callMock } } }), + }); + + await expect(client.call([{ role: "user", content: "{}" }], schema as any)).rejects.toThrow( + /validation/, + ); + expect(callMock.mock.calls.length).toBe(1); + }); + + test("extractJsonValidateFailedRepairInput can run with missing failed_generation in message", async () => { + const callMock = vi + .fn() + .mockRejectedValueOnce(new Error('400 {"error":{"code":"json_validate_failed"}}')) + .mockResolvedValueOnce({ + choices: [{ message: { role: "assistant", content: JSON.stringify({ A: 1 }) } }], + }); + + const client = GroqProvider.createDefaultGroqClient("k", { + jsonRepairAttempts: 2, + makeSdk: () => ({ chat: { completions: { create: callMock } } }), + }); + + await client.call([{ role: "user", content: "{}" }], schema as any); + expect(callMock.mock.calls.length).toBe(2); + }); + + test("extractJsonValidateFailedRepairInput returns undefined when error is nested object but missing code", async () => { + const callMock = vi.fn().mockRejectedValueOnce( + Object.assign(new Error("no-code"), { + error: { error: { message: "details" } }, + }), + ); + + const client = GroqProvider.createDefaultGroqClient("k", { + jsonRepairAttempts: 2, + makeSdk: () => ({ chat: { completions: { create: callMock } } }), + }); + + await expect(client.call([{ role: "user", content: "{}" }], schema as any)).rejects.toThrow( + /no-code/, + ); + expect(callMock.mock.calls.length).toBe(1); + }); + + test("getNestedErrorObject returns undefined for null", async () => { + const callMock = vi.fn().mockRejectedValueOnce(null); + const client = GroqProvider.createDefaultGroqClient("k", { + makeSdk: () => ({ chat: { completions: { create: callMock } } }), + jsonRepairAttempts: 2, + }); + + await expect(client.call([{ role: "user", content: "{}" }], schema as any)).rejects.toBeNull(); + expect(callMock.mock.calls.length).toBe(1); + }); + + test("getNestedErrorObject returns undefined for object with no error property", async () => { + const callMock = vi.fn().mockRejectedValueOnce({ foo: "bar" }); + const client = GroqProvider.createDefaultGroqClient("k", { + makeSdk: () => ({ chat: { completions: { create: callMock } } }), + jsonRepairAttempts: 2, + }); + + await expect(client.call([{ role: "user", content: "{}" }], schema as any)).rejects.toEqual({ + foo: "bar", + }); + expect(callMock.mock.calls.length).toBe(1); + }); + + test("getNestedErrorObject returns undefined for non-object error field", async () => { + const callMock = vi.fn().mockRejectedValueOnce( + Object.assign(new Error("boom"), { + error: "nope", + }), + ); + const client = GroqProvider.createDefaultGroqClient("k", { + makeSdk: () => ({ chat: { completions: { create: callMock } } }), + jsonRepairAttempts: 2, + }); + + await expect(client.call([{ role: "user", content: "{}" }], schema as any)).rejects.toThrow( + /boom/, + ); + expect(callMock.mock.calls.length).toBe(1); + }); + test("timeout forwarded to groq-sdk", async () => { - vi.resetModules(); - vi.clearAllMocks(); const callMock = vi.fn(async (_req: any, opts?: any) => { expect(opts?.timeout).toBe(1); return { choices: [{ message: { role: "assistant", content: JSON.stringify({ A: 1 }) } }], }; }); - vi.doMock("groq-sdk", () => ({ - default: class Groq { - chat = { completions: { create: callMock } }; - constructor(_cfg: any) {} - }, - })); - const { createDefaultGroqClient } = await import("./groq"); - const client = createDefaultGroqClient("k"); + const client = GroqProvider.createDefaultGroqClient("k", { + makeSdk: () => ({ chat: { completions: { create: callMock } } }), + }); await client.call([{ role: "user", content: "{}" }], schema as any, { timeoutMs: 1 }); expect(callMock.mock.calls.length).toBe(1); }); test("does not pass timeout when undefined", async () => { - vi.resetModules(); - vi.clearAllMocks(); const callMock = vi.fn(async (req: any, opts?: any) => { expect(opts).toBeUndefined(); // ensure assistant role mapping works too @@ -130,70 +574,32 @@ describe("groq provider", () => { choices: [{ message: { role: "assistant", content: JSON.stringify({ A: 1 }) } }], }; }); - vi.doMock("groq-sdk", () => ({ - default: class Groq { - chat = { completions: { create: callMock } }; - constructor(_cfg: any) {} - }, - })); - const { createDefaultGroqClient } = await import("./groq"); - const client = createDefaultGroqClient("k"); + const client = GroqProvider.createDefaultGroqClient("k", { + makeSdk: () => ({ chat: { completions: { create: callMock } } }), + }); await client.call([{ role: "system", content: "hi" }], schema as any); const [req] = callMock.mock.calls[0]; expect(req.messages[0].role).toBe("system"); }); test("throws on unsupported 'tool' role", async () => { - vi.resetModules(); - vi.clearAllMocks(); - vi.doMock("groq-sdk", () => ({ - default: class Groq { - chat = { completions: { create: vi.fn() } }; - constructor(_cfg: any) {} - }, - })); - const { createDefaultGroqClient } = await import("./groq"); - const client = createDefaultGroqClient("k"); + const client = GroqProvider.createDefaultGroqClient("k", { + makeSdk: () => ({ chat: { completions: { create: vi.fn() } } }), + }); await expect( client.call([{ role: "tool", content: "" } as any], schema as any), ).rejects.toThrow(/not supported/); }); - test("forwards userId to request", async () => { - vi.resetModules(); - vi.clearAllMocks(); - const callMock = vi.fn(async (req: any) => { - expect(req.user).toBe("u1"); - return { - choices: [{ message: { role: "assistant", content: JSON.stringify({ A: 1 }) } }], - }; - }); - vi.doMock("groq-sdk", () => ({ - default: class Groq { - chat = { completions: { create: callMock } }; - constructor(_cfg: any) {} - }, - })); - const { createDefaultGroqClient } = await import("./groq"); - const client = createDefaultGroqClient("k"); - await client.call([{ role: "user", content: "{}" }], schema as any, {}, "u1"); - expect(callMock.mock.calls.length).toBe(1); - }); + // (userId forwarding covered above) test("accepts assistant role messages", async () => { - vi.resetModules(); - vi.clearAllMocks(); const callMock = vi.fn(async (_req: any) => ({ choices: [{ message: { role: "assistant", content: JSON.stringify({ A: 2 }) } }], })); - vi.doMock("groq-sdk", () => ({ - default: class Groq { - chat = { completions: { create: callMock } }; - constructor(_cfg: any) {} - }, - })); - const { createDefaultGroqClient } = await import("./groq"); - const client = createDefaultGroqClient("k"); + const client = GroqProvider.createDefaultGroqClient("k", { + makeSdk: () => ({ chat: { completions: { create: callMock } } }), + }); const res = await client.call([{ role: "assistant", content: "hello" }] as any, schema as any); expect(res.data).toEqual({ A: 2 }); expect(callMock.mock.calls.length).toBe(1); diff --git a/src/reranker.groq-default.unit.test.ts b/src/reranker.groq-default.unit.test.ts deleted file mode 100644 index 0aece9c..0000000 --- a/src/reranker.groq-default.unit.test.ts +++ /dev/null @@ -1,46 +0,0 @@ -import { describe, expect, test, vi } from "vitest"; - -// Mock groq-sdk before importing the reranker -const callMock = vi.fn(async (_req: any) => ({ - choices: [{ message: { role: "assistant", content: JSON.stringify({ A: 10, B: 0 }) } }], -})); -vi.mock("groq-sdk", () => ({ - default: class Groq { - chat = { completions: { create: callMock } }; - constructor(_cfg: any) {} - }, -})); - -// Mock CONFIG to have a test key -vi.mock("./config", async (importOriginal) => { - const original = await importOriginal(); - return { - ...original, - CONFIG: { - ...original.CONFIG, - GROQ: { - ...original.CONFIG.GROQ, - API_KEY: "test-key", - }, - }, - }; -}); - -const { Reranker } = await import("./reranker"); - -describe("Reranker (default Groq) ", () => { - test("uses groq-sdk when GROQ_API_KEY is set", async () => { - const reranker = new Reranker<{ key: string; summary: string }>( - { - /* no llm */ - }, - { key: (x) => x.key, summary: (x) => x.summary }, - ); - const out = await reranker.rerank("q", [ - { key: "A", summary: "" }, - { key: "B", summary: "" }, - ]); - expect(out.map((c) => c.key)).toEqual(["A"]); - expect(callMock.mock.calls.length).toBe(1); - }); -}); diff --git a/src/reranker.int.test.ts b/src/reranker.int.test.ts deleted file mode 100644 index 84a73df..0000000 --- a/src/reranker.int.test.ts +++ /dev/null @@ -1,26 +0,0 @@ -import { describe, expect, test } from "vitest"; - -import { CONFIG } from "./config"; -import { Reranker } from "./reranker"; - -const hasKey = Boolean(CONFIG.GROQ.API_KEY); - -describe.skipIf(!hasKey)("reranker integration", () => { - test.concurrent( - "reranker end-to-end", - async () => { - const reranker = new Reranker<{ key: string; summary: string }>( - {}, - { key: (x) => x.key, summary: (x) => x.summary }, - { timeoutMs: 5000 }, - ); - const out = await reranker.rerank("select the best doc", [ - { key: "Alpha", summary: "Doc about alpha" }, - { key: "Beta", summary: "Doc about beta" }, - ]); - expect(out.length).toBeGreaterThanOrEqual(0); - expect(out.length).toBeLessThanOrEqual(2); - }, - 20000, - ); -}); diff --git a/src/reranker.ts b/src/reranker.ts deleted file mode 100644 index 1498172..0000000 --- a/src/reranker.ts +++ /dev/null @@ -1,281 +0,0 @@ -import { batchProcess } from "./batches"; -import { CONFIG } from "./config"; -import { clamp } from "./lib/number"; -import { selectLlmClient } from "./llm_client"; -import { buildMessages } from "./messages"; -import { buildRelevancySchema } from "./schema"; - -import type { JSONObject, RerankerExtractors, IntentContext, LlmClient } from "./types"; - -type RerankerConfig = typeof CONFIG.RERANKER; - -/** - * LLM-based reranker for arbitrary items. - * - * Uses a listwise LLM approach to score candidates 0-10 based on relevance to a query, - * then filters by threshold and returns results sorted by score with stable ordering. - * - * @template T - The type of items to rerank - * - * @example - * ```typescript - * type Document = { id: string; title: string; content: string }; - * - * const reranker = new Reranker( - * { llm: myClient, userId: "user-123" }, - * { - * key: doc => doc.title, - * summary: doc => doc.content.slice(0, 200) - * }, - * { RELEVANCY_THRESHOLD: 5, BATCH_SIZE: 20 } - * ); - * - * const ranked = await reranker.rerank("find expense reports", documents); - * ``` - */ -export class Reranker { - private readonly cfg: RerankerConfig; - private readonly llm: LlmClient; - - /** - * Creates a new Reranker instance. - * - * @param ctx - Context containing LLM client, optional logger, and user ID - * @param ctx.llm - Optional LLM client. If omitted, will use Groq client if GROQ_API_KEY is set - * @param ctx.logger - Optional logger for warnings and errors - * @param ctx.userId - Optional user identifier passed to LLM provider for abuse monitoring - * @param extractors - Functions to extract key and summary from items - * @param extractors.key - Required function returning a short human-readable identifier - * @param extractors.summary - Optional function returning a short description for LLM reasoning - * @param overrides - Optional config overrides for model, timeout, threshold, and batch settings - * @throws {Error} If no LLM client is provided and GROQ_API_KEY is not set - */ - constructor( - private readonly ctx: IntentContext, - private readonly extractors: RerankerExtractors, - overrides: Partial = {}, - ) { - this.cfg = { ...CONFIG.RERANKER, ...overrides }; - - // Validate threshold is in valid range - if (this.cfg.RELEVANCY_THRESHOLD < 0 || this.cfg.RELEVANCY_THRESHOLD > 10) { - throw new Error( - `intent: RELEVANCY_THRESHOLD must be between 0 and 10, got ${this.cfg.RELEVANCY_THRESHOLD}`, - ); - } - - const selectedClient = selectLlmClient(ctx); - if (!selectedClient) { - throw new Error( - "intent: No LLM client provided and GROQ_API_KEY not set. Provide ctx.llm or set GROQ_API_KEY.", - ); - } - this.llm = selectedClient; - } - - /** - * Rerank candidates based on relevance to a query. - * - * Calls the LLM to score each candidate 0-10 based on relevance to the query, - * filters results by the configured threshold, and returns items sorted by score - * (highest first) with ties preserving original input order. - * - * Fast-path optimizations: - * - Returns empty array for 0 candidates without LLM call - * - Returns single candidate unchanged without LLM call - * - * Error handling: - * - On any batch error, returns that batch's items in original order - * - On top-level error, returns all items in original order - * - All errors are logged via the configured logger - * - * @param query - The search query or user intent to rank against - * @param candidates - Array of items to rerank - * @param options - Optional per-call configuration - * @param options.userId - Optional user ID for this specific call, overrides ctx.userId - * @returns Filtered and sorted array of items, or original order on any error - * - * @example - * ```typescript - * const results = await reranker.rerank( - * "quarterly expense reports from 2024", - * allDocuments, - * { userId: "session-abc" } - * ); - * // Returns only documents with score > threshold, sorted by relevance - * ``` - */ - public async rerank(query: string, candidates: T[], options?: { userId?: string }): Promise { - try { - if (candidates.length === 0) return []; - if (candidates.length === 1) return candidates; - - const prepared = this.prepareCandidates(candidates); - return await batchProcess( - prepared, - this.cfg.BATCH_SIZE, - this.cfg.TINY_BATCH_FRACTION, - (batch) => this.processBatch(query, batch, options?.userId), - this.ctx.logger, - (batch) => batch.map(({ item }) => item), - ); - } catch (error) { - this.ctx.logger?.warn?.("intent reranker failed, using fallback", { - error: (error as Error)?.message, - }); - return candidates; - } - } - - /** - * Normalize incoming items into a consistent shape for downstream processing. - * - * Extracts the key and summary from each item using the configured extractors, - * and attaches the original input index for stable sorting later. - * - * @param candidates - Raw items to prepare - * @returns Array of prepared candidates with extracted metadata and original index - * @private - */ - private prepareCandidates(candidates: T[]): Array<{ - item: T; - idx: number; - baseKey: string; - summary: string; - }> { - return candidates.map((item, idx) => ({ - item, - idx, - baseKey: this.extractors.key(item), - summary: this.extractors.summary?.(item) ?? "", - })); - } - - /** - * Ensure keys are unique by suffixing duplicates with their input index. - * - * When multiple items share the same key, subsequent occurrences are renamed - * to "Key (idx)" where idx is the original input index. This prevents JSON - * schema validation errors and ensures the LLM can score each item independently. - * - * @param itemsBase - Prepared candidates with potentially duplicate keys - * @returns Candidates with guaranteed unique keys - * @private - */ - private ensureUniqueKeys( - itemsBase: Array<{ item: T; idx: number; baseKey: string; summary: string }>, - ): Array<{ item: T; idx: number; key: string; summary: string }> { - const counts = new Map(); - return itemsBase.map(({ item, baseKey, summary, idx }) => { - const n = (counts.get(baseKey) ?? 0) + 1; - counts.set(baseKey, n); - const key = n === 1 ? baseKey : `${baseKey} (${idx})`; - return { item, idx, key, summary }; - }); - } - - /** - * Build the JSON schema and chat messages payload for the LLM. - * - * Creates a strict JSON schema requiring one integer property (0-10) per candidate key, - * and constructs system + user messages instructing the LLM to score relevance. - * - * @param query - The search query to evaluate candidates against - * @param items - Candidates with unique keys and summaries - * @returns Object containing JSON schema and chat messages array - * @private - */ - private buildRequest( - query: string, - items: Array<{ key: string; summary: string }>, - ): { schema: JSONObject; messages: any[] } { - const keys = items.map((x) => x.key); - const schema: JSONObject = buildRelevancySchema(keys); - const messages = buildMessages(query, items); - return { schema, messages }; - } - - /** - * Invoke the LLM and return the parsed map of candidate scores. - * - * Calls the configured LLM client with the messages, JSON schema, model config, - * and user ID. Returns null if the response is invalid or missing. - * - * @param messages - Chat messages (system + user) to send to LLM - * @param schema - Strict JSON schema defining expected response structure - * @param userId - Optional user identifier for provider abuse monitoring - * @returns Map of candidate keys to numeric scores, or null if response invalid - * @private - */ - private async fetchScores( - messages: any[], - schema: JSONObject, - userId?: string, - ): Promise | null> { - const { data } = await this.llm.call>( - messages, - schema, - { model: this.cfg.MODEL, temperature: 0, timeoutMs: this.cfg.TIMEOUT_MS }, - userId ?? this.ctx.userId, - ); - - if (data == null || typeof data !== "object") return null; - return data as Record; - } - - /** - * Apply relevancy threshold filtering and stable sorting. - * - * Scores are clamped to 0-10 range, then filtered to keep only items with - * score > threshold. Results are sorted by score descending, with ties - * preserving original input order for deterministic results. - * - * @param items - Candidates with unique keys - * @param scores - Map of candidate keys to LLM-assigned scores - * @returns Filtered and sorted array of original items - * @private - */ - private rankAndFilter( - items: Array<{ item: T; idx: number; key: string; summary: string }>, - scores: Record, - ): T[] { - const threshold = this.cfg.RELEVANCY_THRESHOLD; - const scored = items.map(({ item, idx, key }) => ({ - item, - idx, - score: clamp((scores as any)[key], 0, 10), - })); - - const filtered = scored.filter(({ score }) => score > threshold); - const sorted = filtered.sort((a, b) => { - if (b.score !== a.score) return b.score - a.score; - return a.idx - b.idx; - }); - return sorted.map(({ item }) => item); - } - - /** - * Process a single batch of candidates through the LLM. - * - * Ensures unique keys, builds the request payload, fetches scores from the LLM, - * and returns filtered and sorted results. On any error or null response from - * the LLM, returns items in their original order as a fallback. - * - * @param query - The search query to evaluate candidates against - * @param batch - Batch of prepared candidates to process - * @param userId - Optional user identifier for provider abuse monitoring - * @returns Ranked and filtered items, or original order on error - * @private - */ - private async processBatch( - query: string, - batch: Array<{ item: T; idx: number; baseKey: string; summary: string }>, - userId?: string, - ): Promise { - const keyed = this.ensureUniqueKeys(batch); - const { schema, messages } = this.buildRequest(query, keyed); - const scores = await this.fetchScores(messages, schema, userId); - if (scores == null) return keyed.map(({ item }) => item); - return this.rankAndFilter(keyed, scores); - } -} diff --git a/src/reranker.unit.test.ts b/src/reranker.unit.test.ts deleted file mode 100644 index 8f63c8c..0000000 --- a/src/reranker.unit.test.ts +++ /dev/null @@ -1,346 +0,0 @@ -import { describe, expect, test, vi } from "vitest"; - -import { Reranker } from "./reranker"; - -import type { LlmClient, LoggerLike, RerankerCandidate, IntentContext } from "./types"; - -function makeCtx(overrides: Partial = {}): IntentContext & { - llm: LlmClient & { call: ReturnType }; - logger: Required & { warn: ReturnType }; -} { - const logger = { - info: vi.fn(() => {}), - warn: vi.fn(() => {}), - error: vi.fn(() => {}), - } as any; - const llm = { - call: vi.fn(async () => ({ data: {} })), - } as any; - return { - llm, - logger, - userId: undefined, - ...overrides, - } as any; -} - -describe("Reranker.rerank", () => { - test("throws when no llm and no GROQ_API_KEY", async () => { - const { CONFIG } = await import("./config"); - - try { - // Mock CONFIG to return empty API key - vi.doMock("./config", () => ({ - CONFIG: { - ...CONFIG, - GROQ: { - ...CONFIG.GROQ, - API_KEY: "", - }, - }, - })); - - // Re-import modules to get mocked config - vi.resetModules(); - const { Reranker: RerankerWithMock } = await import("./reranker"); - - expect( - () => new RerankerWithMock({} as any, { key: (c) => c.key }), - ).toThrow(/No LLM client provided/); - } finally { - vi.doUnmock("./config"); - vi.resetModules(); - } - }); - - test("throws when threshold is below 0", async () => { - const ctx = makeCtx(); - expect( - () => - new Reranker(ctx, { key: (c) => c.key }, { RELEVANCY_THRESHOLD: -1 }), - ).toThrow(/RELEVANCY_THRESHOLD must be between 0 and 10/); - }); - - test("throws when threshold is above 10", async () => { - const ctx = makeCtx(); - expect( - () => - new Reranker(ctx, { key: (c) => c.key }, { RELEVANCY_THRESHOLD: 11 }), - ).toThrow(/RELEVANCY_THRESHOLD must be between 0 and 10/); - }); - test("returns empty list for zero candidates", async () => { - const ctx = makeCtx(); - const reranker = new Reranker(ctx, { - key: (c) => c.key, - summary: (c) => c.summary, - }); - const res = await reranker.rerank("query", []); - expect(res).toEqual([]); - expect(ctx.llm.call).not.toHaveBeenCalled(); - }); - - test("returns input unchanged for single candidate (no LLM call)", async () => { - const ctx = makeCtx(); - const reranker = new Reranker(ctx, { - key: (c) => c.key, - summary: (c) => c.summary, - }); - const input = [{ key: "Only", summary: "s" }]; - const res = await reranker.rerank("query", input); - expect(res).toEqual(input); - expect(ctx.llm.call).not.toHaveBeenCalled(); - }); - - test("rounds scores, filters zeros, orders by score", async () => { - const ctx = makeCtx(); - (ctx.llm.call as any).mockResolvedValueOnce({ - data: { A: 10, B: 6.8, C: 0 }, - }); - const reranker = new Reranker(ctx, { - key: (c) => c.key, - summary: (c) => c.summary, - }); - const input = ["A", "B", "C"].map((k) => ({ key: k, summary: k })); - const res = await reranker.rerank("query", input); - expect(res.map((c) => c.key)).toEqual(["A", "B"]); - const call = (ctx.llm.call as any).mock.calls[0]; - expect(call[2].timeoutMs).toBe(3000); // default - }); - - test("handles non-numeric or missing scores by clamping to 0", async () => { - const ctx = makeCtx(); - (ctx.llm.call as any).mockResolvedValueOnce({ data: { X: "nope" as any } }); - const reranker = new Reranker(ctx, { - key: (c) => c.key, - summary: (c) => c.summary, - }); - const input = ["X", "Y"].map((k) => ({ key: k, summary: k })); - const res = await reranker.rerank("query", input); - expect(res).toEqual([]); - }); - - test("normalizes out-of-range and infinite values", async () => { - const ctx = makeCtx(); - (ctx.llm.call as any).mockResolvedValueOnce({ - data: { A: -3, B: 11, C: 9.6, D: Number.POSITIVE_INFINITY }, - }); - const reranker = new Reranker(ctx, { - key: (c) => c.key, - summary: (c) => c.summary, - }); - const input = ["A", "B", "C", "D"].map((k) => ({ key: k, summary: k })); - const res = await reranker.rerank("query", input); - // B clamps to 10, D rounds to 10, keep input order for tie: B before D - expect(res.map((c) => c.key)).toEqual(["B", "D", "C"]); - }); - - test("returns original list on error and logs warning", async () => { - const ctx = makeCtx(); - (ctx.llm.call as any).mockRejectedValueOnce(new Error("boom")); - const reranker = new Reranker(ctx, { - key: (c) => c.key, - summary: (c) => c.summary, - }); - const input = [ - { key: "A", summary: "" }, - { key: "B", summary: "" }, - ]; - const res = await reranker.rerank("query", input); - expect(res).toEqual(input); - expect(ctx.logger.warn).toHaveBeenCalled(); - }); - - test("uses empty summary when extractor is missing", async () => { - const ctx = makeCtx(); - // Return some scores so rerank proceeds - (ctx.llm.call as any).mockResolvedValueOnce({ data: { A: 1, B: 0 } }); - const reranker = new Reranker( - ctx, - { - key: (c) => c.key, - // no summary extractor - }, - { BATCH_SIZE: 10 }, - ); - const input: RerankerCandidate[] = [ - { key: "A", summary: "original" }, - { key: "B", summary: "ignored" }, - ]; - await reranker.rerank("query", input); - const call = (ctx.llm.call as any).mock.calls[0]; - const messages = call[0]; - const userPayload = JSON.parse(messages[1].content); - const summaries = userPayload.candidate_search_results.map((c: any) => c.summary); - expect(summaries).toEqual(["", ""]); - }); - - test("top-level rerank catch: logs and returns input on unexpected error", async () => { - const ctx = makeCtx(); - const reranker = new Reranker(ctx, { - key: (c) => c.key, - summary: (c) => c.summary, - }); - // @ts-ignore override private method to throw to trigger top-level catch - reranker.prepareCandidates = () => { - throw new Error("oops"); - }; - const input = [ - { key: "A", summary: "" }, - { key: "B", summary: "" }, - { key: "C", summary: "" }, - ]; - const res = await reranker.rerank("query", input); - expect(res).toEqual(input); - expect(ctx.logger.warn).toHaveBeenCalled(); - }); - - test("passes userId from ctx and allows method override", async () => { - const ctx = makeCtx({ userId: "ctx-user" }); - (ctx.llm.call as any).mockResolvedValueOnce({ data: { A: 5, B: 5 } }); - const reranker = new Reranker(ctx, { - key: (c) => c.key, - summary: (c) => c.summary, - }); - const input = [ - { key: "A", summary: "" }, - { key: "B", summary: "" }, - ]; - await reranker.rerank("query", input, { userId: "call-user" }); - const calls = (ctx.llm.call as any).mock.calls; - expect(calls[0][3]).toBe("call-user"); // override wins - }); - - test("timeout config is forwarded to client", async () => { - const ctx = makeCtx(); - (ctx.llm.call as any).mockResolvedValueOnce({ data: { A: 5, B: 5 } }); - const reranker = new Reranker( - ctx, - { key: (c) => c.key, summary: (c) => c.summary }, - { TIMEOUT_MS: 5 }, - ); - const input = [ - { key: "A", summary: "" }, - { key: "B", summary: "" }, - ]; - await reranker.rerank("query", input); - const calls = (ctx.llm.call as any).mock.calls; - expect(calls[0][2].timeoutMs).toBe(5); - }); - - test("splits long lists into batches and combines results", async () => { - const ctx = makeCtx(); - (ctx.llm.call as any) - .mockResolvedValueOnce({ data: { K2: 10, K0: 7 } }) - .mockResolvedValueOnce({ data: { K7: 10, K6: 9 } }); - - const reranker = new Reranker( - ctx, - { key: (c) => c.key, summary: (c) => c.summary }, - { BATCH_SIZE: 5 }, - ); - - const input: RerankerCandidate[] = Array.from({ length: 10 }).map((_, i) => ({ - key: `K${i}`, - summary: `S${i}`, - })); - - const out = await reranker.rerank("query", input); - expect(out.map((c) => c.key)).toEqual(["K2", "K0", "K7", "K6"]); - }); - - test("merges tiny final batch", async () => { - const ctx = makeCtx(); - (ctx.llm.call as any).mockResolvedValueOnce({ data: {} }).mockResolvedValueOnce({ data: {} }); - - const reranker = new Reranker( - ctx, - { key: (c) => c.key, summary: (c) => c.summary }, - { BATCH_SIZE: 5, TINY_BATCH_FRACTION: 0.2 }, - ); - - const input: RerankerCandidate[] = Array.from({ length: 7 }).map((_, i) => ({ - key: `K${i}`, - summary: `S${i}`, - })); - - await reranker.rerank("query", input); - expect((ctx.llm.call as any).mock.calls.length).toBe(2); - }); - - test("merges really tiny final batch (<= threshold)", async () => { - const ctx = makeCtx(); - (ctx.llm.call as any).mockResolvedValueOnce({ data: {} }); - - const reranker = new Reranker( - ctx, - { key: (c) => c.key, summary: (c) => c.summary }, - { BATCH_SIZE: 5, TINY_BATCH_FRACTION: 0.2 }, - ); - - const input: RerankerCandidate[] = Array.from({ length: 6 }).map((_, i) => ({ - key: `K${i}`, - summary: `S${i}`, - })); - - await reranker.rerank("query", input); - expect((ctx.llm.call as any).mock.calls.length).toBe(1); - }); - - test("one batch fails while others succeed (partial fallback)", async () => { - const ctx = makeCtx(); - (ctx.llm.call as any) - .mockResolvedValueOnce({ data: { K1: 10, K0: 9 } }) - .mockRejectedValueOnce(new Error("boom")); - - const reranker = new Reranker( - ctx, - { key: (c) => c.key, summary: (c) => c.summary }, - { BATCH_SIZE: 3 }, - ); - - const input: RerankerCandidate[] = Array.from({ length: 6 }).map((_, i) => ({ - key: `K${i}`, - summary: `S${i}`, - })); - - const out = await reranker.rerank("query", input); - expect(out.map((c) => c.key)).toEqual(["K1", "K0", "K3", "K4", "K5"]); - expect((ctx.llm.call as any).mock.calls.length).toBe(2); - }); - - test("returns original list when scores payload is null", async () => { - const ctx = makeCtx(); - (ctx.llm.call as any).mockResolvedValueOnce({ data: null }); - const reranker = new Reranker(ctx, { - key: (c) => c.key, - summary: (c) => c.summary, - }); - const input: RerankerCandidate[] = [ - { key: "A", summary: "S" }, - { key: "B", summary: "S" }, - ]; - - const out = await reranker.rerank("query", input); - expect(out).toEqual(input); - }); - - test("stable order for ties within a batch (duplicate keys)", async () => { - const ctx = makeCtx(); - (ctx.llm.call as any).mockResolvedValueOnce({ - data: { Same: 5, "Same (1)": 5 }, - }); - - const reranker = new Reranker(ctx, { - key: (c) => c.key, - summary: (c) => c.summary, - }); - - const input: RerankerCandidate[] = [ - { key: "Same", summary: "S0" }, - { key: "Same", summary: "S1" }, - ]; - - const out = await reranker.rerank("query", input); - expect(out.map((c) => c.summary)).toEqual(["S0", "S1"]); - }); -}); diff --git a/src/schema.ts b/src/schema.ts index aa22899..df3e5ef 100644 --- a/src/schema.ts +++ b/src/schema.ts @@ -1,24 +1,162 @@ import type { JSONObject } from "./types"; +type IntegerSchema = { type: "integer" }; +type StringSchema = { type: "string" }; +type BooleanSchema = { type: "boolean" }; + +type CandidateEvaluationSchema = { + type: "object"; + properties: { + explanation: StringSchema; + score: IntegerSchema; + }; + required: ["explanation", "score"]; + additionalProperties: false; +}; + +type CandidateFilterSchema = { + type: "object"; + properties: { + explanation: StringSchema; + isRelevant: BooleanSchema; + }; + required: ["explanation", "isRelevant"]; + additionalProperties: false; +}; + +type ChoiceSchema = { + type: "object"; + properties: { + explanation: StringSchema; + selectedKey: { type: "string"; enum: string[] }; + }; + required: ["explanation", "selectedKey"]; + additionalProperties: false; +}; + +/** + * Build the schema used for a single candidate's evaluation. + * + * Property order is intentional: `explanation` is defined before `score` to + * encourage structured-output models to generate explanations first. + * + * @returns JSON schema for a single candidate evaluation + */ +export function buildCandidateEvaluationSchema(): CandidateEvaluationSchema { + return { + type: "object", + properties: { + explanation: { type: "string" }, + score: { type: "integer" }, + }, + required: ["explanation", "score"], + additionalProperties: false, + }; +} + /** - * Build a strict JSON schema mapping candidate keys to integer 0-10. + * Build the schema used for a single candidate's filter decision. * - * Creates a JSON schema object with one required integer property per candidate key. - * Sets additionalProperties to false to prevent the LLM from adding extra keys. - * Used with LLM providers that support structured output (e.g., Groq's json_schema mode). + * Property order is intentional: `explanation` is defined before `isRelevant`. + * + * @returns JSON schema for a single candidate filter decision + */ +export function buildCandidateFilterSchema(): CandidateFilterSchema { + return { + type: "object", + properties: { + explanation: { type: "string" }, + isRelevant: { type: "boolean" }, + }, + required: ["explanation", "isRelevant"], + additionalProperties: false, + }; +} + +/** + * Build a strict JSON schema mapping candidate keys to evaluation objects. + * + * Each candidate key maps to an object containing: + * - explanation: a short justification of the score + * - score: an integer within the configured range + * + * Property order is intentional (`explanation` then `score`) to encourage the + * model to generate explanations before scores in structured-output modes. * * @param keys - Array of unique candidate keys * @returns JSON schema object enforcing exact structure of response */ -export function buildRelevancySchema(keys: string[]): JSONObject { - const properties: Record = {}; - for (const k of keys) properties[k] = { type: "integer" }; +export function buildRelevancySchema( + keys: string[], + minScore: number, + maxScore: number, +): JSONObject { + const evaluationSchema = buildCandidateEvaluationSchema(); + + const properties: Record = {}; + for (const k of keys) { + properties[k] = evaluationSchema; + } + return { title: "Query / Candidate Relevancy Assessment", - description: "Map candidate results for a search query to relevancy scores (0-10).", + description: `Map candidate results for a search query to relevancy scores (${minScore}-${maxScore}) with explanations.`, type: "object", properties, required: keys, additionalProperties: false, } as JSONObject; } + +/** + * Build a strict JSON schema mapping candidate keys to boolean relevancy decisions. + * + * Each candidate key maps to an object containing: + * - explanation: a short justification + * - isRelevant: boolean + * + * @param keys - Array of unique candidate keys + * @returns JSON schema object enforcing exact structure of response + */ +export function buildFilterSchema(keys: string[]): JSONObject { + const decisionSchema = buildCandidateFilterSchema(); + + const properties: Record = {}; + for (const k of keys) { + properties[k] = decisionSchema; + } + + return { + title: "Query / Candidate Relevancy Filter", + description: + "Map candidate results for a search query to boolean relevancy decisions with explanations.", + type: "object", + properties, + required: keys, + additionalProperties: false, + } as JSONObject; +} + +/** + * Build a strict JSON schema for choosing a single candidate key. + * + * The model must return: + * - explanation: string + * - selectedKey: one of the provided candidate keys (enum) + * + * @param keys - Array of unique candidate keys + * @returns JSON schema enforcing a single selected key + */ +export function buildChoiceSchema(keys: string[]): JSONObject { + return { + title: "Query / Candidate Single Choice", + description: "Choose exactly one candidate key for the query and explain why.", + type: "object", + properties: { + explanation: { type: "string" }, + selectedKey: { type: "string", enum: keys }, + }, + required: ["explanation", "selectedKey"], + additionalProperties: false, + } as ChoiceSchema as JSONObject; +} diff --git a/src/types.ts b/src/types.ts index ea86ad8..36b8af5 100644 --- a/src/types.ts +++ b/src/types.ts @@ -12,7 +12,7 @@ export type ChatMessage = { export type LlmCallConfig = { model?: string; - temperature?: number; + reasoningEffort?: "low" | "medium" | "high"; timeoutMs?: number; }; @@ -31,13 +31,13 @@ export interface LoggerLike { error?(msg: string, meta?: unknown): void; } -export type RerankerCandidate = { +export type IntentCandidate = { key: string; summary: string; }; -export type RerankerExtractors = { - key: (item: T) => string; +export type IntentExtractors = { + key?: (item: T) => string; summary?: (item: T) => string; }; @@ -46,3 +46,47 @@ export type IntentContext = { logger?: LoggerLike; userId?: string; // optional per-instance user id used for provider abuse monitoring }; + +// Utility types for key-case transformations +export type CamelCase = S extends `${infer H}_${infer T}` + ? `${Lowercase}${Capitalize>}` + : S extends `${infer H}-${infer T}` + ? `${Lowercase}${Capitalize>}` + : Lowercase; + +export type CamelCasedProps = { + [K in keyof T as K extends string ? CamelCase : K]: T[K]; +}; + +/** + * Configuration options for Intent. + * + * This is a camelCase version of the INTENT config object from config.ts. + */ +export type IntentConfig = { + provider?: "GROQ"; + timeoutMs?: number; + relevancyThreshold?: number; + batchSize?: number; + tinyBatchFraction?: number; + minScore?: number; + maxScore?: number; +}; + +/** + * Complete options object for Intent constructor. + * + * Merges all configuration into a single, fully optional object: + * - LLM client and runtime context (llm, logger, userId) + * - Item extractors (key, summary) + * - Intent configuration (model, timeoutMs, relevancyThreshold, batchSize, tinyBatchFraction) + * + * All fields are optional with sensible defaults: + * - llm: Auto-detected from GROQ_API_KEY if available + * - key: Hash-based string from JSON representation + * - summary: Pretty-printed JSON of the item (2-space indentation for LLM readability) + * - Config values: From environment variables or built-in defaults + * + * @template T - The type of items to rerank + */ +export type IntentOptions = IntentContext & IntentExtractors & IntentConfig;