|
# retoor <retoor@molodetz.nl>
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import logging
|
|
import math
|
|
import re
|
|
from dataclasses import dataclass
|
|
|
|
import httpx
|
|
|
|
from devplacepy.config import INTERNAL_EMBED_MODEL, INTERNAL_EMBED_URL
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
EMBED_TIMEOUT_SECONDS = 60.0
|
|
LOCAL_EMBED_DIMS = 256
|
|
TOKEN_PATTERN = re.compile(r"[a-z0-9]+")
|
|
|
|
|
|
@dataclass
|
|
class EmbedResult:
|
|
vectors: list[list[float]]
|
|
backend: str
|
|
|
|
|
|
def _local_vector(text: str) -> list[float]:
|
|
bucket = [0.0] * LOCAL_EMBED_DIMS
|
|
tokens = TOKEN_PATTERN.findall((text or "").lower())
|
|
if not tokens:
|
|
return bucket
|
|
for token in tokens:
|
|
digest = hashlib.sha1(token.encode("utf-8")).digest()
|
|
index = int.from_bytes(digest[:4], "big") % LOCAL_EMBED_DIMS
|
|
sign = 1.0 if digest[4] % 2 == 0 else -1.0
|
|
bucket[index] += sign
|
|
norm = math.sqrt(sum(value * value for value in bucket))
|
|
if norm == 0.0:
|
|
return bucket
|
|
return [value / norm for value in bucket]
|
|
|
|
|
|
def local_embed(texts: list[str]) -> EmbedResult:
|
|
return EmbedResult(vectors=[_local_vector(text) for text in texts], backend="local")
|
|
|
|
|
|
async def embed_texts(
|
|
texts: list[str], api_key: str, *, gateway_url: str = INTERNAL_EMBED_URL
|
|
) -> EmbedResult:
|
|
if not texts:
|
|
return EmbedResult(vectors=[], backend="empty")
|
|
headers = {
|
|
"Authorization": f"Bearer {api_key}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
payload = {"model": INTERNAL_EMBED_MODEL, "input": texts}
|
|
try:
|
|
async with httpx.AsyncClient(timeout=EMBED_TIMEOUT_SECONDS) as client:
|
|
response = await client.post(gateway_url, json=payload, headers=headers)
|
|
if response.status_code >= 400:
|
|
raise RuntimeError(f"embed gateway returned {response.status_code}")
|
|
data = response.json()
|
|
rows = data.get("data") or []
|
|
vectors = [row.get("embedding") or [] for row in rows]
|
|
if len(vectors) != len(texts) or any(not vector for vector in vectors):
|
|
raise RuntimeError("embed gateway returned an incomplete response")
|
|
return EmbedResult(vectors=vectors, backend="gateway")
|
|
except Exception as exc:
|
|
logger.warning("deepsearch embedding gateway failed, using local: %s", exc)
|
|
return local_embed(texts)
|