Add semantic search support to server store

Author: javiermtorres

server/backend/pyproject.toml

@@ -16,6 +16,13 @@ dependencies = [
     "cq-sdk~=0.6.1",
 ]
 
+[project.optional-dependencies]
+semsearch = [
+    "sqlite_vec",
+    "numpy",
+    "httpx",
+]
+
 [project.scripts]
 cq-server = "cq_server.app:main"
 

server/backend/src/cq_server/app.py

@@ -89,19 +89,19 @@ def health() -> dict[str, str]:
 
 
 @api_router.get("/query")
-def query_units(
+async def query_units(
     domains: Annotated[list[str], Query()],
     languages: Annotated[list[str] | None, Query()] = None,
     frameworks: Annotated[list[str] | None, Query()] = None,
     limit: Annotated[int, Query(gt=0)] = 5,
 ) -> list[KnowledgeUnit]:
     """Search knowledge units by domain tags with relevance ranking."""
     store = _get_store()
-    return store.query(domains, languages=languages, frameworks=frameworks, limit=limit)
+    return await store.query(domains, languages=languages, frameworks=frameworks, limit=limit)
 
 
 @api_router.post("/propose", status_code=201)
-def propose_unit(request: ProposeRequest) -> KnowledgeUnit:
+async def propose_unit(request: ProposeRequest) -> KnowledgeUnit:
     """Submit a new knowledge unit."""
     store = _get_store()
     normalized = normalize_domains(request.domains)
@@ -114,42 +114,42 @@ def propose_unit(request: ProposeRequest) -> KnowledgeUnit:
         tier=Tier.PRIVATE,
         created_by=request.created_by,
     )
-    store.insert(unit)
+    await store.insert(unit)
     return unit
 
 
 @api_router.post("/confirm/{unit_id}")
-def confirm_unit(unit_id: str) -> KnowledgeUnit:
+async def confirm_unit(unit_id: str) -> KnowledgeUnit:
     """Confirm a knowledge unit, boosting its confidence."""
     store = _get_store()
-    unit = store.get(unit_id)
+    unit = await store.get(unit_id)
     if unit is None:
         raise HTTPException(status_code=404, detail="Knowledge unit not found")
     confirmed = apply_confirmation(unit)
-    store.update(confirmed)
+    await store.update(confirmed)
     return confirmed
 
 
 @api_router.post("/flag/{unit_id}")
-def flag_unit(unit_id: str, request: FlagRequest) -> KnowledgeUnit:
+async def flag_unit(unit_id: str, request: FlagRequest) -> KnowledgeUnit:
     """Flag a knowledge unit, reducing its confidence."""
     store = _get_store()
-    unit = store.get(unit_id)
+    unit = await store.get(unit_id)
     if unit is None:
         raise HTTPException(status_code=404, detail="Knowledge unit not found")
     flagged = apply_flag(unit, request.reason)
-    store.update(flagged)
+    await store.update(flagged)
     return flagged
 
 
 @api_router.get("/stats")
-def stats() -> StatsResponse:
+async def stats() -> StatsResponse:
     """Return store statistics."""
     store = _get_store()
     return StatsResponse(
-        total_units=store.count(),
-        tiers=store.counts_by_tier(),
-        domains=store.domain_counts(),
+        total_units=await store.count(),
+        tiers=await store.counts_by_tier(),
+        domains=await store.domain_counts(),
     )
 
 

server/backend/src/cq_server/auth.py

@@ -103,7 +103,7 @@ def get_current_user(request: Request) -> str:
 
 
 @router.post("/login")
-def login(request: LoginRequest, store: RemoteStore = Depends(get_store)) -> LoginResponse:
+async def login(request: LoginRequest, store: RemoteStore = Depends(get_store)) -> LoginResponse:
     """Authenticate a user and return a JWT token.
 
     Args:
@@ -116,15 +116,15 @@ def login(request: LoginRequest, store: RemoteStore = Depends(get_store)) -> Log
     Raises:
         HTTPException: With status 401 if credentials are invalid.
     """
-    user = store.get_user(request.username)
+    user = await store.get_user(request.username)
     if user is None or not verify_password(request.password, user["password_hash"]):
         raise HTTPException(status_code=401, detail="Invalid username or password")
     token = create_token(request.username, secret=_get_jwt_secret())
     return LoginResponse(token=token, username=request.username)
 
 
 @router.get("/me")
-def me(username: str = Depends(get_current_user), store: RemoteStore = Depends(get_store)) -> MeResponse:
+async def me(username: str = Depends(get_current_user), store: RemoteStore = Depends(get_store)) -> MeResponse:
     """Return the current user's info.
 
     Args:
@@ -137,7 +137,7 @@ def me(username: str = Depends(get_current_user), store: RemoteStore = Depends(g
     Raises:
         HTTPException: With status 404 if the user no longer exists.
     """
-    user = store.get_user(username)
+    user = await store.get_user(username)
     if user is None:
         raise HTTPException(status_code=404, detail="User not found")
     return MeResponse(username=user["username"], created_at=user["created_at"])

server/backend/src/cq_server/review.py

@@ -85,7 +85,7 @@ def _build_decision(unit_id: str, row: dict[str, str | None]) -> ReviewDecisionR
 
 
 @router.get("/queue")
-def review_queue(
+async def review_queue(
     limit: int = 20,
     offset: int = 0,
     _user: str = Depends(get_current_user),
@@ -102,8 +102,8 @@ def review_queue(
     Returns:
         A paginated list of pending knowledge units with review metadata.
     """
-    items = store.pending_queue(limit=limit, offset=offset)
-    total = store.pending_count()
+    items = await store.pending_queue(limit=limit, offset=offset)
+    total = await store.pending_count()
     return ReviewQueueResponse(
         items=[
             ReviewItem(
@@ -121,7 +121,7 @@ def review_queue(
 
 
 @router.post("/{unit_id}/approve")
-def approve_unit(
+async def approve_unit(
     unit_id: str,
     username: str = Depends(get_current_user),
     store: RemoteStore = Depends(get_store),
@@ -140,19 +140,19 @@ def approve_unit(
         HTTPException: With status 404 if the unit does not exist.
         HTTPException: With status 409 if the unit has already been reviewed.
     """
-    status = store.get_review_status(unit_id)
+    status = await store.get_review_status(unit_id)
     if status is None:
         raise HTTPException(status_code=404, detail="Knowledge unit not found")
     if status["status"] != "pending":
         raise HTTPException(status_code=409, detail=f"Knowledge unit already {status['status']}")
-    store.set_review_status(unit_id, "approved", username)
-    updated = store.get_review_status(unit_id)
+    await store.set_review_status(unit_id, "approved", username)
+    updated = await store.get_review_status(unit_id)
     assert updated is not None  # Unit exists; we just wrote to it.
     return _build_decision(unit_id, updated)
 
 
 @router.post("/{unit_id}/reject")
-def reject_unit(
+async def reject_unit(
     unit_id: str,
     username: str = Depends(get_current_user),
     store: RemoteStore = Depends(get_store),
@@ -171,19 +171,19 @@ def reject_unit(
         HTTPException: With status 404 if the unit does not exist.
         HTTPException: With status 409 if the unit has already been reviewed.
     """
-    status = store.get_review_status(unit_id)
+    status = await store.get_review_status(unit_id)
     if status is None:
         raise HTTPException(status_code=404, detail="Knowledge unit not found")
     if status["status"] != "pending":
         raise HTTPException(status_code=409, detail=f"Knowledge unit already {status['status']}")
-    store.set_review_status(unit_id, "rejected", username)
-    updated = store.get_review_status(unit_id)
+    await store.set_review_status(unit_id, "rejected", username)
+    updated = await store.get_review_status(unit_id)
     assert updated is not None  # Unit exists; we just wrote to it.
     return _build_decision(unit_id, updated)
 
 
 @router.get("/stats")
-def review_stats(
+async def review_stats(
     _user: str = Depends(get_current_user),
     store: RemoteStore = Depends(get_store),
 ) -> ReviewStatsResponse:
@@ -197,24 +197,24 @@ def review_stats(
         Aggregated counts by status, domain distribution, confidence
         distribution, recent activity, and daily trend data.
     """
-    counts = store.counts_by_status()
+    counts = await store.counts_by_status()
     return ReviewStatsResponse(
         counts={
             "pending": counts.get("pending", 0),
             "approved": counts.get("approved", 0),
             "rejected": counts.get("rejected", 0),
         },
-        domains=store.domain_counts(),
-        confidence_distribution=store.confidence_distribution(),
-        recent_activity=store.recent_activity(),
+        domains=await store.domain_counts(),
+        confidence_distribution=await store.confidence_distribution(),
+        recent_activity=await store.recent_activity(),
         trends=TrendsResponse(
-            daily=[DailyCount(**d) for d in store.daily_counts()],
+            daily=[DailyCount(**d) for d in await store.daily_counts()],
         ),
     )
 
 
 @router.get("/units")
-def list_units(
+async def list_units(
     domain: str | None = None,
     confidence_min: float | None = None,
     confidence_max: float | None = None,
@@ -238,7 +238,7 @@ def list_units(
     Returns:
         List of knowledge units with review metadata.
     """
-    items = store.list_units(
+    items = await store.list_units(
         domain=domain,
         confidence_min=confidence_min,
         confidence_max=confidence_max,
@@ -257,7 +257,7 @@ def list_units(
 
 
 @router.get("/{unit_id}")
-def get_unit(
+async def get_unit(
     unit_id: str,
     _user: str = Depends(get_current_user),
     store: RemoteStore = Depends(get_store),
@@ -275,10 +275,10 @@ def get_unit(
     Raises:
         HTTPException: With status 404 if the unit does not exist.
     """
-    ku = store.get_any(unit_id)
+    ku = await store.get_any(unit_id)
     if ku is None:
         raise HTTPException(status_code=404, detail="Knowledge unit not found")
-    review = store.get_review_status(unit_id)
+    review = await store.get_review_status(unit_id)
     assert review is not None  # Unit exists; get_any just returned it.
     return ReviewItem(
         knowledge_unit=ku,

server/backend/src/cq_server/semsearch.py

@@ -0,0 +1,120 @@
+import logging
+import os
+import sqlite3
+import numpy as np
+
+from cq.models import KnowledgeUnit
+
+logger = logging.getLogger(__name__)
+
+_ENABLED = False
+_DIM = int(os.environ.get("SEMSEARCH_EMBEDDING_DIM", 768))
+
+
+_TOKEN_EMBEDDING_URL = os.environ.get("TOKEN_EMBEDDING_URL", "http://localhost:8080")
+if _TOKEN_EMBEDDING_URL:
+    try:
+        import sqlite_vec
+        from httpx import AsyncClient
+        _ENABLED = True
+
+        logger.info(f"Token embedding enabled using encoderfile endpoint at {_TOKEN_EMBEDDING_URL}")
+    except ImportError:
+        logger.warning(
+            "TOKEN_EMBEDDING_URL is set but required packages are not installed; "
+            "semantic search will be unavailable. To enable, install cq with "
+            "the 'embedding' extra: pip install cq-sdk[embedding]",
+            exc_info=True,
+        )
+
+
+_VEC_SCHEMA_SQL = """
+CREATE VIRTUAL TABLE IF NOT EXISTS knowledge_units_vec
+    USING vec0(
+        id TEXT PRIMARY KEY,
+        embedding float[{dim}]
+    );
+"""
+
+_VEC_SEARCH_SQL = """
+    SELECT ku.data
+    FROM knowledge_units_vec vec
+    JOIN knowledge_units ku ON ku.id = vec.id
+    WHERE ku.status = 'approved'
+            AND vec.embedding MATCH ?
+            AND k = ?
+        ORDER BY vec.distance
+"""
+
+_VEC_DELETE_SQL = "DELETE FROM knowledge_units_vec WHERE id = ?"
+_VEC_INSERT_SQL = "INSERT INTO knowledge_units_vec (id, embedding) VALUES (?, ?)"
+
+
+async def _get_embeddings(wordlist: list[str]) -> list[np.array]:
+    """Get embeddings for a list of words using the embedding API."""
+    if not _ENABLED:
+        raise RuntimeError("Semantic search is not enabled. Set TOKEN_EMBEDDING_URL and install required packages to enable.")
+    async with AsyncClient(base_url=_TOKEN_EMBEDDING_URL) as client:
+        request_data = {"inputs": wordlist}
+        response = await client.post("/predict", json=request_data)
+        response.raise_for_status()
+        results = response.json().get("results")
+        if not results:
+            raise RuntimeError(f"Embedding API returned no embeddings for input: {request_data}")
+        return [np.average(np.array([embedding.get("embedding") for embedding in embeddings["embeddings"]]), axis=0) for embeddings in results]
+
+
+def is_enabled() -> bool:
+    """Return whether semantic search dependencies are available."""
+    return _ENABLED
+
+
+def _serialize_embedding(vec: np.ndarray) -> bytes:
+    arr = np.asarray(vec, dtype=np.float32)
+    if hasattr(sqlite_vec, "serialize_float32"):
+        return sqlite_vec.serialize_float32(arr)
+    return arr.tobytes()
+
+def load(conn: sqlite3.Connection) -> None:
+    if not _ENABLED:
+        return
+    conn.enable_load_extension(True)
+    sqlite_vec.load(conn)
+    conn.enable_load_extension(False)
+
+
+def ensure_schema(conn: sqlite3.Connection) -> None:
+    """Create semantic search virtual table if embedding is enabled."""
+    if not _ENABLED:
+        return
+    conn.executescript(_VEC_SCHEMA_SQL.format(dim=_DIM))
+
+
+async def upsert_unit(conn: sqlite3.Connection, unit: KnowledgeUnit) -> None:
+    """Insert or update a knowledge unit embedding row."""
+    if not _ENABLED:
+        return
+    text = " ".join([unit.insight.summary, unit.insight.detail, unit.insight.action]).strip()
+    if not text:
+        return
+    embedding = (await _get_embeddings([text]))[0]
+    serialized = _serialize_embedding(embedding)
+    conn.execute(_VEC_DELETE_SQL, (unit.id,))
+    conn.execute(_VEC_INSERT_SQL, (unit.id, serialized))
+
+
+async def query(conn: sqlite3.Connection, domains: list[str], *, limit: int = 5) -> list[KnowledgeUnit]:
+    """Semantic search implementation."""
+    if not _ENABLED:
+        return []
+    if not domains:
+        return []
+    if limit <= 0:
+        return []
+    vec_emb_search = await _get_embeddings(domains)
+    search_embedding = _serialize_embedding(vec_emb_search[0])
+    try:
+        vec_rows = conn.execute(_VEC_SEARCH_SQL, (search_embedding, limit)).fetchall()
+    except sqlite3.OperationalError:
+        vec_rows = []
+    return [KnowledgeUnit.model_validate_json(row[0]) for row in vec_rows]