business-stack 0.1.0

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

@@ -0,0 +1,224 @@
+from __future__ import annotations
+
+import json
+import logging
+
+from sqlalchemy import text
+from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
+
+from app.config import Settings
+from app.services.embeddings.build_inputs import (
+    chunk_to_multimodal_parts,
+    chunk_to_text_parts,
+    load_document_chunks,
+)
+from app.services.embeddings.dlq import (
+    cleaned_ingest_meta_dict,
+    clear_embedding_dlq,
+    merge_ingest_meta,
+    record_embedding_dlq_failure,
+)
+from app.services.embeddings.gemini_api import batch_embed_contents
+from app.services.embeddings.persist import (
+    delete_gemini_embeddings_for_document,
+    persist_chunk_embeddings,
+)
+from app.services.embeddings.types import EmbeddingPart, parts_to_modality
+from app.services.entities.pipeline import extract_and_store_entities_for_document
+from app.services.hooks.post_ingest import dispatch_post_ingest_hooks
+from app.services.integrations_remote import resolve_gemini_api_key
+from app.storage.blobs import BlobStore
+from app.vectorstore import SqliteVecStore
+
+logger = logging.getLogger(__name__)
+
+
+async def embed_document_gemini(
+    session: AsyncSession,
+    *,
+    document_id: str,
+    settings: Settings,
+    store: SqliteVecStore,
+    multimodal: bool = False,
+) -> int:
+    """
+    Idempotent: embed first, then replace vec rows + Gemini ``embeddings`` rows.
+
+    Prior embedding rows for this model are removed only after a successful
+    ``batch_embed_contents`` call so failures do not leave the document without
+    vectors while a retry is pending.
+    """
+    api_key = await resolve_gemini_api_key(settings)
+    if not api_key:
+        msg = "Gemini API key not configured (GEMINI_API_KEY or gateway store)"
+        raise ValueError(msg)
+
+    row = await session.execute(
+        text("SELECT source_id, ingest_meta FROM documents WHERE id = :id LIMIT 1"),
+        {"id": document_id},
+    )
+    doc_row = row.first()
+    if doc_row is None:
+        msg = "document not found"
+        raise ValueError(msg)
+    source_id = int(doc_row[0])
+    prior_meta = doc_row[1]
+
+    chunks = await load_document_chunks(session, document_id)
+    if not chunks:
+        err = "no document_chunks; run POST .../chunks first"
+        raise ValueError(err)
+
+    blob_store = BlobStore(settings.data_dir / "blobs")
+    contents: list[list[EmbeddingPart]] = []
+    modalities: list[str] = []
+    for ch in chunks:
+        if multimodal:
+            parts = await chunk_to_multimodal_parts(
+                session,
+                document_id,
+                ch,
+                blob_store,
+            )
+        else:
+            parts = chunk_to_text_parts(ch)
+        contents.append(parts)
+        modalities.append(parts_to_modality(parts))
+
+    vectors = await batch_embed_contents(
+        api_key=api_key,
+        model=settings.gemini_embedding_model,
+        contents=contents,
+        settings=settings,
+        task_type=settings.gemini_embed_task_type,
+    )
+
+    await delete_gemini_embeddings_for_document(
+        session,
+        document_id=document_id,
+        model=settings.gemini_embedding_model,
+    )
+    # Use session connection: a second pool connection + open txn causes SQLite
+    # "database is locked" on kb_vec_embeddings (single-writer).
+    await store.delete_document_for_session(session, document_id)
+
+    await persist_chunk_embeddings(
+        session,
+        store,
+        document_id=document_id,
+        source_id=source_id,
+        model=settings.gemini_embedding_model,
+        chunks=chunks,
+        vectors=vectors,
+        modalities=modalities,
+        settings=settings,
+    )
+
+    await clear_embedding_dlq(session, document_id=document_id)
+
+    clean = cleaned_ingest_meta_dict(prior_meta)
+    meta_emb = merge_ingest_meta(
+        json.dumps(clean, default=str),
+        {
+            "embedding_model": settings.gemini_embedding_model,
+            "embedding_dim": settings.vector_embedding_dim,
+        },
+    )
+
+    await session.execute(
+        text(
+            "UPDATE documents SET status = 'ok', ingest_meta = :im WHERE id = :id",
+        ),
+        {"id": document_id, "im": meta_emb},
+    )
+
+    try:
+        n_mentions, n_co = await extract_and_store_entities_for_document(
+            session,
+            document_id=document_id,
+            settings=settings,
+        )
+        await session.execute(
+            text("UPDATE documents SET ingest_meta = :im WHERE id = :id"),
+            {
+                "id": document_id,
+                "im": merge_ingest_meta(
+                    meta_emb,
+                    {
+                        "entity_mentions": n_mentions,
+                        "entity_cooccurrence_pairs": n_co,
+                    },
+                ),
+            },
+        )
+    except Exception as ex:
+        logger.exception("Entity extraction failed for %s", document_id)
+        err_meta = merge_ingest_meta(
+            meta_emb,
+            {"entity_extraction_error": str(ex)},
+        )
+        await session.execute(
+            text(
+                "UPDATE documents SET ingest_meta = :im WHERE id = :id",
+            ),
+            {"id": document_id, "im": err_meta},
+        )
+        if settings.entity_extract_strict:
+            await session.execute(
+                text("UPDATE documents SET status = 'failed' WHERE id = :id"),
+                {"id": document_id},
+            )
+            raise
+    return len(chunks)
+
+
+async def run_embed_document_job(
+    *,
+    document_id: str,
+    multimodal: bool,
+    session_factory: async_sessionmaker[AsyncSession],
+    settings: Settings,
+    store: SqliteVecStore,
+) -> None:
+    async with session_factory() as session:
+        try:
+            await embed_document_gemini(
+                session,
+                document_id=document_id,
+                settings=settings,
+                store=store,
+                multimodal=multimodal,
+            )
+            await session.commit()
+            try:
+                await dispatch_post_ingest_hooks(
+                    session_factory,
+                    document_id=document_id,
+                    settings=settings,
+                )
+            except Exception:
+                logger.exception(
+                    "post-ingest hooks failed (document already ok): %s",
+                    document_id,
+                )
+        except Exception as e:
+            logger.exception("Gemini embed failed for %s", document_id)
+            await session.rollback()
+            async with session_factory() as s2:
+                async with s2.begin():
+                    row = await s2.execute(
+                        text(
+                            "SELECT ingest_meta FROM documents WHERE id = :id LIMIT 1",
+                        ),
+                        {"id": document_id},
+                    )
+                    r = row.first()
+                    im = r[0] if r else None
+                    await record_embedding_dlq_failure(
+                        s2,
+                        document_id=document_id,
+                        error_message=str(e),
+                        multimodal=multimodal,
+                        settings=settings,
+                        prior_ingest_meta=im,
+                    )

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

@@ -0,0 +1,12 @@
+from app.services.entities.pipeline import (
+    extract_and_store_entities_for_document,
+    run_entity_extraction_job,
+)
+from app.services.entities.types import EntityType, ExtractedMention
+
+__all__ = [
+    "EntityType",
+    "ExtractedMention",
+    "extract_and_store_entities_for_document",
+    "run_entity_extraction_job",
+]

package/backend/app/services/entities/gliner_extract.py

@@ -0,0 +1,63 @@
+from __future__ import annotations
+
+import logging
+import re
+
+from app.services.entities.types import EntityType, ExtractedMention
+
+logger = logging.getLogger(__name__)
+
+_GLINER_LABELS = [
+    "person",
+    "company",
+    "organization",
+    "product",
+    "location",
+    "concept",
+]
+
+_LABEL_MAP: dict[str, EntityType] = {
+    "person": "person",
+    "company": "company",
+    "organization": "company",
+    "product": "product",
+    "location": "location",
+    "concept": "concept",
+}
+
+
+def _normalize_name(s: str) -> str:
+    return re.sub(r"\s+", " ", s.strip())
+
+
+def extract_with_gliner(text: str, *, model_id: str) -> list[ExtractedMention]:
+    """Optional GLiNER zero-shot NER (requires ``gliner`` + a checkpoint)."""
+    if not text.strip():
+        return []
+    try:
+        from gliner import GLiNER  # type: ignore[import-not-found]
+    except ImportError:
+        logger.debug("gliner not installed; skipping GLiNER entity extraction")
+        return []
+
+    try:
+        model = GLiNER.from_pretrained(model_id)
+    except Exception:
+        logger.warning("GLiNER model %r could not load", model_id, exc_info=True)
+        return []
+
+    ents = model.predict_entities(text[:8000], _GLINER_LABELS, threshold=0.35)
+    out: list[ExtractedMention] = []
+    for e in ents:
+        label = str(e.get("label", "")).lower()
+        et = _LABEL_MAP.get(label)
+        if et is None:
+            continue
+        name = _normalize_name(str(e.get("text", "")))
+        if len(name) < 2:
+            continue
+        score = float(e.get("score", 0.6))
+        out.append(
+            ExtractedMention(name=name, type=et, confidence=min(1.0, max(0.0, score)))
+        )
+    return out

package/backend/app/services/entities/llm_extract.py

@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+import json
+import logging
+import re
+
+import httpx
+
+from app.config import Settings
+from app.services.entities.types import ExtractedMention
+
+logger = logging.getLogger(__name__)
+
+_ALLOWED = frozenset({"person", "company", "product", "concept", "location"})
+
+_SYSTEM = """Extract named entities from the user text. Return JSON only:
+{"entities":[{"name":"string","type":"person|company|product|concept|location","confidence":0.0-1.0}]}
+Skip duplicates. Prefer high-precision extractions."""
+
+
+def _normalize_name(s: str) -> str:
+    return re.sub(r"\s+", " ", s.strip())
+
+
+async def extract_with_ollama(
+    text: str,
+    *,
+    settings: Settings,
+    client: httpx.AsyncClient | None = None,
+) -> list[ExtractedMention]:
+    """Optional local LLM (Ollama) with JSON-shaped output for harder cases."""
+    if not text.strip() or not settings.entity_llm_enabled:
+        return []
+    base = str(settings.ollama_base_url).rstrip("/")
+    url = f"{base}/api/chat"
+    body = {
+        "model": settings.ollama_entity_model,
+        "stream": False,
+        "format": "json",
+        "messages": [
+            {"role": "system", "content": _SYSTEM},
+            {
+                "role": "user",
+                "content": text[:16_000],
+            },
+        ],
+    }
+    close = False
+    if client is None:
+        client = httpx.AsyncClient(timeout=120.0)
+        close = True
+    try:
+        r = await client.post(url, json=body)
+        r.raise_for_status()
+        data = r.json()
+        raw = data.get("message", {}).get("content", "")
+        parsed = json.loads(raw)
+        items = parsed.get("entities")
+        if not isinstance(items, list):
+            return []
+        out: list[ExtractedMention] = []
+        for it in items:
+            if not isinstance(it, dict):
+                continue
+            name = it.get("name")
+            typ = it.get("type")
+            conf = it.get("confidence", 0.75)
+            if not isinstance(name, str) or not isinstance(typ, str):
+                continue
+            typ_l = typ.lower().strip()
+            if typ_l not in _ALLOWED:
+                continue
+            try:
+                c = float(conf)
+            except (TypeError, ValueError):
+                c = 0.75
+            c = max(0.0, min(1.0, c))
+            nn = _normalize_name(name)
+            if len(nn) < 2:
+                continue
+            out.append(
+                ExtractedMention(
+                    name=nn,
+                    type=typ_l,  # type: ignore[arg-type]
+                    confidence=c,
+                ),
+            )
+        return out
+    except Exception:
+        logger.exception("Ollama entity extraction failed")
+        return []
+    finally:
+        if close:
+            await client.aclose()

package/backend/app/services/entities/pipeline.py

@@ -0,0 +1,179 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+from collections.abc import Iterable
+
+from sqlalchemy import text
+from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
+
+from app.config import Settings
+from app.services.embeddings.build_inputs import load_document_chunks
+from app.services.entities.gliner_extract import extract_with_gliner
+from app.services.entities.llm_extract import extract_with_ollama
+from app.services.entities.spacy_extract import extract_with_spacy
+from app.services.entities.types import ExtractedMention
+
+logger = logging.getLogger(__name__)
+
+
+def _merge_mentions(mentions: Iterable[ExtractedMention]) -> list[ExtractedMention]:
+    """Dedupe by (normalized name, type); keep max confidence."""
+    best: dict[tuple[str, str], ExtractedMention] = {}
+    for m in mentions:
+        key = (m.name.strip().lower(), m.type)
+        prev = best.get(key)
+        if prev is None or m.confidence > prev.confidence:
+            best[key] = m
+    return list(best.values())
+
+
+def _baseline_sync(text: str, settings: Settings) -> list[ExtractedMention]:
+    out: list[ExtractedMention] = []
+    if settings.entity_use_spacy:
+        out.extend(extract_with_spacy(text, model_name=settings.spacy_model))
+    if settings.entity_use_gliner:
+        out.extend(
+            extract_with_gliner(text, model_id=settings.gliner_model_id),
+        )
+    return _merge_mentions(out)
+
+
+async def extract_mentions_for_chunk_text(
+    text: str,
+    settings: Settings,
+) -> list[ExtractedMention]:
+    baseline = await asyncio.to_thread(_baseline_sync, text, settings)
+    merged = baseline
+    sparse = len(merged) < settings.entity_llm_min_mentions
+    if settings.entity_llm_enabled and (sparse or settings.entity_llm_on_all_chunks):
+        llm_ents = await extract_with_ollama(text, settings=settings)
+        merged = _merge_mentions([*merged, *llm_ents])
+    return merged
+
+
+async def get_or_create_entity_id(
+    session: AsyncSession,
+    *,
+    name: str,
+    entity_type: str,
+) -> int:
+    row = await session.execute(
+        text(
+            "SELECT id FROM entities WHERE lower(trim(name)) = lower(trim(:n)) "
+            "AND type = :t LIMIT 1",
+        ),
+        {"n": name, "t": entity_type},
+    )
+    found = row.first()
+    if found is not None:
+        return int(found[0])
+    ins = await session.execute(
+        text(
+            "INSERT INTO entities (name, type, meta) VALUES (:n, :t, NULL) "
+            "RETURNING id",
+        ),
+        {"n": name.strip(), "t": entity_type},
+    )
+    return int(ins.scalar_one())
+
+
+async def extract_and_store_entities_for_document(
+    session: AsyncSession,
+    *,
+    document_id: str,
+    settings: Settings,
+) -> tuple[int, int]:
+    """
+    Replace prior mentions/co-occurrence for this document, then insert fresh rows.
+    Caller should commit with the same session as embedding writes.
+    Returns (mention_rows, cooccurrence_rows).
+    """
+    chunks = await load_document_chunks(session, document_id)
+    if not chunks:
+        return 0, 0
+
+    await session.execute(
+        text("DELETE FROM entity_mentions WHERE document_id = :d"),
+        {"d": document_id},
+    )
+    await session.execute(
+        text("DELETE FROM entity_cooccurrence WHERE document_id = :d"),
+        {"d": document_id},
+    )
+
+    doc_entity_ids: set[int] = set()
+    mention_count = 0
+
+    for ch in chunks:
+        mentions = await extract_mentions_for_chunk_text(ch.text, settings)
+        chunk_entity_ids: set[int] = set()
+        for m in mentions:
+            eid = await get_or_create_entity_id(
+                session,
+                name=m.name,
+                entity_type=m.type,
+            )
+            await session.execute(
+                text(
+                    "INSERT INTO entity_mentions "
+                    "(document_id, entity_id, document_chunk_id, confidence) "
+                    "VALUES (:doc, :eid, :cid, :conf)",
+                ),
+                {
+                    "doc": document_id,
+                    "eid": eid,
+                    "cid": ch.id,
+                    "conf": m.confidence,
+                },
+            )
+            mention_count += 1
+            doc_entity_ids.add(eid)
+            chunk_entity_ids.add(eid)
+
+        logger.debug(
+            "chunk %s document %s: %s mentions, %s unique entities",
+            ch.id,
+            document_id,
+            len(mentions),
+            len(chunk_entity_ids),
+        )
+
+    sorted_ids = sorted(doc_entity_ids)
+    co_pairs = 0
+    for i in range(len(sorted_ids)):
+        for j in range(i + 1, len(sorted_ids)):
+            lo, hi = sorted_ids[i], sorted_ids[j]
+            if lo > hi:
+                lo, hi = hi, lo
+            await session.execute(
+                text(
+                    "INSERT INTO entity_cooccurrence "
+                    "(entity_low_id, entity_high_id, document_id, weight) "
+                    "VALUES (:lo, :hi, :doc, 1.0)",
+                ),
+                {"lo": lo, "hi": hi, "doc": document_id},
+            )
+            co_pairs += 1
+
+    return mention_count, co_pairs
+
+
+async def run_entity_extraction_job(
+    *,
+    document_id: str,
+    session_factory: async_sessionmaker[AsyncSession],
+    settings: Settings,
+) -> None:
+    """Background entrypoint (same pattern as embedding worker)."""
+    async with session_factory() as session:
+        try:
+            await extract_and_store_entities_for_document(
+                session,
+                document_id=document_id,
+                settings=settings,
+            )
+            await session.commit()
+        except Exception:
+            logger.exception("entity extraction failed for %s", document_id)
+            await session.rollback()

package/backend/app/services/entities/spacy_extract.py

@@ -0,0 +1,63 @@
+from __future__ import annotations
+
+import logging
+import re
+from functools import lru_cache
+
+from app.services.entities.types import EntityType, ExtractedMention
+
+logger = logging.getLogger(__name__)
+
+_SPACY_LABEL_TO_TYPE: dict[str, EntityType] = {
+    "PERSON": "person",
+    "ORG": "company",
+    "GPE": "location",
+    "LOC": "location",
+    "FAC": "location",
+    "PRODUCT": "product",
+    "EVENT": "concept",
+    "WORK_OF_ART": "concept",
+    "LAW": "concept",
+    "NORP": "concept",
+    "LANGUAGE": "concept",
+}
+
+
+def _normalize_name(s: str) -> str:
+    t = re.sub(r"\s+", " ", s.strip())
+    return t
+
+
+@lru_cache(maxsize=1)
+def _load_nlp(model_name: str):
+    try:
+        import spacy
+
+        return spacy.load(model_name)
+    except Exception:
+        logger.warning(
+            "spaCy model %r not available; install optional [entities] deps",
+            model_name,
+            exc_info=True,
+        )
+        return None
+
+
+def extract_with_spacy(text: str, *, model_name: str) -> list[ExtractedMention]:
+    """Baseline NER: spaCy labels mapped to person/company/product/concept/location."""
+    if not text or not text.strip():
+        return []
+    nlp = _load_nlp(model_name)
+    if nlp is None:
+        return []
+    doc = nlp(text[:1_000_000])
+    out: list[ExtractedMention] = []
+    for ent in doc.ents:
+        et = _SPACY_LABEL_TO_TYPE.get(ent.label_)
+        if et is None:
+            continue
+        name = _normalize_name(ent.text)
+        if len(name) < 2:
+            continue
+        out.append(ExtractedMention(name=name, type=et, confidence=0.82))
+    return out

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

@@ -0,0 +1,15 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Literal
+
+EntityType = Literal["person", "company", "product", "concept", "location"]
+
+
+@dataclass(frozen=True, slots=True)
+class ExtractedMention:
+    """One surface-form mention with coarse type and confidence."""
+
+    name: str
+    type: EntityType
+    confidence: float