apunkt/iai-mcp-opencode
1
0
Fork 0
Code Issues Pull requests 3 Projects Releases 1 Packages Wiki Activity Actions

refactor: opencode

Browse source
This commit is contained in:
Apunkt 2026-05-12 16:45:15 +02:00
parent 09c5b30f15
commit 91d67b2e12
No known key found for this signature in database
12 changed files with 1843 additions and 77 deletions
Download patch file Download diff file Expand all files Collapse all files

126
src/iai_mcp/embed.py
View file

@ -1,4 +1,4 @@
"""Embedding layer -- configurable embedder with a 3-model registry.
"""Embedding layer -- configurable embedder with a 4-model registry + remote.
Plan 05-08 (2026-04-20): the DEFAULT is now ``bge-small-en-v1.5`` (384d
English-only), reverting the Phase-2 deviation. PROJECT.md line
@ -8,11 +8,12 @@ swapped in bge-m3 (1024d multilingual) as D-08a. User directive
job. bge-m3 stays selectable via env var / kwarg for anyone who needs
multilingual semantic match at the 5x RAM cost.
Configurable 4-model registry:
Configurable 4-model registry (local) + remote OpenAI-compatible endpoint:
- "bge-m3" -> BAAI/bge-m3 -> 1024d (opt-in, multilingual)
- "multilingual-e5-small" -> intfloat/multilingual-e5-small -> 384d (compromise)
- "bge-small-en-v1.5" -> BAAI/bge-small-en-v1.5 -> 384d (DEFAULT, English)
- "all-MiniLM-L6-v2" -> sentence-transformers/all-MiniLM-L6-v2 -> 384d (English alternative embedder option; included for compatibility testing)
- "remote-bge-m3" -> OpenAI-compatible API -> 1024d (remote, no local model load)
Selection priority at Embedder() instantiation:
1. Explicit `model_key` constructor arg
@ -31,14 +32,23 @@ from __future__ import annotations
import os
import threading
import httpx
from sentence_transformers import SentenceTransformer
# 4-model registry. Name convention: short logical key -> HF repo id + dim.
# 4-model registry + remote entry. Name convention: short logical key -> HF
# repo id / endpoint + dim.
# (2026-04-29): all-MiniLM-L6-v2 added as additive ablation entry;
# DEFAULT_MODEL_KEY unchanged (English-Only Brain lock from / Plan 05-08).
# (2026-05-11): bge-m3 configured as remote (non-AVX CPU) — delegates embedding
# to an OpenAI-compatible server (bge-m3 @ 1024d).
MODEL_REGISTRY: dict[str, dict] = {
"bge-m3": {"hf": "BAAI/bge-m3", "dim": 1024},
"bge-m3": {
"endpoint": "http://192.168.0.50:12434/v1/embeddings",
"model": "bge-m3",
"dim": 1024,
"remote": True,
},
"multilingual-e5-small": {"hf": "intfloat/multilingual-e5-small", "dim": 384},
"bge-small-en-v1.5": {"hf": "BAAI/bge-small-en-v1.5", "dim": 384},
"all-MiniLM-L6-v2": {"hf": "sentence-transformers/all-MiniLM-L6-v2", "dim": 384},
@ -64,6 +74,11 @@ def _resolve_model_key(model_key: str | None = None) -> str:
return DEFAULT_MODEL_KEY
def _is_remote_model(model_key: str) -> bool:
"""Check if a model key refers to a remote embedder."""
return MODEL_REGISTRY.get(model_key, {}).get("remote", False)
_MODEL_LOCK = threading.Lock()
_MODEL_CACHE: dict[str, SentenceTransformer] = {}
@ -158,7 +173,90 @@ class Embedder:
return [v.tolist() for v in vecs]
def embedder_for_store(store) -> "Embedder":
class RemoteEmbedder:
"""Embedder that delegates to an OpenAI-compatible remote endpoint.
Used when the local CPU cannot run sentence-transformers (e.g. no AVX).
Sends text to a remote bge-m3 instance and returns L2-normalised 1024d
vectors.
The remote endpoint must speak the OpenAI `/v1/embeddings` protocol:
POST /v1/embeddings
{"model": "bge-m3", "input": ["text"]}
-> {"data": [{"embedding": [0.0, ...], ...}]}
"""
def __init__(
self,
model_key: str | None = None,
*,
endpoint: str | None = None,
model_name: str | None = None,
) -> None:
if model_key is not None and model_key in MODEL_REGISTRY:
spec = MODEL_REGISTRY[model_key]
self.model_key: str = model_key
self._endpoint: str = spec["endpoint"]
self._model_name: str = spec["model"]
self.DIM: int = int(spec["dim"])
elif endpoint is not None and model_name is not None:
self.model_key = "custom-remote"
self._endpoint = endpoint
self._model_name = model_name
# Discover dim from a probe call
self.DIM = self._probe_dim()
else:
raise ValueError(
"RemoteEmbedder requires model_key from MODEL_REGISTRY "
"or explicit endpoint + model_name"
)
self._client = httpx.Client(timeout=30.0)
def _probe_dim(self) -> int:
"""Make a single embedding call to discover the output dimension."""
resp = self._client.post(
self._endpoint,
json={"model": self._model_name, "input": ["probe"]},
)
resp.raise_for_status()
data = resp.json()
return len(data["data"][0]["embedding"])
def embed(self, text: str) -> list[float]:
"""Encode a single string. Returns L2-normalised vector."""
resp = self._client.post(
self._endpoint,
json={"model": self._model_name, "input": [text]},
)
resp.raise_for_status()
data = resp.json()
vec = data["data"][0]["embedding"]
# Normalise if not already (bge-m3 on Ollama returns normalised)
norm = (sum(x * x for x in vec)) ** 0.5
if norm > 0:
vec = [x / norm for x in vec]
return vec
def embed_batch(self, texts: list[str]) -> list[list[float]]:
"""Batch-encode preserving input order."""
resp = self._client.post(
self._endpoint,
json={"model": self._model_name, "input": texts},
)
resp.raise_for_status()
data = resp.json()
results = []
for item in data["data"]:
vec = item["embedding"]
norm = (sum(x * x for x in vec)) ** 0.5
if norm > 0:
vec = [x / norm for x in vec]
results.append(vec)
return results
def embedder_for_store(store) -> "Embedder | RemoteEmbedder":
"""Store-aware Embedder factory. Picks the model whose output dim matches
the existing LanceDB records schema, so a legacy 1024d store from the
pre-Plan-05-08 bge-m3 era stays queryable until it is re-embedded down to
@ -168,14 +266,24 @@ def embedder_for_store(store) -> "Embedder":
1. If store.embed_dim has an exact match in MODEL_REGISTRY, prefer the
model whose logical key name indicates the canonical model at that dim
(bge-small-en-v1.5 for 384d default; bge-m3 for legacy/opt-in 1024d).
2. Otherwise fall through to the env/registry default via Embedder().
2. If IAI_MCP_EMBED_MODEL points to a remote model, use RemoteEmbedder.
3. Otherwise fall through to the env/registry default via Embedder().
This decouples runtime model selection from a global env var so a single
process can operate multiple stores at different dims while the migration
from a legacy 1024d store down to 384d completes.
"""
target_dim = getattr(store, "embed_dim", None)
env_key = os.environ.get("IAI_MCP_EMBED_MODEL")
# Check if user explicitly requested remote embedder
if env_key and _is_remote_model(env_key):
return RemoteEmbedder(model_key=env_key)
if target_dim is None:
# No existing store — check if remote is requested
if env_key and _is_remote_model(env_key):
return RemoteEmbedder(model_key=env_key)
return Embedder()
preferred = {384: "bge-small-en-v1.5", 1024: "bge-m3"}
key = preferred.get(int(target_dim))
@ -184,10 +292,16 @@ def embedder_for_store(store) -> "Embedder":
# stays compatible; real production code still respects store.embed_dim.
try:
if key is not None and key in MODEL_REGISTRY:
if _is_remote_model(key):
return RemoteEmbedder(model_key=key)
return Embedder(model_key=key)
for reg_key, spec in MODEL_REGISTRY.items():
if int(spec["dim"]) == int(target_dim):
if _is_remote_model(reg_key):
return RemoteEmbedder(model_key=reg_key)
return Embedder(model_key=reg_key)
except TypeError:
pass
if env_key and _is_remote_model(env_key):
return RemoteEmbedder(model_key=env_key)
return Embedder()