Дмитрий
2a109455a5
embed.py: get_embedder(provider) — aitunnel (живой /v1/embeddings) по умолчанию, fake (детерминированный без сети) для тестов, local — заглушка NotImplementedError. batch() применяет embed_fn инъекцией — сеть не затрагивается в тестах.
51 lines
2.0 KiB
Python
51 lines
2.0 KiB
Python
"""Эмбеддинги: провайдер AITunnel (по умолчанию) с fallback на локальную модель.
|
||
Живой источник инкапсулирован — вся остальная логика получает embed_fn инъекцией."""
|
||
import os
|
||
import json
|
||
import urllib.request
|
||
|
||
|
||
def batch(texts, embed_fn):
|
||
return [embed_fn(t) for t in texts]
|
||
|
||
|
||
def _aitunnel_embedder():
|
||
key_path = os.environ.get("ALEX_AITUNNEL_KEY", "aitunnel.key")
|
||
key = open(key_path, encoding="utf-8").read().strip()
|
||
model = os.environ.get("ALEX_EMBED_MODEL", "text-embedding-3-small")
|
||
|
||
def embed(text):
|
||
body = {"model": model, "input": text}
|
||
req = urllib.request.Request(
|
||
"https://api.aitunnel.ru/v1/embeddings",
|
||
data=json.dumps(body).encode(), method="POST")
|
||
req.add_header("Authorization", f"Bearer {key}")
|
||
req.add_header("Content-Type", "application/json")
|
||
r = json.load(urllib.request.urlopen(req, timeout=30))
|
||
return r["data"][0]["embedding"]
|
||
|
||
return embed
|
||
|
||
|
||
def _fake_embedder():
|
||
# детерминированный «отпечаток» без сети — для тестов/сухих прогонов
|
||
def embed(text):
|
||
t = text.lower()
|
||
return [
|
||
float(sum(ord(c) for c in t) % 97),
|
||
float(len(t) % 13),
|
||
float(t.count("а") + t.count("о")),
|
||
]
|
||
return embed
|
||
|
||
|
||
def get_embedder(provider=None):
|
||
provider = provider or os.environ.get("ALEX_EMBED", "aitunnel")
|
||
if provider == "aitunnel":
|
||
return _aitunnel_embedder()
|
||
if provider == "local":
|
||
# fallback: локальная модель на GPU-сервере (sentence-transformers).
|
||
# Реализовать при переключении; сигнатура та же: embed(text) -> list[float].
|
||
raise NotImplementedError("local-эмбеддер подключается на GPU-сервере при провале AITunnel")
|
||
return _fake_embedder()
|