business-stack 0.1.0

package/backend/app/services/embeddings/__init__.py

@@ -0,0 +1,20 @@
+"""
+Gemini embedding generation (single- or multi-modal, batched).
+
+Public entrypoints:
+
+- ``batch_embed_contents`` — low-level API batch (list of interleaved part lists).
+
+Import ``run_embed_document_job`` / ``embed_document_gemini`` from
+``app.services.embeddings.worker`` to avoid circular imports with the entity pipeline.
+"""
+
+from app.services.embeddings.gemini_api import batch_embed_contents
+from app.services.embeddings.types import EmbeddingPart, InlineDataPart, TextPart
+
+__all__ = [
+    "EmbeddingPart",
+    "InlineDataPart",
+    "TextPart",
+    "batch_embed_contents",
+]

package/backend/app/services/embeddings/build_inputs.py

@@ -0,0 +1,140 @@
+from __future__ import annotations
+
+import json
+import logging
+from dataclasses import dataclass
+
+from sqlalchemy import text
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from app.services.embeddings.types import EmbeddingPart, InlineDataPart, TextPart
+from app.services.normalization.normalizer import parse_blob_sha256
+from app.storage.blobs import BlobStore
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True, slots=True)
+class ChunkRow:
+    id: int
+    ordinal: int
+    text: str
+    start_block_ordinal: int
+    end_block_ordinal: int
+    meta_raw: str | None
+
+
+@dataclass(frozen=True, slots=True)
+class BlockRow:
+    ordinal: int
+    type: str
+    meta: dict
+    storage_uri: str | None
+    sha256: str | None
+    mime: str | None
+
+
+async def load_document_chunks(
+    session: AsyncSession,
+    document_id: str,
+) -> list[ChunkRow]:
+    r = await session.execute(
+        text(
+            "SELECT id, ordinal, text, start_block_ordinal, end_block_ordinal, meta "
+            "FROM document_chunks WHERE document_id = :d ORDER BY ordinal",
+        ),
+        {"d": document_id},
+    )
+    rows = r.fetchall()
+    return [
+        ChunkRow(
+            id=int(row[0]),
+            ordinal=int(row[1]),
+            text=str(row[2]),
+            start_block_ordinal=int(row[3]),
+            end_block_ordinal=int(row[4]),
+            meta_raw=row[5] if row[5] else None,
+        )
+        for row in rows
+    ]
+
+
+async def load_blocks_span(
+    session: AsyncSession,
+    document_id: str,
+    start_ord: int,
+    end_ord: int,
+) -> list[BlockRow]:
+    r = await session.execute(
+        text(
+            "SELECT ordinal, type, meta, storage_uri, sha256, mime "
+            "FROM content_blocks WHERE document_id = :d "
+            "AND ordinal >= :s AND ordinal <= :e "
+            "ORDER BY ordinal",
+        ),
+        {"d": document_id, "s": start_ord, "e": end_ord},
+    )
+    out: list[BlockRow] = []
+    for row in r.fetchall():
+        meta: dict = {}
+        if row[2]:
+            try:
+                p = json.loads(row[2])
+                if isinstance(p, dict):
+                    meta = p
+            except json.JSONDecodeError:
+                pass
+        out.append(
+            BlockRow(
+                ordinal=int(row[0]),
+                type=str(row[1]),
+                meta=meta,
+                storage_uri=row[3] if row[3] else None,
+                sha256=row[4] if row[4] else None,
+                mime=row[5] if row[5] else None,
+            ),
+        )
+    return out
+
+
+def chunk_to_text_parts(chunk: ChunkRow) -> list[EmbeddingPart]:
+    """Single-modality text from stored chunk row."""
+    if not chunk.text.strip():
+        return [TextPart(" ")]
+    return [TextPart(chunk.text)]
+
+
+async def chunk_to_multimodal_parts(
+    session: AsyncSession,
+    document_id: str,
+    chunk: ChunkRow,
+    blob_store: BlobStore,
+) -> list[EmbeddingPart]:
+    """Interleaved parts from content_blocks in the chunk's ordinal span."""
+    blocks = await load_blocks_span(
+        session,
+        document_id,
+        chunk.start_block_ordinal,
+        chunk.end_block_ordinal,
+    )
+    parts: list[EmbeddingPart] = []
+    for b in blocks:
+        if b.type == "text":
+            t = b.meta.get("text")
+            if isinstance(t, str) and t.strip():
+                parts.append(TextPart(t.strip()))
+            continue
+        if b.type == "image" and b.storage_uri:
+            sha = b.sha256 or parse_blob_sha256(b.storage_uri)
+            if not sha:
+                continue
+            try:
+                data = blob_store.read_bytes(sha)
+            except OSError:
+                logger.warning("missing blob %s for document %s", sha, document_id)
+                continue
+            mime = b.mime or "image/png"
+            parts.append(InlineDataPart(mime_type=mime, data=data))
+    if not parts:
+        return chunk_to_text_parts(chunk)
+    return parts

package/backend/app/services/embeddings/dlq.py

@@ -0,0 +1,128 @@
+from __future__ import annotations
+
+import json
+from datetime import UTC, datetime, timedelta
+from typing import Any
+
+from sqlalchemy import text
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from app.config import Settings
+
+
+def merge_ingest_meta(existing: str | None, patch: dict[str, Any]) -> str:
+    base: dict[str, Any] = {}
+    if existing:
+        try:
+            parsed = json.loads(existing)
+            if isinstance(parsed, dict):
+                base = parsed
+        except json.JSONDecodeError:
+            base = {"_raw_ingest_meta": existing}
+    base.update(patch)
+    return json.dumps(base, default=str)
+
+
+def cleaned_ingest_meta_dict(ingest_meta: str | None) -> dict[str, Any]:
+    """Drop embedding failure / DLQ keys so a successful embed can overwrite status."""
+    base: dict[str, Any] = {}
+    if ingest_meta:
+        try:
+            parsed = json.loads(ingest_meta)
+            if isinstance(parsed, dict):
+                base = parsed
+        except json.JSONDecodeError:
+            base = {}
+    for k in ("embedding_error", "embedding_dlq_attempts", "embedding_dlq_state"):
+        base.pop(k, None)
+    return base
+
+
+async def clear_embedding_dlq(session: AsyncSession, *, document_id: str) -> None:
+    await session.execute(
+        text("DELETE FROM embedding_dlq WHERE document_id = :d"),
+        {"d": document_id},
+    )
+
+
+async def record_embedding_dlq_failure(
+    session: AsyncSession,
+    *,
+    document_id: str,
+    error_message: str,
+    multimodal: bool,
+    settings: Settings,
+    prior_ingest_meta: str | None,
+) -> None:
+    """
+    Increment attempts, set exponential ``next_retry_at``, or mark ``dead`` when
+    attempts exceed ``embedding_dlq_max_attempts``.
+    """
+    row = await session.execute(
+        text(
+            "SELECT attempt_count FROM embedding_dlq WHERE document_id = :d LIMIT 1",
+        ),
+        {"d": document_id},
+    )
+    found = row.first()
+    attempts = int(found[0]) + 1 if found is not None else 1
+    max_a = settings.embedding_dlq_max_attempts
+    base_delay = settings.embedding_dlq_base_delay_s
+
+    if attempts >= max_a:
+        state = "dead"
+        next_retry_at: str | None = None
+        doc_status = "failed"
+    else:
+        state = "pending_retry"
+        delay_s = base_delay * (2 ** (attempts - 1))
+        cap = settings.embedding_dlq_max_backoff_s
+        delay_s = min(delay_s, cap)
+        next_at = datetime.now(UTC) + timedelta(seconds=delay_s)
+        next_retry_at = next_at.strftime("%Y-%m-%dT%H:%M:%SZ")
+        doc_status = "partial"
+
+    now_iso = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ")
+    await session.execute(
+        text(
+            """
+            INSERT INTO embedding_dlq (
+              document_id, last_error, attempt_count, next_retry_at, state,
+              multimodal, updated_at
+            ) VALUES (
+              :doc, :err, :ac, :nra, :st, :mm, :up
+            )
+            ON CONFLICT(document_id) DO UPDATE SET
+              last_error = excluded.last_error,
+              attempt_count = excluded.attempt_count,
+              next_retry_at = excluded.next_retry_at,
+              state = excluded.state,
+              multimodal = excluded.multimodal,
+              updated_at = excluded.updated_at
+            """,
+        ),
+        {
+            "doc": document_id,
+            "err": error_message[:8000],
+            "ac": attempts,
+            "nra": next_retry_at,
+            "st": state,
+            "mm": 1 if multimodal else 0,
+            "up": now_iso,
+        },
+    )
+
+    im = merge_ingest_meta(
+        prior_ingest_meta,
+        {
+            "embedding_error": error_message[:2000],
+            "embedding_dlq_attempts": attempts,
+            "embedding_dlq_state": state,
+        },
+    )
+    await session.execute(
+        text(
+            "UPDATE documents SET status = :st, ingest_meta = :im WHERE id = :id",
+        ),
+        {"st": doc_status, "im": im, "id": document_id},
+    )

package/backend/app/services/embeddings/gemini_api.py

@@ -0,0 +1,207 @@
+from __future__ import annotations
+
+import asyncio
+import base64
+import logging
+from typing import Any
+
+import httpx
+
+from app.config import Settings
+from app.services.embeddings.types import EmbeddingPart, InlineDataPart, TextPart
+
+logger = logging.getLogger(__name__)
+
+_GEMINI_BASE = "https://generativelanguage.googleapis.com/v1beta"
+
+
+def _part_to_api_dict(part: EmbeddingPart) -> dict[str, Any]:
+    if isinstance(part, TextPart):
+        return {"text": part.text}
+    if isinstance(part, InlineDataPart):
+        b64 = base64.standard_b64encode(part.data).decode("ascii")
+        return {
+            "inlineData": {
+                "mimeType": part.mime_type,
+                "data": b64,
+            },
+        }
+    msg = f"unknown part type: {type(part)}"
+    raise TypeError(msg)
+
+
+def _build_request_item(
+    *,
+    model: str,
+    parts: list[EmbeddingPart],
+    task_type: str,
+    output_dimensionality: int | None,
+) -> dict[str, Any]:
+    item: dict[str, Any] = {
+        "model": f"models/{model}",
+        "content": {"parts": [_part_to_api_dict(p) for p in parts]},
+        "taskType": task_type,
+    }
+    if output_dimensionality is not None:
+        item["outputDimensionality"] = output_dimensionality
+    return item
+
+
+def _normalize_vector(values: list[float], target_dim: int) -> list[float]:
+    if len(values) == target_dim:
+        return values
+    if len(values) > target_dim:
+        return values[:target_dim]
+    msg = f"embedding dim {len(values)} < target {target_dim}"
+    raise ValueError(msg)
+
+
+async def _post_batch_once(
+    client: httpx.AsyncClient,
+    url: str,
+    headers: dict[str, str],
+    payload: dict[str, Any],
+    target_dim: int,
+) -> tuple[list[list[float]] | None, bool]:
+    """
+    Returns (vectors, reject_output_dimensionality).
+    Second flag True if caller should retry without outputDimensionality.
+    """
+    r = await client.post(url, headers=headers, json=payload)
+    if r.status_code == 400 and "outputDimensionality" in r.text:
+        return None, True
+    if r.status_code in (429, 500, 502, 503):
+        msg = f"retryable HTTP {r.status_code}"
+        raise RuntimeError(msg)
+    r.raise_for_status()
+    data = r.json()
+    embs = data.get("embeddings")
+    if not isinstance(embs, list):
+        msg = f"unexpected response: {data!r}"[:500]
+        raise RuntimeError(msg)
+    out: list[list[float]] = []
+    for emb in embs:
+        vals = emb.get("values")
+        if not isinstance(vals, list):
+            msg = f"bad embedding entry: {emb!r}"[:300]
+            raise RuntimeError(msg)
+        floats = [float(x) for x in vals]
+        out.append(_normalize_vector(floats, target_dim))
+    return out, False
+
+
+async def _post_batch_with_retry(
+    client: httpx.AsyncClient,
+    url: str,
+    headers: dict[str, str],
+    payload: dict[str, Any],
+    target_dim: int,
+    settings: Settings,
+) -> list[list[float]] | None:
+    delay = settings.gemini_embed_base_delay_s
+    attempts = max(1, settings.gemini_embed_max_retries)
+    for attempt in range(attempts):
+        try:
+            vecs, reject_od = await _post_batch_once(
+                client,
+                url,
+                headers,
+                payload,
+                target_dim,
+            )
+            if reject_od:
+                return None
+            return vecs
+        except (httpx.HTTPError, RuntimeError, ValueError, KeyError, TypeError) as e:
+            logger.warning(
+                "Gemini embed batch attempt %s/%s: %s",
+                attempt + 1,
+                attempts,
+                e,
+            )
+            if attempt == attempts - 1:
+                logger.exception("Gemini embed batch failed")
+                return None
+            await asyncio.sleep(delay * (2**attempt))
+    return None
+
+
+async def batch_embed_contents(
+    *,
+    api_key: str,
+    model: str,
+    contents: list[list[EmbeddingPart]],
+    settings: Settings,
+    task_type: str = "RETRIEVAL_DOCUMENT",
+    client: httpx.AsyncClient | None = None,
+) -> list[list[float]]:
+    """
+    Batch Gemini ``batchEmbedContents`` calls.
+
+    Each inner list is one embed request: a single ``TextPart`` or an
+    interleaved list of ``TextPart`` / ``InlineDataPart`` (multimodal).
+    """
+    if not contents:
+        return []
+    target_dim = settings.vector_embedding_dim
+    url = f"{_GEMINI_BASE}/models/{model}:batchEmbedContents"
+    headers = {
+        "x-goog-api-key": api_key,
+        "Content-Type": "application/json",
+    }
+
+    use_od: int | None = target_dim
+    close_client = False
+    if client is None:
+        client = httpx.AsyncClient(timeout=120.0)
+        close_client = True
+    try:
+        all_out: list[list[float]] = []
+        batch_size = max(1, settings.gemini_embed_batch_size)
+        for start in range(0, len(contents), batch_size):
+            batch = contents[start : start + batch_size]
+            requests_body = [
+                _build_request_item(
+                    model=model,
+                    parts=parts,
+                    task_type=task_type,
+                    output_dimensionality=use_od,
+                )
+                for parts in batch
+            ]
+            payload: dict[str, Any] = {"requests": requests_body}
+            vecs = await _post_batch_with_retry(
+                client,
+                url,
+                headers,
+                payload,
+                target_dim,
+                settings,
+            )
+            if vecs is None and use_od is not None:
+                use_od = None
+                requests_body = [
+                    _build_request_item(
+                        model=model,
+                        parts=parts,
+                        task_type=task_type,
+                        output_dimensionality=None,
+                    )
+                    for parts in batch
+                ]
+                vecs = await _post_batch_with_retry(
+                    client,
+                    url,
+                    headers,
+                    {"requests": requests_body},
+                    target_dim,
+                    settings,
+                )
+            if vecs is None:
+                msg = "Gemini batchEmbedContents failed after retries"
+                raise RuntimeError(msg)
+            all_out.extend(vecs)
+        return all_out
+    finally:
+        if close_client:
+            await client.aclose()

package/backend/app/services/embeddings/persist.py

@@ -0,0 +1,74 @@
+from __future__ import annotations
+
+from datetime import UTC, datetime
+
+from sqlalchemy import text
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from app.config import Settings
+from app.services.embeddings.build_inputs import ChunkRow
+from app.vectorstore import SqliteVecStore, VectorMeta
+
+
+def vector_store_pointer(document_id: str, chunk_pk: int) -> str:
+    return f"sqlite-vec:kb_vec_embeddings:{document_id}:{chunk_pk}"
+
+
+async def delete_gemini_embeddings_for_document(
+    session: AsyncSession,
+    *,
+    document_id: str,
+    model: str,
+) -> None:
+    await session.execute(
+        text("DELETE FROM embeddings WHERE document_id = :d AND model = :m"),
+        {"d": document_id, "m": model},
+    )
+
+
+async def persist_chunk_embeddings(
+    session: AsyncSession,
+    store: SqliteVecStore,
+    *,
+    document_id: str,
+    source_id: int,
+    model: str,
+    chunks: list[ChunkRow],
+    vectors: list[list[float]],
+    modalities: list[str],
+    settings: Settings,
+) -> None:
+    if len(chunks) != len(vectors) or len(chunks) != len(modalities):
+        msg = "chunks, vectors, and modalities must align"
+        raise ValueError(msg)
+    ts = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ")
+    dim = settings.vector_embedding_dim
+
+    metas = [
+        VectorMeta(
+            document_id=document_id,
+            chunk_id=ch.id,
+            source_id=source_id,
+            modality=mod,
+            ingested_at=ts,
+        )
+        for ch, mod in zip(chunks, modalities, strict=True)
+    ]
+
+    await store.upsert_for_session(session, vectors, metas)
+
+    for ch in chunks:
+        ptr = vector_store_pointer(document_id, ch.id)
+        await session.execute(
+            text(
+                "INSERT INTO embeddings "
+                "(document_id, chunk_id, model, dim, vector_store_ref) "
+                "VALUES (:doc, NULL, :model, :dim, :ref)",
+            ),
+            {
+                "doc": document_id,
+                "model": model,
+                "dim": dim,
+                "ref": ptr,
+            },
+        )

package/backend/app/services/embeddings/types.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True, slots=True)
+class TextPart:
+    text: str
+
+
+@dataclass(frozen=True, slots=True)
+class InlineDataPart:
+    """Binary payload as base64 in API (image/video/audio snippet)."""
+
+    mime_type: str
+    data: bytes
+
+
+EmbeddingPart = TextPart | InlineDataPart
+
+
+def parts_to_modality(parts: list[EmbeddingPart]) -> str:
+    if not parts:
+        return "text"
+    if len(parts) == 1 and isinstance(parts[0], TextPart):
+        return "text"
+    kinds = {type(p).__name__ for p in parts}
+    if "InlineDataPart" in kinds and "TextPart" in kinds:
+        return "multimodal"
+    if "InlineDataPart" in kinds:
+        return "image"
+    return "text"