feat: optional sqlite-vec embedding pipeline for Paid-tier RAG

This commit is contained in:
pyr0ball 2026-05-13 16:32:57 -07:00
parent d8c3eba0f8
commit e0bfa11642
3 changed files with 131 additions and 0 deletions

View file

@ -17,6 +17,20 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
# sqlite-vec: optional vector search extension for context embedding (Paid tier)
RUN set -eux; \
SVEC_VER=0.1.6; \
ARCH=$(uname -m); \
case "$ARCH" in \
x86_64) SVEC_ARCH="x86_64-linux-gnu" ;; \
aarch64) SVEC_ARCH="aarch64-linux-gnu" ;; \
*) echo "sqlite-vec: unsupported arch $ARCH — skipping" && exit 0 ;; \
esac; \
wget -q -O /tmp/sqlite_vec.tar.gz \
"https://github.com/asg017/sqlite-vec/releases/download/v${SVEC_VER}/sqlite-vec-${SVEC_VER}-loadable-linux-${SVEC_ARCH}.tar.gz"; \
tar -xz -C /usr/lib/python3/ -f /tmp/sqlite_vec.tar.gz --wildcards '*.so' || true; \
rm /tmp/sqlite_vec.tar.gz
COPY app/ ./app/
COPY patterns/ ./patterns/
COPY scripts/ ./scripts/

64
app/context/embedder.py Normal file
View file

@ -0,0 +1,64 @@
"""Ollama embedding client with sqlite-vec storage — BSL licensed."""
from __future__ import annotations
import logging
import sqlite3
import struct
from pathlib import Path
import httpx
logger = logging.getLogger(__name__)
EMBEDDING_AVAILABLE: bool = False
try:
import sqlite_vec # type: ignore[import] # noqa: F401
EMBEDDING_AVAILABLE = True
logger.debug("sqlite-vec loaded — embedding pipeline enabled")
except ImportError:
logger.debug("sqlite-vec not available — embedding pipeline disabled")
def embed_chunks(
db_path: Path,
document_id: str,
llm_url: str,
model: str = "nomic-embed-text",
timeout: float = 60.0,
) -> int:
"""Embed all unembedded chunks for a document. Returns count embedded. No-op when EMBEDDING_AVAILABLE is False."""
if not EMBEDDING_AVAILABLE:
return 0
conn = sqlite3.connect(str(db_path))
conn.execute("PRAGMA journal_mode=WAL")
conn.row_factory = sqlite3.Row
rows = conn.execute(
"SELECT id, text FROM context_chunks WHERE document_id=? AND embedding IS NULL",
(document_id,),
).fetchall()
count = 0
for row in rows:
try:
resp = httpx.post(
f"{llm_url.rstrip('/')}/api/embeddings",
json={"model": model, "prompt": row["text"]},
timeout=timeout,
)
resp.raise_for_status()
vector: list[float] = resp.json().get("embedding") or []
if vector:
blob = struct.pack(f"{len(vector)}f", *vector)
conn.execute(
"UPDATE context_chunks SET embedding=? WHERE id=?",
(blob, row["id"]),
)
count += 1
except Exception as exc:
logger.warning("Embedding chunk %s failed: %s", row["id"], exc)
conn.commit()
conn.close()
return count

View file

@ -0,0 +1,53 @@
"""Tests for app/context/embedder.py — graceful no-op without sqlite-vec."""
import sqlite3
from pathlib import Path
from unittest.mock import patch
import pytest
from app.context import embedder as emb_mod
@pytest.fixture
def db(tmp_path):
db_path = tmp_path / "t.db"
conn = sqlite3.connect(str(db_path))
conn.executescript("""
CREATE TABLE context_documents (
id TEXT PRIMARY KEY, filename TEXT NOT NULL, doc_type TEXT NOT NULL,
full_text TEXT NOT NULL, file_size INTEGER, uploaded_at TEXT NOT NULL
);
CREATE TABLE context_chunks (
id TEXT PRIMARY KEY, document_id TEXT NOT NULL
REFERENCES context_documents(id) ON DELETE CASCADE,
chunk_index INTEGER NOT NULL, text TEXT NOT NULL, embedding BLOB
);
INSERT INTO context_documents VALUES ('d1','test.md','markdown','hello',5,'2026-01-01T00:00:00+00:00');
INSERT INTO context_chunks VALUES ('c1','d1',0,'hello world',NULL);
""")
conn.commit()
conn.close()
return db_path
def test_embed_skipped_when_extension_absent(db):
with patch.object(emb_mod, "EMBEDDING_AVAILABLE", False):
count = emb_mod.embed_chunks(db, "d1", "http://localhost:11434")
assert count == 0
def test_embed_calls_ollama_when_available(db):
import httpx
class FakeResponse:
status_code = 200
def raise_for_status(self): pass
def json(self): return {"embedding": [0.1, 0.2, 0.3]}
with patch.object(emb_mod, "EMBEDDING_AVAILABLE", True), \
patch("app.context.embedder.httpx.post", return_value=FakeResponse()):
count = emb_mod.embed_chunks(db, "d1", "http://localhost:11434")
assert count == 1
# Verify blob was written
conn = sqlite3.connect(str(db))
row = conn.execute("SELECT embedding FROM context_chunks WHERE id='c1'").fetchone()
conn.close()
assert row[0] is not None