@framers/agentos-ext-topicality 0.1.0

package/src/TopicDriftTracker.ts

@@ -0,0 +1,307 @@
+/**
+ * @fileoverview TopicDriftTracker — session-level EMA drift detection for topicality guardrails.
+ *
+ * This module tracks whether a conversation session is gradually drifting away
+ * from its allowed topics by maintaining a per-session **running embedding**
+ * that is updated with each new message using an Exponential Moving Average
+ * (EMA).
+ *
+ * ### Why EMA?
+ * A simple "last-message" check is too noisy: a single off-topic message in an
+ * otherwise on-topic conversation should not trigger a hard block.  EMA
+ * smooths the signal so that sustained drift is detected while brief tangents
+ * are tolerated.
+ *
+ * The update formula is:
+ * ```
+ * running[i] = alpha * message[i] + (1 - alpha) * running[i]
+ * ```
+ * A smaller `alpha` means the running vector changes slowly (long memory);
+ * a larger `alpha` means it reacts quickly to each new message.
+ *
+ * ### Drift decision
+ * After each EMA update the tracker checks whether the running embedding is
+ * "on-topic" by calling {@link TopicEmbeddingIndex.isOnTopicByVector}.  If the
+ * check fails the `driftStreak` counter is incremented; if it passes the
+ * streak resets to zero.  When `driftStreak >= driftStreakLimit` the result
+ * `driftLimitExceeded` is set to `true`.
+ *
+ * ### Session management
+ * Sessions are stored in an in-memory `Map<sessionId, TopicState>`.  To
+ * prevent unbounded memory growth:
+ * - Stale sessions (inactive for > `sessionTimeoutMs`) are pruned lazily
+ *   whenever `map.size > maxSessions` at the start of an `update()` call.
+ * - Callers can force a full clear via {@link clear}.
+ *
+ * @module topicality/TopicDriftTracker
+ */
+
+import type { TopicEmbeddingIndex } from './TopicEmbeddingIndex';
+import {
+  DEFAULT_DRIFT_CONFIG,
+  type DriftConfig,
+  type DriftResult,
+  type TopicMatch,
+  type TopicState,
+} from './types';
+
+// ---------------------------------------------------------------------------
+// TopicDriftTracker
+// ---------------------------------------------------------------------------
+
+/**
+ * Tracks per-session topic drift using EMA-blended running embeddings.
+ *
+ * Instantiate once per agent process (not per conversation) since the tracker
+ * manages many concurrent sessions internally.
+ *
+ * @example
+ * ```ts
+ * const tracker = new TopicDriftTracker({ alpha: 0.4, driftStreakLimit: 2 });
+ *
+ * // In your message handler:
+ * const embedding = await embed(userMessage);
+ * const result = tracker.update('session-abc', embedding, allowedIndex);
+ *
+ * if (result.driftLimitExceeded) {
+ *   // Take configured action: redirect, warn, or block.
+ * }
+ * ```
+ */
+export class TopicDriftTracker {
+  /** Fully resolved drift configuration (defaults merged with caller overrides). */
+  private readonly config: DriftConfig;
+
+  /**
+   * In-memory session store.
+   * Key: caller-supplied session ID (e.g. conversation UUID).
+   * Value: mutable {@link TopicState} for that session.
+   */
+  private readonly sessions: Map<string, TopicState> = new Map();
+
+  // -------------------------------------------------------------------------
+  // Constructor
+  // -------------------------------------------------------------------------
+
+  /**
+   * Creates a new `TopicDriftTracker`.
+   *
+   * @param config - Optional partial override of {@link DEFAULT_DRIFT_CONFIG}.
+   *   Any fields not provided fall back to their default values.  Pass an
+   *   empty object `{}` or omit entirely to use all defaults.
+   *
+   * @example
+   * ```ts
+   * // Use defaults
+   * const tracker = new TopicDriftTracker();
+   *
+   * // Override only alpha and streakLimit
+   * const strictTracker = new TopicDriftTracker({ alpha: 0.5, driftStreakLimit: 2 });
+   * ```
+   */
+  constructor(config?: Partial<DriftConfig>) {
+    // Merge caller overrides with defaults — undefined fields are taken from
+    // DEFAULT_DRIFT_CONFIG, preserving all caller-supplied values exactly.
+    this.config = { ...DEFAULT_DRIFT_CONFIG, ...(config ?? {}) };
+  }
+
+  // -------------------------------------------------------------------------
+  // Public API — update
+  // -------------------------------------------------------------------------
+
+  /**
+   * Processes a new message embedding for the given session and returns the
+   * current drift assessment.
+   *
+   * ### Steps performed
+   * 1. **Retrieve or create** the session state.  On the very first message
+   *    the running embedding is initialised to a shallow copy of
+   *    `messageEmbedding` (no EMA applied yet).
+   * 2. **Apply EMA** (from the second message onwards):
+   *    `running[i] = alpha * message[i] + (1 - alpha) * running[i]`
+   * 3. **Check topic alignment** using
+   *    `allowedIndex.isOnTopicByVector(running, driftThreshold)`.
+   * 4. **Update streak** — increment `driftStreak` on off-topic, reset to 0
+   *    on on-topic.
+   * 5. **Lazy-prune** stale sessions when `map.size > maxSessions` before
+   *    creating a new session (never during updates of existing sessions to
+   *    avoid deleting the session we are currently updating).
+   * 6. **Persist** the updated state and return a {@link DriftResult}.
+   *
+   * @param sessionId        - Unique identifier for the conversation session.
+   *   Typically a UUID or user ID. Must be consistent across messages in the
+   *   same conversation.
+   * @param messageEmbedding - Pre-computed numeric embedding of the current
+   *   message.  Must have the same dimensionality as the topic centroids used
+   *   by `allowedIndex`.
+   * @param allowedIndex     - Built {@link TopicEmbeddingIndex} containing the
+   *   allowed topics to check against.  The tracker calls only
+   *   `isOnTopicByVector` (no async operations, no extra embedding calls).
+   * @returns A {@link DriftResult} describing whether the session is currently
+   *   drifting and by how much.
+   */
+  update(
+    sessionId: string,
+    messageEmbedding: number[],
+    allowedIndex: TopicEmbeddingIndex,
+  ): DriftResult {
+    const now = Date.now();
+    const isNewSession = !this.sessions.has(sessionId);
+
+    // Lazy prune: only trigger when a NEW session would push us over the limit.
+    // We do not prune during updates of existing sessions to avoid accidentally
+    // deleting a session that is currently being processed.
+    if (isNewSession && this.sessions.size >= this.config.maxSessions) {
+      this.pruneStale();
+    }
+
+    let state: TopicState;
+
+    if (isNewSession) {
+      // First message in this session — initialise running embedding to a copy
+      // of the current message embedding.  A copy prevents external mutation
+      // of the array from silently corrupting the tracker state.
+      state = {
+        runningEmbedding: [...messageEmbedding],
+        messageCount: 0, // will be incremented below
+        lastTopicScore: 0,
+        driftStreak: 0,
+        lastSeenAt: now,
+      };
+    } else {
+      // Retrieve existing state — guaranteed non-null by the `has()` check above.
+      state = this.sessions.get(sessionId)!;
+
+      // Apply the EMA update in-place.
+      // running[i] = alpha * message[i] + (1 - alpha) * running[i]
+      const alpha = this.config.alpha;
+      const oneMinusAlpha = 1 - alpha;
+
+      for (let i = 0; i < state.runningEmbedding.length; i++) {
+        state.runningEmbedding[i] =
+          alpha * messageEmbedding[i] + oneMinusAlpha * state.runningEmbedding[i];
+      }
+    }
+
+    // Increment message counter and timestamp.
+    state.messageCount += 1;
+    state.lastSeenAt = now;
+
+    // -----------------------------------------------------------------------
+    // Topic alignment check
+    // -----------------------------------------------------------------------
+
+    // Check whether the (now-updated) running embedding is on-topic.
+    const onTopic = allowedIndex.isOnTopicByVector(
+      state.runningEmbedding,
+      this.config.driftThreshold,
+    );
+
+    // Retrieve the best-match details from the index for the result payload.
+    // matchByVector is synchronous and does not re-embed anything.
+    const topMatches = allowedIndex.matchByVector(state.runningEmbedding);
+    const nearestTopic: TopicMatch | null = topMatches.length > 0 ? topMatches[0] : null;
+    const currentSimilarity = nearestTopic?.similarity ?? 0;
+
+    // Store the latest similarity score for observability.
+    state.lastTopicScore = currentSimilarity;
+
+    // -----------------------------------------------------------------------
+    // Drift streak management
+    // -----------------------------------------------------------------------
+
+    if (onTopic) {
+      // Good message — reset the drift counter.
+      state.driftStreak = 0;
+    } else {
+      // Off-topic message — accumulate the streak.
+      state.driftStreak += 1;
+    }
+
+    const driftLimitExceeded = state.driftStreak >= this.config.driftStreakLimit;
+
+    // -----------------------------------------------------------------------
+    // Persist state and return result
+    // -----------------------------------------------------------------------
+
+    // Always write back (even for existing sessions, since we mutated in-place
+    // for the EMA; for new sessions we need to insert).
+    this.sessions.set(sessionId, state);
+
+    return {
+      onTopic,
+      currentSimilarity,
+      nearestTopic,
+      driftStreak: state.driftStreak,
+      driftLimitExceeded,
+    };
+  }
+
+  // -------------------------------------------------------------------------
+  // Public API — pruneStale
+  // -------------------------------------------------------------------------
+
+  /**
+   * Removes sessions that have been inactive for longer than `sessionTimeoutMs`.
+   *
+   * This is called lazily inside {@link update} when the session map exceeds
+   * `maxSessions`, but callers may invoke it directly to trigger an immediate
+   * cleanup (e.g. in a scheduled maintenance job).
+   *
+   * Pruning is O(n) in the number of active sessions.
+   */
+  pruneStale(): void {
+    const now = Date.now();
+    const timeoutMs = this.config.sessionTimeoutMs;
+
+    for (const [id, state] of this.sessions) {
+      if (now - state.lastSeenAt > timeoutMs) {
+        this.sessions.delete(id);
+      }
+    }
+  }
+
+  // -------------------------------------------------------------------------
+  // Public API — clear
+  // -------------------------------------------------------------------------
+
+  /**
+   * Removes all sessions from the tracker unconditionally.
+   *
+   * Useful for graceful shutdown, testing teardown, or resetting the agent
+   * context between evaluation runs.
+   */
+  clear(): void {
+    this.sessions.clear();
+  }
+
+  // -------------------------------------------------------------------------
+  // Internal helpers (exposed for testing via package-private pattern)
+  // -------------------------------------------------------------------------
+
+  /**
+   * Returns the current number of active sessions in the internal map.
+   * Useful for observability and testing.
+   *
+   * @internal
+   */
+  get sessionCount(): number {
+    return this.sessions.size;
+  }
+
+  /**
+   * Returns a copy of the {@link TopicState} for the given session, or
+   * `undefined` if the session does not exist.
+   *
+   * Exposed for unit-testing state inspection.  The returned object is a
+   * shallow copy — mutating it does not affect the tracker's internal state.
+   *
+   * @internal
+   */
+  getState(sessionId: string): TopicState | undefined {
+    const state = this.sessions.get(sessionId);
+    if (!state) return undefined;
+    // Shallow copy to prevent callers from accidentally mutating tracker state.
+    return { ...state, runningEmbedding: [...state.runningEmbedding] };
+  }
+}

package/src/TopicEmbeddingIndex.ts

@@ -0,0 +1,346 @@
+/**
+ * @fileoverview TopicEmbeddingIndex — semantic similarity lookup for topic guardrails.
+ *
+ * This module implements a lightweight in-memory embedding index that:
+ *
+ * 1. **Builds** per-topic centroid embeddings from descriptions + examples.
+ * 2. **Matches** an arbitrary embedding or text string against all topic centroids
+ *    using cosine similarity.
+ * 3. **Answers** boolean on-topic queries at a configurable similarity threshold.
+ *
+ * ### How centroids are built
+ * For each {@link TopicDescriptor} the index concatenates:
+ * ```
+ * texts = [descriptor.description, ...descriptor.examples]
+ * ```
+ * All topics are embedded in a single batch call to `embeddingFn` to minimise
+ * round-trips.  The centroid for a topic is the component-wise average (mean)
+ * of all its embedding vectors.
+ *
+ * ### Similarity scoring
+ * Raw cosine similarity can be negative when vectors point in opposite directions.
+ * `matchByVector` clamps scores to `Math.max(0, similarity)` so that all
+ * {@link TopicMatch} values represent non-negative relevance scores.
+ *
+ * @module topicality/TopicEmbeddingIndex
+ */
+
+import { cosineSimilarity } from '@framers/agentos/core/utils/text-utils';
+import type { TopicDescriptor, TopicMatch } from './types';
+
+// ---------------------------------------------------------------------------
+// Internal storage shape
+// ---------------------------------------------------------------------------
+
+/**
+ * Private per-topic record stored in the index map.
+ *
+ * @internal
+ */
+interface TopicEntry {
+  /** Original descriptor, kept for metadata retrieval. */
+  descriptor: TopicDescriptor;
+  /**
+   * Component-wise average of all embeddings derived from this topic's
+   * description and example strings.
+   */
+  centroid: number[];
+}
+
+// ---------------------------------------------------------------------------
+// TopicEmbeddingIndex
+// ---------------------------------------------------------------------------
+
+/**
+ * Semantic embedding index for topicality guardrail matching.
+ *
+ * The index is intentionally **lazy** — it holds no embeddings until
+ * {@link build} is called.  This makes instantiation cheap and lets the
+ * caller defer the (potentially expensive) batch embedding call until the
+ * agent's first message.
+ *
+ * @example
+ * ```ts
+ * const index = new TopicEmbeddingIndex(async (texts) => {
+ *   const res = await openai.embeddings.create({ model: 'text-embedding-3-small', input: texts });
+ *   return res.data.map(d => d.embedding);
+ * });
+ *
+ * await index.build(TOPIC_PRESETS.customerSupport);
+ *
+ * const matches = await index.match('How do I cancel my subscription?');
+ * // → [{ topicId: 'billing', topicName: 'Billing & Payments', similarity: 0.82 }, ...]
+ *
+ * const onTopic = await index.isOnTopic('Tell me a joke', 0.35);
+ * // → false (a joke doesn't match any customer-support topic)
+ * ```
+ */
+export class TopicEmbeddingIndex {
+  /**
+   * Caller-supplied batch embedding function.
+   * Invoked once during {@link build} with all topic texts concatenated.
+   */
+  private readonly embeddingFn: (texts: string[]) => Promise<number[][]>;
+
+  /**
+   * Internal store mapping `topicId → TopicEntry`.
+   * Populated by {@link build}; empty until then.
+   */
+  private readonly entries: Map<string, TopicEntry> = new Map();
+
+  /** Whether {@link build} has been called and completed successfully. */
+  private built: boolean = false;
+
+  // -------------------------------------------------------------------------
+  // Constructor
+  // -------------------------------------------------------------------------
+
+  /**
+   * Creates a new `TopicEmbeddingIndex`.
+   *
+   * @param embeddingFn - Async function that converts an array of text strings
+   *   into corresponding numeric embedding vectors.  All returned vectors must
+   *   share the same dimensionality.  The function is called exactly **once**
+   *   per {@link build} invocation with all texts for all topics batched
+   *   together.
+   */
+  constructor(embeddingFn: (texts: string[]) => Promise<number[][]>) {
+    this.embeddingFn = embeddingFn;
+  }
+
+  // -------------------------------------------------------------------------
+  // Public API — build
+  // -------------------------------------------------------------------------
+
+  /**
+   * Embeds all topic descriptions and examples, computes per-topic centroid
+   * embeddings, and stores them in the internal index.
+   *
+   * Calling `build()` a second time replaces the existing index entirely,
+   * allowing hot-reloading of topic configurations without recreating the
+   * instance.
+   *
+   * ### Centroid computation
+   * For each topic we collect `[description, ...examples]` as a list of
+   * strings, embed them all in one batch, then average the resulting vectors
+   * component-wise to produce a single representative centroid.
+   *
+   * All topics are embedded in a **single batch call** to minimise latency.
+   *
+   * @param topics - Array of {@link TopicDescriptor} objects to index.
+   *   An empty array is valid — the index will simply return no matches.
+   * @returns A promise that resolves once all embeddings are computed and
+   *   stored.  Rejects if `embeddingFn` throws or returns vectors of
+   *   mismatched length.
+   */
+  async build(topics: TopicDescriptor[]): Promise<void> {
+    // Reset state before (re)building so a failed build leaves the index empty
+    // rather than in a partial state.
+    this.entries.clear();
+    this.built = false;
+
+    if (topics.length === 0) {
+      // Nothing to embed — mark as built so isBuilt returns true.
+      this.built = true;
+      return;
+    }
+
+    // Collect, per topic, the list of texts to embed and the range of
+    // indices they will occupy in the flat batch array.
+    //
+    // Layout: [topic0_desc, topic0_ex0, topic0_ex1, …, topic1_desc, …]
+    const allTexts: string[] = [];
+    const topicRanges: Array<{ topic: TopicDescriptor; start: number; end: number }> = [];
+
+    for (const topic of topics) {
+      const start = allTexts.length;
+      // Always include the description as the first text for this topic.
+      allTexts.push(topic.description);
+      // Then all examples (may be empty — the centroid will just be the description).
+      for (const example of topic.examples) {
+        allTexts.push(example);
+      }
+      const end = allTexts.length; // exclusive
+      topicRanges.push({ topic, start, end });
+    }
+
+    // Single batch embedding call — one round-trip regardless of how many
+    // topics or examples are configured.
+    const allEmbeddings = await this.embeddingFn(allTexts);
+
+    // Validate that the embedding function returned the right number of vectors.
+    if (allEmbeddings.length !== allTexts.length) {
+      throw new Error(
+        `TopicEmbeddingIndex.build: embeddingFn returned ${allEmbeddings.length} vectors ` +
+          `but ${allTexts.length} texts were provided.`,
+      );
+    }
+
+    // Compute centroid for each topic from its slice of the batch result.
+    for (const { topic, start, end } of topicRanges) {
+      const slice = allEmbeddings.slice(start, end);
+      const centroid = computeCentroid(slice);
+      this.entries.set(topic.id, { descriptor: topic, centroid });
+    }
+
+    this.built = true;
+  }
+
+  // -------------------------------------------------------------------------
+  // Public API — matchByVector
+  // -------------------------------------------------------------------------
+
+  /**
+   * Computes similarity between a pre-computed embedding vector and all topic
+   * centroids **without** making any additional embedding calls.
+   *
+   * This is the hot path invoked by {@link TopicDriftTracker}, which maintains
+   * its own running embedding and never needs to re-embed.
+   *
+   * Results are clamped to `[0, 1]` (negative cosine → 0) and sorted
+   * descending by similarity.
+   *
+   * @param embedding - A numeric vector with the same dimensionality as the
+   *   centroids produced during {@link build}.
+   * @returns Array of {@link TopicMatch} objects sorted by similarity
+   *   descending.  Returns an empty array if the index was not yet built or
+   *   contains no topics.
+   */
+  matchByVector(embedding: number[]): TopicMatch[] {
+    if (!this.built || this.entries.size === 0) {
+      // Return empty rather than throwing — callers can treat no match as off-topic.
+      return [];
+    }
+
+    const matches: TopicMatch[] = [];
+
+    for (const [topicId, entry] of this.entries) {
+      const raw = cosineSimilarity(embedding, entry.centroid);
+      // Clamp to [0, 1] — negative similarity means "opposite direction" which
+      // is no more useful than "unrelated" for topic matching.
+      const similarity = Math.max(0, raw);
+
+      matches.push({
+        topicId,
+        topicName: entry.descriptor.name,
+        similarity,
+      });
+    }
+
+    // Sort descending so the best match is first.
+    matches.sort((a, b) => b.similarity - a.similarity);
+    return matches;
+  }
+
+  // -------------------------------------------------------------------------
+  // Public API — match
+  // -------------------------------------------------------------------------
+
+  /**
+   * Embeds `text` and returns similarity scores against all topic centroids.
+   *
+   * This is a convenience wrapper that handles the embedding step.  If you
+   * already have an embedding (e.g. from the drift tracker's running vector)
+   * prefer {@link matchByVector} to avoid a redundant embedding call.
+   *
+   * @param text - The user message or assistant output to evaluate.
+   * @returns A promise resolving to {@link TopicMatch}[] sorted descending.
+   */
+  async match(text: string): Promise<TopicMatch[]> {
+    // Embed the single query text.
+    const [embedding] = await this.embeddingFn([text]);
+    return this.matchByVector(embedding);
+  }
+
+  // -------------------------------------------------------------------------
+  // Public API — isOnTopicByVector
+  // -------------------------------------------------------------------------
+
+  /**
+   * Returns `true` if the given embedding vector scores above `threshold`
+   * against **at least one** topic in the index.
+   *
+   * Uses {@link matchByVector} internally so no additional embedding call is
+   * made.
+   *
+   * @param embedding - Pre-computed numeric vector.
+   * @param threshold - Minimum similarity (in `[0, 1]`) for a topic to count
+   *   as a match.
+   * @returns `true` if any topic centroid has similarity > threshold; otherwise `false`.
+   */
+  isOnTopicByVector(embedding: number[], threshold: number): boolean {
+    const matches = this.matchByVector(embedding);
+    // The list is sorted descending, so we only need to check the first entry.
+    return matches.length > 0 && matches[0].similarity > threshold;
+  }
+
+  // -------------------------------------------------------------------------
+  // Public API — isOnTopic
+  // -------------------------------------------------------------------------
+
+  /**
+   * Embeds `text` and returns `true` if it scores above `threshold` against
+   * at least one allowed topic.
+   *
+   * @param text      - The text to evaluate.
+   * @param threshold - Minimum cosine similarity for the text to be considered on-topic.
+   * @returns A promise resolving to `true` if on-topic, `false` otherwise.
+   */
+  async isOnTopic(text: string, threshold: number): Promise<boolean> {
+    const [embedding] = await this.embeddingFn([text]);
+    return this.isOnTopicByVector(embedding, threshold);
+  }
+
+  // -------------------------------------------------------------------------
+  // Getter
+  // -------------------------------------------------------------------------
+
+  /**
+   * Whether {@link build} has been called and completed successfully.
+   *
+   * Use this to guard against calling {@link match} or {@link matchByVector}
+   * before the index is ready.
+   *
+   * @example
+   * ```ts
+   * if (!index.isBuilt) await index.build(topics);
+   * ```
+   */
+  get isBuilt(): boolean {
+    return this.built;
+  }
+}
+
+// ---------------------------------------------------------------------------
+// Internal helpers
+// ---------------------------------------------------------------------------
+
+/**
+ * Computes the component-wise average (centroid) of an array of embedding
+ * vectors.
+ *
+ * All input vectors are assumed to have the same dimensionality.  If the
+ * input array is empty an empty array is returned (safe no-op).
+ *
+ * @param vectors - One or more numeric vectors of equal length.
+ * @returns A single vector whose i-th element is the mean of i-th elements
+ *   across all input vectors.
+ *
+ * @internal
+ */
+function computeCentroid(vectors: number[][]): number[] {
+  if (vectors.length === 0) return [];
+
+  const dim = vectors[0].length;
+  // Initialise accumulator to all zeros.
+  const sum = new Array<number>(dim).fill(0);
+
+  for (const vec of vectors) {
+    for (let i = 0; i < dim; i++) {
+      sum[i] += vec[i];
+    }
+  }
+
+  // Divide each component by the number of vectors to get the mean.
+  return sum.map((v) => v / vectors.length);
+}