diff --git a/python/fi/evals/feedback/store.py b/python/fi/evals/feedback/store.py index da1eb9a8..843743f5 100644 --- a/python/fi/evals/feedback/store.py +++ b/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, diff --git a/python/fi/evals/feedback/torchembed_encoder.py b/python/fi/evals/feedback/torchembed_encoder.py new file mode 100644 index 00000000..357160bc --- /dev/null +++ b/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|``). 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}|" diff --git a/python/pyproject.toml b/python/pyproject.toml index 2d921c66..31d21a6e 100644 --- a/python/pyproject.toml +++ b/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", diff --git a/python/tests/evals/test_torchembed_encoder.py b/python/tests/evals/test_torchembed_encoder.py new file mode 100644 index 00000000..b2d1815f --- /dev/null +++ b/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"