diff --git a/docker-compose.yml b/docker-compose.yml index 6f704d3c..28c898ec 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -82,6 +82,12 @@ services: MINIO_SECURE: ${MINIO_SECURE:-false} MINIO_BUCKET: ${MINIO_BUCKET:-openrunner} REDIS_URL: ${REDIS_URL:-redis://redis:6379/0} + # Speech-catalog semantic search: TEI bge-m3 embed server on the host + # (docker-published, firewalled to internal only). host.docker.internal -> + # host gateway -> tei :8081. Used to embed the NL query for /catalog/semantic. + # Literal (not ${...}) so it wins over a stale value in env_file/.env. + CATALOG_EMBED_URL: "http://host.docker.internal:8081/v1/embeddings" + CATALOG_EMBED_MODEL: ${CATALOG_EMBED_MODEL:-BAAI/bge-m3} # SECRET_KEY / JWT_*_SECRET: validators in app/core/config.py refuse # weak/default values (>=32 chars, no `change-me*` placeholders) — the # entrypoint auto-generates and persists strong values to @@ -125,6 +131,11 @@ services: - ./secrets:/secrets:ro ports: - "127.0.0.1:8000:8000" + extra_hosts: + # host.docker.internal -> host gateway, so the API can reach the TEI embed + # server published on the host (:8081) for /catalog/semantic. Not automatic + # on Linux compose. + - "host.docker.internal:host-gateway" depends_on: postgres: condition: service_healthy diff --git a/docs/datasets/derived-datasets.md b/docs/datasets/derived-datasets.md new file mode 100644 index 00000000..33a5bc55 --- /dev/null +++ b/docs/datasets/derived-datasets.md @@ -0,0 +1,185 @@ +# Derived datasets (query- & search-derived) + +A **derived dataset** is a saved selection over a parent dataset's LanceDB table. +You run a **word**, **full-text**, or **semantic** query and OpenRunner records the +rows that matched — as a new dataset, or appended into an existing one. + +It is built for scale and correctness: + +- **No replication.** Only the matching row keys (`debug_id`) are stored, never a + copy of the rows. A derived dataset over 5.8M clips is a few thousand IDs. +- **Inherits every parent field dynamically.** At read time the parent rows are + fetched live, so any column the parent gains later shows up automatically. +- **Extendable.** Add your own columns on top (labels, notes, review status) — + stored only in the derived dataset, the parent is never touched. +- **Auto de-duplication.** The selection is keyed on `debug_id`; re-running a + query or appending a second query never double-counts a row. +- **Auto provenance.** Every derived row carries `derived_from_request` (the query + that pulled it in) and `derived_from_request_date` (when). + +A canonical example: derive a **free-user medical** dataset from the speech +catalog by semantically searching *"a medical consultation between a doctor and a +patient about symptoms, diagnosis and treatment"*. + +--- + +## How it works + +``` +parent dataset (LanceDB table, e.g. transcriptions — 5.8M rows) + │ word / fts / semantic query + ▼ + matching debug_ids ──► members table (debug_id + provenance + your extra cols) + │ │ merge_insert(debug_id) → dedup + ▼ ▼ + read a page ── join by debug_id ──► live parent fields + your extra columns +``` + +The members table lives in the **same** LanceDB as the parent (no new storage +bucket). Only keys + your extra columns are written. Reads page the member keys +and point-look-up exactly that page from the parent (BTREE on `debug_id`), so a +derived dataset over millions of rows still lists instantly. + +--- + +## Create a derived dataset + +=== "UI" + + Open the parent dataset → **Indexes** tab → **🧬 Derive a dataset from this + query**. Pick a mode, type the query, name the dataset, **Create**. + +=== "API" + + ```bash + curl -X POST /api/v1/datasets/{parent_id}/derive \ + -H 'Authorization: Bearer $TOKEN' -H 'Content-Type: application/json' \ + -d '{ + "name": "free-user-medical", + "mode": "semantic", + "q": "a medical consultation between a doctor and a patient about symptoms, diagnosis and treatment", + "limit": 2000 + }' + ``` + + Returns the new `DatasetRead`. `num_examples` is the de-duplicated member count. + +### Query modes + +| Mode | Backend | Use it for | +|------|---------|-----------| +| `word` | inverted `lexical_map` (O(1)-ish) | one exact term, cheapest | +| `fts` | BM25 full-text | keyword queries, ranked | +| `exact` | phrase full-text | adjacent-word phrase | +| `semantic` | cosine ANN over embeddings | meaning / cross-lingual (needs a vector index) | + +`lang` (optional) restricts to a language; `limit` caps the matches (default 2000). + +--- + +## Append another query (same or new dataset) + +Pass `target_dataset_id` to add a second query's matches **into an existing +derived dataset**. De-dup is automatic — rows already present are left as-is +(keeping their original `derived_from_request`). + +```bash +curl -X POST /api/v1/datasets/{parent_id}/derive \ + -H 'Authorization: Bearer $TOKEN' -H 'Content-Type: application/json' \ + -d '{ "target_dataset_id": "", "mode": "word", "q": "prescription", "limit": 5000 }' +``` + +The derived dataset's metadata keeps a `queries[]` log (mode, query, matched, +added, timestamp) for full provenance. + +--- + +## Read derived rows + +```bash +curl '/api/v1/datasets/{derived_id}/derived-rows?offset=0&limit=100' +``` + +```jsonc +{ + "rows": [ + { + "debug_id": "G-f0ab…", + "language": "en", + "full_transcript": "…doctor: what symptoms…", // inherited from parent, live + "duration_seconds": 12.4, // inherited + "derived_from_request": "semantic:a medical consultation…", // auto + "derived_from_request_date": "2026-07-22T06:37:19", // auto + "my_label": "confirmed" // your extra column + } + ], + "total": 1873, + "extra_columns": ["my_label"], + "queries": [ { "mode": "semantic", "q": "…", "matched": 1873, "added": 1873, "at": "…" } ] +} +``` + +Every parent column is present without being copied; add a column to the parent +later and it simply appears here on the next read. + +--- + +## Add & edit your own columns + +Extra columns live only in the derived dataset. + +```bash +# add a column (all rows default to the given value / NULL) +curl -X POST /api/v1/datasets/{derived_id}/columns \ + -H 'Content-Type: application/json' -d '{ "name": "my_label", "default": "unlabeled" }' + +# set one row's value (creates the column if new) +curl -X PATCH /api/v1/datasets/{derived_id}/cell \ + -H 'Content-Type: application/json' \ + -d '{ "debug_id": "G-f0ab…", "column": "my_label", "value": "confirmed" }' +``` + +--- + +## API reference + +| Method & path | Purpose | +|---------------|---------| +| `POST /datasets/{parent_id}/derive` | Create (with `name`) or append (with `target_dataset_id`) from a query | +| `GET /datasets/{derived_id}/derived-rows` | Page of inherited parent fields + auto + extra columns | +| `POST /datasets/{derived_id}/columns` | Add an extra column | +| `PATCH /datasets/{derived_id}/cell` | Set one row's extra-column value | + +Lineage: a derived dataset records a `derived-from` edge to its parent, so it +appears in the dataset lineage graph. + +## Enabling semantic search (optimize embedding search) + +Semantic mode (and `GET /datasets/{id}/semantic`) needs an **embedding column + +ANN index** on the table. Any indexed dataset can be optimized for it: + +```bash +# embed rows missing a vector + build/rebuild the IVF index (background job) +curl -X POST /api/v1/datasets/{id}/optimize-embedding-search \ + -H 'Content-Type: application/json' -d '{ "text_column": "full_transcript", "max_rows": 200000 }' +# poll dataset_metadata.lance.semantic.status -> running | ready | error + +# then search by meaning +curl '/api/v1/datasets/{id}/semantic?q=a doctor discussing symptoms&limit=20' +``` + +Or in the UI: dataset → **Indexes** tab → **⚡ Optimize embedding search**. + +It embeds only rows whose text column is non-empty, and — critically — uses a +**stage-then-flush** write (embed to a staging table via cheap appends, then a +single `merge_insert`). Per-chunk `merge_insert` on a wide table rewrites the +whole embedding column each call (minutes per chunk), so it is never used. + +## Notes & limits + +- The parent must be **indexed** (have a LanceDB table). Semantic mode also needs + a vector index on the embedding column (run *optimize embedding search*) and a + reachable embed server. +- Provenance columns `derived_from_request` / `derived_from_request_date` are + reserved and written automatically; they are not counted as your extra columns. +- Extra columns are stored as strings. diff --git a/mkdocs.yml b/mkdocs.yml index b45643e3..dbacadd7 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -65,7 +65,9 @@ nav: - Quick Start: getting-started/quickstart.md - Troubleshooting (First Run): troubleshooting-first-run.md - Skills: skills/index.md - - Datasets: datasets/index.md + - Datasets: + - Overview: datasets/index.md + - Derived datasets: datasets/derived-datasets.md - SDK Reference: - init(): sdk/init.md - Logging: sdk/logging.md diff --git a/src/api/app/api/v1/datasets.py b/src/api/app/api/v1/datasets.py index f543da91..ff7e00b7 100644 --- a/src/api/app/api/v1/datasets.py +++ b/src/api/app/api/v1/datasets.py @@ -282,6 +282,40 @@ def _exactly_one_parent(self) -> "DerivationRequest": return self +class DeriveRequest(BaseModel): + """Create (or, with ``target_dataset_id``, append to) a query-derived dataset.""" + name: str | None = None # required when creating a new derived dataset + target_dataset_id: uuid.UUID | None = None # append into this existing derived dataset instead + mode: str = Field("semantic", pattern="^(word|fts|exact|semantic)$") + q: str = Field(..., min_length=1, max_length=500) + lang: str | None = None + limit: int = Field(2000, ge=1, le=100000) + project_id: uuid.UUID | None = None # optional: attach the new dataset to a project + + @field_validator("name") + @classmethod + def _trim(cls, v): + return v.strip() if v else v + + +class OptimizeEmbedRequest(BaseModel): + """Kick off embedding-search optimization (embed missing rows + build ANN index).""" + text_column: str | None = None # defaults to the first FTS column / full_transcript + max_rows: int = Field(200000, ge=1, le=5000000) + model: str | None = None # embed model override (default: server config) + + +class DeriveColumnRequest(BaseModel): + name: str = Field(..., min_length=1, max_length=64) + default: str | None = None + + +class DeriveCellRequest(BaseModel): + debug_id: str = Field(..., min_length=1, max_length=128) + column: str = Field(..., min_length=1, max_length=64) + value: str | None = None + + # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- @@ -2966,3 +3000,411 @@ async def remove_hf_derivation( ) ) await db.commit() + + +# --------------------------------------------------------------------------- +# Query-/search-derived datasets (no replication, dedup, inherit + extend) +# --------------------------------------------------------------------------- + + +def _embed_query_sync(q: str) -> list[float]: + """Embed a query string via the configured catalog embed server (bge-m3).""" + import json as _json + import urllib.request as _u + + from app.core.config import get_settings + + s = get_settings() + body = _json.dumps({"model": s.catalog_embed_model, "input": [q]}).encode() + req = _u.Request(s.catalog_embed_url, data=body, headers={"Content-Type": "application/json"}) + with _u.urlopen(req, timeout=30) as r: + return _json.loads(r.read())["data"][0]["embedding"] + + +def _parent_lance_or_400(parent: Dataset) -> dict: + lance = (parent.dataset_metadata or {}).get("lance") + if not lance or parent.connection_id is None or not parent.storage_prefix: + raise HTTPException( + status_code=400, + detail="Parent dataset is not indexed (no LanceDB table to query)", + ) + return lance + + +async def _derived_meta_or_404(ds: Dataset) -> dict: + meta = (ds.dataset_metadata or {}).get("derived") + if not meta: + raise HTTPException(status_code=400, detail="Not a query-derived dataset") + return meta + + +@router.post("/datasets/{dataset_id}/derive", response_model=DatasetRead) +async def derive_dataset( + dataset_id: uuid.UUID, + data: DeriveRequest, + current_user: User = Depends(get_current_user), + db: AsyncSession = Depends(get_db), +): + """Create — or append to — a dataset **derived from a query** over this dataset. + + Runs a ``word`` / ``fts`` / ``exact`` / ``semantic`` search on the parent's + LanceDB table and records only the matching primary keys (``debug_id``): no + row data is replicated. De-duplication is automatic, so re-running a query or + appending a second query never double-counts a row. The derived dataset + inherits every parent field dynamically and can carry extra columns of its own. + + - New dataset: pass ``name``. + - Append into an existing derived dataset: pass ``target_dataset_id`` (must + derive from this same parent). + """ + from starlette.concurrency import run_in_threadpool + + from app.services import dataset_derive as dd + + parent = await _get_dataset(db, dataset_id) + await _verify_org_membership(db, parent.org_id, current_user.id) + lance = _parent_lance_or_400(parent) + + conn = ( + await db.execute(select(StorageConnection).where(StorageConnection.id == parent.connection_id)) + ).scalar_one_or_none() + if conn is None: + raise HTTPException(status_code=409, detail="Parent storage connection no longer exists") + + uri = dd.resolve_uri(conn, parent.storage_prefix, lance.get("uri")) + parent_table = lance["table"] + map_table = lance.get("map_table", "lexical_map") + pk = lance.get("primary_key", "debug_id") + + # embed the query up-front for semantic mode (blocking HTTP -> threadpool) + vec = None + if data.mode == "semantic": + try: + vec = await run_in_threadpool(_embed_query_sync, data.q) + except Exception: + raise HTTPException(status_code=502, detail="embedding service unavailable for semantic derive") + + await _release_row_locks(db) + ids = await run_in_threadpool( + dd.derive_member_ids_sync, conn, uri, parent_table, + data.mode, data.q, data.lang, data.limit, vec, map_table, pk, + ) + + # --- resolve the target derived dataset (append vs create) --- + if data.target_dataset_id is not None: + target = await _get_dataset(db, data.target_dataset_id) + await _verify_org_membership(db, target.org_id, current_user.id) + dmeta = await _derived_meta_or_404(target) + if str(dmeta.get("parent_dataset_id")) != str(dataset_id): + raise HTTPException(status_code=400, detail="Target does not derive from this parent") + members_table = dmeta["members_table"] + is_new = False + else: + if not data.name: + raise HTTPException(status_code=400, detail="name is required to create a derived dataset") + dup = ( + await db.execute(select(Dataset).where(Dataset.org_id == parent.org_id, Dataset.name == data.name)) + ).scalar_one_or_none() + if dup is not None: + raise HTTPException(status_code=409, detail=f"Dataset '{data.name}' already exists") + members_table = dd.new_members_table_name() + target = Dataset( + id=uuid.uuid4(), + org_id=parent.org_id, + name=data.name, + description=f"Derived from {parent.name} via {data.mode} query: {data.q!r}", + source_uri=parent.source_uri, + connection_id=parent.connection_id, + storage_prefix=parent.storage_prefix, + modality=parent.modality, + source_type="derived", + dataset_type=parent.dataset_type, + languages=parent.languages, + license=parent.license, + license_use=parent.license_use, + dataset_metadata={ + "derived": { + "parent_dataset_id": str(dataset_id), + "parent_table": parent_table, + "uri": uri, + "members_table": members_table, + "primary_key": pk, + "map_table": map_table, + "queries": [], + }, + # so /search, /row, /facets keep working against the parent table + "lance": lance, + }, + ) + db.add(target) + dmeta = target.dataset_metadata["derived"] + is_new = True + + # --- write the members (dedup) --- + result = await run_in_threadpool( + dd.upsert_members_sync, conn, uri, members_table, ids, f"{data.mode}:{data.q}", pk, + ) + + # --- provenance + counts --- + meta = dict(target.dataset_metadata or {}) + d = dict(meta.get("derived") or dmeta) + d["queries"] = list(d.get("queries") or []) + [{ + "mode": data.mode, "q": data.q, "lang": data.lang or None, + "limit": data.limit, "matched": len(ids), "added": result["added"], + "at": datetime.utcnow().isoformat(), + }] + meta["derived"] = d + target.dataset_metadata = meta + target.num_examples = result["total"] + target.updated_at = datetime.utcnow() + + if is_new: + # lineage edge parent -> derived (reuse HF-style derivation table) + db.add(DatasetDerivation(dataset_id=target.id, parent_dataset_id=dataset_id)) + if data.project_id is not None: + db.add(DatasetProjectLink(dataset_id=target.id, project_id=data.project_id)) + + await db.commit() + await db.refresh(target) + return await _to_read(db, target) + + +@router.get("/datasets/{dataset_id}/derived-rows") +async def derived_dataset_rows( + dataset_id: uuid.UUID, + offset: int = 0, + limit: int = 100, + current_user: User | None = Depends(get_current_user_optional), + db: AsyncSession = Depends(get_db), +): + """A page of a derived dataset's rows: **live parent fields (inherited) + the + derived dataset's own extra columns**. Bounded — pages the member keys and + point-looks-up only that page from the parent.""" + from starlette.concurrency import run_in_threadpool + + from app.services import dataset_derive as dd + + ds = await _get_dataset(db, dataset_id) + await enforce_read_visibility( + db, visibility=getattr(ds, "visibility", "private") or "private", + org_id=ds.org_id, resource_type="dataset", resource_id=ds.id, user=current_user, + ) + meta = await _derived_meta_or_404(ds) + conn = ( + await db.execute(select(StorageConnection).where(StorageConnection.id == ds.connection_id)) + ).scalar_one_or_none() + if conn is None: + raise HTTPException(status_code=409, detail="Bound storage connection no longer exists") + offset = max(0, offset) + limit = max(1, min(limit, 500)) + await _release_row_locks(db) + rows, total = await run_in_threadpool( + dd.derived_rows_sync, conn, meta["uri"], meta["parent_table"], + meta["members_table"], offset, limit, meta.get("primary_key", "debug_id"), + ) + extra = await run_in_threadpool( + dd.extra_columns_sync, conn, meta["uri"], meta["members_table"], + meta.get("primary_key", "debug_id"), + ) + return {"rows": rows, "total": total, "offset": offset, "limit": limit, + "extra_columns": extra, "queries": meta.get("queries", []), "source": "derived"} + + +@router.post("/datasets/{dataset_id}/columns") +async def add_derived_column( + dataset_id: uuid.UUID, + data: DeriveColumnRequest, + current_user: User = Depends(get_current_user), + db: AsyncSession = Depends(get_db), +): + """Add a new (string) column to a derived dataset — stored only in the derived + members table, never touching the parent.""" + from starlette.concurrency import run_in_threadpool + + from app.services import dataset_derive as dd + + ds = await _get_dataset(db, dataset_id) + await _verify_org_membership(db, ds.org_id, current_user.id) + meta = await _derived_meta_or_404(ds) + conn = ( + await db.execute(select(StorageConnection).where(StorageConnection.id == ds.connection_id)) + ).scalar_one_or_none() + if conn is None: + raise HTTPException(status_code=409, detail="Bound storage connection no longer exists") + await _release_row_locks(db) + try: + cols = await run_in_threadpool( + dd.add_member_column_sync, conn, meta["uri"], meta["members_table"], + data.name, data.default, + ) + except ValueError as e: + raise HTTPException(status_code=400, detail=str(e)) + return {"columns": cols} + + +@router.patch("/datasets/{dataset_id}/cell") +async def set_derived_cell( + dataset_id: uuid.UUID, + data: DeriveCellRequest, + current_user: User = Depends(get_current_user), + db: AsyncSession = Depends(get_db), +): + """Set one derived row's extra-column value (creates the column if new).""" + from starlette.concurrency import run_in_threadpool + + from app.services import dataset_derive as dd + + ds = await _get_dataset(db, dataset_id) + await _verify_org_membership(db, ds.org_id, current_user.id) + meta = await _derived_meta_or_404(ds) + conn = ( + await db.execute(select(StorageConnection).where(StorageConnection.id == ds.connection_id)) + ).scalar_one_or_none() + if conn is None: + raise HTTPException(status_code=409, detail="Bound storage connection no longer exists") + await _release_row_locks(db) + try: + await run_in_threadpool( + dd.set_member_cell_sync, conn, meta["uri"], meta["members_table"], + data.debug_id, data.column, data.value, meta.get("primary_key", "debug_id"), + ) + except ValueError as e: + raise HTTPException(status_code=400, detail=str(e)) + return {"ok": True, "debug_id": data.debug_id, "column": data.column} + + +# --------------------------------------------------------------------------- +# Optimize embedding search (generic: embed missing rows + build ANN index) +# --------------------------------------------------------------------------- + + +def _default_text_column(lance: dict) -> str: + fts = lance.get("fts_columns") or [] + return fts[0] if fts else "full_transcript" + + +async def _run_optimize_embedding(dataset_id: uuid.UUID, text_column: str, + max_rows: int, model: str) -> None: + """Background worker: embed missing vectors + build the ANN index, then write + the result into dataset_metadata.lance.semantic. Runs in its own DB session + + engine so it outlives the request.""" + from starlette.concurrency import run_in_threadpool + + from app.core.config import get_settings + from app.core.database import async_session_factory, create_engine + from app.services.dataset_embed_optimize import optimize_embedding_search_sync + + settings = get_settings() + eng = create_engine(settings) + try: + async with async_session_factory(eng)() as db: + ds = await db.get(Dataset, dataset_id) + if ds is None: + return + lance = (ds.dataset_metadata or {}).get("lance") or {} + conn = await db.get(StorageConnection, ds.connection_id) + status: dict + try: + summary = await run_in_threadpool( + optimize_embedding_search_sync, conn, ds.storage_prefix, + lance["table"], text_column, settings.catalog_embed_url, + model, 1024, lance.get("uri"), + lance.get("primary_key", "debug_id"), max_rows, + ) + status = {"status": "ready", **summary} + except Exception as e: # noqa: BLE001 + status = {"status": "error", "error": str(e)[:300]} + ds = await db.get(Dataset, dataset_id) + meta = dict(ds.dataset_metadata or {}) + lc = dict(meta.get("lance") or {}) + lc["semantic"] = {**status, "at": datetime.utcnow().isoformat()} + if status.get("status") == "ready": + lc["embedded"] = True + meta["lance"] = lc + ds.dataset_metadata = meta + ds.updated_at = datetime.utcnow() + await db.commit() + finally: + await eng.dispose() + + +@router.post("/datasets/{dataset_id}/optimize-embedding-search", status_code=202) +async def optimize_embedding_search( + dataset_id: uuid.UUID, + data: OptimizeEmbedRequest, + current_user: User = Depends(get_current_user), + db: AsyncSession = Depends(get_db), +): + """Optimize this dataset for semantic search: embed the rows that still lack a + vector (stage-then-flush, never per-chunk merge) and build/rebuild an IVF ANN + index. Long-running -> runs in the background; poll + ``dataset_metadata.lance.semantic.status`` (running/ready/error).""" + import asyncio + + ds = await _get_dataset(db, dataset_id) + await _verify_org_membership(db, ds.org_id, current_user.id) + lance = (ds.dataset_metadata or {}).get("lance") + if not lance or ds.connection_id is None or not ds.storage_prefix: + raise HTTPException(status_code=400, detail="Dataset is not indexed. POST /index first.") + if (lance.get("semantic") or {}).get("status") == "running": + raise HTTPException(status_code=409, detail="Optimization already running") + text_column = (data.text_column or _default_text_column(lance)).strip() + model = data.model or "BAAI/bge-m3" + + meta = dict(ds.dataset_metadata or {}) + lc = dict(meta.get("lance") or {}) + lc["semantic"] = {"status": "running", "at": datetime.utcnow().isoformat(), + "text_column": text_column} + meta["lance"] = lc + ds.dataset_metadata = meta + await db.commit() + + asyncio.create_task( + _run_optimize_embedding(ds.id, text_column, data.max_rows, model) + ) + return {"status": "running", "dataset_id": str(ds.id), "text_column": text_column} + + +@router.get("/datasets/{dataset_id}/semantic") +async def dataset_semantic_search( + dataset_id: uuid.UUID, + q: str, + limit: int = 20, + lang: str = "", + current_user: User = Depends(get_current_user), + db: AsyncSession = Depends(get_db), +): + """Semantic (embedding) search over the dataset's LanceDB table. Requires the + dataset to have been optimized (an embedding column + ideally an ANN index).""" + from starlette.concurrency import run_in_threadpool + + from app.services.dataset_embed_optimize import semantic_search_sync + + ds = await _get_dataset(db, dataset_id) + await _verify_org_membership(db, ds.org_id, current_user.id) + lance = (ds.dataset_metadata or {}).get("lance") + if not lance or ds.connection_id is None or not ds.storage_prefix: + raise HTTPException(status_code=400, detail="Dataset is not indexed") + conn = ( + await db.execute(select(StorageConnection).where(StorageConnection.id == ds.connection_id)) + ).scalar_one_or_none() + if conn is None: + raise HTTPException(status_code=409, detail="Bound storage connection no longer exists") + q = (q or "").strip()[:300] + if not q: + raise HTTPException(status_code=400, detail="empty query") + limit = max(1, min(limit, 200)) + lang = lang.strip() if lang else "" + try: + vec = await run_in_threadpool(_embed_query_sync, q) + except Exception: + raise HTTPException(status_code=502, detail="embedding service unavailable") + await _release_row_locks(db) + try: + rows = await run_in_threadpool( + semantic_search_sync, conn, ds.storage_prefix, lance["table"], vec, + limit, lance.get("uri"), lang or None, + ) + except Exception as e: + raise HTTPException(status_code=502, detail=f"Semantic search failed: {e}") + return {"query": q, "table": lance["table"], "rows": rows, "count": len(rows), "semantic": True} diff --git a/src/api/app/core/config.py b/src/api/app/core/config.py index 822afc50..f172a70e 100644 --- a/src/api/app/core/config.py +++ b/src/api/app/core/config.py @@ -333,7 +333,7 @@ def minio_external_endpoint(self) -> str: # Speech catalog rich-search: embedding + NL-query services (optional; # endpoints degrade gracefully if unreachable). The LanceDB corpus itself # comes from the org's "catalog" storage connection. - catalog_embed_url: str = "http://127.0.0.1:8001/v1/embeddings" + catalog_embed_url: str = "http://host.docker.internal:8081/v1/embeddings" catalog_embed_model: str = "BAAI/bge-m3" catalog_ollama_url: str = "http://127.0.0.1:11434/api/generate" catalog_nl_model: str = "gemma3:4b" diff --git a/src/api/app/services/dataset_derive.py b/src/api/app/services/dataset_derive.py new file mode 100644 index 00000000..450de059 --- /dev/null +++ b/src/api/app/services/dataset_derive.py @@ -0,0 +1,234 @@ +"""Query-/search-derived datasets — a saved selection over a parent LanceDB table. + +A derived dataset is NOT a copy. It stores only the primary keys (``debug_id``) of +the rows a word/FTS/semantic query matched on the parent table, plus any *extra* +columns the user adds on top. At read time it fetches the live parent rows for +that page and merges the extra columns — so the derived dataset **inherits every +parent field dynamically** (new parent columns appear automatically) while adding +its own. De-duplication is intrinsic: the members table is keyed on the primary +key and appended with ``merge_insert(...).when_not_matched_insert_all()``. + +All heavy work is blocking LanceDB/S3 I/O — call these from a threadpool. +""" +from __future__ import annotations + +import unicodedata +import uuid +from datetime import datetime + +import pyarrow as pa + +from app.models.storage_connection import StorageConnection +from app.services.dataset_index import _lance_storage_options, _lance_uri, _scalar + +# Columns never surfaced in a derived row (huge / internal). +_HEAVY = {"embedding", "vector", "_text", "raw_input_json", "raw_output_uri", + "_rowid", "_rowaddr"} +# Auto-provenance columns written on every derived row (named, user-visible, but +# not hand-editable "extra" columns): which request pulled the row in, and when. +COL_REQUEST = "derived_from_request" +COL_REQUEST_DATE = "derived_from_request_date" +_RESERVED = {COL_REQUEST, COL_REQUEST_DATE} + + +def _connect(conn: StorageConnection, uri: str): + import lancedb + + return lancedb.connect(uri, storage_options=_lance_storage_options(conn)) + + +def new_members_table_name() -> str: + return f"derived_{uuid.uuid4().hex[:12]}" + + +def resolve_uri(conn: StorageConnection, storage_prefix: str, uri: str | None) -> str: + return uri or _lance_uri(conn, storage_prefix) + + +def derive_member_ids_sync( + conn: StorageConnection, + uri: str, + table: str, + mode: str, + q: str, + lang: str | None, + limit: int, + vec: list[float] | None = None, + map_table: str = "lexical_map", + pk: str = "debug_id", +) -> list[str]: + """Run a query on the PARENT table and return the matching primary keys. + + ``mode``: ``word`` (exact term via the lexical_map, cheapest), ``fts`` / + ``exact`` (BM25 / phrase full-text), ``semantic`` (cosine ANN — needs ``vec``). + Only the primary-key column is projected: no row data leaves the parent. + """ + db = _connect(conn, uri) + tbl = db.open_table(table) + names = set(tbl.schema.names) + + if mode == "word": + norm = unicodedata.normalize("NFC", (q or "").strip()).lower().replace("'", "") + if not norm: + return [] + if map_table in db.table_names(): + rows = (db.open_table(map_table).search() + .where(f"term = '{norm}'").limit(1).to_list()) + ids = list(rows[0].get("debug_ids") or []) if rows else [] + return [_scalar(i) for i in ids[:limit]] + mode = "fts" # no map built -> fall back to full-text + + if mode == "semantic": + if vec is None: + raise ValueError("semantic derive requires an embedded query vector") + qq = tbl.search(vec).metric("cosine").select([pk]).limit(limit) + if lang and "language" in names: + qq = qq.where(f"language = '{lang}'") + return [_scalar(r[pk]) for r in qq.to_list() if r.get(pk)] + + # full-text (fts = ranked tokens, exact = adjacent phrase) + q2 = f'"{q}"' if mode == "exact" else (q or "") + qq = tbl.search(q2, query_type="fts").select([pk]).limit(limit) + if lang and "language" in names: + qq = qq.where(f"language = '{lang}'") + return [_scalar(r[pk]) for r in qq.to_list() if r.get(pk)] + + +def upsert_members_sync( + conn: StorageConnection, + uri: str, + members_table: str, + ids: list[str], + source_q: str, + pk: str = "debug_id", +) -> dict: + """Create-or-append the members table with only ``ids`` + provenance. + + De-dup is automatic: existing primary keys are left untouched + (``when_not_matched_insert_all``), so re-deriving or appending a second query + never double-counts a row. Returns ``{added, total}``. + """ + db = _connect(conn, uri) + ids = [i for i in dict.fromkeys(ids) if i] # de-dup within the batch, drop empties + now = datetime.utcnow().isoformat() + existed = members_table in db.table_names() + + if not ids: + total = db.open_table(members_table).count_rows() if existed else 0 + return {"added": 0, "total": total} + + arrow = pa.table({ + pk: pa.array(ids, pa.string()), + COL_REQUEST: pa.array([source_q[:500]] * len(ids), pa.string()), + COL_REQUEST_DATE: pa.array([now] * len(ids), pa.string()), + }) + if not existed: + t = db.create_table(members_table, data=arrow) + try: + t.create_scalar_index(pk, index_type="BTREE", replace=True) + except Exception: + pass + return {"added": len(ids), "total": t.count_rows()} + + t = db.open_table(members_table) + before = t.count_rows() + # Align to the table schema so extra user columns (added later) stay NULL for + # new rows instead of erroring the insert. + (t.merge_insert(pk).when_not_matched_insert_all().execute(arrow)) + total = t.count_rows() + return {"added": total - before, "total": total} + + +def member_count_sync(conn: StorageConnection, uri: str, members_table: str) -> int: + db = _connect(conn, uri) + return db.open_table(members_table).count_rows() if members_table in db.table_names() else 0 + + +def derived_rows_sync( + conn: StorageConnection, + uri: str, + parent_table: str, + members_table: str, + offset: int, + limit: int, + pk: str = "debug_id", +) -> tuple[list[dict], int]: + """A page of derived rows: parent fields (live, inherited) + extra member cols. + + Pages the members table, fetches exactly that page's parent rows by primary + key (BTREE point lookups), then overlays the derived extra columns. Bounded — + never scans the whole parent. + """ + db = _connect(conn, uri) + if members_table not in db.table_names(): + return [], 0 + m = db.open_table(members_table) + total = m.count_rows() + mrows = m.search().limit(offset + limit).to_list() + page = mrows[offset:offset + limit] + if not page: + return [], total + ids = [r[pk] for r in page if r.get(pk)] + extra_by_id = { + r[pk]: {k: _scalar(v) for k, v in r.items() + if k != pk and k not in _RESERVED and not k.startswith("_")} + for r in page + } + prov_by_id = {r[pk]: {k: _scalar(v) for k, v in r.items() if k in _RESERVED} for r in page} + + ptbl = db.open_table(parent_table) + inlist = ",".join("'" + str(i).replace("'", "") + "'" for i in ids) + prows = ptbl.search().where(f"{pk} IN ({inlist})").limit(len(ids)).to_list() + pby = {r.get(pk): r for r in prows} + + out = [] + for i in ids: + base = {k: _scalar(v) for k, v in (pby.get(i) or {}).items() if k not in _HEAVY} + base.update(prov_by_id.get(i, {})) # _added_at / _source_q provenance + base.update(extra_by_id.get(i, {})) # user extra columns (win on conflict) + base[pk] = i + out.append(base) + return out, total + + +def add_member_column_sync( + conn: StorageConnection, uri: str, members_table: str, name: str, default=None, +) -> list[str]: + """Add a new (string) column to the derived dataset. Parent stays untouched.""" + db = _connect(conn, uri) + t = db.open_table(members_table) + safe = "".join(c for c in name if c.isalnum() or c in "_-").strip("_-") + if not safe: + raise ValueError("invalid column name") + if safe in t.schema.names: + return list(t.schema.names) + lit = "NULL" if default in (None, "") else "'" + str(default).replace("'", "''") + "'" + t.add_columns({safe: f"cast({lit} as string)"}) + return list(db.open_table(members_table).schema.names) + + +def set_member_cell_sync( + conn: StorageConnection, uri: str, members_table: str, + pk_value: str, col: str, value, pk: str = "debug_id", +) -> bool: + """Set an extra-column value for one derived row (creates the column if new).""" + db = _connect(conn, uri) + t = db.open_table(members_table) + safe = "".join(c for c in col if c.isalnum() or c in "_-").strip("_-") + if not safe or safe == pk: + raise ValueError("invalid column name") + if safe not in t.schema.names: + add_member_column_sync(conn, uri, members_table, safe) + t = db.open_table(members_table) + arrow = pa.table({pk: pa.array([pk_value], pa.string()), + safe: pa.array([None if value is None else str(value)], pa.string())}) + t.merge_insert(pk).when_matched_update_all().execute(arrow) + return True + + +def extra_columns_sync(conn: StorageConnection, uri: str, members_table: str, pk: str = "debug_id") -> list[str]: + db = _connect(conn, uri) + if members_table not in db.table_names(): + return [] + return [c for c in db.open_table(members_table).schema.names + if c != pk and c not in _RESERVED and not c.startswith("_")] diff --git a/src/api/app/services/dataset_embed_optimize.py b/src/api/app/services/dataset_embed_optimize.py new file mode 100644 index 00000000..2d275831 --- /dev/null +++ b/src/api/app/services/dataset_embed_optimize.py @@ -0,0 +1,177 @@ +"""Optimize a dataset's LanceDB table for fast semantic (embedding) search. + +Generalises the speech-catalog backfill into a reusable dataset action: embed the +rows that still lack a vector, then build/rebuild an IVF ANN index so semantic +search is sub-second instead of a brute-force scan. + +Two hard-won rules baked in (see the catalog backfill): +- **Never per-chunk `merge_insert`.** On a wide many-column table it rewrites the + whole embedding column every call (O(column size), not O(rows)) — minutes per + chunk. Instead embed to a staging table via cheap `add` (append), then ONE + `merge_insert` at the end (a single column rewrite). +- **Only embed rows with real text.** Rows with an empty text column are skipped + (nothing to embed) so they don't churn the scan. + +Blocking LanceDB/HTTP I/O — call from a threadpool. +""" +from __future__ import annotations + +import json +import time +import urllib.request +from concurrent.futures import ThreadPoolExecutor + +import pyarrow as pa + +from app.models.storage_connection import StorageConnection +from app.services.dataset_index import _lance_storage_options, _lance_uri + +EMB_COL = "embedding" + + +def _embed(embed_url: str, model: str, texts: list[str]) -> list[list[float]]: + body = json.dumps({"model": model, "input": texts}).encode() + req = urllib.request.Request(embed_url, data=body, headers={"Content-Type": "application/json"}) + for attempt in range(5): + try: + with urllib.request.urlopen(req, timeout=180) as r: + d = json.loads(r.read())["data"] + return [e["embedding"] for e in sorted(d, key=lambda x: x["index"])] + except Exception: + if attempt == 4: + raise + time.sleep(1.0 * (attempt + 1)) + + +def optimize_embedding_search_sync( + conn: StorageConnection, + storage_prefix: str, + table: str, + text_column: str, + embed_url: str, + model: str = "BAAI/bge-m3", + dim: int = 1024, + uri: str | None = None, + pk: str = "debug_id", + max_rows: int = 200000, + batch: int = 16, + workers: int = 16, + maxchars: int = 4000, + num_partitions: int = 256, + num_sub_vectors: int = 128, +) -> dict: + """Embed rows missing a vector (stage-then-flush) + build an IVF index. + + Returns ``{embedded, remaining_null, indexed, total, seconds}``. + """ + import lancedb + + if dim % num_sub_vectors != 0: + num_sub_vectors = 128 if dim % 128 == 0 else 64 if dim % 64 == 0 else 16 + uri = uri or _lance_uri(conn, storage_prefix) + db = lancedb.connect(uri, storage_options=_lance_storage_options(conn)) + t = db.open_table(table) + names = set(t.schema.names) + if pk not in names: + raise ValueError(f"table has no primary key column {pk!r}") + if text_column not in names: + raise ValueError(f"table has no text column {text_column!r}") + has_emb = EMB_COL in names + + start = time.time() + staging = f"_embopt_{table}"[:60] + embedded = 0 + + if has_emb: + # embed only the rows still missing a vector (with real text) + where = (f"{EMB_COL} IS NULL AND {text_column} IS NOT NULL " + f"AND {text_column} != ''") + else: + # brand-new column: embed every row that has text + where = f"{text_column} IS NOT NULL AND {text_column} != ''" + + rows = (t.search().where(where).select([pk, text_column]) + .limit(max_rows).to_list()) + + if rows: + pool = ThreadPoolExecutor(workers) + if staging in db.table_names(): + db.drop_table(staging) + stg = None + CHUNK = 20000 + for c0 in range(0, len(rows), CHUNK): + chunk = rows[c0:c0 + CHUNK] + ids = [r[pk] for r in chunk] + txts = [(r.get(text_column) or "")[:maxchars] for r in chunk] + subs = [txts[j:j + batch] for j in range(0, len(txts), batch)] + vecs = [v for part in pool.map(lambda b: _embed(embed_url, model, b), subs) for v in part] + arrow = pa.table({pk: pa.array(ids, pa.string()), + EMB_COL: pa.array(vecs, pa.list_(pa.float32(), dim))}) + if stg is None: + stg = db.create_table(staging, data=arrow) + else: + stg.add(arrow) + embedded += len(ids) + pool.shutdown() + + # ONE merge (single column rewrite). For a brand-new column, merge adds it. + data = stg.search().select([pk, EMB_COL]).limit(10 ** 9).to_arrow() + if has_emb: + t.merge_insert(pk).when_matched_update_all().execute(data) + else: + t.merge_insert(pk).when_matched_update_all().when_not_matched_insert_all().execute(data) + db.drop_table(staging) + + # build / rebuild the ANN index + try: + t.create_index(metric="cosine", vector_column_name=EMB_COL, + index_type="IVF_PQ", num_partitions=num_partitions, + num_sub_vectors=num_sub_vectors, replace=True) + except Exception as e: + # tables below the training floor can't build IVF_PQ; brute-force still works + if "not enough" not in str(e).lower() and "training" not in str(e).lower(): + raise + + total = t.count_rows() + remaining = t.count_rows(f"{EMB_COL} IS NULL") if EMB_COL in t.schema.names else total + indexed = 0 + for i in t.list_indices(): + if getattr(i, "vector_column_name", None) == EMB_COL or i.name == f"{EMB_COL}_idx": + indexed = i.num_indexed_rows + return {"embedded": embedded, "remaining_null": remaining, "indexed": indexed, + "total": total, "seconds": round(time.time() - start, 1)} + + +def semantic_search_sync( + conn: StorageConnection, + storage_prefix: str, + table: str, + vec: list[float], + limit: int, + uri: str | None = None, + lang: str | None = None, + light_columns: tuple[str, ...] = ("debug_id", "language", "date", + "duration_seconds", "number_of_speakers", + "word_count", "audio_uri", "full_transcript"), +) -> list[dict]: + """Cosine ANN search over the table's ``embedding`` column. Returns light hits.""" + import lancedb + + from app.services.dataset_index import _fts_snippet, _scalar + + uri = uri or _lance_uri(conn, storage_prefix) + db = lancedb.connect(uri, storage_options=_lance_storage_options(conn)) + t = db.open_table(table) + cols = [c for c in light_columns if c in t.schema.names] + q = t.search(vec).metric("cosine").select(cols).limit(limit) + if lang and "language" in t.schema.names: + q = q.where(f"language = '{lang}'") + out = [] + for r in q.to_list(): + d = {k: _scalar(v) for k, v in r.items() if k not in ("full_transcript", "_distance")} + d["has_audio"] = bool(r.get("audio_uri")) + d.pop("audio_uri", None) + d["score"] = round(1 - r.get("_distance", 0), 3) + d["snippet"] = _fts_snippet(r.get("full_transcript") or "", "") + out.append(d) + return out diff --git a/src/web/src/pages/DatasetDetail.tsx b/src/web/src/pages/DatasetDetail.tsx index 5095451f..a7da2040 100644 --- a/src/web/src/pages/DatasetDetail.tsx +++ b/src/web/src/pages/DatasetDetail.tsx @@ -1000,6 +1000,61 @@ function IndexesTab({ datasetId, managed, columns, lance }: { finally { setWloading(false); } }; + // Derive a new dataset from a word / full-text / semantic query — stores only + // matching debug_ids (no replication), inherits all parent fields, auto-dedup. + const [dmode, setDmode] = useState<"word" | "fts" | "semantic">("semantic"); + const [dname, setDname] = useState(""); + const [dq, setDq] = useState(""); + const [dlimit, setDlimit] = useState(2000); + const [dbusy, setDbusy] = useState(false); + const [derr, setDerr] = useState(null); + const [dresult, setDresult] = useState<{ id: string; name: string; num_examples: number | null } | null>(null); + const runDerive = async () => { + if (!dq.trim() || !dname.trim()) { setDerr("name and query are required"); return; } + setDbusy(true); setDerr(null); setDresult(null); + try { + const res = await apiFetch(`/datasets/${datasetId}/derive`, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ name: dname.trim(), mode: dmode, q: dq.trim(), limit: dlimit }), + }); + if (!res.ok) { setDerr(`error ${res.status}: ${(await res.text()).slice(0, 200)}`); return; } + const ds = await res.json(); + setDresult({ id: ds.id, name: ds.name, num_examples: ds.num_examples }); + } catch (e) { setDerr(e instanceof Error ? e.message : "derive failed"); } + finally { setDbusy(false); } + }; + + // Optimize embedding search: embed missing rows + build an ANN index so + // semantic (meaning) search is sub-second. Generic — works on any indexed dataset. + const sem = (lance as unknown as { semantic?: { status?: string; embedded?: number; indexed?: number; remaining_null?: number } } | null)?.semantic; + const [optStatus, setOptStatus] = useState(sem?.status ?? null); + const [optErr, setOptErr] = useState(null); + const [sq, setSq] = useState(""); + const [sres, setSres] = useState[] | null>(null); + const [sloading, setSloading] = useState(false); + const [serr, setSerr] = useState(null); + const runOptimize = async () => { + setOptErr(null); + try { + const res = await apiFetch(`/datasets/${datasetId}/optimize-embedding-search`, { + method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify({}), + }); + if (!res.ok) { setOptErr(`error ${res.status}: ${(await res.text()).slice(0, 200)}`); return; } + setOptStatus("running"); + } catch (e) { setOptErr(e instanceof Error ? e.message : "optimize failed"); } + }; + const runSemantic = async () => { + if (!sq.trim()) return; + setSloading(true); setSerr(null); setSres(null); + try { + const res = await apiFetch(`/datasets/${datasetId}/semantic?q=${encodeURIComponent(sq.trim())}&limit=20`); + if (!res.ok) { setSerr(`error ${res.status}: ${(await res.text()).slice(0, 160)}`); return; } + setSres((await res.json()).rows); + } catch (e) { setSerr(e instanceof Error ? e.message : "semantic search failed"); } + finally { setSloading(false); } + }; + if (!managed) return

Indexes are available for managed datasets (bound to a storage connection). Upload files through a connection to enable indexing.

; @@ -1097,6 +1152,43 @@ function IndexesTab({ datasetId, managed, columns, lance }: { )} +
+
+ ⚡ Optimize embedding search + {sem?.indexed != null && {sem.indexed.toLocaleString()} vectors indexed} +
+

+ Embed the rows that still lack a vector and build an ANN index so + semantic (meaning) search is sub-second. Runs in the background — + works on any indexed dataset. Also powers semantic-mode dataset derivation. +

+
+ + {optStatus && + status: {optStatus}{sem?.embedded != null ? ` · ${sem.embedded.toLocaleString()} embedded` : ""}{sem?.remaining_null != null ? ` · ${sem.remaining_null.toLocaleString()} without text` : ""} + } +
+ {optErr &&
{optErr}
} +
+ setSq(e.target.value)} placeholder="semantic query (meaning)…" + disabled={!lance} style={{ ...s.input, flex: 1 }} onKeyDown={(e) => e.key === "Enter" && runSemantic()} /> + +
+ {serr &&
{serr}
} + {sres && (sres.length === 0 ?

No matches.

: ( +
+ {sres.map((r, i) => ( +
+ {typeof r.score === "number" ? r.score.toFixed(3) : ""} + {" "}[{String(r.language ?? "")}] {String(r.snippet ?? "").slice(0, 180)} +
+ ))} +
+ ))} +
+
Search
@@ -1124,6 +1216,42 @@ function IndexesTab({ datasetId, managed, columns, lance }: {
))}
+ +
+
+ 🧬 Derive a dataset from this query + no replication · auto-dedup · inherits all fields +
+

+ Run a word / full-text / semantic query and save the matches as a new dataset. + It stores only the matching row IDs (not a copy), inherits every parent column + dynamically, auto-tags each row with derived_from_request +{" "} + derived_from_request_date, and de-duplicates automatically. +

+
+ + setDq(e.target.value)} placeholder={dmode === "word" ? "a single word…" : "query…"} + disabled={!lance} style={{ ...s.input, flex: 1, minWidth: 180 }} /> + setDlimit(Math.max(1, Number(e.target.value) || 1))} + title="max rows" style={{ ...s.input, flex: "0 0 90px" }} /> +
+
+ setDname(e.target.value)} placeholder="new dataset name (e.g. free-user-medical)" + disabled={!lance} style={{ ...s.input, flex: 1, minWidth: 180 }} onKeyDown={(e) => e.key === "Enter" && runDerive()} /> + +
+ {derr &&
{derr}
} + {dresult && ( +
+ Created {dresult.name} + {" — "}{(dresult.num_examples ?? 0).toLocaleString()} rows (deduplicated). Open it to add columns or append more queries. +
+ )} +
); }