137 lines
4.8 KiB
Python
137 lines
4.8 KiB
Python
"""Embedding- und Rerank-Clients sowie Skill-Index."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import math
|
|
import re
|
|
from collections.abc import Sequence
|
|
from typing import Protocol
|
|
|
|
import httpx
|
|
|
|
Vector = list[float]
|
|
|
|
|
|
class Embedder(Protocol):
|
|
def embed(self, texts: Sequence[str]) -> list[Vector]: ...
|
|
|
|
|
|
class Reranker(Protocol):
|
|
def rerank(self, query: str, documents: Sequence[str]) -> list[float]: ...
|
|
|
|
|
|
def cosine(a: Sequence[float], b: Sequence[float]) -> float:
|
|
dot = sum(x * y for x, y in zip(a, b, strict=True))
|
|
na = math.sqrt(sum(x * x for x in a))
|
|
nb = math.sqrt(sum(y * y for y in b))
|
|
if na == 0 or nb == 0:
|
|
return 0.0
|
|
return dot / (na * nb)
|
|
|
|
|
|
class HashingEmbedder:
|
|
"""Deterministischer Bag-of-Words-Embedder für Tests und als Notfall-Fallback ohne GPU/Netz."""
|
|
|
|
_WORD = re.compile(r"\w+", re.UNICODE)
|
|
|
|
def __init__(self, dim: int = 512):
|
|
self.dim = dim
|
|
|
|
def _vector(self, text: str) -> Vector:
|
|
vec = [0.0] * self.dim
|
|
for word in self._WORD.findall(text.lower()):
|
|
h = int.from_bytes(hashlib.sha1(word.encode()).digest()[:4], "big")
|
|
vec[h % self.dim] += 1.0 if (h >> 31) & 1 else -1.0
|
|
norm = math.sqrt(sum(v * v for v in vec)) or 1.0
|
|
return [v / norm for v in vec]
|
|
|
|
def embed(self, texts: Sequence[str]) -> list[Vector]:
|
|
return [self._vector(t) for t in texts]
|
|
|
|
|
|
class OpenAIEmbedder:
|
|
"""Embeddings über ein OpenAI-kompatibles API (LiteLLM → llama.cpp/TEI)."""
|
|
|
|
def __init__(self, base_url: str, model: str, api_key: str = "", timeout: float = 30.0,
|
|
client: httpx.Client | None = None):
|
|
self.base_url = base_url.rstrip("/")
|
|
self.model = model
|
|
self.headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
|
|
self.client = client or httpx.Client(timeout=timeout)
|
|
|
|
def embed(self, texts: Sequence[str]) -> list[Vector]:
|
|
resp = self.client.post(
|
|
f"{self.base_url}/v1/embeddings",
|
|
json={"model": self.model, "input": list(texts)},
|
|
headers=self.headers,
|
|
)
|
|
resp.raise_for_status()
|
|
data = sorted(resp.json()["data"], key=lambda d: d["index"])
|
|
return [d["embedding"] for d in data]
|
|
|
|
|
|
class LlamaCppReranker:
|
|
"""Reranking über den llama.cpp-Server (`--reranking`, Endpoint /v1/rerank)."""
|
|
|
|
def __init__(self, base_url: str, model: str = "rerank-default", timeout: float = 30.0,
|
|
client: httpx.Client | None = None):
|
|
self.base_url = base_url.rstrip("/")
|
|
self.model = model
|
|
self.client = client or httpx.Client(timeout=timeout)
|
|
|
|
def rerank(self, query: str, documents: Sequence[str]) -> list[float]:
|
|
resp = self.client.post(
|
|
f"{self.base_url}/v1/rerank",
|
|
json={"model": self.model, "query": query, "documents": list(documents)},
|
|
)
|
|
resp.raise_for_status()
|
|
scores = [0.0] * len(documents)
|
|
for item in resp.json()["results"]:
|
|
scores[item["index"]] = float(item["relevance_score"])
|
|
return scores
|
|
|
|
|
|
class SkillIndex:
|
|
"""Hält pro Skill die Vektoren von Beschreibung und Triggern (wenige hundert Skills → im RAM)."""
|
|
|
|
def __init__(self, embedder: Embedder):
|
|
self.embedder = embedder
|
|
self._cache: dict[str, Vector] = {}
|
|
self._entries: dict[str, list[Vector]] = {}
|
|
|
|
def _embed_cached(self, texts: list[str]) -> list[Vector]:
|
|
missing = [t for t in texts if t not in self._cache]
|
|
if missing:
|
|
for text, vec in zip(missing, self.embedder.embed(missing), strict=True):
|
|
self._cache[text] = vec
|
|
return [self._cache[t] for t in texts]
|
|
|
|
def rebuild(self, skills: dict[str, list[str]]) -> None:
|
|
"""skills: name -> Texte (Beschreibung + Trigger)."""
|
|
all_texts = sorted({t for texts in skills.values() for t in texts})
|
|
self._embed_cached(all_texts)
|
|
self._entries = {name: [self._cache[t] for t in texts] for name, texts in skills.items()}
|
|
|
|
def vector(self, name: str) -> Vector:
|
|
return self._entries[name][0]
|
|
|
|
def search(self, query_vec: Vector, k: int = 5, only: set[str] | None = None) -> list[tuple[str, float]]:
|
|
scored = []
|
|
for name, vecs in self._entries.items():
|
|
if only is not None and name not in only:
|
|
continue
|
|
scored.append((name, max(cosine(query_vec, v) for v in vecs)))
|
|
scored.sort(key=lambda x: (-x[1], x[0]))
|
|
return scored[:k]
|
|
|
|
def pairwise_similar(self, threshold: float) -> list[tuple[str, str, float]]:
|
|
names = sorted(self._entries)
|
|
pairs = []
|
|
for i, a in enumerate(names):
|
|
for b in names[i + 1:]:
|
|
sim = cosine(self.vector(a), self.vector(b))
|
|
if sim >= threshold:
|
|
pairs.append((a, b, sim))
|
|
return pairs
|