Files
oai-web/server/brain/embeddings.py
2026-04-08 12:43:24 +02:00

45 lines
1.2 KiB
Python

"""
brain/embeddings.py — OpenRouter embedding generation.
Uses text-embedding-3-small (1536 dims) via the OpenAI-compatible OpenRouter API.
Falls back gracefully if OpenRouter is not configured.
"""
from __future__ import annotations
import logging
logger = logging.getLogger(__name__)
_MODEL = "text-embedding-3-small"
async def get_embedding(text: str) -> list[float]:
"""
Generate a 1536-dimensional embedding for text using OpenRouter.
Returns a list of floats suitable for pgvector storage.
"""
from openai import AsyncOpenAI
from ..database import credential_store
api_key = await credential_store.get("system:openrouter_api_key")
if not api_key:
raise RuntimeError(
"OpenRouter API key is not configured — required for brain embeddings. "
"Set it via Settings → Credentials → OpenRouter API Key."
)
client = AsyncOpenAI(
api_key=api_key,
base_url="https://openrouter.ai/api/v1",
default_headers={
"HTTP-Referer": "https://mac.oai.pm",
"X-Title": "oAI-Web",
},
)
response = await client.embeddings.create(
model=_MODEL,
input=text.replace("\n", " "),
)
return response.data[0].embedding