orchestrator-core 4.4.1__py3-none-any.whl → 4.5.0a3__py3-none-any.whl

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,122 @@
+# 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."""
+        pas_dec = Decimal(str(score_value)).quantize(Decimal("0.000000000001"))
+        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

orchestrator/search/retrieval/retrievers/hybrid.py

@@ -0,0 +1,188 @@
+# 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 BindParameter, Select, and_, bindparam, case, 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 RrfHybridRetriever(Retriever):
+    """Reciprocal Rank Fusion of semantic and fuzzy ranking with parent-child retrieval."""
+
+    def __init__(
+        self,
+        q_vec: list[float],
+        fuzzy_term: str,
+        pagination_params: PaginationParams,
+        k: int = 60,
+        field_candidates_limit: int = 100,
+    ) -> None:
+        self.q_vec = q_vec
+        self.fuzzy_term = fuzzy_term
+        self.page_after_score = pagination_params.page_after_score
+        self.page_after_id = pagination_params.page_after_id
+        self.k = k
+        self.field_candidates_limit = field_candidates_limit
+
+    def apply(self, candidate_query: Select) -> Select:
+        cand = candidate_query.subquery()
+        q_param: BindParameter[list[float]] = bindparam("q_vec", self.q_vec, type_=AiSearchIndex.embedding.type)
+
+        best_similarity = func.word_similarity(self.fuzzy_term, AiSearchIndex.value)
+        sem_expr = case(
+            (AiSearchIndex.embedding.is_(None), None),
+            else_=AiSearchIndex.embedding.op("<->")(q_param),
+        )
+        sem_val = func.coalesce(sem_expr, literal(1.0)).label("semantic_distance")
+
+        filter_condition = literal(self.fuzzy_term).op("<%")(AiSearchIndex.value)
+
+        field_candidates = (
+            select(
+                AiSearchIndex.entity_id,
+                AiSearchIndex.path,
+                AiSearchIndex.value,
+                sem_val,
+                best_similarity.label("fuzzy_score"),
+            )
+            .select_from(AiSearchIndex)
+            .join(cand, cand.c.entity_id == AiSearchIndex.entity_id)
+            .where(
+                and_(
+                    AiSearchIndex.value_type.in_(self.SEARCHABLE_FIELD_TYPES),
+                    filter_condition,
+                )
+            )
+            .order_by(
+                best_similarity.desc().nulls_last(),
+                sem_expr.asc().nulls_last(),
+                AiSearchIndex.entity_id.asc(),
+            )
+            .limit(self.field_candidates_limit)
+        ).cte("field_candidates")
+
+        entity_scores = (
+            select(
+                field_candidates.c.entity_id,
+                func.avg(field_candidates.c.semantic_distance).label("avg_semantic_distance"),
+                func.avg(field_candidates.c.fuzzy_score).label("avg_fuzzy_score"),
+            ).group_by(field_candidates.c.entity_id)
+        ).cte("entity_scores")
+
+        entity_highlights = (
+            select(
+                field_candidates.c.entity_id,
+                func.first_value(field_candidates.c.value)
+                .over(
+                    partition_by=field_candidates.c.entity_id,
+                    order_by=[field_candidates.c.fuzzy_score.desc(), field_candidates.c.path.asc()],
+                )
+                .label(self.HIGHLIGHT_TEXT_LABEL),
+                func.first_value(field_candidates.c.path)
+                .over(
+                    partition_by=field_candidates.c.entity_id,
+                    order_by=[field_candidates.c.fuzzy_score.desc(), field_candidates.c.path.asc()],
+                )
+                .label(self.HIGHLIGHT_PATH_LABEL),
+            ).distinct(field_candidates.c.entity_id)
+        ).cte("entity_highlights")
+
+        ranked = (
+            select(
+                entity_scores.c.entity_id,
+                entity_scores.c.avg_semantic_distance,
+                entity_scores.c.avg_fuzzy_score,
+                entity_highlights.c.highlight_text,
+                entity_highlights.c.highlight_path,
+                func.dense_rank()
+                .over(
+                    order_by=[entity_scores.c.avg_semantic_distance.asc().nulls_last(), entity_scores.c.entity_id.asc()]
+                )
+                .label("sem_rank"),
+                func.dense_rank()
+                .over(order_by=[entity_scores.c.avg_fuzzy_score.desc().nulls_last(), entity_scores.c.entity_id.asc()])
+                .label("fuzzy_rank"),
+            ).select_from(
+                entity_scores.join(entity_highlights, entity_scores.c.entity_id == entity_highlights.c.entity_id)
+            )
+        ).cte("ranked_results")
+
+        # RRF (rank-based)
+        rrf_raw = (1.0 / (self.k + ranked.c.sem_rank)) + (1.0 / (self.k + ranked.c.fuzzy_rank))
+        rrf_num = cast(rrf_raw, self.SCORE_NUMERIC_TYPE)
+
+        # Perfect flag to boost near perfect fuzzy matches as this most likely indicates the desired record.
+        perfect = case((ranked.c.avg_fuzzy_score >= 0.9, 1), else_=0).label("perfect_match")
+
+        # Dynamic beta based on k (and number of sources)
+        # rrf_max = n_sources / (k + 1)
+        k_num = literal(float(self.k), type_=self.SCORE_NUMERIC_TYPE)
+        n_sources = literal(2.0, type_=self.SCORE_NUMERIC_TYPE)  # semantic + fuzzy
+        rrf_max = n_sources / (k_num + literal(1.0, type_=self.SCORE_NUMERIC_TYPE))
+
+        # Choose a small positive margin above rrf_max to ensure strict separation
+        # Keep it small to avoid compressing perfects near 1 after normalization
+        margin = rrf_max * literal(0.05, type_=self.SCORE_NUMERIC_TYPE)  # 5% above bound
+        beta = rrf_max + margin
+
+        fused_num = rrf_num + beta * cast(perfect, self.SCORE_NUMERIC_TYPE)
+
+        # Normalize to [0,1] via the theoretical max (beta + rrf_max)
+        norm_den = beta + rrf_max
+        normalized_score = fused_num / norm_den
+
+        score = cast(
+            func.round(cast(normalized_score, self.SCORE_NUMERIC_TYPE), self.SCORE_PRECISION),
+            self.SCORE_NUMERIC_TYPE,
+        ).label(self.SCORE_LABEL)
+
+        stmt = select(
+            ranked.c.entity_id,
+            score,
+            ranked.c.highlight_text,
+            ranked.c.highlight_path,
+            perfect.label("perfect_match"),
+        ).select_from(ranked)
+
+        stmt = self._apply_fused_pagination(stmt, score, ranked.c.entity_id)
+
+        return stmt.order_by(
+            score.desc().nulls_last(),
+            ranked.c.entity_id.asc(),
+        ).params(q_vec=self.q_vec)
+
+    def _apply_fused_pagination(
+        self,
+        stmt: Select,
+        score_column: ColumnElement,
+        entity_id_column: ColumnElement,
+    ) -> Select:
+        """Keyset paginate by fused score + id."""
+        if self.page_after_score is not None and self.page_after_id is not None:
+            score_param = self._quantize_score_for_pagination(self.page_after_score)
+            stmt = stmt.where(
+                or_(
+                    score_column < score_param,
+                    and_(score_column == score_param, entity_id_column > self.page_after_id),
+                )
+            )
+        return stmt
+
+    @property
+    def metadata(self) -> SearchMetadata:
+        return SearchMetadata.hybrid()