56fdc57389
- Fix double /v1 in TTS audio/speech URL when LITELLM_API_BASE includes /v1 - Fix double /v1 in embedding service and vector service URLs - Clean up docs: remove second-person language in deployment, frontend, migrations Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
103 lines
3.8 KiB
Python
103 lines
3.8 KiB
Python
"""Embedding generation for semantic search via pgvector.
|
|
|
|
Priority:
|
|
1. LiteLLM proxy — model from Redis settings (overrides env)
|
|
2. LiteLLM proxy — model from LITELLM_EMBEDDING_MODEL env
|
|
3. AWS Bedrock Titan Embed V2 (direct fallback)
|
|
"""
|
|
import logging
|
|
|
|
from app.config import settings
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _get_embedding_model() -> str:
|
|
"""Return the active embedding model: Redis setting takes precedence over env."""
|
|
try:
|
|
import redis as redis_lib
|
|
r = redis_lib.from_url(settings.REDIS_URL, decode_responses=True)
|
|
model = r.get("settings:embedding_model")
|
|
if model:
|
|
return model
|
|
except Exception:
|
|
pass
|
|
return settings.LITELLM_EMBEDDING_MODEL
|
|
|
|
|
|
def _text_for_question(question_text: str, options: list[str] | None) -> str:
|
|
"""Build the text to embed for a question — stem + options, no explanation."""
|
|
parts = [question_text]
|
|
if options:
|
|
parts.extend(options)
|
|
return " ".join(parts)[:4000]
|
|
|
|
|
|
def generate_embedding(text: str) -> list[float] | None:
|
|
"""Generate a 1024-dim embedding.
|
|
|
|
Priority:
|
|
1. LiteLLM proxy (openai/titan-embed-v2) — scores ~0.71 cosine similarity
|
|
2. AWS Bedrock direct — fallback, scores ~0.48
|
|
"""
|
|
clean = " ".join(text.split())[:4000]
|
|
if not clean:
|
|
return None
|
|
|
|
# ── 1. LiteLLM proxy (direct httpx — avoids LiteLLM param validation) ──
|
|
embedding_model = _get_embedding_model()
|
|
api_base = (settings.LITELLM_API_BASE or "").rstrip("/").removesuffix("/v1")
|
|
if embedding_model and settings.LITELLM_API_KEY and api_base:
|
|
try:
|
|
import httpx, json as _json
|
|
body: dict = {"model": embedding_model, "input": [clean], "dimensions": settings.EMBEDDING_DIMENSIONS}
|
|
resp = httpx.post(
|
|
f"{api_base}/v1/embeddings",
|
|
headers={"Authorization": f"Bearer {settings.LITELLM_API_KEY}", "Content-Type": "application/json"},
|
|
content=_json.dumps(body),
|
|
timeout=30,
|
|
)
|
|
resp.raise_for_status()
|
|
emb = resp.json()["data"][0]["embedding"]
|
|
if len(emb) == settings.EMBEDDING_DIMENSIONS:
|
|
return emb
|
|
logger.warning(f"Embedding dim mismatch: got {len(emb)}, expected {settings.EMBEDDING_DIMENSIONS}")
|
|
except Exception as e:
|
|
logger.warning(f"LiteLLM embedding failed: {e}")
|
|
|
|
# ── 2. AWS Bedrock Titan direct (fallback) ──────────────────
|
|
if settings.AWS_ACCESS_KEY_ID and settings.AWS_SECRET_ACCESS_KEY:
|
|
try:
|
|
import boto3, json
|
|
client = boto3.client(
|
|
"bedrock-runtime",
|
|
aws_access_key_id=settings.AWS_ACCESS_KEY_ID,
|
|
aws_secret_access_key=settings.AWS_SECRET_ACCESS_KEY,
|
|
region_name=settings.AWS_BEDROCK_REGION or "us-east-1",
|
|
)
|
|
body = json.dumps({
|
|
"inputText": clean,
|
|
"dimensions": settings.EMBEDDING_DIMENSIONS,
|
|
"normalize": True,
|
|
})
|
|
resp = client.invoke_model(
|
|
modelId="amazon.titan-embed-text-v2:0",
|
|
body=body,
|
|
contentType="application/json",
|
|
accept="application/json",
|
|
)
|
|
return json.loads(resp["body"].read())["embedding"]
|
|
except Exception as e:
|
|
logger.warning(f"Bedrock embedding failed: {e}")
|
|
|
|
return None
|
|
|
|
|
|
def embed_question(question) -> bool:
|
|
"""Generate and store embedding for a Question ORM object. Returns True on success."""
|
|
text = _text_for_question(question.question_text, question.options)
|
|
emb = generate_embedding(text)
|
|
if emb:
|
|
question.embedding = emb
|
|
return True
|
|
return False
|