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

orchestrator/search/indexing/tasks.py

@@ -0,0 +1,53 @@
+import structlog
+from sqlalchemy.orm import Query
+
+from orchestrator.db import db
+from orchestrator.search.core.types import EntityType
+from orchestrator.search.indexing.indexer import Indexer
+from orchestrator.search.indexing.registry import ENTITY_CONFIG_REGISTRY
+
+logger = structlog.get_logger(__name__)
+
+
+def run_indexing_for_entity(
+    entity_kind: EntityType,
+    entity_id: str | None = None,
+    dry_run: bool = False,
+    force_index: bool = False,
+    chunk_size: int = 1000,
+) -> None:
+    """Stream and index entities for the given kind.
+
+    Builds a streaming query via the entity's registry config, disables ORM eager
+    loads when applicable and delegates processing to `Indexer`.
+
+    Args:
+        entity_kind (EntityType): The entity type to index (must exist in
+            `ENTITY_CONFIG_REGISTRY`).
+        entity_id (Optional[str]): If provided, restricts indexing to a single
+            entity (UUID string).
+        dry_run (bool): When True, runs the full pipeline without performing
+            writes or external embedding calls.
+        force_index (bool): When True, re-indexes all fields regardless of
+            existing hashes.
+        chunk_size (int): Number of rows fetched per round-trip and passed to
+            the indexer per batch.
+
+    Returns:
+        None
+    """
+    config = ENTITY_CONFIG_REGISTRY[entity_kind]
+
+    q = config.get_all_query(entity_id)
+
+    if isinstance(q, Query):
+        q = q.enable_eagerloads(False)
+        stmt = q.statement
+    else:
+        stmt = q
+
+    stmt = stmt.execution_options(stream_results=True, yield_per=chunk_size)
+    entities = db.session.execute(stmt).scalars()
+
+    indexer = Indexer(config=config, dry_run=dry_run, force_index=force_index, chunk_size=chunk_size)
+    indexer.run(entities)

orchestrator/search/indexing/traverse.py

@@ -0,0 +1,322 @@
+import re
+from abc import ABC, abstractmethod
+from collections.abc import Iterable
+from enum import Enum
+from typing import Any, cast, get_args
+from uuid import uuid4
+
+import structlog
+
+from orchestrator.db import ProcessTable, ProductTable, SubscriptionTable, WorkflowTable
+from orchestrator.domain import (
+    SUBSCRIPTION_MODEL_REGISTRY,
+    SubscriptionModel,
+)
+from orchestrator.domain.base import ProductBlockModel, ProductModel
+from orchestrator.domain.lifecycle import (
+    lookup_specialized_type,
+)
+from orchestrator.schemas.process import ProcessSchema
+from orchestrator.schemas.workflow import WorkflowSchema
+from orchestrator.search.core.exceptions import ModelLoadError, ProductNotInRegistryError
+from orchestrator.search.core.types import ExtractedField, FieldType
+from orchestrator.types import SubscriptionLifecycle
+
+logger = structlog.get_logger(__name__)
+
+DatabaseEntity = SubscriptionTable | ProductTable | ProcessTable | WorkflowTable
+
+
+class BaseTraverser(ABC):
+    """Base class for traversing database models and extracting searchable fields."""
+
+    _LTREE_SEPARATOR = "."
+    _MAX_DEPTH = 40
+
+    @classmethod
+    def get_fields(cls, entity: DatabaseEntity, pk_name: str, root_name: str) -> list[ExtractedField]:
+        """Main entry point for extracting fields from an entity. Default implementation delegates to _load_model."""
+        try:
+            model = cls._load_model(entity)
+            if model is None:
+                return []
+            return sorted(cls.traverse(model, root_name), key=lambda f: f.path)
+
+        except (ProductNotInRegistryError, ModelLoadError) as e:
+            entity_id = getattr(entity, pk_name, "unknown")
+            logger.error(f"Failed to extract fields from {entity.__class__.__name__}", id=str(entity_id), error=str(e))
+            return []
+
+    @classmethod
+    def traverse(cls, instance: Any, path: str = "") -> Iterable[ExtractedField]:
+        """Walks the fields of a Pydantic model, dispatching each to a field handler."""
+        model_class = type(instance)
+
+        # Handle both standard and computed fields from the Pydantic model
+        all_fields = model_class.model_fields.copy()
+        all_fields.update(getattr(model_class, "__pydantic_computed_fields__", {}))
+
+        for name, field in all_fields.items():
+            try:
+                value = getattr(instance, name, None)
+            except Exception as e:
+                logger.error(f"Failed to access field '{name}' on {model_class.__name__}", error=str(e))
+                continue
+            new_path = f"{path}{cls._LTREE_SEPARATOR}{name}" if path else name
+            annotation = field.annotation if hasattr(field, "annotation") else field.return_type
+            yield from cls._yield_fields_for_value(value, new_path, annotation)
+
+    @classmethod
+    def _yield_fields_for_value(cls, value: Any, path: str, annotation: Any) -> Iterable[ExtractedField]:
+        """Yields fields for a given value based on its type (model, list, or scalar)."""
+        if value is None:
+            return
+
+        # If the value is a list, pass it to the list traverser
+        if isinstance(value, list):
+            if element_annotation := get_args(annotation):
+                yield from cls._traverse_list(value, path, element_annotation[0])
+            return
+
+        # If the value is another Pydantic model, recurse into it
+        if hasattr(type(value), "model_fields"):
+            yield from cls.traverse(value, path)
+            return
+
+        ftype = FieldType.from_type_hint(annotation)
+
+        if isinstance(value, Enum):
+            yield ExtractedField(path, str(value.value), ftype)
+        else:
+            yield ExtractedField(path, str(value), ftype)
+
+    @classmethod
+    def _traverse_list(cls, items: list[Any], path: str, element_annotation: Any) -> Iterable[ExtractedField]:
+        """Recursively traverses items in a list."""
+        for i, item in enumerate(items):
+            item_path = f"{path}.{i}"
+            yield from cls._yield_fields_for_value(item, item_path, element_annotation)
+
+    @classmethod
+    def _load_model_with_schema(cls, entity: Any, schema_class: type[Any], pk_name: str) -> Any:
+        """Generic helper for loading models using Pydantic schema validation."""
+        try:
+            return schema_class.model_validate(entity)
+        except Exception as e:
+            entity_id = getattr(entity, pk_name, "unknown")
+            raise ModelLoadError(f"Failed to load {schema_class.__name__} for {pk_name} '{entity_id}'") from e
+
+    @classmethod
+    @abstractmethod
+    def _load_model(cls, entity: Any) -> Any: ...
+
+
+class SubscriptionTraverser(BaseTraverser):
+    """Traverser for subscription entities using full Pydantic model extraction."""
+
+    @classmethod
+    def _load_model(cls, sub: SubscriptionTable) -> SubscriptionModel | None:
+        base_model_cls = SUBSCRIPTION_MODEL_REGISTRY.get(sub.product.name)
+        if not base_model_cls:
+            raise ProductNotInRegistryError(f"Product '{sub.product.name}' not in registry.")
+
+        specialized_model_cls = cast(type[SubscriptionModel], lookup_specialized_type(base_model_cls, sub.status))
+
+        try:
+            return specialized_model_cls.from_subscription(sub.subscription_id)
+        except Exception as e:
+            raise ModelLoadError(f"Failed to load model for subscription_id '{sub.subscription_id}'") from e
+
+
+class ProductTraverser(BaseTraverser):
+    """Traverser for product entities using a template SubscriptionModel instance."""
+
+    @classmethod
+    def _sanitize_for_ltree(cls, name: str) -> str:
+        """Sanitizes a string to be a valid ltree path label."""
+        # Convert to lowercase
+        sanitized = name.lower()
+
+        # Replace all non-alphanumeric (and non-underscore) characters with an underscore
+        sanitized = re.sub(r"[^a-z0-9_]", "_", sanitized)
+
+        # Collapse multiple underscores into a single one
+        sanitized = re.sub(r"__+", "_", sanitized)
+
+        # Remove leading or trailing underscores
+        sanitized = sanitized.strip("_")
+
+        # Handle cases where the name was only invalid characters
+        if not sanitized:
+            return "unnamed_product"
+
+        return sanitized
+
+    @classmethod
+    def get_fields(cls, entity: ProductTable, pk_name: str, root_name: str) -> list[ExtractedField]:  # type: ignore[override]
+        """Extracts fields by creating a template SubscriptionModel instance for the product.
+
+        Extracts product metadata and block schema structure.
+        """
+        try:
+            model = cls._load_model(entity)
+
+            if not model:
+                return []
+
+            fields: list[ExtractedField] = []
+
+            product_fields = cls.traverse(model.product, root_name)
+            fields.extend(product_fields)
+
+            product_name = cls._sanitize_for_ltree(model.product.name)
+
+            product_block_root = f"{root_name}.{product_name}.product_block"
+
+            # Extract product block schema structure
+            model_class = type(model)
+            product_block_fields = getattr(model_class, "_product_block_fields_", {})
+
+            for field_name in product_block_fields:
+                block_value = getattr(model, field_name, None)
+                if block_value is not None:
+                    block_path = f"{product_block_root}.{field_name}"
+                    schema_fields = cls._extract_block_schema(block_value, block_path)
+                    fields.extend(schema_fields)
+
+            return sorted(fields, key=lambda f: f.path)
+
+        except (ProductNotInRegistryError, ModelLoadError) as e:
+            entity_id = getattr(entity, pk_name, "unknown")
+            logger.error(f"Failed to extract fields from {entity.__class__.__name__}", id=str(entity_id), error=str(e))
+            return []
+
+    @classmethod
+    def _extract_block_schema(cls, block_instance: ProductBlockModel, block_path: str) -> list[ExtractedField]:
+        """Extract schema information from a block instance, returning field names as RESOURCE_TYPE."""
+        fields = []
+
+        # Add the block itself as a BLOCK type
+        block_name = block_path.split(cls._LTREE_SEPARATOR)[-1]
+        fields.append(ExtractedField(path=block_path, value=block_name, value_type=FieldType.BLOCK))
+
+        # Extract all field names from the block as RESOURCE_TYPE
+        if hasattr(type(block_instance), "model_fields"):
+            all_fields = type(block_instance).model_fields
+            computed_fields = getattr(block_instance, "__pydantic_computed_fields__", None)
+            if computed_fields:
+                all_fields.update(computed_fields)
+
+            for field_name in all_fields:
+                field_value = getattr(block_instance, field_name, None)
+                field_path = f"{block_path}.{field_name}"
+
+                # If it's a nested block, recurse
+                if field_value is not None and isinstance(field_value, ProductBlockModel):
+                    nested_fields = cls._extract_block_schema(field_value, field_path)
+                    fields.extend(nested_fields)
+                elif field_value is not None and isinstance(field_value, list):
+                    # Handle list of blocks
+                    if field_value and isinstance(field_value[0], ProductBlockModel):
+                        # For lists, we still add the list field as a resource type
+                        fields.append(
+                            ExtractedField(path=field_path, value=field_name, value_type=FieldType.RESOURCE_TYPE)
+                        )
+                        # And potentially traverse the first item for schema
+                        first_item_path = f"{field_path}{cls._LTREE_SEPARATOR}0"
+                        nested_fields = cls._extract_block_schema(field_value[0], first_item_path)
+                        fields.extend(nested_fields)
+                    else:
+                        fields.append(
+                            ExtractedField(path=field_path, value=field_name, value_type=FieldType.RESOURCE_TYPE)
+                        )
+                else:
+                    # Regular fields are resource types
+                    fields.append(ExtractedField(path=field_path, value=field_name, value_type=FieldType.RESOURCE_TYPE))
+
+        return fields
+
+    @classmethod
+    def _load_model(cls, product: ProductTable) -> SubscriptionModel | None:
+        """Creates a template instance of a SubscriptionModel for a given product.
+
+        This allows us to traverse the product's defined block structure, even
+        without a real subscription instance in the database.
+        """
+        # Find the SubscriptionModel class associated with this product's name.
+        domain_model_cls = SUBSCRIPTION_MODEL_REGISTRY.get(product.name)
+        if not domain_model_cls:
+            raise ProductNotInRegistryError(f"Product '{product.name}' not in registry.")
+
+        # Get the initial lifecycle version of that class, as it represents the base structure.
+        try:
+            subscription_model_cls = cast(
+                type[SubscriptionModel], lookup_specialized_type(domain_model_cls, SubscriptionLifecycle.INITIAL)
+            )
+        except Exception:
+            subscription_model_cls = domain_model_cls
+
+        try:
+            product_model = ProductModel(
+                product_id=product.product_id,
+                name=product.name,
+                description=product.description,
+                product_type=product.product_type,
+                tag=product.tag,
+                status=product.status,
+            )
+
+            # Generate a fake subscription ID for the template
+            subscription_id = uuid4()
+
+            # Get fixed inputs for the product
+            fixed_inputs = {fi.name: fi.value for fi in product.fixed_inputs}
+
+            # Initialize product blocks
+            instances = subscription_model_cls._init_instances(subscription_id)
+
+            return subscription_model_cls(
+                product=product_model,
+                customer_id="traverser_template",
+                subscription_id=subscription_id,
+                description="Template for schema traversal",
+                status=SubscriptionLifecycle.INITIAL,
+                insync=False,
+                start_date=None,
+                end_date=None,
+                note=None,
+                version=1,
+                **fixed_inputs,
+                **instances,
+            )
+        except Exception:
+            logger.exception("Failed to instantiate template model for product", product_name=product.name)
+            return None
+
+
+class ProcessTraverser(BaseTraverser):
+    """Traverser for process entities using ProcessSchema model.
+
+    Note: Currently extracts only top-level process fields. Could be extended to include:
+    - Related subscriptions (entity.subscriptions)
+    - Related workflow information beyond workflow_name
+    """
+
+    @classmethod
+    def _load_model(cls, process: ProcessTable) -> ProcessSchema:
+        """Load process model using ProcessSchema."""
+        return cls._load_model_with_schema(process, ProcessSchema, "process_id")
+
+
+class WorkflowTraverser(BaseTraverser):
+    """Traverser for workflow entities using WorkflowSchema model.
+
+    Note: Currently extracts only top-level workflow fields. Could be extended to include:
+    - Related products (entity.products) - each with their own block structures
+    - Related processes (entity.processes) - each with their own process data
+    """
+
+    @classmethod
+    def _load_model(cls, workflow: WorkflowTable) -> WorkflowSchema:
+        """Load workflow model using WorkflowSchema."""
+        return cls._load_model_with_schema(workflow, WorkflowSchema, "workflow_id")

orchestrator/search/retrieval/__init__.py

@@ -0,0 +1,3 @@
+from .engine import execute_search
+
+__all__ = ["execute_search"]

orchestrator/search/retrieval/builder.py

@@ -0,0 +1,108 @@
+from collections import defaultdict
+from typing import Sequence
+
+from sqlalchemy import Select, String, cast, func, select
+from sqlalchemy.engine import Row
+
+from orchestrator.db.models import AiSearchIndex
+from orchestrator.search.core.types import EntityType, FieldType, FilterOp, UIType
+from orchestrator.search.filters import LtreeFilter
+from orchestrator.search.schemas.parameters import BaseSearchParameters
+from orchestrator.search.schemas.results import ComponentInfo, LeafInfo
+
+
+def create_path_autocomplete_lquery(prefix: str) -> str:
+    """Create the lquery pattern for a multi-level path autocomplete search."""
+    return f"{prefix}*.*"
+
+
+def build_candidate_query(params: BaseSearchParameters) -> Select:
+    """Build the base query for retrieving candidate entities.
+
+    Constructs a `SELECT` statement that retrieves distinct `entity_id` values
+    from the index table for the given entity type, applying any structured
+    filters from the provided search parameters.
+
+    Parameters
+    ----------
+    params : BaseSearchParameters
+        The search parameters containing the entity type and optional filters.
+
+    Returns:
+    -------
+    Select
+        The SQLAlchemy `Select` object representing the query.
+    """
+    stmt = select(AiSearchIndex.entity_id).where(AiSearchIndex.entity_type == params.entity_type.value).distinct()
+
+    if params.filters is not None:
+        entity_id_col = AiSearchIndex.entity_id
+        stmt = stmt.where(
+            params.filters.to_expression(
+                entity_id_col,
+                entity_type_value=params.entity_type.value,
+            )
+        )
+
+    return stmt
+
+
+def build_paths_query(entity_type: EntityType, prefix: str | None = None, q: str | None = None) -> Select:
+    """Build the query for retrieving paths and their value types for leaves/components processing."""
+    stmt = select(AiSearchIndex.path, AiSearchIndex.value_type).where(AiSearchIndex.entity_type == entity_type.value)
+
+    if prefix:
+        lquery_pattern = create_path_autocomplete_lquery(prefix)
+        ltree_filter = LtreeFilter(op=FilterOp.MATCHES_LQUERY, value=lquery_pattern)
+        stmt = stmt.where(ltree_filter.to_expression(AiSearchIndex.path, path=""))
+
+    stmt = stmt.group_by(AiSearchIndex.path, AiSearchIndex.value_type)
+
+    if q:
+        score = func.similarity(cast(AiSearchIndex.path, String), q)
+        stmt = stmt.order_by(score.desc(), AiSearchIndex.path)
+    else:
+        stmt = stmt.order_by(AiSearchIndex.path)
+
+    return stmt
+
+
+def process_path_rows(rows: Sequence[Row]) -> tuple[list[LeafInfo], list[ComponentInfo]]:
+    """Process query results to extract leaves and components information.
+
+    Parameters
+    ----------
+    rows : Sequence[Row]
+        Database rows containing path and value_type information
+
+    Returns:
+    -------
+    tuple[list[LeafInfo], list[ComponentInfo]]
+        Processed leaves and components
+    """
+    leaves_dict: dict[str, set[UIType]] = defaultdict(set)
+    components_set: set[str] = set()
+
+    for row in rows:
+        path, value_type = row
+
+        path_str = str(path)
+        path_segments = path_str.split(".")
+
+        # Remove numeric segments
+        clean_segments = [seg for seg in path_segments if not seg.isdigit()]
+
+        if clean_segments:
+            # Last segment is a leaf
+            leaf_name = clean_segments[-1]
+            ui_type = UIType.from_field_type(FieldType(value_type))
+            leaves_dict[leaf_name].add(ui_type)
+
+            # All segments except the first/last are components
+            for component in clean_segments[1:-1]:
+                components_set.add(component)
+
+    leaves = [LeafInfo(name=leaf, ui_types=list(types)) for leaf, types in leaves_dict.items()]
+    components = [ComponentInfo(name=component, ui_types=[UIType.COMPONENT]) for component in sorted(components_set)]
+
+    return leaves, components

orchestrator/search/retrieval/engine.py

@@ -0,0 +1,152 @@
+from collections.abc import Sequence
+
+import structlog
+from sqlalchemy.engine.row import RowMapping
+from sqlalchemy.orm import Session
+
+from orchestrator.search.core.types import FilterOp, SearchMetadata
+from orchestrator.search.filters import FilterTree, LtreeFilter
+from orchestrator.search.schemas.parameters import BaseSearchParameters
+from orchestrator.search.schemas.results import MatchingField, SearchResponse, SearchResult
+
+from .builder import build_candidate_query
+from .pagination import PaginationParams
+from .retriever import Retriever
+from .utils import generate_highlight_indices
+
+logger = structlog.get_logger(__name__)
+
+
+def _format_response(
+    db_rows: Sequence[RowMapping], search_params: BaseSearchParameters, metadata: SearchMetadata
+) -> SearchResponse:
+    """Format database query results into a `SearchResponse`.
+
+    Converts raw SQLAlchemy `RowMapping` objects into `SearchResult` instances,
+    including highlight metadata if present in the database results.
+
+    Parameters
+    ----------
+    db_rows : Sequence[RowMapping]
+        The rows returned from the executed SQLAlchemy query.
+
+    Returns:
+    -------
+    SearchResponse
+        A list of `SearchResult` objects containing entity IDs, scores, and
+        optional highlight information.
+    """
+
+    if not db_rows:
+        return SearchResponse(results=[], metadata=metadata)
+
+    user_query = search_params.query
+
+    results = []
+    for row in db_rows:
+        matching_field = None
+
+        if user_query and row.get("highlight_text") and row.get("highlight_path"):
+            # Text/semantic searches
+            text = row.highlight_text
+            path = row.highlight_path
+
+            if not isinstance(text, str):
+                text = str(text)
+            if not isinstance(path, str):
+                path = str(path)
+
+            highlight_indices = generate_highlight_indices(text, user_query) or None
+            matching_field = MatchingField(text=text, path=path, highlight_indices=highlight_indices)
+
+        elif not user_query and search_params.filters and metadata.search_type == "structured":
+            # Structured search (filter-only)
+            matching_field = _extract_matching_field_from_filters(search_params.filters)
+
+        results.append(
+            SearchResult(
+                entity_id=str(row.entity_id),
+                score=row.score,
+                perfect_match=row.get("perfect_match", 0),
+                matching_field=matching_field,
+            )
+        )
+    return SearchResponse(results=results, metadata=metadata)
+
+
+def _extract_matching_field_from_filters(filters: FilterTree) -> MatchingField | None:
+    """Extract the first path filter to use as matching field for structured searches.
+
+    TODO: Should we allow a list of matched fields in the MatchingField model?
+    We need a different approach, probably a cross join in StructuredRetriever.
+    """
+    leaves = filters.get_all_leaves()
+    if len(leaves) != 1:
+        return None
+
+    pf = leaves[0]
+
+    if isinstance(pf.condition, LtreeFilter):
+        op = pf.condition.op
+        # Prefer the original component/pattern (validator may set path="*" and move the value)
+        display = str(getattr(pf.condition, "value", "") or pf.path)
+
+        # There can be no match for abscence.
+        if op == FilterOp.NOT_HAS_COMPONENT:
+            return None
+
+        return MatchingField(text=display, path=display, highlight_indices=[(0, len(display))])
+
+    # Everything thats not Ltree
+    val = getattr(pf.condition, "value", "")
+    text = "" if val is None else str(val)
+    return MatchingField(text=text, path=pf.path, highlight_indices=[(0, len(text))])
+
+
+async def execute_search(
+    search_params: BaseSearchParameters,
+    db_session: Session,
+    pagination_params: PaginationParams | None = None,
+) -> SearchResponse:
+    """Execute a hybrid search and return ranked results.
+
+    Builds a candidate entity query based on the given search parameters,
+    applies the appropriate ranking strategy, and executes the final ranked
+    query to retrieve results.
+
+    Parameters
+    ----------
+    search_params : BaseSearchParameters
+        The search parameters specifying vector, fuzzy, or filter criteria.
+    db_session : Session
+        The active SQLAlchemy session for executing the query.
+    limit : int, optional
+        The maximum number of search results to return, by default 5.
+
+    Returns:
+    -------
+    SearchResponse
+        A list of `SearchResult` objects containing entity IDs, scores, and
+        optional highlight metadata.
+
+    Notes:
+    -----
+    If no vector query, filters, or fuzzy term are provided, a warning is logged
+    and an empty result set is returned.
+    """
+    if not search_params.vector_query and not search_params.filters and not search_params.fuzzy_term:
+        logger.warning("No search criteria provided (vector_query, fuzzy_term, or filters).")
+        return SearchResponse(results=[], metadata=SearchMetadata.empty())
+
+    candidate_query = build_candidate_query(search_params)
+
+    pagination_params = pagination_params or PaginationParams()
+    retriever = await Retriever.from_params(search_params, pagination_params)
+    logger.debug("Using retriever", retriever_type=retriever.__class__.__name__)
+
+    final_stmt = retriever.apply(candidate_query)
+    final_stmt = final_stmt.limit(search_params.limit)
+    logger.debug(final_stmt)
+    result = db_session.execute(final_stmt).mappings().all()
+
+    return _format_response(result, search_params, retriever.metadata)