feat: optional sqlite-vec embedding pipeline for Paid-tier RAG
This commit is contained in:
parent
d8c3eba0f8
commit
e0bfa11642
3 changed files with 131 additions and 0 deletions
14
Dockerfile
14
Dockerfile
|
|
@ -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
64
app/context/embedder.py
Normal 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
|
||||
53
tests/context/test_embedder.py
Normal file
53
tests/context/test_embedder.py
Normal 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
|
||||
Loading…
Reference in a new issue