orchestrator-core 4.4.0rc3__py3-none-any.whl → 4.5.1a1__py3-none-any.whl

orchestrator/search/filters/date_filters.py

@@ -0,0 +1,75 @@
+from datetime import date, datetime
+from typing import Annotated, Any, Literal
+
+from dateutil.parser import parse as dt_parse
+from pydantic import BaseModel, BeforeValidator, Field, model_validator
+from sqlalchemy import TIMESTAMP, and_
+from sqlalchemy import cast as sa_cast
+from sqlalchemy.sql.elements import ColumnElement
+
+from orchestrator.search.core.types import FilterOp, SQLAColumn
+
+
+def _validate_date_string(v: Any) -> Any:
+    if not isinstance(v, str):
+        return v
+    try:
+        dt_parse(v)
+        return v
+    except Exception as exc:
+        raise ValueError("is not a valid date or datetime string") from exc
+
+
+DateValue = datetime | date | str
+ValidatedDateValue = Annotated[DateValue, BeforeValidator(_validate_date_string)]
+
+
+class DateRange(BaseModel):
+
+    start: ValidatedDateValue
+    end: ValidatedDateValue
+
+    @model_validator(mode="after")
+    def _order(self) -> "DateRange":
+        to_datetime = dt_parse(str(self.end))
+        from_datetime = dt_parse(str(self.start))
+        if to_datetime <= from_datetime:
+            raise ValueError("'to' must be after 'from'")
+        return self
+
+
+class DateValueFilter(BaseModel):
+    """A filter that operates on a single date value."""
+
+    op: Literal[FilterOp.EQ, FilterOp.NEQ, FilterOp.LT, FilterOp.LTE, FilterOp.GT, FilterOp.GTE]
+    value: ValidatedDateValue
+
+    def to_expression(self, column: SQLAColumn, path: str) -> ColumnElement[bool]:
+        date_column = sa_cast(column, TIMESTAMP(timezone=True))
+        match self.op:
+            case FilterOp.EQ:
+                return date_column == self.value
+            case FilterOp.NEQ:
+                return date_column != self.value
+            case FilterOp.LT:
+                return date_column < self.value
+            case FilterOp.LTE:
+                return date_column <= self.value
+            case FilterOp.GT:
+                return date_column > self.value
+            case FilterOp.GTE:
+                return date_column >= self.value
+
+
+class DateRangeFilter(BaseModel):
+    """A filter that operates on a range of dates."""
+
+    op: Literal[FilterOp.BETWEEN]
+    value: DateRange
+
+    def to_expression(self, column: SQLAColumn, path: str) -> ColumnElement[bool]:
+        date_column = sa_cast(column, TIMESTAMP(timezone=True))
+        return and_(date_column >= self.value.start, date_column < self.value.end)
+
+
+DateFilter = Annotated[DateValueFilter | DateRangeFilter, Field(discriminator="op")]

orchestrator/search/filters/definitions.py

@@ -0,0 +1,93 @@
+from orchestrator.search.core.types import FieldType, FilterOp, UIType
+from orchestrator.search.schemas.results import TypeDefinition, ValueSchema
+
+
+def operators_for(ft: FieldType) -> list[FilterOp]:
+    """Return the list of valid operators for a given FieldType."""
+    return list(value_schema_for(ft).keys())
+
+
+def component_operators() -> dict[FilterOp, ValueSchema]:
+    """Return operators available for path components."""
+    return {
+        FilterOp.HAS_COMPONENT: ValueSchema(kind=UIType.COMPONENT),
+        FilterOp.NOT_HAS_COMPONENT: ValueSchema(kind=UIType.COMPONENT),
+    }
+
+
+def value_schema_for(ft: FieldType) -> dict[FilterOp, ValueSchema]:
+    """Return the value schema map for a given FieldType."""
+    if ft in (FieldType.INTEGER, FieldType.FLOAT):
+        return {
+            FilterOp.EQ: ValueSchema(kind=UIType.NUMBER),
+            FilterOp.NEQ: ValueSchema(kind=UIType.NUMBER),
+            FilterOp.LT: ValueSchema(kind=UIType.NUMBER),
+            FilterOp.LTE: ValueSchema(kind=UIType.NUMBER),
+            FilterOp.GT: ValueSchema(kind=UIType.NUMBER),
+            FilterOp.GTE: ValueSchema(kind=UIType.NUMBER),
+            FilterOp.BETWEEN: ValueSchema(
+                kind="object",
+                fields={
+                    "start": ValueSchema(kind=UIType.NUMBER),
+                    "end": ValueSchema(kind=UIType.NUMBER),
+                },
+            ),
+        }
+
+    if ft == FieldType.BOOLEAN:
+        return {
+            FilterOp.EQ: ValueSchema(kind=UIType.BOOLEAN),
+            FilterOp.NEQ: ValueSchema(kind=UIType.BOOLEAN),
+        }
+
+    if ft == FieldType.DATETIME:
+        return {
+            FilterOp.EQ: ValueSchema(kind=UIType.DATETIME),
+            FilterOp.NEQ: ValueSchema(kind=UIType.DATETIME),
+            FilterOp.LT: ValueSchema(kind=UIType.DATETIME),
+            FilterOp.LTE: ValueSchema(kind=UIType.DATETIME),
+            FilterOp.GT: ValueSchema(kind=UIType.DATETIME),
+            FilterOp.GTE: ValueSchema(kind=UIType.DATETIME),
+            FilterOp.BETWEEN: ValueSchema(
+                kind="object",
+                fields={
+                    "start": ValueSchema(kind=UIType.DATETIME),
+                    "end": ValueSchema(kind=UIType.DATETIME),
+                },
+            ),
+        }
+
+    return {
+        FilterOp.EQ: ValueSchema(kind=UIType.STRING),
+        FilterOp.NEQ: ValueSchema(kind=UIType.STRING),
+    }
+
+
+def generate_definitions() -> dict[UIType, TypeDefinition]:
+    """Generate the full definitions dictionary for all UI types."""
+    definitions: dict[UIType, TypeDefinition] = {}
+
+    for ui_type in UIType:
+        if ui_type == UIType.COMPONENT:
+            # Special case for component filtering
+            comp_ops = component_operators()
+            definitions[ui_type] = TypeDefinition(
+                operators=list(comp_ops.keys()),
+                valueSchema=comp_ops,
+            )
+        else:
+            # Regular field types
+            if ui_type == UIType.NUMBER:
+                rep_ft = FieldType.INTEGER
+            elif ui_type == UIType.DATETIME:
+                rep_ft = FieldType.DATETIME
+            elif ui_type == UIType.BOOLEAN:
+                rep_ft = FieldType.BOOLEAN
+            else:
+                rep_ft = FieldType.STRING
+
+            definitions[ui_type] = TypeDefinition(
+                operators=operators_for(rep_ft),
+                valueSchema=value_schema_for(rep_ft),
+            )
+    return definitions

orchestrator/search/filters/ltree_filters.py

@@ -0,0 +1,43 @@
+from typing import Literal
+
+from pydantic import BaseModel, Field
+from sqlalchemy import TEXT, bindparam
+from sqlalchemy.sql.elements import ColumnElement
+from sqlalchemy_utils.types.ltree import Ltree
+
+from orchestrator.search.core.types import FilterOp, SQLAColumn
+
+
+class LtreeFilter(BaseModel):
+    """Filter for ltree path operations."""
+
+    op: Literal[
+        FilterOp.MATCHES_LQUERY,
+        FilterOp.IS_ANCESTOR,
+        FilterOp.IS_DESCENDANT,
+        FilterOp.PATH_MATCH,
+        FilterOp.HAS_COMPONENT,
+        FilterOp.NOT_HAS_COMPONENT,
+        FilterOp.ENDS_WITH,
+    ]
+    value: str = Field(description="The ltree path or lquery pattern to compare against.")
+
+    def to_expression(self, column: SQLAColumn, path: str) -> ColumnElement[bool]:
+        """Converts the filter condition into a SQLAlchemy expression."""
+        match self.op:
+            case FilterOp.IS_DESCENDANT:
+                ltree_value = Ltree(self.value)
+                return column.op("<@")(ltree_value)
+            case FilterOp.IS_ANCESTOR:
+                ltree_value = Ltree(self.value)
+                return column.op("@>")(ltree_value)
+            case FilterOp.MATCHES_LQUERY:
+                param = bindparam(None, self.value, type_=TEXT)
+                return column.op("~")(param)
+            case FilterOp.PATH_MATCH:
+                ltree_value = Ltree(path)
+                return column == ltree_value
+            case FilterOp.HAS_COMPONENT | FilterOp.NOT_HAS_COMPONENT:
+                return column.op("~")(bindparam(None, f"*.{self.value}.*", type_=TEXT))
+            case FilterOp.ENDS_WITH:
+                return column.op("~")(bindparam(None, f"*.{self.value}", type_=TEXT))

orchestrator/search/filters/numeric_filter.py

@@ -0,0 +1,60 @@
+from typing import Annotated, Any, Literal
+
+from pydantic import BaseModel, Field, model_validator
+from sqlalchemy import DOUBLE_PRECISION, INTEGER, and_
+from sqlalchemy import cast as sa_cast
+from sqlalchemy.sql.elements import ColumnElement
+from typing_extensions import Self
+
+from orchestrator.search.core.types import FilterOp, SQLAColumn
+
+
+class NumericRange(BaseModel):
+    start: int | float
+    end: int | float
+
+    @model_validator(mode="after")
+    def validate_order(self) -> Self:
+        if self.end <= self.start:
+            raise ValueError("'end' must be greater than 'start'")
+        return self
+
+
+class NumericValueFilter(BaseModel):
+    """A filter for single numeric value comparisons (int or float)."""
+
+    op: Literal[FilterOp.EQ, FilterOp.NEQ, FilterOp.LT, FilterOp.LTE, FilterOp.GT, FilterOp.GTE]
+    value: int | float
+
+    def to_expression(self, column: SQLAColumn, path: str) -> ColumnElement[bool]:
+        cast_type = INTEGER if isinstance(self.value, int) else DOUBLE_PRECISION
+        numeric_column: ColumnElement[Any] = sa_cast(column, cast_type)
+        match self.op:
+
+            case FilterOp.EQ:
+                return numeric_column == self.value
+            case FilterOp.NEQ:
+                return numeric_column != self.value
+            case FilterOp.LT:
+                return numeric_column < self.value
+            case FilterOp.LTE:
+                return numeric_column <= self.value
+            case FilterOp.GT:
+                return numeric_column > self.value
+            case FilterOp.GTE:
+                return numeric_column >= self.value
+
+
+class NumericRangeFilter(BaseModel):
+    """A filter for a range of numeric values (int or float)."""
+
+    op: Literal[FilterOp.BETWEEN]
+    value: NumericRange
+
+    def to_expression(self, column: SQLAColumn, path: str) -> ColumnElement[bool]:
+        cast_type = INTEGER if isinstance(self.value.start, int) else DOUBLE_PRECISION
+        numeric_column: ColumnElement[Any] = sa_cast(column, cast_type)
+        return and_(numeric_column >= self.value.start, numeric_column <= self.value.end)
+
+
+NumericFilter = Annotated[NumericValueFilter | NumericRangeFilter, Field(discriminator="op")]

orchestrator/search/indexing/__init__.py

@@ -0,0 +1,3 @@
+from .tasks import run_indexing_for_entity
+
+__all__ = ["run_indexing_for_entity"]

orchestrator/search/indexing/indexer.py

@@ -0,0 +1,323 @@
+import hashlib
+from collections.abc import Generator, Iterable, Iterator
+from contextlib import contextmanager, nullcontext
+from functools import lru_cache
+from typing import Any
+
+import structlog
+from litellm.utils import encode, get_max_tokens
+from sqlalchemy import delete, tuple_
+from sqlalchemy.dialects.postgresql import insert
+from sqlalchemy.dialects.postgresql.dml import Insert
+from sqlalchemy.orm import Session
+from sqlalchemy_utils.types.ltree import Ltree
+
+from orchestrator.db import db
+from orchestrator.db.models import AiSearchIndex
+from orchestrator.llm_settings import llm_settings
+from orchestrator.search.core.embedding import EmbeddingIndexer
+from orchestrator.search.core.types import ExtractedField, IndexableRecord
+from orchestrator.search.indexing.registry import EntityConfig
+from orchestrator.search.indexing.traverse import DatabaseEntity
+
+logger = structlog.get_logger(__name__)
+
+
+@contextmanager
+def _maybe_begin(session: Session | None) -> Iterator[None]:
+    if session is None:
+        yield
+    else:
+        with session.begin():
+            yield
+
+
+class Indexer:
+    """Index entities into `AiSearchIndex` using streaming reads and batched writes.
+
+    Entities are read from a streaming iterator and accumulated into chunks of
+    size `chunk_size`. For each chunk, the indexer extracts fields, diffs via
+    content hashes, deletes stale paths, and prepares upserts using a two-list
+    buffer:
+    - Embeddable list (STRING fields): maintains a running token count against a
+        token budget (model context window minus a safety margin) and flushes when
+        adding the next item would exceed the budget.
+    - Non-embeddable list: accumulated in parallel and does not contribute to the
+        flush condition.
+    Each flush (or end-of-chunk) emits a single combined UPSERT batch from both
+    lists (wrapped in a per-chunk transaction in non-dry-runs).
+
+    Args:
+        config (EntityConfig): Registry config describing the entity kind,
+            ORM table, and traverser.
+        dry_run (bool): If True, skip DELETE/UPSERT statements and external
+            embedding calls.
+        force_index (bool): If True, ignore existing hashes and reindex all
+            fields for each entity.
+        chunk_size (int): Number of entities to process per batch. Defaults to 1000.
+
+    Notes:
+        - Non-dry-run runs open a write session and wrap each processed chunk in
+          a transaction (`Session.begin()`).
+        - Read queries use the passed session when available, otherwise the
+          generic `db.session`.
+
+    Workflow:
+        1) Stream entities (yield_per=chunk_size) and accumulate into a chunk.
+        2) Begin transaction for the chunk.
+        3) determine_changes() → fields_to_upsert, paths_to_delete.
+        4) Delete stale paths.
+        5) Build UPSERT batches with a two-list buffer:
+        - Embeddable list (STRING): track running token count; flush when next item
+            would exceed the token budget (model max context - safety margin).
+        - Non-embeddable list: accumulate in parallel; does not affect flushing.
+        6) Execute UPSERT for each batch (skip in dry_run).
+        7) Commit transaction (auto on context exit).
+        8) Repeat until the stream is exhausted.
+    """
+
+    def __init__(self, config: EntityConfig, dry_run: bool, force_index: bool, chunk_size: int = 1000) -> None:
+        self.config = config
+        self.dry_run = dry_run
+        self.force_index = force_index
+        self.chunk_size = chunk_size
+        self.embedding_model = llm_settings.EMBEDDING_MODEL
+        self.logger = logger.bind(entity_kind=config.entity_kind.value)
+
+    def run(self, entities: Iterable[DatabaseEntity]) -> int:
+        """Orchestrates the entire indexing process."""
+        chunk: list[DatabaseEntity] = []
+        total_records_processed = 0
+        total_identical_records = 0
+
+        write_scope = db.database_scope() if not self.dry_run else nullcontext()
+
+        def flush() -> None:
+            nonlocal total_records_processed, total_identical_records
+            with _maybe_begin(session):
+                processed_in_chunk, identical_in_chunk = self._process_chunk(chunk, session)
+                total_records_processed += processed_in_chunk
+                total_identical_records += identical_in_chunk
+            chunk.clear()
+
+        with write_scope as database:
+            session: Session | None = getattr(database, "session", None)
+            for entity in entities:
+                chunk.append(entity)
+                if len(chunk) >= self.chunk_size:
+                    flush()
+
+            if chunk:
+                flush()
+
+        final_log_message = (
+            f"processed {total_records_processed} records and skipped {total_identical_records} identical records."
+        )
+        self.logger.info(
+            f"Dry run, would have indexed {final_log_message}"
+            if self.dry_run
+            else f"Indexing done, {final_log_message}"
+        )
+        return total_records_processed
+
+    def _process_chunk(self, entity_chunk: list[DatabaseEntity], session: Session | None = None) -> tuple[int, int]:
+        """Process a chunk of entities."""
+        if not entity_chunk:
+            return 0, 0
+
+        fields_to_upsert, paths_to_delete, identical_count = self._determine_changes(entity_chunk, session)
+
+        if paths_to_delete and session is not None:
+            self.logger.debug(f"Deleting {len(paths_to_delete)} stale records in chunk.")
+            self._execute_batched_deletes(paths_to_delete, session)
+
+        if fields_to_upsert:
+            upsert_stmt = self._get_upsert_statement()
+            batch_generator = self._generate_upsert_batches(fields_to_upsert)
+
+            for batch in batch_generator:
+                if self.dry_run:
+                    self.logger.debug(f"Dry Run: Would upsert {len(batch)} records.")
+                elif batch and session:
+                    session.execute(upsert_stmt, batch)
+
+        return len(fields_to_upsert), identical_count
+
+    def _determine_changes(
+        self, entities: list[DatabaseEntity], session: Session | None = None
+    ) -> tuple[list[tuple[str, ExtractedField]], list[tuple[str, Ltree]], int]:
+        """Identifies all changes across all entities using pre-fetched data."""
+        entity_ids = [str(getattr(e, self.config.pk_name)) for e in entities]
+        read_session = session or db.session
+        existing_hashes = {} if self.force_index else self._get_all_existing_hashes(entity_ids, read_session)
+
+        fields_to_upsert: list[tuple[str, ExtractedField]] = []
+        paths_to_delete: list[tuple[str, Ltree]] = []
+        identical_records_count = 0
+
+        for entity in entities:
+            entity_id = str(getattr(entity, self.config.pk_name))
+            current_fields = self.config.traverser.get_fields(
+                entity, pk_name=self.config.pk_name, root_name=self.config.root_name
+            )
+
+            entity_hashes = existing_hashes.get(entity_id, {})
+            current_paths = set()
+
+            for field in current_fields:
+                current_paths.add(field.path)
+                current_hash = self._compute_content_hash(field.path, field.value, field.value_type)
+                if field.path not in entity_hashes or entity_hashes[field.path] != current_hash:
+                    fields_to_upsert.append((entity_id, field))
+                else:
+                    identical_records_count += 1
+
+            stale_paths = set(entity_hashes.keys()) - current_paths
+            paths_to_delete.extend([(entity_id, Ltree(p)) for p in stale_paths])
+
+        return fields_to_upsert, paths_to_delete, identical_records_count
+
+    def _execute_batched_deletes(self, paths_to_delete: list[tuple[str, Ltree]], session: Session) -> None:
+        """Execute delete operations in batches to avoid PostgreSQL stack depth limits."""
+        for i in range(0, len(paths_to_delete), self.chunk_size):
+            batch = paths_to_delete[i : i + self.chunk_size]
+            delete_stmt = delete(AiSearchIndex).where(tuple_(AiSearchIndex.entity_id, AiSearchIndex.path).in_(batch))
+            session.execute(delete_stmt)
+            self.logger.debug(f"Deleted batch of {len(batch)} records.")
+
+    def _get_all_existing_hashes(self, entity_ids: list[str], session: Session) -> dict[str, dict[str, str]]:
+        """Fetches all existing hashes for a list of entity IDs in a single query."""
+        if not entity_ids:
+            return {}
+
+        results = (
+            session.query(AiSearchIndex.entity_id, AiSearchIndex.path, AiSearchIndex.content_hash)
+            .filter(AiSearchIndex.entity_id.in_(entity_ids))
+            .all()
+        )
+
+        hashes_by_entity: dict[str, dict[str, str]] = {eid: {} for eid in entity_ids}
+        for entity_id, path, content_hash in results:
+            hashes_by_entity[str(entity_id)][str(path)] = content_hash
+        return hashes_by_entity
+
+    def _generate_upsert_batches(
+        self, fields_to_upsert: Iterable[tuple[str, ExtractedField]]
+    ) -> Generator[list[IndexableRecord], None, None]:
+        """Streams through fields, buffers them by token count, and yields batches."""
+        embeddable_buffer: list[tuple[str, ExtractedField]] = []
+        non_embeddable_records: list[IndexableRecord] = []
+        current_tokens = 0
+
+        max_ctx = self._get_max_tokens()
+        safe_margin = int(max_ctx * llm_settings.EMBEDDING_SAFE_MARGIN_PERCENT)
+        token_budget = max(1, max_ctx - safe_margin)
+
+        max_batch_size = None
+        if llm_settings.OPENAI_BASE_URL:  # We are using a local model
+            max_batch_size = llm_settings.EMBEDDING_MAX_BATCH_SIZE
+
+        for entity_id, field in fields_to_upsert:
+            if field.value_type.is_embeddable(field.value):
+                text = self._prepare_text_for_embedding(field)
+                try:
+                    item_tokens = len(encode(model=self.embedding_model, text=text))
+                except Exception as e:
+                    self.logger.warning("Tokenization failed; skipping.", path=field.path, err=str(e))
+                    continue
+
+                if item_tokens > max_ctx:
+                    self.logger.warning(
+                        "Field exceeds context; skipping.", path=field.path, tokens=item_tokens, max_ctx=max_ctx
+                    )
+                    continue
+
+                should_flush = embeddable_buffer and (
+                    current_tokens + item_tokens > token_budget
+                    or (max_batch_size and len(embeddable_buffer) >= max_batch_size)
+                )
+
+                if should_flush:
+                    yield self._flush_buffer(embeddable_buffer, non_embeddable_records)
+                    embeddable_buffer.clear()
+                    non_embeddable_records.clear()
+                    current_tokens = 0
+
+                embeddable_buffer.append((entity_id, field))
+                current_tokens += item_tokens
+            else:
+                record = self._make_indexable_record(field, entity_id, embedding=None)
+                non_embeddable_records.append(record)
+
+        if embeddable_buffer or non_embeddable_records:
+            yield self._flush_buffer(embeddable_buffer, non_embeddable_records)
+
+    def _flush_buffer(self, embeddable_buffer: list, non_embeddable_records: list) -> list[IndexableRecord]:
+        """Processes and combines buffers into a single batch."""
+        if not embeddable_buffer:
+            return non_embeddable_records
+
+        texts_to_embed = [self._prepare_text_for_embedding(f) for _, f in embeddable_buffer]
+        embeddings = EmbeddingIndexer.get_embeddings_from_api_batch(texts_to_embed, self.dry_run)
+
+        if len(embeddable_buffer) != len(embeddings):
+            raise ValueError(f"Embedding mismatch: sent {len(embeddable_buffer)}, received {len(embeddings)}")
+
+        with_embeddings = [
+            self._make_indexable_record(field, entity_id, embedding)
+            for (entity_id, field), embedding in zip(embeddable_buffer, embeddings)
+        ]
+        return non_embeddable_records + with_embeddings
+
+    def _get_max_tokens(self) -> int:
+        """Gets max tokens, using a fallback from settings if necessary."""
+        try:
+            max_ctx = get_max_tokens(self.embedding_model)
+            if isinstance(max_ctx, int):
+                return max_ctx
+        except Exception:
+            # Allow local(unknown) models to fall back.
+            self.logger.warning("Could not auto-detect max tokens.", model=self.embedding_model)
+
+        max_ctx = llm_settings.EMBEDDING_FALLBACK_MAX_TOKENS
+        if not isinstance(max_ctx, int):
+            raise RuntimeError("Model not recognized and EMBEDDING_FALLBACK_MAX_TOKENS not set.")
+        self.logger.warning("Using configured fallback token limit.", fallback=max_ctx)
+        return max_ctx
+
+    @staticmethod
+    def _prepare_text_for_embedding(field: ExtractedField) -> str:
+        return f"{field.path}: {str(field.value)}"
+
+    @staticmethod
+    def _compute_content_hash(path: str, value: Any, value_type: Any) -> str:
+        v = "" if value is None else str(value)
+        content = f"{path}:{v}:{value_type}"
+        return hashlib.sha256(content.encode("utf-8")).hexdigest()
+
+    def _make_indexable_record(
+        self, field: ExtractedField, entity_id: str, embedding: list[float] | None
+    ) -> IndexableRecord:
+        return IndexableRecord(
+            entity_id=entity_id,
+            entity_type=self.config.entity_kind.value,
+            path=Ltree(field.path),
+            value=field.value,
+            value_type=field.value_type,
+            content_hash=self._compute_content_hash(field.path, field.value, field.value_type),
+            embedding=embedding if embedding else None,
+        )
+
+    @staticmethod
+    @lru_cache(maxsize=1)
+    def _get_upsert_statement() -> Insert:
+        stmt = insert(AiSearchIndex)
+        return stmt.on_conflict_do_update(
+            index_elements=[AiSearchIndex.entity_id, AiSearchIndex.path],
+            set_={
+                AiSearchIndex.value: stmt.excluded.value,
+                AiSearchIndex.value_type: stmt.excluded.value_type,
+                AiSearchIndex.content_hash: stmt.excluded.content_hash,
+                AiSearchIndex.embedding: stmt.excluded.embedding,
+            },
+        )

orchestrator/search/indexing/registry.py

@@ -0,0 +1,88 @@
+from dataclasses import dataclass
+from typing import Generic, TypeVar
+from uuid import UUID
+
+from sqlalchemy.orm import Query
+from sqlalchemy.sql import Select
+
+from orchestrator.db import (
+    ProcessTable,
+    ProductTable,
+    SubscriptionTable,
+    WorkflowTable,
+)
+from orchestrator.db.database import BaseModel
+from orchestrator.search.core.types import EntityType
+
+from .traverse import (
+    BaseTraverser,
+    ProcessTraverser,
+    ProductTraverser,
+    SubscriptionTraverser,
+    WorkflowTraverser,
+)
+
+ModelT = TypeVar("ModelT", bound=BaseModel)
+
+
+@dataclass(frozen=True)
+class EntityConfig(Generic[ModelT]):
+    """A container for all configuration related to a specific entity type."""
+
+    entity_kind: EntityType
+    table: type[ModelT]
+
+    traverser: "type[BaseTraverser]"
+    pk_name: str
+    root_name: str
+
+    def get_all_query(self, entity_id: str | None = None) -> Query | Select:
+        query = self.table.query
+        if entity_id:
+            pk_column = getattr(self.table, self.pk_name)
+            query = query.filter(pk_column == UUID(entity_id))
+        return query
+
+
+@dataclass(frozen=True)
+class WorkflowConfig(EntityConfig[WorkflowTable]):
+    """Workflows have a custom select() function that filters out deleted workflows."""
+
+    def get_all_query(self, entity_id: str | None = None) -> Select:
+        query = self.table.select()
+        if entity_id:
+            pk_column = getattr(self.table, self.pk_name)
+            query = query.where(pk_column == UUID(entity_id))
+        return query
+
+
+ENTITY_CONFIG_REGISTRY: dict[EntityType, EntityConfig] = {
+    EntityType.SUBSCRIPTION: EntityConfig(
+        entity_kind=EntityType.SUBSCRIPTION,
+        table=SubscriptionTable,
+        traverser=SubscriptionTraverser,
+        pk_name="subscription_id",
+        root_name="subscription",
+    ),
+    EntityType.PRODUCT: EntityConfig(
+        entity_kind=EntityType.PRODUCT,
+        table=ProductTable,
+        traverser=ProductTraverser,
+        pk_name="product_id",
+        root_name="product",
+    ),
+    EntityType.PROCESS: EntityConfig(
+        entity_kind=EntityType.PROCESS,
+        table=ProcessTable,
+        traverser=ProcessTraverser,
+        pk_name="process_id",
+        root_name="process",
+    ),
+    EntityType.WORKFLOW: WorkflowConfig(
+        entity_kind=EntityType.WORKFLOW,
+        table=WorkflowTable,
+        traverser=WorkflowTraverser,
+        pk_name="workflow_id",
+        root_name="workflow",
+    ),
+}