orchestrator-core 4.4.2__py3-none-any.whl → 4.5.0__py3-none-any.whl

orchestrator/search/retrieval/engine.py

@@ -0,0 +1,154 @@
+# Copyright 2019-2025 SURF, GÉANT.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+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 .retrievers 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.
+
+    Args:
+        db_rows (Sequence[RowMapping]): The rows returned from the executed SQLAlchemy query.
+        search_params (BaseSearchParameters): The search parameters, including query text and filters.
+        metadata (SearchMetadata): Metadata about the search execution.
+
+    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 (text := row.get(Retriever.HIGHLIGHT_TEXT_LABEL))
+            and (path := row.get(Retriever.HIGHLIGHT_PATH_LABEL))
+        ):
+            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."""
+    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.
+
+    Args:
+        search_params (BaseSearchParameters): The search parameters specifying vector, fuzzy, or filter criteria.
+        db_session (Session): The active SQLAlchemy session for executing the query.
+        pagination_params (PaginationParams): Parameters controlling pagination of the search results.
+        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)

orchestrator/search/retrieval/exceptions.py

@@ -0,0 +1,90 @@
+# Copyright 2019-2025 SURF, GÉANT.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from orchestrator.search.core.types import FilterOp
+
+
+class FilterValidationError(Exception):
+    """Base exception for filter validation errors."""
+
+    pass
+
+
+class InvalidLtreePatternError(FilterValidationError):
+    """Raised when an ltree pattern has invalid ltree query syntax."""
+
+    def __init__(self, pattern: str) -> None:
+        message = f"Ltree pattern '{pattern}' has invalid syntax. Use valid PostgreSQL ltree lquery syntax."
+        super().__init__(message)
+
+
+class EmptyFilterPathError(FilterValidationError):
+    """Raised when a filter path is empty or contains only whitespace."""
+
+    def __init__(self) -> None:
+        message = (
+            "Filter path cannot be empty. Provide a valid path like 'subscription.product.name' or 'workflow.name'."
+        )
+        super().__init__(message)
+
+
+class PathNotFoundError(FilterValidationError):
+    """Raised when a filter path doesn't exist in the database schema.
+
+    Examples:
+        Using a non-existent filter path:
+
+        >>> print(PathNotFoundError('subscription.nonexistent.field'))
+        Path 'subscription.nonexistent.field' does not exist in the database.
+    """
+
+    def __init__(self, path: str) -> None:
+        message = f"Path '{path}' does not exist in the database."
+        super().__init__(message)
+
+
+class IncompatibleFilterTypeError(FilterValidationError):
+    """Raised when a filter operator is incompatible with the field's data type.
+
+    Examples:
+        Using a numeric comparison operator on a string field:
+
+        >>> print(IncompatibleFilterTypeError(
+        ...     operator='gt',
+        ...     field_type='string',
+        ...     path='subscription.customer_name',
+        ...     expected_operators=[FilterOp.EQ, FilterOp.NEQ, FilterOp.LIKE],
+        ... ))
+        Operator 'gt' is not compatible with field type 'string' for path 'subscription.customer_name'. Valid operators for 'string': [eq, neq, like]
+    """
+
+    def __init__(self, operator: str, field_type: str, path: str, expected_operators: list[FilterOp]) -> None:
+        valid_ops_str = ", ".join([op.value for op in expected_operators])
+        message = f"Operator '{operator}' is not compatible with field type '{field_type}' for path '{path}'. Valid operators for '{field_type}': [{valid_ops_str}]"
+
+        super().__init__(message)
+
+
+class InvalidEntityPrefixError(FilterValidationError):
+    """Raised when a filter path doesn't have the correct entity type prefix.
+
+    Examples:
+        Using wrong entity prefix in filter path:
+
+        >>> print(InvalidEntityPrefixError('workflow.name', 'subscription.', 'SUBSCRIPTION'))
+        Filter path 'workflow.name' must start with 'subscription.' for SUBSCRIPTION searches, or use '*' for wildcard paths.
+    """
+
+    def __init__(self, path: str, expected_prefix: str, entity_type: str) -> None:
+        message = f"Filter path '{path}' must start with '{expected_prefix}' for {entity_type} searches, or use '*' for wildcard paths."
+        super().__init__(message)

orchestrator/search/retrieval/pagination.py

@@ -0,0 +1,96 @@
+# Copyright 2019-2025 SURF, GÉANT.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import array
+import base64
+from dataclasses import dataclass
+
+from pydantic import BaseModel
+
+from orchestrator.search.core.exceptions import InvalidCursorError
+from orchestrator.search.schemas.parameters import BaseSearchParameters
+from orchestrator.search.schemas.results import SearchResult
+
+
+@dataclass
+class PaginationParams:
+    """Parameters for pagination in search queries."""
+
+    page_after_score: float | None = None
+    page_after_id: str | None = None
+    q_vec_override: list[float] | None = None
+
+
+def floats_to_b64(v: list[float]) -> str:
+    a = array.array("f", v)
+    return base64.urlsafe_b64encode(a.tobytes()).decode("ascii")
+
+
+def b64_to_floats(s: str) -> list[float]:
+    raw = base64.urlsafe_b64decode(s.encode("ascii"))
+    a = array.array("f")
+    a.frombytes(raw)
+    return list(a)
+
+
+class PageCursor(BaseModel):
+    score: float
+    id: str
+    q_vec_b64: str
+
+    def encode(self) -> str:
+        """Encode the cursor data into a URL-safe Base64 string."""
+        json_str = self.model_dump_json()
+        return base64.urlsafe_b64encode(json_str.encode("utf-8")).decode("utf-8")
+
+    @classmethod
+    def decode(cls, cursor: str) -> "PageCursor":
+        """Decode a Base64 string back into a PageCursor instance."""
+        try:
+            decoded_str = base64.urlsafe_b64decode(cursor).decode("utf-8")
+            return cls.model_validate_json(decoded_str)
+        except Exception as e:
+            raise InvalidCursorError("Invalid pagination cursor") from e
+
+
+async def process_pagination_cursor(cursor: str | None, search_params: BaseSearchParameters) -> PaginationParams:
+    """Process pagination cursor and return pagination parameters."""
+    if cursor:
+        c = PageCursor.decode(cursor)
+        return PaginationParams(
+            page_after_score=c.score,
+            page_after_id=c.id,
+            q_vec_override=b64_to_floats(c.q_vec_b64),
+        )
+    if search_params.vector_query:
+        from orchestrator.search.core.embedding import QueryEmbedder
+
+        q_vec_override = await QueryEmbedder.generate_for_text_async(search_params.vector_query)
+        return PaginationParams(q_vec_override=q_vec_override)
+    return PaginationParams()
+
+
+def create_next_page_cursor(
+    search_results: list[SearchResult], pagination_params: PaginationParams, limit: int
+) -> str | None:
+    """Create next page cursor if there are more results."""
+    has_next_page = len(search_results) == limit and limit > 0
+    if has_next_page:
+        last_item = search_results[-1]
+        cursor_data = PageCursor(
+            score=float(last_item.score),
+            id=last_item.entity_id,
+            q_vec_b64=floats_to_b64(pagination_params.q_vec_override or []),
+        )
+        return cursor_data.encode()
+    return None

orchestrator/search/retrieval/retrievers/__init__.py

@@ -0,0 +1,26 @@
+# Copyright 2019-2025 SURF, GÉANT.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .base import Retriever
+from .fuzzy import FuzzyRetriever
+from .hybrid import RrfHybridRetriever
+from .semantic import SemanticRetriever
+from .structured import StructuredRetriever
+
+__all__ = [
+    "Retriever",
+    "FuzzyRetriever",
+    "RrfHybridRetriever",
+    "SemanticRetriever",
+    "StructuredRetriever",
+]

orchestrator/search/retrieval/retrievers/base.py

@@ -0,0 +1,123 @@
+# Copyright 2019-2025 SURF, GÉANT.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from abc import ABC, abstractmethod
+from decimal import Decimal
+
+import structlog
+from sqlalchemy import BindParameter, Numeric, Select, literal
+
+from orchestrator.search.core.types import FieldType, SearchMetadata
+from orchestrator.search.schemas.parameters import BaseSearchParameters
+
+from ..pagination import PaginationParams
+
+logger = structlog.get_logger(__name__)
+
+
+class Retriever(ABC):
+    """Abstract base class for applying a ranking strategy to a search query."""
+
+    SCORE_PRECISION = 12
+    SCORE_NUMERIC_TYPE = Numeric(38, 12)
+    HIGHLIGHT_TEXT_LABEL = "highlight_text"
+    HIGHLIGHT_PATH_LABEL = "highlight_path"
+    SCORE_LABEL = "score"
+    SEARCHABLE_FIELD_TYPES = [
+        FieldType.STRING.value,
+        FieldType.UUID.value,
+        FieldType.BLOCK.value,
+        FieldType.RESOURCE_TYPE.value,
+    ]
+
+    @classmethod
+    async def from_params(
+        cls,
+        params: BaseSearchParameters,
+        pagination_params: PaginationParams,
+    ) -> "Retriever":
+        """Create the appropriate retriever instance from search parameters.
+
+        Args:
+            params (BaseSearchParameters): Search parameters including vector queries, fuzzy terms, and filters.
+            pagination_params (PaginationParams): Pagination parameters for cursor-based paging.
+
+        Returns:
+            Retriever: A concrete retriever instance (semantic, fuzzy, hybrid, or structured).
+        """
+
+        from .fuzzy import FuzzyRetriever
+        from .hybrid import RrfHybridRetriever
+        from .semantic import SemanticRetriever
+        from .structured import StructuredRetriever
+
+        fuzzy_term = params.fuzzy_term
+        q_vec = await cls._get_query_vector(params.vector_query, pagination_params.q_vec_override)
+
+        # If semantic search was attempted but failed, fall back to fuzzy with the full query
+        fallback_fuzzy_term = fuzzy_term
+        if q_vec is None and params.vector_query is not None and params.query is not None:
+            fallback_fuzzy_term = params.query
+
+        if q_vec is not None and fallback_fuzzy_term is not None:
+            return RrfHybridRetriever(q_vec, fallback_fuzzy_term, pagination_params)
+        if q_vec is not None:
+            return SemanticRetriever(q_vec, pagination_params)
+        if fallback_fuzzy_term is not None:
+            return FuzzyRetriever(fallback_fuzzy_term, pagination_params)
+
+        return StructuredRetriever(pagination_params)
+
+    @classmethod
+    async def _get_query_vector(
+        cls, vector_query: str | None, q_vec_override: list[float] | None
+    ) -> list[float] | None:
+        """Get query vector either from override or by generating from text."""
+        if q_vec_override:
+            return q_vec_override
+
+        if not vector_query:
+            return None
+
+        from orchestrator.search.core.embedding import QueryEmbedder
+
+        q_vec = await QueryEmbedder.generate_for_text_async(vector_query)
+        if not q_vec:
+            logger.warning("Embedding generation failed; using non-semantic retriever")
+            return None
+
+        return q_vec
+
+    @abstractmethod
+    def apply(self, candidate_query: Select) -> Select:
+        """Apply the ranking logic to the given candidate query.
+
+        Args:
+            candidate_query (Select): A SQLAlchemy `Select` statement returning candidate entity IDs.
+
+        Returns:
+            Select: A new `Select` statement with ranking expressions applied.
+        """
+        ...
+
+    def _quantize_score_for_pagination(self, score_value: float) -> BindParameter[Decimal]:
+        """Convert score value to properly quantized Decimal parameter for pagination."""
+        quantizer = Decimal(1).scaleb(-self.SCORE_PRECISION)
+        pas_dec = Decimal(str(score_value)).quantize(quantizer)
+        return literal(pas_dec, type_=self.SCORE_NUMERIC_TYPE)
+
+    @property
+    @abstractmethod
+    def metadata(self) -> SearchMetadata:
+        """Return metadata describing this search strategy."""
+        ...

orchestrator/search/retrieval/retrievers/fuzzy.py

@@ -0,0 +1,94 @@
+# Copyright 2019-2025 SURF, GÉANT.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from sqlalchemy import Select, and_, cast, func, literal, or_, select
+from sqlalchemy.sql.expression import ColumnElement
+
+from orchestrator.db.models import AiSearchIndex
+from orchestrator.search.core.types import SearchMetadata
+
+from ..pagination import PaginationParams
+from .base import Retriever
+
+
+class FuzzyRetriever(Retriever):
+    """Ranks results based on the max of fuzzy text similarity scores."""
+
+    def __init__(self, fuzzy_term: str, pagination_params: PaginationParams) -> None:
+        self.fuzzy_term = fuzzy_term
+        self.page_after_score = pagination_params.page_after_score
+        self.page_after_id = pagination_params.page_after_id
+
+    def apply(self, candidate_query: Select) -> Select:
+        cand = candidate_query.subquery()
+
+        similarity_expr = func.word_similarity(self.fuzzy_term, AiSearchIndex.value)
+
+        raw_max = func.max(similarity_expr).over(partition_by=AiSearchIndex.entity_id)
+        score = cast(
+            func.round(cast(raw_max, self.SCORE_NUMERIC_TYPE), self.SCORE_PRECISION), self.SCORE_NUMERIC_TYPE
+        ).label(self.SCORE_LABEL)
+
+        combined_query = (
+            select(
+                AiSearchIndex.entity_id,
+                score,
+                func.first_value(AiSearchIndex.value)
+                .over(partition_by=AiSearchIndex.entity_id, order_by=[similarity_expr.desc(), AiSearchIndex.path.asc()])
+                .label(self.HIGHLIGHT_TEXT_LABEL),
+                func.first_value(AiSearchIndex.path)
+                .over(partition_by=AiSearchIndex.entity_id, order_by=[similarity_expr.desc(), AiSearchIndex.path.asc()])
+                .label(self.HIGHLIGHT_PATH_LABEL),
+            )
+            .select_from(AiSearchIndex)
+            .join(cand, cand.c.entity_id == AiSearchIndex.entity_id)
+            .where(
+                and_(
+                    AiSearchIndex.value_type.in_(self.SEARCHABLE_FIELD_TYPES),
+                    literal(self.fuzzy_term).op("<%")(AiSearchIndex.value),
+                )
+            )
+            .distinct(AiSearchIndex.entity_id)
+        )
+        final_query = combined_query.subquery("ranked_fuzzy")
+
+        stmt = select(
+            final_query.c.entity_id,
+            final_query.c.score,
+            final_query.c.highlight_text,
+            final_query.c.highlight_path,
+        ).select_from(final_query)
+
+        stmt = self._apply_score_pagination(stmt, final_query.c.score, final_query.c.entity_id)
+
+        return stmt.order_by(final_query.c.score.desc().nulls_last(), final_query.c.entity_id.asc())
+
+    @property
+    def metadata(self) -> SearchMetadata:
+        return SearchMetadata.fuzzy()
+
+    def _apply_score_pagination(
+        self, stmt: Select, score_column: ColumnElement, entity_id_column: ColumnElement
+    ) -> Select:
+        """Apply standard score + entity_id pagination."""
+        if self.page_after_score is not None and self.page_after_id is not None:
+            stmt = stmt.where(
+                or_(
+                    score_column < self.page_after_score,
+                    and_(
+                        score_column == self.page_after_score,
+                        entity_id_column > self.page_after_id,
+                    ),
+                )
+            )
+        return stmt