diff --git a/.env.example b/.env.example index a4ee12c..7bae72a 100644 --- a/.env.example +++ b/.env.example @@ -27,6 +27,49 @@ DOCUMENTDB_DB=copilot_memory ## Logging level: trace, debug, info, warn, error, fatal. MEMORY_LOG_LEVEL=info +## ───────────────────────────────────────────────────────────────────────── +## DocumentDB Search (embedding-backed vector search) +## ───────────────────────────────────────────────────────────────────────── +## Adds semantic recall to the knowledge-graph `search_nodes` tool: entity +## text is embedded and indexed as a vector in DocumentDB, then blended with +## the existing keyword (`$text`) search via reciprocal-rank fusion. Entirely +## optional — if the provider is `none` or the embedding backend is +## unreachable, search transparently falls back to keyword-only. + +## Embedding provider: local | ollama | openai | azure-openai | none +## `local` and `ollama` both talk to a local Ollama server. Default: local. +MEMORY_EMBEDDING_PROVIDER=local + +## Embedding model. Defaults per provider: +## local / ollama -> nomic-embed-text +## openai / azure-openai -> text-embedding-3-small +# MEMORY_EMBEDDING_MODEL=nomic-embed-text + +## Base URL override for local/ollama (default http://localhost:11434) +## or openai (default https://api.openai.com/v1). +# MEMORY_EMBEDDING_BASE_URL=http://localhost:11434 + +## Credentials for hosted providers. +# MEMORY_EMBEDDING_API_KEY= +## Azure OpenAI: full resource endpoint, e.g. https://my-res.openai.azure.com +# MEMORY_EMBEDDING_ENDPOINT= +# MEMORY_EMBEDDING_API_VERSION=2024-02-01 + +## Optional: force the vector dimensionality. Normally probed from the model +## at startup; set only if you need to override (mismatch logs a warning). +# MEMORY_EMBEDDING_DIMENSIONS= + +## Vector index tuning (DocumentDB cosmosSearch). +## Index kind: vector-ivf (default) | vector-hnsw +# MEMORY_VECTOR_INDEX_KIND=vector-ivf +## Similarity metric: COS (default) | L2 | IP — must match the model's vectors. +# MEMORY_VECTOR_SIMILARITY=COS +## IVF only: number of lists (default 100). +# MEMORY_VECTOR_NUM_LISTS=100 +## HNSW only: graph degree M (default 16) and construction ef (default 64). +# MEMORY_VECTOR_HNSW_M=16 +# MEMORY_VECTOR_HNSW_EF_CONSTRUCTION=64 + ## Sync daemon polling interval (used by `documentdb-memory sessions sync --watch`). ## Accepts durations like 10s, 30s, 1m, 5m. SYNC_INTERVAL=30s diff --git a/README.md b/README.md index b06d738..3d429bc 100644 --- a/README.md +++ b/README.md @@ -150,6 +150,7 @@ ones that come up daily: | `COPILOT_SESSION_STORE` | `~/.copilot/session-store.db` | SQLite source for the sync daemon. | | `MEMORY_LOG_LEVEL` | `info` | Pino log level (`debug` is useful when things misbehave). | | `SYNC_INTERVAL` | `30s` | How often `sessions sync --watch` polls. | +| `MEMORY_EMBEDDING_PROVIDER` | `local` | DocumentDB Search backend (`local`/`ollama`/`openai`/`azure-openai`/`none`). | The compose-managed DocumentDB enforces TLS server-side with a self-signed cert, so the internal URI looks like @@ -158,6 +159,80 @@ The same database reached from the host (where TLS isn't required) is `mongodb://localadmin:Admin100@localhost:10260/?tls=false`. The compose stack uses the first; the deploy templates default to the second. +## DocumentDB Search — semantic recall for the knowledge graph + +By default the `search_nodes` tool matches entities by keyword (a Mongo +`$text` index). **DocumentDB Search** adds an embedding-backed vector index +on top and blends the two, so a query recalls entities that are *semantically* +related even when they share no words with the query — e.g. searching `k8s` +surfaces an entity named `Kubernetes`. This follows the hybrid dense-vector + +knowledge-graph retrieval model popularized by +[OmniMem](https://omnimem.org). + +It is **entirely optional and degrades gracefully**: if the provider is +`none`, has no credentials, or the embedding backend is unreachable, search +silently falls back to keyword-only and entities are still stored. This keeps +the server a drop-in replacement for the upstream memory server. + +### How it works + +- On `create_entities` / `add_observations`, the entity's name + type + + observations are embedded and stored alongside the document. +- A DocumentDB vector index (`cosmosSearch`) is created automatically at + startup when an embedder is available. +- `search_nodes` runs the keyword search and the vector kNN search in + parallel, then fuses the two ranked lists with reciprocal-rank fusion (RRF). + The relation-containment rule (returned relations only reference returned + entities) is preserved. + +### Enabling it + +Pick an embedding provider via `MEMORY_EMBEDDING_PROVIDER` (default `local`): + +| Provider | Backend | Needs | +| -------------- | ----------------------------------------- | ---------------------------------------- | +| `local` / `ollama` | Local [Ollama](https://ollama.com) server | Ollama running (default `nomic-embed-text`) | +| `openai` | OpenAI embeddings API | `MEMORY_EMBEDDING_API_KEY` | +| `azure-openai` | Azure OpenAI deployment | `MEMORY_EMBEDDING_ENDPOINT` + `MEMORY_EMBEDDING_API_KEY` | +| `none` | disabled (keyword-only) | — | + +For the default local setup, the `compose.full.yml` stack already bundles an +`ollama` service and pulls `nomic-embed-text` automatically before the MCP +server starts — so DocumentDB Search works out of the box: + +```bash +docker compose -f compose.full.yml up -d +``` + +The MCP container reaches it over the compose network at +`http://ollama:11434` (not `localhost`, which inside a container would point +at the container itself). Models persist in the `ollama-models` volume. + +Running the MCP server as a **host binary** instead? Point it at a local +Ollama yourself: + +```bash +ollama pull nomic-embed-text # one-time +# MEMORY_EMBEDDING_PROVIDER=local is already the default (localhost:11434) +documentdb-memory-mcp +``` + +All embedding/index knobs (`MEMORY_EMBEDDING_*`, `MEMORY_VECTOR_*`) are +documented inline in [`.env.example`](./.env.example). + +### Backfilling existing entities + +Entities created before DocumentDB Search was enabled (or before you switched +models) have no vector yet. Backfill them with the CLI: + +```bash +# embed only entities missing / stale vectors +documentdb-memory graph reembed + +# force re-embed every entity (e.g. after changing model) +documentdb-memory graph reembed --all +``` + ## What's where ``` @@ -165,7 +240,7 @@ stack uses the first; the deploy templates default to the second. ├── README.md ← you are here ├── .env.example ← every config knob, commented ├── compose.yml ← DocumentDB + DocumentDBFUSE -├── compose.full.yml ← + MCP server + sync daemon +├── compose.full.yml ← + MCP server + sync daemon + ollama ├── compose.fuse-host-bind.yml ← overlay to bind FUSE to host ├── compose.dev.yml ← live-reload variant for hacking on src/ ├── deploy/ ← launchd + systemd templates for host install diff --git a/compose.full.yml b/compose.full.yml index 01bce85..cf15870 100644 --- a/compose.full.yml +++ b/compose.full.yml @@ -15,6 +15,16 @@ # ~/.copilot/session-store.db into # DocumentDB on a poll interval. # +# ollama — local embedding backend for DocumentDB +# Search. The MCP server's default +# `local` provider talks to it over the +# compose network at http://ollama:11434. +# +# ollama-pull — one-shot init container that pulls the +# embedding model into the shared ollama +# volume, then exits. The MCP server waits +# for it to finish before starting. +# # Bring up: # # docker compose -f compose.full.yml up -d @@ -55,6 +65,11 @@ services: depends_on: documentdb: condition: service_started + # Block startup until the embedding model has been pulled, so the + # first search already benefits from DocumentDB Search instead of + # silently falling back to keyword-only on a cold model cache. + ollama-pull: + condition: service_completed_successfully environment: # Hardcoded to the in-network service name so this works even when # the user's .env has a localhost-based DOCUMENTDB_URI for @@ -63,6 +78,12 @@ services: DOCUMENTDB_URI: mongodb://localadmin:Admin100@documentdb:10260/?directConnection=true&tls=true&tlsInsecure=true DOCUMENTDB_DB: ${DOCUMENTDB_DB:-copilot_memory} MEMORY_LOG_LEVEL: ${MEMORY_LOG_LEVEL:-info} + # DocumentDB Search: the `local` provider points at the in-network + # ollama service (NOT localhost — that would resolve to this + # container). Override MEMORY_EMBEDDING_PROVIDER=none to disable. + MEMORY_EMBEDDING_PROVIDER: ${MEMORY_EMBEDDING_PROVIDER:-local} + MEMORY_EMBEDDING_MODEL: ${MEMORY_EMBEDDING_MODEL:-nomic-embed-text} + MEMORY_EMBEDDING_BASE_URL: ${MEMORY_EMBEDDING_BASE_URL:-http://ollama:11434} # No `ports:` on purpose — the MCP protocol speaks stdio, not TCP. # External clients attach via `docker exec -i`. stdin_open: true @@ -114,3 +135,52 @@ services: # source path (e.g. to point at a non-default Copilot home). - ${COPILOT_DIR:-${HOME}/.copilot}:/copilot:ro restart: unless-stopped + + ollama: + # Local embedding backend for DocumentDB Search. CPU-only by default; + # nomic-embed-text is tiny (~275MB) so CPU inference is plenty fast + # for the short entity strings we embed. To use a host GPU, add a + # `deploy.resources.reservations.devices` block per the ollama docs. + image: ollama/ollama:latest + container_name: documentdb-agentic-memory-ollama + # No host `ports:` — only the in-network services need it. Uncomment + # to also reach it from the host (e.g. for `ollama` CLI debugging): + # ports: + # - "11434:11434" + volumes: + # Persist pulled models across restarts so we don't re-download on + # every `up`. + - ollama-models:/root/.ollama + healthcheck: + # `ollama list` exits non-zero until the server is accepting API + # calls, which is exactly the readiness signal we want. + test: ["CMD", "ollama", "list"] + interval: 5s + timeout: 5s + retries: 20 + start_period: 10s + restart: unless-stopped + + ollama-pull: + # One-shot: wait for ollama to be healthy, pull the embedding model + # into the shared volume, then exit 0. The MCP server depends on this + # completing successfully (service_completed_successfully), so the + # first request already has a warm model instead of racing the pull. + image: ollama/ollama:latest + container_name: documentdb-agentic-memory-ollama-pull + depends_on: + ollama: + condition: service_healthy + environment: + # Point the ollama CLI at the server container rather than its own + # (unstarted) localhost daemon. + OLLAMA_HOST: http://ollama:11434 + entrypoint: + - /bin/sh + - -c + - ollama pull "${MEMORY_EMBEDDING_MODEL:-nomic-embed-text}" + restart: "no" + +volumes: + ollama-models: + name: documentdb-agentic-memory-ollama-models diff --git a/docs/architecture.md b/docs/architecture.md index fd777e6..66008bb 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -76,7 +76,9 @@ differs. Indexes (declared in `src/storage/graph/schema.ts`): - `graph_entities`: text index over `_id`, `entityType`, `observations.text` - (powers `search_nodes`); `entityType` ascending; `updatedAt` ascending. + (powers keyword `search_nodes`); `entityType` ascending; `updatedAt` + ascending. When DocumentDB Search is enabled, a `cosmosSearch` vector index + over the `embedding` field is also created (see below). - `graph_relations`: `from`, `to`, `relationType`, `createdAt` (each separately, ascending). @@ -91,14 +93,56 @@ delete_entities delete_relations delete_observations read_graph search_nodes open_nodes ``` -`search_nodes` uses the text index; results include any relation whose -**both** endpoints are in the match set — the same containment rule the -upstream server applies. +`search_nodes` matches on the text index by default. When **DocumentDB +Search** is enabled it becomes a hybrid search (keyword + vector, fused with +RRF — see below). Either way, results include any relation whose **both** +endpoints are in the match set — the same containment rule the upstream +server applies. Writes are race-safe: `createEntities` / `createRelations` use `insertMany({ ordered: false })` and swallow E11000 duplicate-key errors, so concurrent sessions adding the same entity converge instead of crashing. +### DocumentDB Search (embedding-backed vector search) + +An optional layer that adds semantic recall to `search_nodes`, modeled on the +hybrid dense-vector + knowledge-graph retrieval used by +[OmniMem](https://omnimem.org). It is **off-by-default-safe**: if no embedder +is configured or the backend is unreachable, everything degrades to the +keyword-only path above and entities are still stored, preserving drop-in +compatibility with the upstream server. + +**Data model.** `EntityDoc` gains three optional fields +(`src/storage/graph/schema.ts`): `embedding` (the vector), `embeddingModel` +(which model produced it, so stale vectors can be detected), and +`embeddingText` (the exact text embedded — name + type + observations joined +by newline, via `entityEmbeddingText()`). + +**Embedders.** A pluggable `Embedder` interface +(`src/shared/embeddings/`) with three HTTP implementations — Ollama +(local, default), OpenAI, and Azure OpenAI — selected by +`MEMORY_EMBEDDING_PROVIDER`. `createEmbedder()` is the factory: it probes the +backend once at startup to learn the vector dimensionality and **returns +`null` (never throws)** when disabled, misconfigured, or offline. The +providers are dependency-free (global `fetch`). + +**Index.** When an embedder is present the store issues a `createIndexes` +command with `key: { embedding: "cosmosSearch" }` and `cosmosSearchOptions` +(`kind`, `similarity`, `dimensions`, plus `numLists` for IVF or +`m`/`efConstruction` for HNSW). This is DocumentDB's native vector index; the +`cosmosSearch` token is a wire-level name only — the feature is called +DocumentDB Search everywhere else. + +**Query.** `searchNodes` runs the keyword search (`$text`) and the vector kNN +search (`$search: { cosmosSearch: { vector, path, k } }`) in parallel, then +fuses the two ranked ID lists with **reciprocal-rank fusion** (`fuseRankings`, +`RRF_K = 60`). A vector-query failure is caught and treated as an empty vector +list, so hybrid search never regresses below keyword-only. + +**Backfill.** `documentdb-memory graph reembed [--all]` re-embeds existing +entities (only stale/missing vectors, or all) after enabling the feature or +switching models. + ## Track 2: session-history (`history_*`) Mirror of the SQLite database that Copilot CLI writes at @@ -166,7 +210,9 @@ start it: 1. Loads `.env` via `dotenv` (best-effort). 2. Opens one shared `MongoClient` (`src/shared/mongo.ts`). -3. Runs the index bootstrap for both stores (idempotent). +3. Runs the index bootstrap for both stores (idempotent). When an embedding + provider is configured and reachable, this also builds the DocumentDB + Search vector index. 4. Registers all 17 tools (9 graph + 8 history) via the MCP SDK. 5. Installs SIGINT/SIGTERM handlers for clean shutdown. @@ -197,7 +243,7 @@ appetite for containerization: docker compose -f compose.full.yml up -d ``` -Brings up four services: +Brings up all services on the `copilot-memory-net` bridge: - `documentdb` — `ghcr.io/documentdb/documentdb/documentdb-local:latest` with `--username localadmin --password Admin100` passed as CLI args @@ -211,6 +257,12 @@ Brings up four services: - `documentdb-memory-sync` — same image, runs `documentdb-memory sessions sync --watch`; mounts your `~/.copilot` directory read-only at `/copilot`. +- `ollama` — local embedding backend for DocumentDB Search, reached at + `http://ollama:11434` inside the network. Models persist in the + `ollama-models` volume. +- `ollama-pull` — one-shot init container that pulls the embedding model + (`MEMORY_EMBEDDING_MODEL`, default `nomic-embed-text`) then exits; the + MCP server waits on its `service_completed_successfully`. Inside the compose network, DocumentDB requires TLS (the image enforces it with a self-signed cert), so the internal URI is @@ -256,6 +308,7 @@ load-bearing: | `COPILOT_DIR` | `$HOME/.copilot` | `compose.full.yml` (mounts read-only) | | `MEMORY_LOG_LEVEL` | `info` | pino (server + sync) | | `SYNC_INTERVAL` | `30s` | sync daemon `--watch` | +| `MEMORY_EMBEDDING_PROVIDER` | `local` | DocumentDB Search (server + CLI) | | `FUSE_HOST_BIND` | (off) | `compose.fuse-host-bind.yml` | | `FUSE_MOUNT_PATH` | (off) | `doctor` only — manual opt-in | diff --git a/docs/cli.md b/docs/cli.md index 1361ac6..b35ec75 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -195,6 +195,25 @@ documentdb-memory graph import [--merge | --replace] observations appended. - `--replace` — wipe the graph first, then import. Use with care. +### `graph reembed` + +Recompute DocumentDB Search embeddings for existing entities. Useful after +enabling the feature for the first time, or after switching embedding models +(existing vectors from a different model are stale). + +```text +documentdb-memory graph reembed [--all] +``` + +- default — embed only entities whose vector is missing or was produced by a + different model. +- `--all` — force re-embed every entity. + +Requires an embedding provider to be configured and reachable +(`MEMORY_EMBEDDING_PROVIDER`, default `local`). If none is available the +command warns and makes no changes. See the +[DocumentDB Search section of the README](../README.md#documentdb-search--semantic-recall-for-the-knowledge-graph). + --- ## `documentdb-memory sessions` diff --git a/src/cli/commands/graph/add.ts b/src/cli/commands/graph/add.ts index a225f8c..0e16715 100644 --- a/src/cli/commands/graph/add.ts +++ b/src/cli/commands/graph/add.ts @@ -33,7 +33,7 @@ export function registerAdd(graph: Command): void { } else { info(`entity "${name}" already exists; no changes`); } - }); + }, { withEmbedder: true }); }); add @@ -71,7 +71,7 @@ export function registerAdd(graph: Command): void { } throw err; } - }); + }, { withEmbedder: true }); }); } diff --git a/src/cli/commands/graph/index.ts b/src/cli/commands/graph/index.ts index e1d6b9e..8a30c72 100644 --- a/src/cli/commands/graph/index.ts +++ b/src/cli/commands/graph/index.ts @@ -9,6 +9,7 @@ import type { Command } from "commander"; import { registerAdd } from "./add.js"; import { registerDelete } from "./delete.js"; import { registerPrune } from "./prune.js"; +import { registerReembed } from "./reembed.js"; import { registerWipe } from "./wipe.js"; import { registerExport, registerImport } from "./io.js"; import { registerDump, registerStats } from "./inspect.js"; @@ -24,6 +25,7 @@ export function registerGraph(program: Command): void { registerAdd(graph); registerDelete(graph); registerPrune(graph); + registerReembed(graph); registerWipe(graph); registerExport(graph); registerImport(graph); diff --git a/src/cli/commands/graph/reembed.ts b/src/cli/commands/graph/reembed.ts new file mode 100644 index 0000000..753cf15 --- /dev/null +++ b/src/cli/commands/graph/reembed.ts @@ -0,0 +1,42 @@ +// `documentdb-memory graph reembed` — (re)compute embedding vectors for the +// knowledge-graph entities so DocumentDB Search can find them semantically. +// +// Use this to backfill vectors for entities created before embeddings were +// enabled, or after switching the embedding model (which changes the vector +// dimensionality and therefore requires a full re-embed + index rebuild). + +import type { Command } from "commander"; +import { info, ok, runWithStore, warn } from "./util.js"; + +interface ReembedOptions { + all?: boolean; +} + +export function registerReembed(graph: Command): void { + graph + .command("reembed") + .description("(Re)compute embedding vectors for entities (DocumentDB Search).") + .option( + "--all", + "Re-embed every entity, not just those missing or stale vectors.", + false, + ) + .action(async function (this: Command, opts: ReembedOptions) { + await runWithStore( + this, + async ({ store }) => { + if (!store.vectorSearchEnabled) { + warn( + "no embedding provider available (set MEMORY_EMBEDDING_PROVIDER and ensure the backend is reachable); nothing to do", + ); + return 1; + } + info(opts.all === true ? "re-embedding all entities…" : "embedding missing/stale entities…"); + const { scanned, embedded } = await store.reembed({ onlyStale: opts.all !== true }); + ok(`reembed complete: scanned ${scanned}, embedded ${embedded}`); + return 0; + }, + { withEmbedder: true }, + ); + }); +} diff --git a/src/cli/commands/graph/util.ts b/src/cli/commands/graph/util.ts index 1e2b583..0a0e8a0 100644 --- a/src/cli/commands/graph/util.ts +++ b/src/cli/commands/graph/util.ts @@ -11,6 +11,8 @@ import pc from "picocolors"; import type { Db } from "mongodb"; import { ConfigError, loadConfig, type AppConfig } from "../../../shared/config.js"; import { closeMongo, getMongo, runIndexBootstrap } from "../../../shared/mongo.js"; +import { createEmbedder } from "../../../shared/embeddings/index.js"; +import { createLogger } from "../../../shared/logging.js"; import { KnowledgeGraphStore } from "../../../storage/graph/index.js"; // Options every leaf graph command inherits from the parent `graph` command. @@ -45,6 +47,9 @@ export function readGlobalOptions(cmd: Command): GraphGlobalOptions { // * Opens the shared Mongo handle. // * Runs `runIndexBootstrap` so indexes exist on first run. Importing the // graph store module already registers its bootstrap as a side effect. +// * When `options.withEmbedder` is set, builds the embedding provider (so +// writes produce vectors) and ensures the vector index exists. Read-only +// admin commands omit this to avoid a network probe on every invocation. // * Calls `body(ctx)` which is expected to return a process exit code. // * Always closes the Mongo client in `finally`. // @@ -53,6 +58,7 @@ export function readGlobalOptions(cmd: Command): GraphGlobalOptions { export async function runWithStore( cmd: Command, body: (ctx: GraphContext) => Promise, + options: { withEmbedder?: boolean } = {}, ): Promise { const opts = readGlobalOptions(cmd); const debug = opts.debug === true || process.env.DEBUG === "1"; @@ -62,7 +68,19 @@ export async function runWithStore( const config = loadConfig({ uri: opts.uri, db: opts.db }); const handle = await getMongo(config); await runIndexBootstrap(handle.db); - const store = new KnowledgeGraphStore(handle.db); + let store: KnowledgeGraphStore; + if (options.withEmbedder === true) { + const log = createLogger(config.logLevel, "documentdb-memory-cli"); + const embedder = await createEmbedder(config.embedding, log); + store = new KnowledgeGraphStore(handle.db, { + embedder, + embeddingConfig: config.embedding, + logger: log, + }); + if (embedder !== null) await store.ensureVectorIndex(); + } else { + store = new KnowledgeGraphStore(handle.db); + } const result = await body({ config, db: handle.db, store }); code = typeof result === "number" ? result : 0; } catch (err) { diff --git a/src/server/index.ts b/src/server/index.ts index 077f5b9..6c69dce 100644 --- a/src/server/index.ts +++ b/src/server/index.ts @@ -15,6 +15,7 @@ import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js" import { loadConfig } from "../shared/config.js"; import { createLogger, type Logger } from "../shared/logging.js"; import { closeMongo, getMongo, runIndexBootstrap } from "../shared/mongo.js"; +import { createEmbedder } from "../shared/embeddings/index.js"; import { KnowledgeGraphStore } from "../storage/graph/index.js"; // Side-effect imports: each storage module registers its `ensureIndexes` // callback via `registerIndexBootstrap` at module load. Importing the history @@ -67,7 +68,17 @@ async function main(): Promise { await runIndexBootstrap(db); log.debug({ db: config.documentdbDb }, "indexes bootstrapped"); - const graphStore = new KnowledgeGraphStore(db); + // DocumentDB Search: build the embedder (null if disabled/unreachable) and + // ensure the vector index exists before serving. + const embedder = await createEmbedder(config.embedding, log); + const graphStore = new KnowledgeGraphStore(db, { + embedder, + embeddingConfig: config.embedding, + logger: log, + }); + if (embedder !== null) { + await graphStore.ensureVectorIndex(); + } const historyStore = new SessionHistoryStore(db); const server = new McpServer( diff --git a/src/shared/config.ts b/src/shared/config.ts index 5eddd1c..5e3b48a 100644 --- a/src/shared/config.ts +++ b/src/shared/config.ts @@ -8,10 +8,75 @@ export interface AppConfig { copilotSessionStore: string; logLevel: LogLevel; syncIntervalMs: number; + embedding: EmbeddingConfig; } export type LogLevel = "trace" | "debug" | "info" | "warn" | "error" | "fatal"; +// -- DocumentDB Search (vector search) configuration ------------------------ +// +// "DocumentDB Search" is our name for the embedding-backed vector search that +// augments the graph `search_nodes` tool. Embeddings are produced by a +// pluggable provider; the vectors are indexed and queried through DocumentDB's +// native vector index (wire token `cosmosSearch`). All fields have defaults so +// an unconfigured deployment still works — it just falls back to text-only +// search when the provider is `none` or unreachable. + +export type EmbeddingProvider = "local" | "ollama" | "openai" | "azure-openai" | "none"; + +// Vector index algorithm + distance metric. `vector-ivf` is the lighter, +// faster-to-build default; `vector-hnsw` trades build cost for recall. The +// similarity metric MUST match how the provider normalises its vectors — +// `COS` (cosine) is the safe default for every embedding model we support. +export type VectorIndexKind = "vector-ivf" | "vector-hnsw"; +export type VectorSimilarity = "COS" | "L2" | "IP"; + +export interface EmbeddingConfig { + // `local` is an alias for `ollama` (a self-hosted Ollama HTTP endpoint), and + // is the default so a self-hosted deployment needs no external credentials. + provider: EmbeddingProvider; + model: string; + // Base URL for `ollama`/`local` (default http://localhost:11434) and an + // optional override for `openai` (default https://api.openai.com/v1). + baseUrl?: string; + // Bearer token (openai) or api-key (azure-openai). + apiKey?: string; + // Azure OpenAI resource endpoint, e.g. https://my-res.openai.azure.com. + endpoint?: string; + // Azure OpenAI REST api-version. + apiVersion: string; + // Explicit embedding dimensionality. When undefined we probe the provider + // once at startup and use the observed length to build the vector index. + dimensions?: number; + // Vector index tuning. + indexKind: VectorIndexKind; + similarity: VectorSimilarity; + numLists: number; // vector-ivf only + m: number; // vector-hnsw only + efConstruction: number; // vector-hnsw only +} + +const VALID_EMBEDDING_PROVIDERS: ReadonlySet = new Set([ + "local", + "ollama", + "openai", + "azure-openai", + "none", +]); + +const VALID_INDEX_KINDS: ReadonlySet = new Set(["vector-ivf", "vector-hnsw"]); +const VALID_SIMILARITIES: ReadonlySet = new Set(["COS", "L2", "IP"]); + +// Sensible default embedding model per provider. Kept small + widely-available +// so a fresh install has a working model name without extra configuration. +const DEFAULT_MODEL: Record = { + local: "nomic-embed-text", + ollama: "nomic-embed-text", + openai: "text-embedding-3-small", + "azure-openai": "text-embedding-3-small", + none: "", +}; + const VALID_LOG_LEVELS: ReadonlySet = new Set([ "trace", "debug", @@ -63,15 +128,81 @@ export function loadConfig(overrides: ConfigOverrides = {}): AppConfig { throw new ConfigError(`SYNC_INTERVAL must be a positive duration (got "${syncIntervalStr}").`); } + const embedding = loadEmbeddingConfig(); + return { documentdbUri: uri, documentdbDb: db, copilotSessionStore, logLevel: logLevelRaw as LogLevel, syncIntervalMs, + embedding, + }; +} + +// Parse the DocumentDB Search / embedding settings from the environment. All +// values have defaults; only genuinely invalid values (bad enum, non-numeric +// dimensions) raise `ConfigError`. Absent credentials are allowed here — the +// embedder factory downgrades to text-only search at runtime instead. +function loadEmbeddingConfig(): EmbeddingConfig { + const providerRaw = (process.env.MEMORY_EMBEDDING_PROVIDER ?? "local").toLowerCase(); + if (!VALID_EMBEDDING_PROVIDERS.has(providerRaw as EmbeddingProvider)) { + throw new ConfigError( + `Invalid MEMORY_EMBEDDING_PROVIDER "${providerRaw}". Expected one of: ${[...VALID_EMBEDDING_PROVIDERS].join(", ")}.`, + ); + } + const provider = providerRaw as EmbeddingProvider; + + const model = + envString("MEMORY_EMBEDDING_MODEL") ?? DEFAULT_MODEL[provider] ?? DEFAULT_MODEL.local; + + const indexKindRaw = (process.env.MEMORY_VECTOR_INDEX_KIND ?? "vector-ivf").toLowerCase(); + if (!VALID_INDEX_KINDS.has(indexKindRaw as VectorIndexKind)) { + throw new ConfigError( + `Invalid MEMORY_VECTOR_INDEX_KIND "${indexKindRaw}". Expected one of: ${[...VALID_INDEX_KINDS].join(", ")}.`, + ); + } + + const similarityRaw = (process.env.MEMORY_VECTOR_SIMILARITY ?? "COS").toUpperCase(); + if (!VALID_SIMILARITIES.has(similarityRaw as VectorSimilarity)) { + throw new ConfigError( + `Invalid MEMORY_VECTOR_SIMILARITY "${similarityRaw}". Expected one of: ${[...VALID_SIMILARITIES].join(", ")}.`, + ); + } + + return { + provider, + model, + baseUrl: envString("MEMORY_EMBEDDING_BASE_URL"), + apiKey: envString("MEMORY_EMBEDDING_API_KEY"), + endpoint: envString("MEMORY_EMBEDDING_ENDPOINT"), + apiVersion: envString("MEMORY_EMBEDDING_API_VERSION") ?? "2024-02-01", + dimensions: envPositiveInt("MEMORY_EMBEDDING_DIMENSIONS"), + indexKind: indexKindRaw as VectorIndexKind, + similarity: similarityRaw as VectorSimilarity, + numLists: envPositiveInt("MEMORY_VECTOR_NUM_LISTS") ?? 100, + m: envPositiveInt("MEMORY_VECTOR_HNSW_M") ?? 16, + efConstruction: envPositiveInt("MEMORY_VECTOR_HNSW_EF_CONSTRUCTION") ?? 64, }; } +function envString(key: string): string | undefined { + const v = process.env[key]; + if (v === undefined) return undefined; + const trimmed = v.trim(); + return trimmed === "" ? undefined : trimmed; +} + +function envPositiveInt(key: string): number | undefined { + const raw = envString(key); + if (raw === undefined) return undefined; + const n = Number(raw); + if (!Number.isInteger(n) || n <= 0) { + throw new ConfigError(`${key} must be a positive integer (got "${raw}").`); + } + return n; +} + export function defaultSessionStorePath(): string { return resolve(homedir(), ".copilot", "session-store.db"); } diff --git a/src/shared/embeddings/index.ts b/src/shared/embeddings/index.ts new file mode 100644 index 0000000..1714997 --- /dev/null +++ b/src/shared/embeddings/index.ts @@ -0,0 +1,106 @@ +// Embedder factory for DocumentDB Search. +// +// `createEmbedder(config, log)` returns a ready, health-checked `Embedder`, or +// `null` when embeddings are disabled (`provider: "none"`) or the backend is +// unreachable. Returning `null` — never throwing — is deliberate: the graph +// store treats a missing embedder as "text-only search", so a misconfigured or +// offline embedding backend degrades gracefully instead of taking the whole +// MCP server down. + +import type { EmbeddingConfig } from "../config.js"; +import type { Logger } from "../logging.js"; +import { AzureOpenAIEmbedder, OllamaEmbedder, OpenAIEmbedder } from "./providers.js"; +import { EmbeddingError, type Embedder } from "./types.js"; + +export type { Embedder } from "./types.js"; +export { EmbeddingError } from "./types.js"; + +const DEFAULT_OLLAMA_URL = "http://localhost:11434"; +const DEFAULT_OPENAI_URL = "https://api.openai.com/v1"; +// Short, cheap string used to probe the backend and learn the vector length. +const PROBE_TEXT = "documentdb search probe"; + +// Build a concrete embedder for the given dimensions. Returns null (with a +// warning) when required credentials/endpoints are missing so the caller can +// fall back to text-only search. +function construct(config: EmbeddingConfig, dimensions: number, log?: Logger): Embedder | null { + switch (config.provider) { + case "local": + case "ollama": + return new OllamaEmbedder(config.model, dimensions, config.baseUrl ?? DEFAULT_OLLAMA_URL); + case "openai": { + if (config.apiKey === undefined) { + log?.warn("embeddings: openai provider selected but MEMORY_EMBEDDING_API_KEY is unset"); + return null; + } + return new OpenAIEmbedder( + config.model, + dimensions, + config.baseUrl ?? DEFAULT_OPENAI_URL, + config.apiKey, + ); + } + case "azure-openai": { + if (config.endpoint === undefined || config.apiKey === undefined) { + log?.warn( + "embeddings: azure-openai provider needs MEMORY_EMBEDDING_ENDPOINT and MEMORY_EMBEDDING_API_KEY", + ); + return null; + } + return new AzureOpenAIEmbedder( + config.model, + dimensions, + config.endpoint, + config.apiKey, + config.apiVersion, + ); + } + case "none": + return null; + } +} + +export async function createEmbedder( + config: EmbeddingConfig, + log?: Logger, +): Promise { + if (config.provider === "none") { + log?.debug("embeddings: disabled (provider=none); using text-only search"); + return null; + } + + // Provisional instance (embedding does not depend on `dimensions`) used only + // to probe. We then rebuild with the resolved dimension count. + const provisional = construct(config, config.dimensions ?? 0, log); + if (provisional === null) return null; + + let probeVector: number[]; + try { + const [vec] = await provisional.embed([PROBE_TEXT]); + if (vec === undefined) throw new EmbeddingError(`${config.provider}: probe returned no vector`); + probeVector = vec; + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + log?.warn( + { err: message, provider: config.provider, model: config.model }, + "embeddings: provider unreachable; falling back to text-only search", + ); + return null; + } + + const observed = probeVector.length; + if (config.dimensions !== undefined && config.dimensions !== observed) { + log?.warn( + { configured: config.dimensions, observed }, + "embeddings: MEMORY_EMBEDDING_DIMENSIONS does not match the model; using observed length", + ); + } + + const embedder = construct(config, observed, log); + if (embedder === null) return null; + log?.info( + { provider: embedder.provider, model: embedder.model, dimensions: embedder.dimensions }, + "embeddings: DocumentDB Search enabled", + ); + return embedder; +} diff --git a/src/shared/embeddings/providers.ts b/src/shared/embeddings/providers.ts new file mode 100644 index 0000000..4b9b111 --- /dev/null +++ b/src/shared/embeddings/providers.ts @@ -0,0 +1,161 @@ +// HTTP-based embedder implementations. +// +// Each class talks to one embedding backend over `fetch` (global since +// Node 18). None of them hold the vector dimensionality up front — the factory +// probes with a short string and reads the observed length. Any network or +// shape error becomes an `EmbeddingError` so the factory can downgrade cleanly. + +import { EmbeddingError, type Embedder } from "./types.js"; + +const REQUEST_TIMEOUT_MS = 30_000; + +// -- shared fetch helper ---------------------------------------------------- + +interface PostJsonArgs { + url: string; + headers: Record; + body: unknown; + provider: string; +} + +async function postJson({ url, headers, body, provider }: PostJsonArgs): Promise { + const controller = new AbortController(); + const timer = setTimeout(() => controller.abort(), REQUEST_TIMEOUT_MS); + let res: Response; + try { + res = await fetch(url, { + method: "POST", + headers: { "content-type": "application/json", ...headers }, + body: JSON.stringify(body), + signal: controller.signal, + }); + } catch (err) { + throw new EmbeddingError(`${provider}: request to ${url} failed`, { cause: err }); + } finally { + clearTimeout(timer); + } + if (!res.ok) { + const text = await res.text().catch(() => ""); + throw new EmbeddingError( + `${provider}: ${res.status} ${res.statusText}${text ? ` — ${text.slice(0, 200)}` : ""}`, + ); + } + try { + return (await res.json()) as T; + } catch (err) { + throw new EmbeddingError(`${provider}: response was not valid JSON`, { cause: err }); + } +} + +function assertVectors(vectors: unknown, provider: string, expected: number): number[][] { + if (!Array.isArray(vectors) || vectors.length !== expected) { + throw new EmbeddingError( + `${provider}: expected ${expected} embedding vector(s), got ${ + Array.isArray(vectors) ? vectors.length : typeof vectors + }`, + ); + } + for (const v of vectors) { + if (!Array.isArray(v) || v.length === 0 || !v.every((n) => typeof n === "number")) { + throw new EmbeddingError(`${provider}: embedding vector had an unexpected shape`); + } + } + return vectors as number[][]; +} + +// -- Ollama (self-hosted / "local") ----------------------------------------- +// +// POST {baseUrl}/api/embed { model, input: string[] } -> { embeddings: number[][] } + +export class OllamaEmbedder implements Embedder { + readonly provider = "ollama"; + constructor( + readonly model: string, + readonly dimensions: number, + private readonly baseUrl: string, + ) {} + + async embed(texts: string[]): Promise { + if (texts.length === 0) return []; + const data = await postJson<{ embeddings?: number[][] }>({ + url: `${this.baseUrl.replace(/\/$/, "")}/api/embed`, + headers: {}, + body: { model: this.model, input: texts }, + provider: this.provider, + }); + return assertVectors(data.embeddings, this.provider, texts.length); + } +} + +// -- OpenAI ----------------------------------------------------------------- +// +// POST {baseUrl}/embeddings { model, input: string[] } +// -> { data: [{ index, embedding: number[] }, ...] } + +interface OpenAIEmbeddingResponse { + data?: { index: number; embedding: number[] }[]; +} + +function orderOpenAIVectors( + data: OpenAIEmbeddingResponse["data"], + provider: string, + expected: number, +): number[][] { + if (!Array.isArray(data)) { + throw new EmbeddingError(`${provider}: response missing "data" array`); + } + const ordered = [...data].sort((a, b) => a.index - b.index).map((d) => d.embedding); + return assertVectors(ordered, provider, expected); +} + +export class OpenAIEmbedder implements Embedder { + readonly provider = "openai"; + constructor( + readonly model: string, + readonly dimensions: number, + private readonly baseUrl: string, + private readonly apiKey: string, + ) {} + + async embed(texts: string[]): Promise { + if (texts.length === 0) return []; + const data = await postJson({ + url: `${this.baseUrl.replace(/\/$/, "")}/embeddings`, + headers: { authorization: `Bearer ${this.apiKey}` }, + body: { model: this.model, input: texts }, + provider: this.provider, + }); + return orderOpenAIVectors(data.data, this.provider, texts.length); + } +} + +// -- Azure OpenAI ----------------------------------------------------------- +// +// POST {endpoint}/openai/deployments/{model}/embeddings?api-version=... +// header: api-key: +// body: { input: string[] } +// -> { data: [{ index, embedding }] } (same shape as OpenAI) + +export class AzureOpenAIEmbedder implements Embedder { + readonly provider = "azure-openai"; + constructor( + readonly model: string, + readonly dimensions: number, + private readonly endpoint: string, + private readonly apiKey: string, + private readonly apiVersion: string, + ) {} + + async embed(texts: string[]): Promise { + if (texts.length === 0) return []; + const base = this.endpoint.replace(/\/$/, ""); + const url = `${base}/openai/deployments/${encodeURIComponent(this.model)}/embeddings?api-version=${encodeURIComponent(this.apiVersion)}`; + const data = await postJson({ + url, + headers: { "api-key": this.apiKey }, + body: { input: texts }, + provider: this.provider, + }); + return orderOpenAIVectors(data.data, this.provider, texts.length); + } +} diff --git a/src/shared/embeddings/types.ts b/src/shared/embeddings/types.ts new file mode 100644 index 0000000..f9339d2 --- /dev/null +++ b/src/shared/embeddings/types.ts @@ -0,0 +1,29 @@ +// Embedder abstraction for DocumentDB Search. +// +// An `Embedder` turns text into dense vectors. Implementations are thin HTTP +// clients (Ollama / OpenAI / Azure OpenAI) so the project stays dependency-free +// — no native model runtime is bundled. The factory in `./index.ts` picks an +// implementation from `EmbeddingConfig` and health-checks it once; if the probe +// fails the whole feature degrades to text-only search rather than erroring. + +export interface Embedder { + // Provider identifier, for logs/diagnostics (e.g. "ollama", "openai"). + readonly provider: string; + // Model name in use. + readonly model: string; + // Vector dimensionality. Known only after a successful `embed`/probe, so it + // is resolved during factory construction and fixed thereafter. + readonly dimensions: number; + // Embed a batch of texts, returning one vector per input in the same order. + embed(texts: string[]): Promise; +} + +// Raised by an embedder when the backend is unreachable or returns an +// unexpected shape. The factory catches this during the startup probe and +// downgrades to text-only search. +export class EmbeddingError extends Error { + constructor(message: string, options?: { cause?: unknown }) { + super(message, options); + this.name = "EmbeddingError"; + } +} diff --git a/src/storage/graph/index.ts b/src/storage/graph/index.ts index 9d152b0..f612817 100644 --- a/src/storage/graph/index.ts +++ b/src/storage/graph/index.ts @@ -2,10 +2,13 @@ // Importing this module also registers the index bootstrap as a side effect // (via `./store.js` → `registerIndexBootstrap`). -export { KnowledgeGraphStore, ensureGraphIndexes } from "./store.js"; +export { KnowledgeGraphStore, ensureGraphIndexes, type GraphStoreOptions } from "./store.js"; export { ENTITIES_COLLECTION, RELATIONS_COLLECTION, + VECTOR_INDEX_NAME, + EMBEDDING_FIELD, + entityEmbeddingText, relationId, type Entity, type EntityDoc, diff --git a/src/storage/graph/schema.ts b/src/storage/graph/schema.ts index 8fcd2b9..c991e2a 100644 --- a/src/storage/graph/schema.ts +++ b/src/storage/graph/schema.ts @@ -26,6 +26,14 @@ export interface EntityDoc { observations: EntityObservation[]; createdAt: Date; updatedAt: Date; + // DocumentDB Search fields (optional; present only when an embedder is + // configured). `embedding` is the dense vector indexed for vector search; + // `embeddingModel` records which model produced it (so a model change can be + // detected and re-embedded); `embeddingText` is the exact text that was + // embedded, used to skip re-embedding when nothing relevant changed. + embedding?: number[]; + embeddingModel?: string; + embeddingText?: string; } export interface RelationDoc { @@ -76,3 +84,23 @@ export interface ObservationDeletion { export function relationId(from: string, relationType: string, to: string): string { return `${from}__${relationType}__${to}`; } + +// Name of the DocumentDB vector index on `graph_entities.embedding`. The wire +// token for the index type is `cosmosSearch`; this project refers to the +// feature as "DocumentDB Search" everywhere else. +export const VECTOR_INDEX_NAME = "graph_entities_embedding"; +export const EMBEDDING_FIELD = "embedding"; + +// Canonical text an entity is embedded from. Combines the name, type, and all +// observations so vector search matches on any of them — mirroring the fields +// the `$text` index already covers. Keeping this in one place guarantees the +// write path and the `reembed` CLI produce identical vectors. +export function entityEmbeddingText(entity: { + name: string; + entityType: string; + observations: string[]; +}): string { + return [entity.name, entity.entityType, ...entity.observations] + .filter((s) => s.length > 0) + .join("\n"); +} diff --git a/src/storage/graph/store.ts b/src/storage/graph/store.ts index a8d9132..b64e3a6 100644 --- a/src/storage/graph/store.ts +++ b/src/storage/graph/store.ts @@ -8,9 +8,15 @@ import { type UpdateFilter, } from "mongodb"; import { registerIndexBootstrap } from "../../shared/mongo.js"; +import type { EmbeddingConfig } from "../../shared/config.js"; +import type { Embedder } from "../../shared/embeddings/index.js"; +import type { Logger } from "../../shared/logging.js"; import { + EMBEDDING_FIELD, ENTITIES_COLLECTION, RELATIONS_COLLECTION, + VECTOR_INDEX_NAME, + entityEmbeddingText, relationId, type Entity, type EntityDoc, @@ -27,6 +33,26 @@ import { // silently skips entities/relations that already exist. const DUPLICATE_KEY_ERROR = 11000; +// Reciprocal Rank Fusion constant. Standard value from the original RRF paper; +// dampens the contribution of low-ranked results so top hits from either the +// vector or text list dominate the fused ordering. +const RRF_K = 60; + +// How many nearest neighbours to request from the vector index in a hybrid +// search. Kept modest — text matches are unioned on top, and RRF re-ranks the +// combined set. +const DEFAULT_VECTOR_K = 20; + +// Options controlling the optional DocumentDB Search (vector) behaviour. When +// `embedder` is absent the store behaves exactly as before: `$text`-only search +// and no `embedding` field written. This keeps the store a drop-in replacement +// for the official memory server when embeddings are not configured. +export interface GraphStoreOptions { + embedder?: Embedder | null; + embeddingConfig?: EmbeddingConfig; + logger?: Logger; +} + // `KnowledgeGraphStore` is the storage-layer counterpart to the official // `KnowledgeGraphManager`. Method names and semantics are kept identical so // the MCP tool layer can map 1:1 onto it. @@ -39,12 +65,26 @@ const DUPLICATE_KEY_ERROR = 11000; // * Errors propagate. The only "expected" miss is `addObservations` against // a non-existent entity, which throws to match official semantics. export class KnowledgeGraphStore { + private readonly db: Db; private readonly entities: Collection; private readonly relations: Collection; + private readonly embedder: Embedder | null; + private readonly embeddingConfig: EmbeddingConfig | undefined; + private readonly log: Logger | undefined; - constructor(db: Db) { + constructor(db: Db, options: GraphStoreOptions = {}) { + this.db = db; this.entities = db.collection(ENTITIES_COLLECTION); this.relations = db.collection(RELATIONS_COLLECTION); + this.embedder = options.embedder ?? null; + this.embeddingConfig = options.embeddingConfig; + this.log = options.logger; + } + + // True when vector search is available. Exposed so callers (e.g. the doctor + // command) can report whether DocumentDB Search is active. + get vectorSearchEnabled(): boolean { + return this.embedder !== null; } // -- create -------------------------------------------------------------- @@ -80,6 +120,15 @@ export class KnowledgeGraphStore { updatedAt: now, })); + // Best-effort embeddings: if an embedder is configured, embed each new + // entity so it is immediately vector-searchable. Failures are swallowed — + // the entity is still stored, just without a vector until the next + // `reembed`, keeping writes resilient to a flaky embedding backend. + await this.attachEmbeddings( + docs, + toInsert.map((e) => entityEmbeddingText(e)), + ); + try { await this.entities.insertMany(docs, { ordered: false }); return toInsert; @@ -138,7 +187,7 @@ export class KnowledgeGraphStore { const names = observations.map((o) => o.entityName); const docs = await this.entities - .find({ _id: { $in: names } }, { projection: { _id: 1, observations: 1 } }) + .find({ _id: { $in: names } }, { projection: { _id: 1, entityType: 1, observations: 1 } }) .toArray(); const byName = new Map(docs.map((d) => [d._id, d] as const)); @@ -154,6 +203,9 @@ export class KnowledgeGraphStore { const now = new Date(); const results: ObservationResult[] = []; const ops: AnyBulkWriteOperation[] = []; + // Entities whose observations changed, paired with the full text to embed. + // Collected here so we can batch a single embed() call after the loop. + const toEmbed: { name: string; text: string }[] = []; for (const input of observations) { const doc = byName.get(input.entityName); @@ -184,6 +236,42 @@ export class KnowledgeGraphStore { update, }, }); + // The entity content changed, so its vector is now stale. Re-embed + // from name + type + the full (existing + new) observation set. + const allTexts = [...doc.observations.map((o) => o.text), ...newTexts]; + toEmbed.push({ + name: input.entityName, + text: entityEmbeddingText({ + name: input.entityName, + entityType: doc.entityType, + observations: allTexts, + }), + }); + } + } + + // Best-effort re-embed of the changed entities as a second set of update + // ops. Failures are swallowed so the observation append still succeeds. + if (this.embedder !== null && toEmbed.length > 0) { + const vectors = await this.embedTexts(toEmbed.map((t) => t.text)); + if (vectors !== null) { + for (let i = 0; i < toEmbed.length; i++) { + const entry = toEmbed[i]; + const vec = vectors[i]; + if (entry === undefined || vec === undefined) continue; + ops.push({ + updateOne: { + filter: { _id: entry.name }, + update: { + $set: { + embedding: vec, + embeddingModel: this.embedder.model, + embeddingText: entry.text, + }, + }, + }, + }); + } } } @@ -233,6 +321,15 @@ export class KnowledgeGraphStore { // Missing entities / missing observation texts both result in zero // modifications, which is the silent-on-missing behaviour we want. await this.entities.bulkWrite(ops, { ordered: false }); + + // Removing observations changes entity content, so refresh their vectors. + // Best-effort: a failure here leaves a slightly stale vector, not an error. + if (this.embedder !== null) { + const names = deletions.filter((d) => d.observations.length > 0).map((d) => d.entityName); + if (names.length > 0) { + await this.reembed({ onlyStale: true, names }); + } + } } async deleteRelations(relations: Relation[]): Promise { @@ -255,23 +352,99 @@ export class KnowledgeGraphStore { } async searchNodes(query: string): Promise { - // Uses the text index declared in `ensureGraphIndexes`. Returns matched - // entities AND the relations whose endpoints are BOTH in the match set — - // the same containment rule the official server applies. - const matched = await this.entities.find({ $text: { $search: query } }).toArray(); - if (matched.length === 0) { + // Hybrid retrieval: combine keyword (`$text`) and, when an embedder is + // configured, semantic (vector) results, fused with Reciprocal Rank + // Fusion. When no embedder is present this collapses to the original + // text-only behaviour, keeping the store wire-compatible with the official + // memory server. In both cases we return matched entities AND the relations + // whose endpoints are BOTH in the match set — the containment rule the + // official server applies. + const trimmed = query.trim(); + if (trimmed.length === 0) { + return { entities: [], relations: [] }; + } + + const [textIds, vectorHits] = await Promise.all([ + this.textSearchIds(trimmed), + this.vectorSearchIds(trimmed, DEFAULT_VECTOR_K), + ]); + + // Fuse the two ranked lists. `vectorHits === null` means vector search was + // unavailable (no embedder or a runtime failure) — fall back to text order. + const rankedNames = + vectorHits === null + ? textIds + : fuseRankings(textIds, vectorHits.map((h) => h.id)); + + if (rankedNames.length === 0) { return { entities: [], relations: [] }; } - const names = matched.map((d) => d._id); + + const entityDocs = await this.entities.find({ _id: { $in: rankedNames } }).toArray(); + // Reorder the fetched docs to match the fused ranking (find() does not + // preserve `$in` order). + const byId = new Map(entityDocs.map((d) => [d._id, d] as const)); + const orderedEntities = rankedNames + .map((name) => byId.get(name)) + .filter((d): d is EntityDoc => d !== undefined); + const relationDocs = await this.relations - .find({ from: { $in: names }, to: { $in: names } }) + .find({ from: { $in: rankedNames }, to: { $in: rankedNames } }) .toArray(); + return { - entities: matched.map((d) => this.mapEntity(d)), + entities: orderedEntities.map((d) => this.mapEntity(d)), relations: relationDocs.map((d) => mapRelation(d)), }; } + // Keyword search via the `$text` index, returned as entity names ranked by + // text score (best first). + private async textSearchIds(query: string): Promise { + const docs = await this.entities + .aggregate<{ _id: string }>([ + { $match: { $text: { $search: query } } }, + { $addFields: { __score: { $meta: "textScore" } } }, + { $sort: { __score: -1 } }, + { $project: { _id: 1 } }, + ]) + .toArray(); + return docs.map((d) => d._id); + } + + // Semantic search via the DocumentDB vector index (`cosmosSearch`). Returns + // null — never throws — when embeddings are disabled or the query embedding / + // vector query fails, so callers transparently fall back to text-only search. + private async vectorSearchIds( + query: string, + k: number, + ): Promise<{ id: string; score: number }[] | null> { + if (this.embedder === null) return null; + const vectors = await this.embedTexts([query]); + const queryVector = vectors?.[0]; + if (queryVector === undefined) return null; + + try { + const docs = await this.entities + .aggregate<{ _id: string; __score: number }>([ + { + $search: { + cosmosSearch: { vector: queryVector, path: EMBEDDING_FIELD, k }, + }, + }, + { $project: { _id: 1, __score: { $meta: "searchScore" } } }, + ]) + .toArray(); + return docs.map((d) => ({ id: d._id, score: d.__score })); + } catch (err) { + this.log?.warn( + { err: err instanceof Error ? err.message : String(err) }, + "vector search failed; falling back to text-only results", + ); + return null; + } + } + async openNodes(names: string[]): Promise { if (names.length === 0) { return { entities: [], relations: [] }; @@ -287,6 +460,161 @@ export class KnowledgeGraphStore { }; } + // -- embeddings (DocumentDB Search) ------------------------------------- + + // Embed a batch of texts, returning one vector per input, or null on any + // failure / when no embedder is configured. Never throws: embeddings are a + // best-effort enhancement, so the calling write/search path stays resilient. + private async embedTexts(texts: string[]): Promise { + if (this.embedder === null) return null; + try { + return await this.embedder.embed(texts); + } catch (err) { + this.log?.warn( + { err: err instanceof Error ? err.message : String(err) }, + "embedding failed; entity stored without a vector", + ); + return null; + } + } + + // Mutate the given docs in place, attaching `embedding`/`embeddingModel`/ + // `embeddingText` computed from the parallel `texts` array. No-op when + // embeddings are unavailable. + private async attachEmbeddings(docs: EntityDoc[], texts: string[]): Promise { + if (this.embedder === null || docs.length === 0) return; + const vectors = await this.embedTexts(texts); + if (vectors === null) return; + for (let i = 0; i < docs.length; i++) { + const doc = docs[i]; + const vec = vectors[i]; + const text = texts[i]; + if (doc === undefined || vec === undefined || text === undefined) continue; + doc.embedding = vec; + doc.embeddingModel = this.embedder.model; + doc.embeddingText = text; + } + } + + // Re-embed entities, refreshing their stored vectors. When `onlyStale` is + // true (default), entities whose `embeddingText` already matches the current + // content AND were embedded by the current model are skipped. Returns the + // number of entities scanned and (re)embedded. Used by the `graph reembed` + // CLI command and after observation deletions. + async reembed( + options: { onlyStale?: boolean; names?: string[] } = {}, + ): Promise<{ scanned: number; embedded: number }> { + if (this.embedder === null) return { scanned: 0, embedded: 0 }; + const onlyStale = options.onlyStale ?? true; + const filter: Filter = + options.names && options.names.length > 0 ? { _id: { $in: options.names } } : {}; + + let scanned = 0; + let embedded = 0; + const batch: EntityDoc[] = []; + + const flush = async (): Promise => { + if (batch.length === 0) return; + const texts = batch.map((d) => + entityEmbeddingText({ + name: d._id, + entityType: d.entityType, + observations: d.observations.map((o) => o.text), + }), + ); + const vectors = await this.embedTexts(texts); + if (vectors !== null) { + const ops: AnyBulkWriteOperation[] = []; + for (let i = 0; i < batch.length; i++) { + const doc = batch[i]; + const vec = vectors[i]; + const text = texts[i]; + if (doc === undefined || vec === undefined || text === undefined) continue; + ops.push({ + updateOne: { + filter: { _id: doc._id }, + update: { + $set: { + embedding: vec, + embeddingModel: this.embedder!.model, + embeddingText: text, + }, + }, + }, + }); + } + if (ops.length > 0) { + await this.entities.bulkWrite(ops, { ordered: false }); + embedded += ops.length; + } + } + batch.length = 0; + }; + + const cursor = this.entities.find(filter); + for await (const doc of cursor) { + scanned++; + if (onlyStale) { + const text = entityEmbeddingText({ + name: doc._id, + entityType: doc.entityType, + observations: doc.observations.map((o) => o.text), + }); + if (doc.embeddingModel === this.embedder.model && doc.embeddingText === text) { + continue; + } + } + batch.push(doc); + if (batch.length >= 32) await flush(); + } + await flush(); + return { scanned, embedded }; + } + + // Create the DocumentDB vector index on `graph_entities.embedding`. Requires + // an embedder (for the resolved dimensionality). Safe to call repeatedly — + // an existing equivalent index makes this a no-op. When embeddings are + // disabled this returns false without touching the database. + async ensureVectorIndex(): Promise { + if (this.embedder === null) return false; + const cfg = this.embeddingConfig; + const kind = cfg?.indexKind ?? "vector-ivf"; + const similarity = cfg?.similarity ?? "COS"; + const cosmosSearchOptions: Record = { + kind, + similarity, + dimensions: this.embedder.dimensions, + }; + if (kind === "vector-ivf") { + cosmosSearchOptions.numLists = cfg?.numLists ?? 100; + } else { + cosmosSearchOptions.m = cfg?.m ?? 16; + cosmosSearchOptions.efConstruction = cfg?.efConstruction ?? 64; + } + + try { + await this.db.command({ + createIndexes: ENTITIES_COLLECTION, + indexes: [ + { + name: VECTOR_INDEX_NAME, + key: { [EMBEDDING_FIELD]: "cosmosSearch" }, + cosmosSearchOptions, + }, + ], + }); + return true; + } catch (err) { + // An index already exists with different options (e.g. dimensions changed + // after switching models). Surface a clear, actionable message. + this.log?.warn( + { err: err instanceof Error ? err.message : String(err), kind, similarity }, + "could not create vector index; DocumentDB Search may be unavailable until it is recreated", + ); + return false; + } + } + // -- helpers ------------------------------------------------------------ private mapEntity(doc: EntityDoc): Entity { @@ -298,6 +626,32 @@ export class KnowledgeGraphStore { } } +// Reciprocal Rank Fusion of two ranked lists of entity names into a single +// ranked, de-duplicated list. Each list contributes 1/(RRF_K + rank) to a +// name's score; names appearing high in both lists rank best. Order within the +// output is by descending fused score, ties broken by first appearance. +// Exported for unit testing — it is a pure function with no I/O. +export function fuseRankings(listA: string[], listB: string[]): string[] { + const score = new Map(); + const firstSeen = new Map(); + let order = 0; + const add = (list: string[]): void => { + for (let rank = 0; rank < list.length; rank++) { + const name = list[rank]; + if (name === undefined) continue; + score.set(name, (score.get(name) ?? 0) + 1 / (RRF_K + rank)); + if (!firstSeen.has(name)) firstSeen.set(name, order++); + } + }; + add(listA); + add(listB); + return [...score.keys()].sort((a, b) => { + const diff = (score.get(b) ?? 0) - (score.get(a) ?? 0); + if (diff !== 0) return diff; + return (firstSeen.get(a) ?? 0) - (firstSeen.get(b) ?? 0); + }); +} + function mapRelation(doc: RelationDoc): Relation { return { from: doc.from, diff --git a/test/storage/graph-store-embeddings.test.ts b/test/storage/graph-store-embeddings.test.ts new file mode 100644 index 0000000..8713b13 --- /dev/null +++ b/test/storage/graph-store-embeddings.test.ts @@ -0,0 +1,194 @@ +import { afterAll, beforeAll, beforeEach, describe, expect, it } from "vitest"; +import { KnowledgeGraphStore, ensureGraphIndexes } from "../../src/storage/graph/store.js"; +import { ENTITIES_COLLECTION } from "../../src/storage/graph/schema.js"; +import type { EntityDoc } from "../../src/storage/graph/schema.js"; +import type { Embedder } from "../../src/shared/embeddings/index.js"; +import type { EmbeddingConfig } from "../../src/shared/config.js"; +import { closeTestClient, getTestDb, type TestDbHandle } from "../setup.js"; + +// Deterministic fake embedder: maps text length + a char sum to a fixed-length +// vector. No network. Records the batches it was asked to embed so tests can +// assert the write path calls it exactly when expected. +class FakeEmbedder implements Embedder { + readonly provider = "fake"; + readonly model = "fake-model-v1"; + readonly dimensions = 4; + readonly calls: string[][] = []; + private failNext = false; + + failOnce(): void { + this.failNext = true; + } + + async embed(texts: string[]): Promise { + if (this.failNext) { + this.failNext = false; + throw new Error("simulated embedding failure"); + } + this.calls.push(texts); + return texts.map((t) => { + const sum = [...t].reduce((acc, c) => acc + c.charCodeAt(0), 0); + return [t.length, sum % 97, (sum * 7) % 101, 1]; + }); + } +} + +const EMBEDDING_CONFIG: EmbeddingConfig = { + provider: "ollama", + model: "fake-model-v1", + apiVersion: "2024-02-01", + indexKind: "vector-ivf", + similarity: "COS", + numLists: 100, + m: 16, + efConstruction: 64, +}; + +describe("KnowledgeGraphStore — DocumentDB Search (embeddings)", () => { + let handle: TestDbHandle; + let embedder: FakeEmbedder; + let store: KnowledgeGraphStore; + + beforeAll(async () => { + handle = await getTestDb("graph_store_embeddings_tests"); + await ensureGraphIndexes(handle.db); + embedder = new FakeEmbedder(); + store = new KnowledgeGraphStore(handle.db, { embedder, embeddingConfig: EMBEDDING_CONFIG }); + }); + + afterAll(async () => { + await closeTestClient(handle); + }); + + beforeEach(async () => { + await handle.db.collection(ENTITIES_COLLECTION).deleteMany({}); + embedder.calls.length = 0; + }); + + const rawDoc = (name: string): Promise => + handle.db.collection(ENTITIES_COLLECTION).findOne({ _id: name }); + + it("writes an embedding vector + metadata on createEntities", async () => { + await store.createEntities([ + { name: "Alice", entityType: "person", observations: ["likes tea"] }, + ]); + const doc = await rawDoc("Alice"); + expect(doc?.embedding).toHaveLength(4); + expect(doc?.embeddingModel).toBe("fake-model-v1"); + expect(doc?.embeddingText).toContain("Alice"); + expect(doc?.embeddingText).toContain("likes tea"); + expect(embedder.calls).toHaveLength(1); + }); + + it("refreshes the embedding when addObservations changes content", async () => { + await store.createEntities([{ name: "Bob", entityType: "person", observations: [] }]); + const before = await rawDoc("Bob"); + await store.addObservations([{ entityName: "Bob", contents: ["now has a hobby"] }]); + const after = await rawDoc("Bob"); + expect(after?.embeddingText).toContain("now has a hobby"); + expect(after?.embedding).not.toEqual(before?.embedding); + }); + + it("does not re-embed when addObservations adds nothing new", async () => { + await store.createEntities([{ name: "Carol", entityType: "person", observations: ["x"] }]); + embedder.calls.length = 0; + await store.addObservations([{ entityName: "Carol", contents: ["x"] }]); // duplicate + expect(embedder.calls).toHaveLength(0); + }); + + it("still stores the entity when the embedder throws (graceful)", async () => { + embedder.failOnce(); + const created = await store.createEntities([ + { name: "Dave", entityType: "person", observations: ["resilient"] }, + ]); + expect(created.map((e) => e.name)).toEqual(["Dave"]); + const doc = await rawDoc("Dave"); + expect(doc).not.toBeNull(); + expect(doc?.embedding).toBeUndefined(); // no vector, but the entity exists + }); + + it("reembed backfills entities missing a vector and skips fresh ones", async () => { + // Insert a doc directly (no embedding), plus one created normally. + await handle.db.collection(ENTITIES_COLLECTION).insertOne({ + _id: "Legacy", + entityType: "thing", + observations: [{ text: "old", createdAt: new Date() }], + createdAt: new Date(), + updatedAt: new Date(), + }); + await store.createEntities([{ name: "Fresh", entityType: "thing", observations: ["new"] }]); + embedder.calls.length = 0; + + const result = await store.reembed({ onlyStale: true }); + expect(result.scanned).toBe(2); + expect(result.embedded).toBe(1); // only "Legacy" needed embedding + + const legacy = await rawDoc("Legacy"); + expect(legacy?.embedding).toHaveLength(4); + }); + + it("reembed({onlyStale:false}) re-embeds everything", async () => { + await store.createEntities([{ name: "One", entityType: "t", observations: ["a"] }]); + await store.createEntities([{ name: "Two", entityType: "t", observations: ["b"] }]); + const result = await store.reembed({ onlyStale: false }); + expect(result.scanned).toBe(2); + expect(result.embedded).toBe(2); + }); + + it("searchNodes falls back to text results when vector search is unavailable", async () => { + // The in-memory Mongo has no cosmosSearch, so vectorSearchIds catches the + // aggregation error and returns null — search must still return text hits. + await store.createEntities([ + { name: "Kubernetes", entityType: "tech", observations: ["container orchestration"] }, + { name: "Postgres", entityType: "tech", observations: ["relational database"] }, + ]); + const graph = await store.searchNodes("orchestration"); + expect(graph.entities.map((e) => e.name)).toContain("Kubernetes"); + }); +}); + +describe("KnowledgeGraphStore — no embedder (text-only, drop-in)", () => { + let handle: TestDbHandle; + let store: KnowledgeGraphStore; + + beforeAll(async () => { + handle = await getTestDb("graph_store_no_embedder_tests"); + await ensureGraphIndexes(handle.db); + store = new KnowledgeGraphStore(handle.db); + }); + + afterAll(async () => { + await closeTestClient(handle); + }); + + beforeEach(async () => { + await handle.db.collection(ENTITIES_COLLECTION).deleteMany({}); + }); + + it("reports vector search disabled", () => { + expect(store.vectorSearchEnabled).toBe(false); + }); + + it("does not write any embedding field", async () => { + await store.createEntities([{ name: "Zoe", entityType: "person", observations: ["hi"] }]); + const doc = await handle.db + .collection(ENTITIES_COLLECTION) + .findOne({ _id: "Zoe" }); + expect(doc?.embedding).toBeUndefined(); + expect(doc?.embeddingModel).toBeUndefined(); + }); + + it("reembed is a no-op returning zero counts", async () => { + await store.createEntities([{ name: "Yan", entityType: "person", observations: [] }]); + const result = await store.reembed(); + expect(result).toEqual({ scanned: 0, embedded: 0 }); + }); + + it("searchNodes returns text matches", async () => { + await store.createEntities([ + { name: "Redis", entityType: "tech", observations: ["in-memory cache"] }, + ]); + const graph = await store.searchNodes("cache"); + expect(graph.entities.map((e) => e.name)).toContain("Redis"); + }); +}); diff --git a/test/unit/embeddings-factory.test.ts b/test/unit/embeddings-factory.test.ts new file mode 100644 index 0000000..79a9051 --- /dev/null +++ b/test/unit/embeddings-factory.test.ts @@ -0,0 +1,66 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import { createEmbedder } from "../../src/shared/embeddings/index.js"; +import type { EmbeddingConfig } from "../../src/shared/config.js"; + +function baseConfig(overrides: Partial = {}): EmbeddingConfig { + return { + provider: "ollama", + model: "nomic-embed-text", + apiVersion: "2024-02-01", + indexKind: "vector-ivf", + similarity: "COS", + numLists: 100, + m: 16, + efConstruction: 64, + ...overrides, + }; +} + +describe("createEmbedder", () => { + afterEach(() => { + vi.unstubAllGlobals(); + }); + + it("returns null when provider is none (no probe)", async () => { + const spy = vi.fn(); + vi.stubGlobal("fetch", spy); + const e = await createEmbedder(baseConfig({ provider: "none" })); + expect(e).toBeNull(); + expect(spy).not.toHaveBeenCalled(); + }); + + it("returns null (graceful) when the backend is unreachable", async () => { + vi.stubGlobal( + "fetch", + vi.fn(async () => Promise.reject(new Error("ECONNREFUSED"))), + ); + const e = await createEmbedder(baseConfig()); + expect(e).toBeNull(); + }); + + it("returns null when openai provider lacks an API key", async () => { + const spy = vi.fn(); + vi.stubGlobal("fetch", spy); + const e = await createEmbedder(baseConfig({ provider: "openai", apiKey: undefined })); + expect(e).toBeNull(); + expect(spy).not.toHaveBeenCalled(); + }); + + it("probes and resolves dimensions from the observed vector length", async () => { + vi.stubGlobal( + "fetch", + vi.fn(async () => + Promise.resolve( + new Response(JSON.stringify({ embeddings: [[0.1, 0.2, 0.3, 0.4]] }), { + status: 200, + headers: { "content-type": "application/json" }, + }), + ), + ), + ); + const e = await createEmbedder(baseConfig()); + expect(e).not.toBeNull(); + expect(e?.dimensions).toBe(4); + expect(e?.provider).toBe("ollama"); + }); +}); diff --git a/test/unit/embeddings-providers.test.ts b/test/unit/embeddings-providers.test.ts new file mode 100644 index 0000000..4360cbd --- /dev/null +++ b/test/unit/embeddings-providers.test.ts @@ -0,0 +1,100 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import { + AzureOpenAIEmbedder, + OllamaEmbedder, + OpenAIEmbedder, +} from "../../src/shared/embeddings/providers.js"; +import { EmbeddingError } from "../../src/shared/embeddings/types.js"; + +// Helper: stub the global fetch with a single JSON response. +function stubFetch(status: number, body: unknown): void { + vi.stubGlobal( + "fetch", + vi.fn(async () => + Promise.resolve( + new Response(JSON.stringify(body), { + status, + headers: { "content-type": "application/json" }, + }), + ), + ), + ); +} + +describe("embedding providers", () => { + afterEach(() => { + vi.unstubAllGlobals(); + }); + + describe("OllamaEmbedder", () => { + it("returns the embeddings array from /api/embed", async () => { + stubFetch(200, { embeddings: [[1, 2, 3]] }); + const e = new OllamaEmbedder("nomic-embed-text", 3, "http://localhost:11434"); + const vecs = await e.embed(["hello"]); + expect(vecs).toEqual([[1, 2, 3]]); + }); + + it("returns [] for an empty input without calling fetch", async () => { + const spy = vi.fn(); + vi.stubGlobal("fetch", spy); + const e = new OllamaEmbedder("m", 3, "http://localhost:11434"); + expect(await e.embed([])).toEqual([]); + expect(spy).not.toHaveBeenCalled(); + }); + + it("throws EmbeddingError when the count does not match", async () => { + stubFetch(200, { embeddings: [[1, 2, 3]] }); + const e = new OllamaEmbedder("m", 3, "http://localhost:11434"); + await expect(e.embed(["a", "b"])).rejects.toBeInstanceOf(EmbeddingError); + }); + + it("throws EmbeddingError on a non-2xx status", async () => { + stubFetch(500, { error: "boom" }); + const e = new OllamaEmbedder("m", 3, "http://localhost:11434"); + await expect(e.embed(["a"])).rejects.toBeInstanceOf(EmbeddingError); + }); + }); + + describe("OpenAIEmbedder", () => { + it("orders vectors by response index", async () => { + stubFetch(200, { + data: [ + { index: 1, embedding: [9, 9] }, + { index: 0, embedding: [1, 1] }, + ], + }); + const e = new OpenAIEmbedder("text-embedding-3-small", 2, "https://api.openai.com/v1", "k"); + const vecs = await e.embed(["first", "second"]); + expect(vecs).toEqual([ + [1, 1], + [9, 9], + ]); + }); + }); + + describe("AzureOpenAIEmbedder", () => { + it("builds the deployment URL and parses OpenAI-shaped data", async () => { + const fetchMock = vi.fn(async () => + Promise.resolve( + new Response(JSON.stringify({ data: [{ index: 0, embedding: [0.5, 0.5] }] }), { + status: 200, + headers: { "content-type": "application/json" }, + }), + ), + ); + vi.stubGlobal("fetch", fetchMock); + const e = new AzureOpenAIEmbedder( + "my-deploy", + 2, + "https://res.openai.azure.com", + "secret", + "2024-02-01", + ); + const vecs = await e.embed(["x"]); + expect(vecs).toEqual([[0.5, 0.5]]); + const calledUrl = fetchMock.mock.calls[0]?.[0] as string; + expect(calledUrl).toContain("/openai/deployments/my-deploy/embeddings"); + expect(calledUrl).toContain("api-version=2024-02-01"); + }); + }); +}); diff --git a/test/unit/rrf.test.ts b/test/unit/rrf.test.ts new file mode 100644 index 0000000..4c1ee4c --- /dev/null +++ b/test/unit/rrf.test.ts @@ -0,0 +1,33 @@ +import { describe, expect, it } from "vitest"; +import { fuseRankings } from "../../src/storage/graph/store.js"; + +describe("fuseRankings (Reciprocal Rank Fusion)", () => { + it("returns the single list unchanged when the other is empty", () => { + expect(fuseRankings(["a", "b", "c"], [])).toEqual(["a", "b", "c"]); + expect(fuseRankings([], ["x", "y"])).toEqual(["x", "y"]); + }); + + it("de-duplicates names appearing in both lists", () => { + const fused = fuseRankings(["a", "b"], ["b", "c"]); + expect([...fused].sort()).toEqual(["a", "b", "c"]); + expect(new Set(fused).size).toBe(fused.length); + }); + + it("ranks a name appearing high in both lists above singletons", () => { + // "shared" is rank 1 in listA and rank 0 in listB → highest fused score. + const fused = fuseRankings(["onlyA", "shared"], ["shared", "onlyB"]); + expect(fused[0]).toBe("shared"); + }); + + it("preserves ordering when a list has a clear top hit", () => { + // Both lists agree "a" is best. + const fused = fuseRankings(["a", "b", "c"], ["a", "c", "b"]); + expect(fused[0]).toBe("a"); + }); + + it("breaks score ties by first appearance", () => { + // Disjoint lists, same rank 0 → equal score; listA's element wins the tie. + const fused = fuseRankings(["a"], ["b"]); + expect(fused).toEqual(["a", "b"]); + }); +});