orchestrator-core 4.4.0rc2__py3-none-any.whl → 5.0.0a1__py3-none-any.whl

orchestrator/search/retrieval/builder.py

@@ -0,0 +1,64 @@
+from sqlalchemy import Select, String, cast, func, select
+
+from orchestrator.db.models import AiSearchIndex
+from orchestrator.search.core.types import EntityType, FilterOp
+from orchestrator.search.filters import LtreeFilter
+from orchestrator.search.schemas.parameters import BaseSearchParameters
+
+
+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."""
+    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=""))
+
+    if q:
+        score = func.similarity(cast(AiSearchIndex.path, String), q).label("score")
+        stmt = (
+            stmt.add_columns(score)
+            .group_by(AiSearchIndex.path, AiSearchIndex.value_type, score)
+            .order_by(score.desc(), AiSearchIndex.path)
+        )
+    else:
+        stmt = stmt.group_by(AiSearchIndex.path, AiSearchIndex.value_type).order_by(AiSearchIndex.path)
+
+    return stmt

orchestrator/search/retrieval/engine.py

@@ -0,0 +1,96 @@
+from collections.abc import Sequence
+
+import structlog
+from sqlalchemy.engine.row import RowMapping
+from sqlalchemy.orm import Session
+
+from orchestrator.search.schemas.parameters import BaseSearchParameters
+from orchestrator.search.schemas.results import Highlight, SearchResponse, SearchResult
+
+from .builder import build_candidate_query
+from .ranker import Ranker
+from .utils import generate_highlight_indices
+
+logger = structlog.get_logger(__name__)
+
+
+def _format_response(db_rows: Sequence[RowMapping], search_params: BaseSearchParameters) -> SearchResponse:
+    """Format database query results into a `SearchResponse`.
+
+    Converts raw SQLAlchemy `RowMapping` objects into `SearchResult` instances,
+    optionally generating highlight metadata if a fuzzy term is present.
+
+    Parameters
+    ----------
+    db_rows : Sequence[RowMapping]
+        The rows returned from the executed SQLAlchemy query.
+    search_params : BaseSearchParameters
+        The parameters used for the search, including any fuzzy term for highlighting.
+
+    Returns:
+    -------
+    SearchResponse
+        A list of `SearchResult` objects containing entity IDs, scores, and
+        optional highlight information.
+    """
+    response: SearchResponse = []
+    for row in db_rows:
+        highlight = None
+        if search_params.fuzzy_term and row.get("highlight_text"):
+            text = row.highlight_text
+            indices = generate_highlight_indices(text, search_params.fuzzy_term)
+            if indices:
+                highlight = Highlight(text=text, indices=indices)
+
+        response.append(
+            SearchResult(
+                entity_id=str(row.entity_id),
+                score=row.score,
+                highlight=highlight,
+            )
+        )
+    return response
+
+
+async def execute_search(search_params: BaseSearchParameters, db_session: Session, limit: int = 5) -> 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 []
+
+    candidate_query = build_candidate_query(search_params)
+
+    ranker = await Ranker.from_params(search_params)
+    logger.debug("Using ranker", ranker_type=ranker.__class__.__name__)
+
+    final_stmt = ranker.apply(candidate_query)
+    final_stmt = final_stmt.limit(limit)
+    logger.debug(final_stmt)
+    result = db_session.execute(final_stmt).mappings().all()
+
+    return _format_response(result, search_params)

orchestrator/search/retrieval/ranker.py

@@ -0,0 +1,202 @@
+from abc import ABC, abstractmethod
+from typing import Any
+
+import structlog
+from sqlalchemy import Select, bindparam, case, func, literal, select
+from sqlalchemy.sql.expression import ColumnElement
+
+from orchestrator.db.models import AiSearchIndex
+from orchestrator.search.core.embedding import QueryEmbedder
+from orchestrator.search.schemas.parameters import BaseSearchParameters
+
+logger = structlog.get_logger(__name__)
+Index = AiSearchIndex
+
+
+class Ranker(ABC):
+    """Abstract base class for applying a ranking strategy to a search query."""
+
+    @classmethod
+    async def from_params(cls, params: BaseSearchParameters, use_rrf: bool = True) -> "Ranker":
+        """Create the appropriate ranker instance from search parameters.
+
+        Parameters
+        ----------
+        params : BaseSearchParameters
+            Search parameters including vector queries, fuzzy terms, and filters.
+        use_rrf : bool, optional
+            Whether to use Reciprocal Rank Fusion for hybrid searches, by default True.
+
+        Returns:
+        -------
+        Ranker
+            A concrete ranker instance (semantic, fuzzy, hybrid, RRF hybrid, or structured).
+        """
+        vq, fq = params.vector_query, params.fuzzy_term
+        q_vec = None
+        if vq:
+            q_vec = await QueryEmbedder.generate_for_text_async(vq)
+            if not q_vec:
+                logger.warning("Embedding generation failed; using non-semantic ranker")
+                vq = None
+
+        if vq and fq and q_vec is not None:
+            return RrfHybridRanker(q_vec, fq) if use_rrf else HybridRanker(q_vec, fq)
+        if vq and q_vec is not None:
+            return SemanticRanker(q_vec)
+        if fq:
+            return FuzzyRanker(fq)
+        return StructuredRanker()
+
+    @abstractmethod
+    def apply(self, candidate_query: Select) -> Select:
+        """Apply the ranking logic to the given candidate query.
+
+        Parameters
+        ----------
+        candidate_query : Select
+            A SQLAlchemy `Select` statement returning candidate entity IDs.
+
+        Returns:
+        -------
+        Select
+            A new `Select` statement with ranking expressions applied.
+        """
+        ...
+
+
+class StructuredRanker(Ranker):
+    """Applies a dummy score for purely structured searches with no text query."""
+
+    def apply(self, candidate_query: Select) -> Select:
+        cand = candidate_query.subquery()
+        return select(cand.c.entity_id, literal(1.0).label("score")).select_from(cand).order_by(cand.c.entity_id.asc())
+
+
+class FuzzyRanker(Ranker):
+    """Ranks results based on the max of fuzzy text similarity scores."""
+
+    def __init__(self, fuzzy_term: str) -> None:
+        self.fuzzy_term = fuzzy_term
+
+    def apply(self, candidate_query: Select) -> Select:
+
+        cand = candidate_query.subquery()
+        score_expr = func.max(func.similarity(Index.value, self.fuzzy_term))
+
+        return (
+            select(Index.entity_id, score_expr.label("score"))
+            .select_from(Index)
+            .join(cand, cand.c.entity_id == Index.entity_id)
+            .group_by(Index.entity_id)
+            .order_by(score_expr.desc().nulls_last(), Index.entity_id.asc())
+        )
+
+
+class SemanticRanker(Ranker):
+    """Ranks results based on the minimum semantic vector distance."""
+
+    def __init__(self, vector_query: list[float]) -> None:
+        self.vector_query = vector_query
+
+    def apply(self, candidate_query: Select) -> Select:
+        cand = candidate_query.subquery()
+
+        dist = Index.embedding.l2_distance(self.vector_query)
+        score_expr = func.min(dist).label("score")
+
+        return (
+            select(Index.entity_id, score_expr)
+            .select_from(Index)
+            .join(cand, cand.c.entity_id == Index.entity_id)
+            .where(Index.embedding.isnot(None))
+            .group_by(Index.entity_id)
+            .order_by(score_expr.asc().nulls_last(), Index.entity_id.asc())
+        )
+
+
+class HybridRanker(Ranker):
+    """Ranks results by combining semantic distance and fuzzy similarity.
+
+    Prioritizes fuzzy score, using semantic score as a tie-breaker.
+    """
+
+    def __init__(self, q_vec: list[float], fuzzy_term: str) -> None:
+        self.q_vec = q_vec
+        self.fuzzy_term = fuzzy_term
+
+    def apply(self, candidate_query: Select) -> Select:
+        cand = candidate_query.subquery()
+
+        dist = Index.embedding.l2_distance(self.q_vec)
+        # Semantic: only consider rows where an embedding exists
+        sem_agg = func.min(dist).filter(Index.embedding.isnot(None))
+        # Fuzzy: consider all rows (strings, uuids, etc.)
+        fuzzy_agg = func.max(func.similarity(Index.value, self.fuzzy_term))
+
+        score = sem_agg.label("score")
+
+        return (
+            select(Index.entity_id, score)
+            .select_from(Index)
+            .join(cand, cand.c.entity_id == Index.entity_id)
+            .group_by(Index.entity_id)
+            .order_by(
+                fuzzy_agg.desc().nulls_last(),
+                sem_agg.asc().nulls_last(),
+                Index.entity_id.asc(),
+            )
+        )
+
+
+class RrfHybridRanker(Ranker):
+    """Reciprocal Rank Fusion of semantic and fuzzy ranking."""
+
+    def __init__(self, q_vec: list[float], fuzzy_term: str, k: int = 60) -> None:
+        self.q_vec = q_vec
+        self.fuzzy_term = fuzzy_term
+        self.k = k
+
+    def apply(self, candidate_query: Select) -> Select:
+        cand = candidate_query.subquery()
+
+        # centroid over rows that have embeddings
+        q_param: ColumnElement[Any] = bindparam("q_vec", self.q_vec, type_=Index.embedding.type)
+        avg_vec = func.avg(Index.embedding).filter(Index.embedding.isnot(None))
+        sem_dist = avg_vec.op("<->")(q_param)
+
+        # fuzzy over ALL rows, substring-friendly for partial UUIDs
+        sim_base = func.similarity(Index.value, self.fuzzy_term)
+        sim_word = func.word_similarity(self.fuzzy_term, Index.value)
+        fuzzy_agg = func.max(func.greatest(sim_base, sim_word))
+
+        scores = (
+            select(
+                Index.entity_id,
+                sem_dist.label("semantic_distance"),
+                fuzzy_agg.label("fuzzy_score"),
+            )
+            .select_from(Index)
+            .join(cand, cand.c.entity_id == Index.entity_id)
+            .group_by(Index.entity_id)
+            .cte("scores")
+        )
+
+        ranked = select(
+            scores.c.entity_id,
+            scores.c.semantic_distance,
+            scores.c.fuzzy_score,
+            func.dense_rank().over(order_by=scores.c.semantic_distance.asc().nulls_last()).label("sem_rank"),
+            func.dense_rank().over(order_by=scores.c.fuzzy_score.desc().nulls_last()).label("fuzzy_rank"),
+        ).cte("ranked_results")
+
+        rrf = (1.0 / (self.k + ranked.c.sem_rank)) + (1.0 / (self.k + ranked.c.fuzzy_rank))
+        score_expr = rrf.label("score")
+        perfect = case((ranked.c.fuzzy_score >= 0.9, 0), else_=1)
+
+        return (
+            select(ranked.c.entity_id, score_expr)
+            .select_from(ranked)
+            .order_by(perfect.asc(), score_expr.desc(), ranked.c.entity_id.asc())
+            .params(q_vec=self.q_vec)
+        )

orchestrator/search/retrieval/utils.py

@@ -0,0 +1,88 @@
+import json
+
+import structlog
+from sqlalchemy import and_
+from sqlalchemy_utils.types.ltree import Ltree
+
+from orchestrator.db.database import WrappedSession
+from orchestrator.db.models import AiSearchIndex
+from orchestrator.search.core.types import EntityType
+from orchestrator.search.indexing.registry import ENTITY_CONFIG_REGISTRY
+from orchestrator.search.schemas.parameters import BaseSearchParameters
+from orchestrator.search.schemas.results import SearchResult
+
+logger = structlog.get_logger(__name__)
+
+
+def generate_highlight_indices(text: str, term: str) -> list[tuple[int, int]]:
+    if not text or not term:
+        return []
+    indices = []
+    start = text.lower().find(term.lower())
+    if start != -1:
+        end = start + len(term)
+        indices.append((start, end))
+    return indices
+
+
+def display_filtered_paths_only(
+    results: list[SearchResult], search_params: BaseSearchParameters, db_session: WrappedSession
+) -> None:
+    """Display only the paths that were searched for in the results."""
+    if not results:
+        logger.info("No results found.")
+        return
+
+    logger.info("--- Search Results ---")
+
+    searched_paths = search_params.filters.get_all_paths() if search_params.filters else []
+    if not searched_paths:
+        return
+
+    for result in results:
+        for path in searched_paths:
+            record: AiSearchIndex | None = (
+                db_session.query(AiSearchIndex)
+                .filter(and_(AiSearchIndex.entity_id == result.entity_id, AiSearchIndex.path == Ltree(path)))
+                .first()
+            )
+
+            if record:
+                logger.info(f"  {record.path}: {record.value}")
+
+        logger.info("-" * 40)
+
+
+def display_results(
+    results: list[SearchResult],
+    db_session: WrappedSession,
+    score_label: str = "Score",
+) -> None:
+    """Finds the original DB record for each search result and logs its traversed fields."""
+    if not results:
+        logger.info("No results found.")
+        return
+
+    logger.info("--- Search Results ---")
+    for result in results:
+        entity_id = result.entity_id
+        score = result.score
+
+        index_records = db_session.query(AiSearchIndex).filter(AiSearchIndex.entity_id == entity_id).all()
+        if not index_records:
+            logger.warning(f"Could not find indexed records for entity_id={entity_id}")
+            continue
+
+        first_record = index_records[0]
+        kind = EntityType(first_record.entity_type)
+        config = ENTITY_CONFIG_REGISTRY[kind]
+
+        db_entity = db_session.get(config.table, entity_id) if config.table else None
+
+        if db_entity and config.traverser:
+            fields = config.traverser.get_fields(db_entity, config.pk_name, config.root_name)
+            result_obj = {p: v for p, v, _ in fields}
+            logger.info(json.dumps(result_obj, indent=2, default=str))
+            logger.info(f"{score_label}: {score:.4f}\n" + "-" * 20)
+        else:
+            logger.warning(f"Could not display entity {kind.value} with id={entity_id}")

orchestrator/search/retrieval/validation.py

@@ -0,0 +1,174 @@
+from typing import assert_never
+
+from sqlalchemy import select, text
+from sqlalchemy.exc import ProgrammingError
+from sqlalchemy_utils import Ltree
+
+from orchestrator.db import db
+from orchestrator.db.database import WrappedSession
+from orchestrator.db.models import AiSearchIndex
+from orchestrator.search.core.types import EntityType, FieldType
+from orchestrator.search.filters import (
+    DateRangeFilter,
+    DateValueFilter,
+    EqualityFilter,
+    FilterCondition,
+    FilterTree,
+    LtreeFilter,
+    NumericRangeFilter,
+    NumericValueFilter,
+    PathFilter,
+    StringFilter,
+)
+
+
+def is_filter_compatible_with_field_type(filter_condition: FilterCondition, field_type: FieldType) -> bool:
+    """Check whether a filter condition is compatible with a given field type.
+
+    Parameters
+    ----------
+    filter_condition : FilterCondition
+        The filter condition instance to check.
+    field_type : FieldType
+        The type of field from the index schema.
+
+    Returns:
+    -------
+    bool
+        True if the filter condition is valid for the given field type, False otherwise.
+    """
+
+    match filter_condition:
+        case LtreeFilter():
+            return True  # Filters for path only
+        case DateRangeFilter() | DateValueFilter():
+            return field_type == FieldType.DATETIME
+        case NumericRangeFilter() | NumericValueFilter():
+            return field_type in {FieldType.INTEGER, FieldType.FLOAT}
+        case StringFilter():
+            return field_type == FieldType.STRING
+        case EqualityFilter():
+            return field_type in {
+                FieldType.BOOLEAN,
+                FieldType.UUID,
+                FieldType.BLOCK,
+                FieldType.RESOURCE_TYPE,
+                FieldType.STRING,
+            }
+        case _:
+            assert_never(filter_condition)
+
+
+def is_lquery_syntactically_valid(pattern: str, db_session: WrappedSession) -> bool:
+    """Validate whether a string is a syntactically correct `lquery` pattern.
+
+    Parameters
+    ----------
+    pattern : str
+        The LTree lquery pattern string to validate.
+    db_session : WrappedSession
+        The database session used to test casting.
+
+    Returns:
+    -------
+    bool
+        True if the pattern is valid, False if it fails to cast in PostgreSQL.
+    """
+    try:
+        with db_session.begin_nested():
+            db_session.execute(text("SELECT CAST(:pattern AS lquery)"), {"pattern": pattern})
+        return True
+    except ProgrammingError:
+        return False
+
+
+def get_structured_filter_schema() -> dict[str, str]:
+    """Retrieve all distinct filterable paths and their field types from the index.
+
+    Returns:
+    -------
+    Dict[str, str]
+        Mapping of path strings to their corresponding field type values.
+    """
+    stmt = select(AiSearchIndex.path, AiSearchIndex.value_type).distinct().order_by(AiSearchIndex.path)
+    result = db.session.execute(stmt)
+    return {str(path): value_type.value for path, value_type in result}
+
+
+def validate_filter_path(path: str) -> str | None:
+    """Check if a given path exists in the index and return its field type.
+
+    Parameters
+    ----------
+    path : str
+        The fully qualified LTree path.
+
+    Returns:
+    -------
+    Optional[str]
+        The value type of the field if found, otherwise None.
+    """
+    stmt = select(AiSearchIndex.value_type).where(AiSearchIndex.path == Ltree(path)).limit(1)
+    result = db.session.execute(stmt).scalar_one_or_none()
+    return result.value if result else None
+
+
+async def complete_filter_validation(filter: PathFilter, entity_type: EntityType) -> None:
+    """Validate a PathFilter against the database schema and entity type.
+
+    Checks performed:
+    1. LTree filter syntax (for LtreeFilter only)
+    2. Non-empty path
+    3. Path exists in the database schema
+    4. Filter type matches the field's value_type
+    5. Path starts with the correct entity type prefix (unless wildcard)
+
+    Parameters
+    ----------
+    filter : PathFilter
+        The filter to validate.
+    entity_type : EntityType
+        The entity type being searched.
+
+    Raises:
+    ------
+    ValueError
+        If any of the validation checks fail.
+    """
+    # Ltree is a special case
+    if isinstance(filter.condition, LtreeFilter):
+        lquery_pattern = filter.condition.value
+        if not is_lquery_syntactically_valid(lquery_pattern, db.session):
+            raise ValueError(f"Ltree pattern '{lquery_pattern}' has invalid syntax.")
+        return
+
+    if not filter.path or not filter.path.strip():
+        raise ValueError("Filter path cannot be empty")
+
+    # 1. Check if path exists in database
+    db_field_type_str = validate_filter_path(filter.path)
+    if db_field_type_str is None:
+        raise ValueError(f"Path '{filter.path}' does not exist in database schema")
+
+    db_field_type = FieldType(db_field_type_str)
+
+    # 2. Check filter compatibility with field type
+    if not is_filter_compatible_with_field_type(filter.condition, db_field_type):
+        raise ValueError(
+            f"Filter '{type(filter.condition).__name__}' not compatible with field type '{db_field_type.value}'"
+        )
+
+    # 3. Check entity type prefix requirements (unless it's a wildcard path)
+    expected_prefix = f"{entity_type.value.lower()}."
+    if not filter.path.startswith(expected_prefix) and not filter.path.startswith("*"):
+        raise ValueError(
+            f"Filter path '{filter.path}' must start with '{expected_prefix}' for {entity_type.value} searches."
+        )
+
+
+async def validate_filter_tree(filters: FilterTree | None, entity_type: EntityType) -> None:
+    """Validate all PathFilter leaves in a FilterTree."""
+    if filters is None:
+        return
+    for leaf in filters.get_all_leaves():
+        await complete_filter_validation(leaf, entity_type)