File size: 5,454 Bytes
bde2f3a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 | """RAG Dual-Dim Embedder (M3 in v3 unfuck plan).
Wraps the FROZEN app/crypto_embeddings.py to support both bge-m3 (1024d, current)
and qwen3-embedding:4b (2048d, target) during zero-downtime migration.
Strategy: reindex one collection per night, smallest first, verify each morning.
If eval harness reports <8% nDCG@10 improvement, defer the rest.
See DESIGN.md §M3 for the full math and migration plan.
"""
from __future__ import annotations
from enum import Enum
from typing import Any
class EmbedderBackend(Enum):
"""Available embedding backends."""
BGE_M3 = "bge-m3" # 1024d, legacy
QWEN3_4B = "qwen3-embedding:4b" # 2048d, target
# ── Per-backend configuration ──────────────────────────────────────────
_BACKEND_DIMS: dict[EmbedderBackend, int] = {
EmbedderBackend.BGE_M3: 1024,
EmbedderBackend.QWEN3_4B: 2048,
}
def backend_dim(backend: EmbedderBackend) -> int:
"""Return the embedding dimension for a backend."""
return _BACKEND_DIMS[backend]
# ── Dual-dim wrapper ────────────────────────────────────────────────────
class DualDimEmbedder:
"""Wraps Ollama embedding for any backend. Same interface regardless of dim."""
def __init__(
self,
backend: EmbedderBackend,
ollama_url: str = "http://ollama:11434",
) -> None:
self.backend = backend
self.ollama_url = ollama_url.rstrip("/")
self.dim = backend_dim(backend)
async def embed(self, texts: list[str]) -> list[list[float]]:
"""Embed a batch of texts. Returns one vector per text.
Uses the Ollama /api/embeddings endpoint with the configured model.
"""
import httpx
if not texts:
return []
async with httpx.AsyncClient(timeout=60.0) as client:
resp = await client.post(
f"{self.ollama_url}/api/embeddings",
json={"model": self.backend.value, "prompt": texts},
)
resp.raise_for_status()
data = resp.json()
# Ollama returns {"embedding": [[...], [...]]} for single, or {"embeddings": [[...]]}
if "embeddings" in data:
vectors = data["embeddings"]
else:
# Single-text fallback — Ollama returns {"embedding": [...]}
vectors = [data["embedding"]] if "embedding" in data else []
# Validate dimensions to catch backend mismatches early.
for i, vec in enumerate(vectors):
if len(vec) != self.dim:
raise RuntimeError(
f"Embedder {self.backend.value} returned {len(vec)}d vector "
f"for text[{i}], expected {self.dim}d. "
f"Check that ollama has the right model pulled."
)
return vectors
async def embed_one(self, text: str) -> list[float]:
"""Convenience: embed a single text."""
result = await self.embed([text])
return result[0] if result else [0.0] * self.dim
# ── Migration helpers ──────────────────────────────────────────────────
async def reindex_collection(
name: str,
source: DualDimEmbedder,
target: DualDimEmbedder,
fetch_docs: Any,
write_collection: Any,
verify_queries: list[str] | None = None,
) -> bool:
"""Re-embed one collection from source backend to target backend.
Args:
name: collection name
source: existing backend (e.g. bge-m3)
target: new backend (e.g. qwen3-embedding:4b)
fetch_docs: async () -> list[(id, text)] callable
write_collection: async (name, vectors) -> None callable
verify_queries: optional list of known queries to verify retrieval
Returns True if migration succeeded (or verification skipped).
"""
if source.dim == target.dim:
# Same dim — no migration needed.
return True
docs = await fetch_docs()
if not docs:
return True
texts = [text for _, text in docs]
new_vectors = await target.embed(texts)
new_name = f"{name}_v2"
await write_collection(new_name, list(zip((id_ for id_, _ in docs), new_vectors, strict=False)))
# Atomic swap — implementation-specific. Caller handles.
# For FAISS: rename .index files. For Qdrant: rename collections.
if verify_queries:
# Verify retrieval on known queries before swapping production traffic.
# Implementation-specific; placeholder for now.
pass
return True
# ── CLI ────────────────────────────────────────────────────────────────
def _main() -> None:
"""CLI: list available backends and their dims."""
import argparse
parser = argparse.ArgumentParser(description="RAG dual-dim embedder")
parser.add_argument(
"--list-backends", action="store_true", help="List available backends"
)
args = parser.parse_args()
if args.list_backends:
for backend in EmbedderBackend:
print(f"{backend.value} {_BACKEND_DIMS[backend]}d")
if __name__ == "__main__":
_main()
|