c762564d40
- Fix 71 invalid-syntax files (class-body newline-broken assignments) - Add from/None chain to 307 B904 raise-without-from sites - Add B008 ignore to ruff.toml (already in pyproject.toml) - Noqa F401 on __init__.py re-exports (137 sites) - Noqa E402 on deferred imports (63 sites) - Bulk-add stdlib/FastAPI/project imports for F821 (127 sites) - Replace ×→x, –→-, …→... in docstrings (4093 chars) - Manual refactor of 5 SIM103/SIM116 patterns Tests: 791 passed (66 deselected due to pre-existing Redis issues in test_rag.py) Co-authored-by: opencode <opencode@rugmunch.io>
107 lines
3.3 KiB
Python
107 lines
3.3 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Supabase vector search helper (RAG integration).
|
|
Uses the existing search_embeddings RPC already installed on the project.
|
|
No pinecone/weaviate needed - SQL + pgvector.
|
|
"""
|
|
|
|
import os
|
|
from typing import Any
|
|
|
|
import httpx
|
|
from dotenv import load_dotenv
|
|
|
|
# Load env vars dynamically - override stale Docker env
|
|
load_dotenv("/app/.env", override=True)
|
|
|
|
|
|
def _get_url():
|
|
return os.environ.get("SUPABASE_URL", "")
|
|
|
|
|
|
def _get_key():
|
|
return os.environ.get("SUPABASE_SERVICE_ROLE_KEY", "") or os.environ.get("SUPABASE_SERVICE_KEY", "")
|
|
|
|
|
|
def _get_headers():
|
|
key = _get_key()
|
|
return {
|
|
"apikey": key,
|
|
"Authorization": f"Bearer {key}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
|
|
|
|
async def search_similar(
|
|
query_embedding: list[float],
|
|
namespace: str = "default",
|
|
match_count: int = 10,
|
|
similarity_threshold: float = 0.7,
|
|
) -> list[dict[str, Any]]:
|
|
"""
|
|
Search for semantically similar documents using pgvector.
|
|
Returns matching document IDs with similarity scores.
|
|
|
|
Args:
|
|
query_embedding: The embedding vector from your model
|
|
namespace: Search namespace to restrict results
|
|
match_count: Number of results to return
|
|
similarity_threshold: Minimum cosine similarity (0-1)
|
|
"""
|
|
# Pad/truncate to match the pgvector table dimension
|
|
from app.supabase_vector import TABLE_DIM, pad_vector
|
|
|
|
padded_embedding = pad_vector(query_embedding, TABLE_DIM)
|
|
|
|
url = f"{_get_url()}/rest/v1/rpc/search_embeddings"
|
|
payload = {
|
|
"query_embedding": padded_embedding,
|
|
"match_count": match_count,
|
|
"namespace": namespace,
|
|
"similarity_threshold": similarity_threshold,
|
|
}
|
|
async with httpx.AsyncClient(timeout=30) as client:
|
|
r = await client.post(url, json=payload, headers=_get_headers())
|
|
if r.status_code == 200:
|
|
return r.json()
|
|
return []
|
|
|
|
|
|
async def store_embedding(
|
|
document_id: str,
|
|
embedding: list[float],
|
|
namespace: str = "default",
|
|
content_hash: str = "",
|
|
metadata: dict | None = None,
|
|
model_name: str = "",
|
|
) -> dict | None:
|
|
"""
|
|
Store an embedding for later retrieval.
|
|
Idempotent - uses ON CONFLICT (document_id) for upsert via REST.
|
|
"""
|
|
url = f"{_get_url()}/rest/v1/embeddings"
|
|
payload = {
|
|
"document_id": document_id,
|
|
"embedding": embedding,
|
|
"namespace": namespace,
|
|
"content_hash": content_hash,
|
|
"metadata": metadata or {},
|
|
"model_name": model_name,
|
|
}
|
|
headers = dict(_get_headers())
|
|
headers["Prefer"] = "resolution=merge-duplicates"
|
|
async with httpx.AsyncClient(timeout=30) as client:
|
|
r = await client.post(url, json=payload, headers=headers)
|
|
return r.json() if r.status_code in (200, 201) else None
|
|
|
|
|
|
async def get_namespace_stats(namespace: str = "default") -> dict:
|
|
"""Get document count and stats for a namespace."""
|
|
url = f"{_get_url()}/rest/v1/embeddings?namespace=eq.{namespace}&select=id"
|
|
async with httpx.AsyncClient(timeout=10) as client:
|
|
r = await client.get(url, headers={**_get_headers(), "Prefer": "count=exact"})
|
|
return {
|
|
"namespace": namespace,
|
|
"count": int(r.headers.get("content-range", "0").split("/")[-1] or 0),
|
|
"status": r.status_code,
|
|
}
|