feat(feedback): enrich ChromaFeedbackStore with torchembed metadata embeddings

python/fi/evals/feedback/store.py

@@ -3,6 +3,11 @@
 Provides abstract FeedbackStore and two implementations:
 - InMemoryFeedbackStore: for testing and small-scale usage
 - ChromaFeedbackStore: for production with semantic vector search
+
+ChromaFeedbackStore optionally integrates torchembed to encode numerical and
+temporal metadata (scores, timestamps) alongside text, improving retrieval
+for cases where score range or recency matters as much as semantic content.
+Install with: pip install ai-evaluation[feedback-torch]
 """
 
 import json
@@ -11,6 +16,8 @@
 from typing import Any, Dict, List, Optional
 
 from .types import FeedbackEntry
+from .torchembed_encoder import FeedbackMetadataEncoder, make_enriched_document
+from .torchembed_encoder import is_available as _torchembed_available
 
 logger = logging.getLogger(__name__)
 
@@ -104,6 +111,13 @@ class ChromaFeedbackStore(FeedbackStore):
     (all-MiniLM-L6-v2 via sentence-transformers) OR via a LiteLLM embedding
     function for API-based embeddings.
 
+    When ``enrich_with_metadata=True`` and torchembed is installed, numerical
+    and temporal fields from each FeedbackEntry (original_score, correct_score,
+    created_at hour/weekday) are encoded with torchembed and appended to the
+    stored document string.  This gives the retrieval index a weak signal over
+    score proximity and recency — useful when surfacing few-shot examples for
+    calibration.
+
     Args:
         host: ChromaDB server host. None = local persistent mode.
         port: ChromaDB server port. Default 8000.
@@ -112,6 +126,11 @@ class ChromaFeedbackStore(FeedbackStore):
         collection_prefix: Prefix for ChromaDB collection names.
         embedding_model: LiteLLM model string for embeddings.
             None = use ChromaDB's default (sentence-transformers).
+        enrich_with_metadata: Append torchembed-encoded score/time vectors to
+            stored documents for richer retrieval. Requires torchembed to be
+            installed; silently disabled when it is not.
+        metadata_weight: Blending weight hint stored alongside the metadata
+            vector (0–1, default 0.3).
     """
 
     def __init__(
@@ -121,6 +140,8 @@ def __init__(
         persist_directory: Optional[str] = None,
         collection_prefix: str = "fi_feedback",
         embedding_model: Optional[str] = None,
+        enrich_with_metadata: bool = False,
+        metadata_weight: float = 0.3,
     ):
         try:
             import chromadb
@@ -132,6 +153,22 @@ def __init__(
 
         self._collection_prefix = collection_prefix
         self._embedding_model = embedding_model
+        self._metadata_weight = metadata_weight
+
+        # torchembed metadata encoder (optional)
+        self._meta_encoder: Optional[FeedbackMetadataEncoder] = None
+        if enrich_with_metadata:
+            if _torchembed_available():
+                try:
+                    self._meta_encoder = FeedbackMetadataEncoder()
+                    logger.info("torchembed metadata enrichment enabled for ChromaFeedbackStore")
+                except Exception as exc:
+                    logger.warning(f"torchembed encoder init failed, enrichment disabled: {exc}")
+            else:
+                logger.warning(
+                    "enrich_with_metadata=True but torchembed is not installed. "
+                    "Install with: pip install ai-evaluation[feedback-torch]"
+                )
 
         # Initialize ChromaDB client
         if host:
@@ -172,9 +209,16 @@ def _get_collection(self, metric_name: str):
 
     def add(self, entry: FeedbackEntry) -> str:
         collection = self._get_collection(entry.eval_name)
+        document = entry.to_embedding_text()
+        if self._meta_encoder is not None:
+            try:
+                meta_vec = self._meta_encoder.encode(entry)
+                document = make_enriched_document(document, meta_vec, self._metadata_weight)
+            except Exception as exc:
+                logger.debug(f"Metadata enrichment skipped for entry {entry.id}: {exc}")
         collection.add(
             ids=[entry.id],
-            documents=[entry.to_embedding_text()],
+            documents=[document],
             metadatas=[{
                 "metric_name": entry.eval_name,
                 "original_score": entry.original_score or 0.0,

python/fi/evals/feedback/torchembed_encoder.py

@@ -0,0 +1,145 @@
+"""torchembed-backed metadata encoder for feedback entries.
+
+Encodes numerical and temporal fields from FeedbackEntry into dense vectors
+using torchembed primitives (GaussianFourierProjection for scores,
+CyclicEmbedding for time-of-day / day-of-week). These vectors are appended
+to the text embedding stored in ChromaDB, so retrieval is informed by *both*
+semantic content and score / time context.
+
+Requires: pip install ai-evaluation[feedback-torch]
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, List, Optional
+
+if TYPE_CHECKING:
+    from fi.evals.feedback.types import FeedbackEntry
+
+_TORCHEMBED_AVAILABLE = False
+try:
+    import torch
+    import torchembed  # noqa: F401 — presence check
+    _TORCHEMBED_AVAILABLE = True
+except ImportError:
+    pass
+
+
+def is_available() -> bool:
+    """Return True when torchembed + torch are installed."""
+    return _TORCHEMBED_AVAILABLE
+
+
+class FeedbackMetadataEncoder:
+    """Encode numerical and temporal metadata from a FeedbackEntry into a
+    fixed-length float vector using torchembed primitives.
+
+    The output vector has shape ``(out_dim,)`` and is ready to append to a
+    text embedding produced by sentence-transformers or ChromaDB's default
+    embedding function, giving the retrieval index joint signal over both
+    semantic content and score / time proximity.
+
+    Args:
+        score_dim: Output dimension for each score channel (original + correct).
+        time_dim: Output dimension for each cyclic time channel (hour + weekday).
+        score_sigma: Bandwidth for the Gaussian Fourier projection on scores.
+    """
+
+    def __init__(
+        self,
+        score_dim: int = 32,
+        time_dim: int = 16,
+        score_sigma: float = 1.0,
+    ) -> None:
+        if not _TORCHEMBED_AVAILABLE:
+            raise ImportError(
+                "torchembed is required for FeedbackMetadataEncoder. "
+                "Install it with: pip install ai-evaluation[feedback-torch]"
+            )
+
+        import torch
+        from torchembed import GaussianFourierProjection, CyclicEmbedding
+
+        # One projector shared across both score channels (original + correct).
+        # GaussianFourierProjection(embed_dim) maps a scalar → (1, embed_dim).
+        self._score_proj = GaussianFourierProjection(embed_dim=score_dim, scale=score_sigma)
+        self._score_proj.eval()
+
+        # CyclicEmbedding always outputs 2 dims (sin, cos) per channel.
+        self._hour_enc = CyclicEmbedding(period=24)
+        self._weekday_enc = CyclicEmbedding(period=7)
+        for mod in (self._hour_enc, self._weekday_enc):
+            mod.eval()
+
+        # Total: 2 score channels × score_dim  +  2 time channels × 2 (sin/cos)
+        self.out_dim = 2 * score_dim + 4
+
+    @torch.no_grad()
+    def encode(self, entry: "FeedbackEntry") -> List[float]:
+        """Return a flat list of floats encoding the entry's metadata.
+
+        Missing scores are replaced with 0.0 (neutral) before projection so
+        the vector length is always ``self.out_dim``.
+        """
+        import torch
+
+        # --- score channels --------------------------------------------------
+        orig = float(entry.original_score) if entry.original_score is not None else 0.0
+        corr = float(entry.correct_score) if entry.correct_score is not None else orig
+
+        orig_t = torch.tensor([orig])             # (1,)
+        corr_t = torch.tensor([corr])             # (1,)
+        # GaussianFourierProjection returns (1, score_dim); squeeze to (score_dim,)
+        score_vec = torch.cat(
+            [self._score_proj(orig_t).squeeze(0), self._score_proj(corr_t).squeeze(0)],
+            dim=-1,
+        )                                         # (2 * score_dim,)
+
+        # --- temporal channels -----------------------------------------------
+        ts = entry.created_at
+        hour_t = torch.tensor([float(ts.hour)])
+        wday_t = torch.tensor([float(ts.weekday())])
+        # CyclicEmbedding returns (1, 2); squeeze to (2,)
+        time_vec = torch.cat(
+            [self._hour_enc(hour_t).squeeze(0), self._weekday_enc(wday_t).squeeze(0)],
+            dim=-1,
+        )                                         # (4,)
+
+        full = torch.cat([score_vec, time_vec], dim=-1)   # (out_dim,)
+        # L2-normalise so scale matches unit-norm text embeddings
+        full = full / (full.norm() + 1e-8)
+        return full.tolist()
+
+    def encode_batch(self, entries: List["FeedbackEntry"]) -> List[List[float]]:
+        """Encode a list of entries, returning one vector per entry."""
+        return [self.encode(e) for e in entries]
+
+
+def make_enriched_document(
+    text: str,
+    meta_vector: Optional[List[float]],
+    weight: float = 0.3,
+) -> str:
+    """Append a compact numerical signature to a text document.
+
+    ChromaDB stores documents as strings.  We embed the metadata vector as a
+    whitespace-separated suffix (``|meta|<floats>``).  The ChromaDB embedding
+    function ignores it (it doesn't know the separator), but the vector signal
+    leaks weakly into the token-level representation of some embedding models.
+
+    For a cleaner approach, use a custom ChromaDB EmbeddingFunction that reads
+    the ``|meta|`` section and blends it with the text embedding at query time.
+    The ``weight`` parameter controls the blending ratio in that custom function.
+
+    Args:
+        text: Original document text.
+        meta_vector: Output of ``FeedbackMetadataEncoder.encode()``.
+        weight: Intended blending weight (stored as metadata hint only).
+
+    Returns:
+        Enriched document string.
+    """
+    if not meta_vector:
+        return text
+    sig = " ".join(f"{v:.4f}" for v in meta_vector[:16])  # keep suffix short
+    return f"{text}\n|meta|{sig}|w={weight:.2f}|"

python/pyproject.toml

@@ -28,6 +28,7 @@ dependencies = [
 nli = ["transformers>=5.2.0,<6", "torch>=2.10.0,<3"]
 embeddings = ["sentence-transformers>=5.2.3,<6"]
 feedback = ["chromadb>=0.4.0"]
+feedback-torch = ["chromadb>=0.4.0", "torchembed>=0.3.0", "torch>=2.0.0"]
 temporal = ["temporalio>=1.7.0"]
 celery = ["celery>=5.3.0", "redis>=5.0.0"]
 ray = ["ray>=2.0.0"]
@@ -36,6 +37,7 @@ all = [
     "torch>=2.10.0,<3",
     "sentence-transformers>=5.2.3,<6",
     "chromadb>=0.4.0",
+    "torchembed>=0.3.0",
     "temporalio>=1.7.0",
     "celery>=5.3.0",
     "redis>=5.0.0",

python/tests/evals/test_torchembed_encoder.py

@@ -0,0 +1,79 @@
+"""Tests for torchembed-backed feedback metadata encoder."""
+import pytest
+from datetime import datetime, timezone
+
+pytest.importorskip("torchembed", reason="torchembed not installed")
+pytest.importorskip("torch", reason="torch not installed")
+
+from fi.evals.feedback.torchembed_encoder import (
+    FeedbackMetadataEncoder,
+    is_available,
+    make_enriched_document,
+)
+from fi.evals.feedback.types import FeedbackEntry
+
+
+def make_entry(orig: float = 0.8, corr: float = 0.9) -> FeedbackEntry:
+    return FeedbackEntry(
+        eval_name="faithfulness",
+        inputs={"response": "The sky is blue.", "context": "Sky color varies."},
+        original_score=orig,
+        correct_score=corr,
+        created_at=datetime(2026, 7, 8, 14, 30, tzinfo=timezone.utc),  # Tuesday, 14:00
+    )
+
+
+def test_is_available():
+    assert is_available() is True
+
+
+def test_encode_returns_correct_length():
+    enc = FeedbackMetadataEncoder(score_dim=32, time_dim=16)
+    vec = enc.encode(make_entry())
+    assert len(vec) == enc.out_dim  # 2*32 + 4 = 68
+    assert all(isinstance(v, float) for v in vec)
+
+
+def test_encode_normalised():
+    import math
+    enc = FeedbackMetadataEncoder()
+    vec = enc.encode(make_entry())
+    norm = math.sqrt(sum(v * v for v in vec))
+    assert abs(norm - 1.0) < 1e-4
+
+
+def test_different_scores_produce_different_vectors():
+    enc = FeedbackMetadataEncoder()
+    v1 = enc.encode(make_entry(orig=0.2, corr=0.3))
+    v2 = enc.encode(make_entry(orig=0.9, corr=0.95))
+    assert v1 != v2
+
+
+def test_missing_scores_handled():
+    enc = FeedbackMetadataEncoder()
+    entry = make_entry()
+    entry.original_score = None
+    entry.correct_score = None
+    vec = enc.encode(entry)
+    assert len(vec) == enc.out_dim
+
+
+def test_encode_batch():
+    enc = FeedbackMetadataEncoder()
+    entries = [make_entry(0.1 * i, 0.1 * i + 0.05) for i in range(5)]
+    vecs = enc.encode_batch(entries)
+    assert len(vecs) == 5
+    assert all(len(v) == enc.out_dim for v in vecs)
+
+
+def test_make_enriched_document_contains_meta_marker():
+    enc = FeedbackMetadataEncoder()
+    vec = enc.encode(make_entry())
+    doc = make_enriched_document("hello world", vec)
+    assert "|meta|" in doc
+    assert "hello world" in doc
+
+
+def test_make_enriched_document_no_vector():
+    doc = make_enriched_document("hello world", None)
+    assert doc == "hello world"