orchestrator-core 4.6.2__py3-none-any.whl → 4.6.3rc1__py3-none-any.whl

orchestrator/search/aggregations/base.py

@@ -0,0 +1,201 @@
+# 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 abstractmethod
+from enum import Enum
+from typing import Annotated, Any, Literal, TypeAlias
+
+from pydantic import BaseModel, ConfigDict, Field
+from sqlalchemy import Integer, cast, func
+from sqlalchemy.sql.elements import ColumnElement, Label
+
+
+class AggregationType(str, Enum):
+    """Types of aggregations that can be computed."""
+
+    COUNT = "count"
+    SUM = "sum"
+    AVG = "avg"
+    MIN = "min"
+    MAX = "max"
+
+
+class TemporalPeriod(str, Enum):
+    """Time periods for temporal grouping."""
+
+    YEAR = "year"
+    QUARTER = "quarter"
+    MONTH = "month"
+    WEEK = "week"
+    DAY = "day"
+    HOUR = "hour"
+
+
+class TemporalGrouping(BaseModel):
+    """Defines temporal grouping for date/time fields.
+
+    Used to group query results by time periods (e.g., group subscriptions by start_date per month).
+    """
+
+    field: str = Field(description="The datetime field path to group by temporally.")
+    period: TemporalPeriod = Field(description="The time period to group by (year, quarter, month, week, day, hour).")
+
+    model_config = ConfigDict(
+        extra="forbid",
+        json_schema_extra={
+            "examples": [
+                {"field": "subscription.start_date", "period": "month"},
+                {"field": "subscription.end_date", "period": "year"},
+                {"field": "process.created_at", "period": "day"},
+            ]
+        },
+    )
+
+    def get_pivot_fields(self) -> list[str]:
+        """Return fields that need to be pivoted for this temporal grouping."""
+        return [self.field]
+
+    def to_expression(self, pivot_cte_columns: Any) -> tuple[Label, Any, str]:
+        """Build SQLAlchemy expression for temporal grouping.
+
+        Args:
+            pivot_cte_columns: The columns object from the pivot CTE
+
+        Returns:
+            tuple: (select_column, group_by_column, column_name)
+                - select_column: Labeled column for SELECT
+                - group_by_column: Column expression for GROUP BY
+                - column_name: The label/name of the column in results
+        """
+        from sqlalchemy import TIMESTAMP, cast, func
+
+        field_alias = BaseAggregation.field_to_alias(self.field)
+        col = getattr(pivot_cte_columns, field_alias)
+        truncated_col = func.date_trunc(self.period.value, cast(col, TIMESTAMP(timezone=True)))
+
+        # Column name without prefix
+        col_name = f"{field_alias}_{self.period.value}"
+        select_col = truncated_col.label(col_name)
+        return select_col, truncated_col, col_name
+
+
+class BaseAggregation(BaseModel):
+    """Base class for all aggregation types."""
+
+    type: AggregationType = Field(description="The type of aggregation to perform.")
+    alias: str = Field(description="The name for this aggregation in the results.")
+
+    @classmethod
+    def create(cls, data: dict) -> "Aggregation":
+        """Create the correct aggregation instance based on type field.
+
+        Args:
+            data: Dictionary with aggregation data including 'type' discriminator
+
+        Returns:
+            Validated aggregation instance (CountAggregation or FieldAggregation)
+
+        Raises:
+            ValidationError: If data is invalid or type is unknown
+        """
+        from pydantic import TypeAdapter
+
+        adapter: TypeAdapter = TypeAdapter(Aggregation)
+        return adapter.validate_python(data)
+
+    @staticmethod
+    def field_to_alias(field_path: str) -> str:
+        """Convert field path to SQL column alias.
+
+        Examples:
+            'subscription.name' -> 'subscription_name'
+            'product.serial-number' -> 'product_serial_number'
+        """
+        return field_path.replace(".", "_").replace("-", "_")
+
+    def get_pivot_fields(self) -> list[str]:
+        """Return fields that need to be pivoted for this aggregation."""
+        return []
+
+    @abstractmethod
+    def to_expression(self, *args: Any, **kwargs: Any) -> Label:
+        """Build SQLAlchemy expression for this aggregation.
+
+        Returns:
+            Label: A labeled SQLAlchemy expression
+        """
+        raise NotImplementedError
+
+
+class CountAggregation(BaseAggregation):
+    """Count aggregation - counts number of entities."""
+
+    type: Literal[AggregationType.COUNT]
+
+    def to_expression(self, entity_id_column: ColumnElement) -> Label:
+        """Build SQLAlchemy expression for count aggregation.
+
+        Args:
+            entity_id_column: The entity_id column from the pivot CTE
+
+        Returns:
+            Label: A labeled SQLAlchemy expression
+        """
+        return func.count(entity_id_column).label(self.alias)
+
+
+class FieldAggregation(BaseAggregation):
+    """Field-based aggregation (sum, avg, min, max)."""
+
+    type: Literal[AggregationType.SUM, AggregationType.AVG, AggregationType.MIN, AggregationType.MAX]
+    field: str = Field(description="The field path to aggregate on.")
+
+    def get_pivot_fields(self) -> list[str]:
+        """Return fields that need to be pivoted for this aggregation."""
+        return [self.field]
+
+    def to_expression(self, pivot_cte_columns: Any) -> Label:
+        """Build SQLAlchemy expression for field-based aggregation.
+
+        Args:
+            pivot_cte_columns: The columns object from the pivot CTE
+
+        Returns:
+            Label: A labeled SQLAlchemy expression
+
+        Raises:
+            ValueError: If the field is not found in the pivot CTE
+        """
+        field_alias = self.field_to_alias(self.field)
+
+        if not hasattr(pivot_cte_columns, field_alias):
+            raise ValueError(f"Field '{self.field}' (alias: '{field_alias}') not found in pivot CTE columns")
+
+        col = getattr(pivot_cte_columns, field_alias)
+
+        numeric_col = cast(col, Integer)
+
+        match self.type:
+            case AggregationType.SUM:
+                return func.sum(numeric_col).label(self.alias)
+            case AggregationType.AVG:
+                return func.avg(numeric_col).label(self.alias)
+            case AggregationType.MIN:
+                return func.min(numeric_col).label(self.alias)
+            case AggregationType.MAX:
+                return func.max(numeric_col).label(self.alias)
+            case _:
+                raise ValueError(f"Unsupported aggregation type: {self.type}")
+
+
+Aggregation: TypeAlias = Annotated[CountAggregation | FieldAggregation, Field(discriminator="type")]

orchestrator/search/core/types.py

@@ -105,7 +105,8 @@ class ActionType(str, Enum):
     """Defines the explicit, safe actions the agent can request."""
 
     SELECT = "select"  # Retrieve a list of matching records.
-    # COUNT = "count"  # For phase1; the agent will not support this yet.
+    COUNT = "count"  # Count matching records, optionally grouped.
+    AGGREGATE = "aggregate"  # Compute aggregations (sum, avg, etc.) over matching records.
 
 
 class UIType(str, Enum):
@@ -261,7 +262,7 @@ class FieldType(str, Enum):
 
     def is_embeddable(self, value: str | None) -> bool:
         """Check if a field should be embedded."""
-        if value is None:
+        if value is None or value == "":
             return False
 
         # If inference suggests it's not actually a string, don't embed it

orchestrator/search/filters/__init__.py

@@ -19,6 +19,7 @@ from .base import (
     StringFilter,
 )
 from .date_filters import DateFilter, DateRangeFilter, DateValueFilter
+from .definitions import TypeDefinition, ValueSchema
 from .ltree_filters import LtreeFilter
 from .numeric_filter import NumericFilter, NumericRangeFilter, NumericValueFilter
 
@@ -37,4 +38,7 @@ __all__ = [
     "DateFilter",
     "LtreeFilter",
     "NumericFilter",
+    # Schema types
+    "TypeDefinition",
+    "ValueSchema",
 ]

orchestrator/search/filters/definitions.py

@@ -11,8 +11,29 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Literal
+
+from pydantic import BaseModel, ConfigDict
+
 from orchestrator.search.core.types import FieldType, FilterOp, UIType
-from orchestrator.search.schemas.results import TypeDefinition, ValueSchema
+
+
+class ValueSchema(BaseModel):
+    """Schema describing the expected value type for a filter operator."""
+
+    kind: UIType | Literal["none", "object"] = UIType.STRING
+    fields: dict[str, "ValueSchema"] | None = None
+
+    model_config = ConfigDict(extra="forbid")
+
+
+class TypeDefinition(BaseModel):
+    """Definition of available operators and their value schemas for a field type."""
+
+    operators: list[FilterOp]
+    value_schema: dict[FilterOp, ValueSchema]
+
+    model_config = ConfigDict(use_enum_values=True)
 
 
 def operators_for(ft: FieldType) -> list[FilterOp]:

orchestrator/search/filters/numeric_filter.py

@@ -14,7 +14,7 @@
 from typing import Annotated, Any, Literal
 
 from pydantic import BaseModel, Field, model_validator
-from sqlalchemy import DOUBLE_PRECISION, INTEGER, and_
+from sqlalchemy import BIGINT, DOUBLE_PRECISION, and_
 from sqlalchemy import cast as sa_cast
 from sqlalchemy.sql.elements import ColumnElement
 from typing_extensions import Self
@@ -40,7 +40,7 @@ class NumericValueFilter(BaseModel):
     value: int | float
 
     def to_expression(self, column: SQLAColumn, path: str) -> ColumnElement[bool]:
-        cast_type = INTEGER if isinstance(self.value, int) else DOUBLE_PRECISION
+        cast_type = BIGINT if isinstance(self.value, int) else DOUBLE_PRECISION
         numeric_column: ColumnElement[Any] = sa_cast(column, cast_type)
         match self.op:
 
@@ -65,7 +65,7 @@ class NumericRangeFilter(BaseModel):
     value: NumericRange
 
     def to_expression(self, column: SQLAColumn, path: str) -> ColumnElement[bool]:
-        cast_type = INTEGER if isinstance(self.value.start, int) else DOUBLE_PRECISION
+        cast_type = BIGINT if isinstance(self.value.start, int) else DOUBLE_PRECISION
         numeric_column: ColumnElement[Any] = sa_cast(column, cast_type)
         return and_(numeric_column >= self.value.start, numeric_column <= self.value.end)
 

orchestrator/search/llm_migration.py

@@ -13,6 +13,7 @@
 
 """Simple search migration function that runs when SEARCH_ENABLED = True."""
 
+import pydantic_ai
 from sqlalchemy import text
 from sqlalchemy.engine import Connection
 from structlog import get_logger
@@ -28,7 +29,7 @@ TARGET_DIM = 1536
 
 def run_migration(connection: Connection) -> None:
     """Run LLM migration with ON CONFLICT DO NOTHING pattern."""
-    logger.info("Running LLM migration")
+    logger.info("Running LLM migration", pydantic_ai_version=pydantic_ai.__version__)
 
     try:
         # Test to see if the extenstion exists and then skip the migration; Needed for certain situations where db user

orchestrator/search/query/__init__.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.
+
+"""Query building and execution module."""
+
+from orchestrator.search.aggregations import TemporalGrouping
+
+from . import engine
+from .builder import (
+    ComponentInfo,
+    LeafInfo,
+    build_aggregation_query,
+    build_candidate_query,
+    build_paths_query,
+    process_path_rows,
+)
+from .exceptions import (
+    EmptyFilterPathError,
+    IncompatibleAggregationTypeError,
+    IncompatibleFilterTypeError,
+    IncompatibleTemporalGroupingTypeError,
+    InvalidEntityPrefixError,
+    InvalidLtreePatternError,
+    PathNotFoundError,
+    QueryValidationError,
+)
+from .queries import AggregateQuery, CountQuery, ExportQuery, Query, SelectQuery
+from .results import (
+    AggregationResponse,
+    AggregationResult,
+    MatchingField,
+    SearchResponse,
+    SearchResult,
+    VisualizationType,
+    format_aggregation_response,
+    format_search_response,
+    generate_highlight_indices,
+)
+from .state import QueryState
+
+__all__ = [
+    # Builder functions
+    "build_aggregation_query",
+    "build_candidate_query",
+    "build_paths_query",
+    "process_path_rows",
+    # Builder metadata
+    "ComponentInfo",
+    "LeafInfo",
+    # Engine
+    "engine",
+    # Exceptions
+    "EmptyFilterPathError",
+    "IncompatibleAggregationTypeError",
+    "IncompatibleFilterTypeError",
+    "IncompatibleTemporalGroupingTypeError",
+    "InvalidEntityPrefixError",
+    "InvalidLtreePatternError",
+    "PathNotFoundError",
+    "QueryValidationError",
+    # Query models
+    "AggregateQuery",
+    "CountQuery",
+    "ExportQuery",
+    "Query",
+    "SelectQuery",
+    "TemporalGrouping",
+    # Results
+    "AggregationResponse",
+    "AggregationResult",
+    "MatchingField",
+    "SearchResponse",
+    "SearchResult",
+    "VisualizationType",
+    "format_aggregation_response",
+    "format_search_response",
+    "generate_highlight_indices",
+    # State
+    "QueryState",
+]

orchestrator/search/query/builder.py

@@ -0,0 +1,285 @@
+# 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 import defaultdict
+from typing import Any, Sequence
+
+from pydantic import BaseModel, ConfigDict
+from sqlalchemy import Select, String, case, cast, func, select
+from sqlalchemy.engine import Row
+from sqlalchemy.sql.elements import Label
+from sqlalchemy.sql.selectable import CTE
+from sqlalchemy_utils.types.ltree import Ltree
+
+from orchestrator.db.models import AiSearchIndex
+from orchestrator.search.aggregations import AggregationType, BaseAggregation, CountAggregation
+from orchestrator.search.core.types import EntityType, FieldType, FilterOp, UIType
+from orchestrator.search.filters import LtreeFilter
+from orchestrator.search.query.queries import AggregateQuery, CountQuery, Query
+
+
+class LeafInfo(BaseModel):
+    """Information about a leaf (terminal field) in the entity schema."""
+
+    name: str
+    ui_types: list[UIType]
+    paths: list[str]
+
+    model_config = ConfigDict(
+        extra="forbid",
+        use_enum_values=True,
+    )
+
+
+class ComponentInfo(BaseModel):
+    """Information about a component (nested object) in the entity schema."""
+
+    name: str
+    ui_types: list[UIType]
+
+    model_config = ConfigDict(
+        extra="forbid",
+        use_enum_values=True,
+    )
+
+
+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(query: Query) -> 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 query plan.
+
+    Args:
+        query: Any query type (SelectQuery, CountQuery, AggregateQuery) containing entity type and optional filters.
+
+    Returns:
+        Select: The SQLAlchemy `Select` object representing the query.
+    """
+
+    stmt = (
+        select(AiSearchIndex.entity_id, AiSearchIndex.entity_title)
+        .where(AiSearchIndex.entity_type == query.entity_type.value)
+        .distinct()
+    )
+
+    if query.filters is not None:
+        entity_id_col = AiSearchIndex.entity_id
+        stmt = stmt.where(
+            query.filters.to_expression(
+                entity_id_col,
+                entity_type_value=query.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)
+    leaves_paths_dict: dict[str, set[str]] = 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)
+            leaves_paths_dict[leaf_name].add(path_str)
+
+            # 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), paths=sorted(leaves_paths_dict[leaf]))
+        for leaf, types in leaves_dict.items()
+    ]
+    components = [ComponentInfo(name=component, ui_types=[UIType.COMPONENT]) for component in sorted(components_set)]
+
+    return leaves, components
+
+
+def _build_pivot_cte(base_query: Select, pivot_fields: list[str]) -> CTE:
+    """Build CTE that pivots EAV rows into columns using CASE WHEN."""
+    from orchestrator.search.aggregations import BaseAggregation
+
+    pivot_columns = [AiSearchIndex.entity_id.label("entity_id")]
+
+    for field_path in pivot_fields:
+        pivot_columns.append(
+            func.max(case((AiSearchIndex.path == Ltree(field_path), AiSearchIndex.value), else_=None)).label(
+                BaseAggregation.field_to_alias(field_path)
+            )
+        )
+
+    return (
+        select(*pivot_columns)
+        .where(
+            AiSearchIndex.entity_id.in_(select(base_query.c.entity_id)),
+            AiSearchIndex.path.in_([Ltree(p) for p in pivot_fields]),
+        )
+        .group_by(AiSearchIndex.entity_id)
+        .cte("pivoted_entities")
+    )
+
+
+def _build_grouping_columns(
+    query: CountQuery | AggregateQuery, pivot_cte: CTE
+) -> tuple[list[Any], list[Any], list[str]]:
+    """Build GROUP BY columns and their SELECT columns.
+
+    Args:
+        query: CountQuery or AggregateQuery with group_by and temporal_group_by fields
+        pivot_cte: The pivoted CTE containing entity fields as columns
+
+    Returns:
+        tuple: (select_columns, group_by_columns, group_column_names)
+            - select_columns: List of labeled columns for SELECT
+            - group_by_columns: List of columns for GROUP BY clause
+            - group_column_names: List of column names (labels) that are grouping columns
+    """
+
+    select_columns = []
+    group_by_columns = []
+    group_column_names = []
+
+    if query.group_by:
+        for group_field in query.group_by:
+            field_alias = BaseAggregation.field_to_alias(group_field)
+            col = getattr(pivot_cte.c, field_alias)
+            select_columns.append(col.label(field_alias))
+            group_by_columns.append(col)
+            group_column_names.append(field_alias)
+
+    if query.temporal_group_by:
+        for temp_group in query.temporal_group_by:
+            select_col, group_col, col_name = temp_group.to_expression(pivot_cte.c)
+            select_columns.append(select_col)
+            group_by_columns.append(group_col)
+            group_column_names.append(col_name)
+
+    return select_columns, group_by_columns, group_column_names
+
+
+def _build_aggregation_columns(query: CountQuery | AggregateQuery, pivot_cte: CTE) -> list[Label]:
+    """Build aggregation columns (COUNT, SUM, AVG, MIN, MAX).
+
+    Args:
+        query: CountQuery or AggregateQuery
+        pivot_cte: The pivoted CTE containing entity fields as columns
+
+    Returns:
+        List of labeled aggregation expressions
+    """
+
+    if isinstance(query, AggregateQuery):
+        # AGGREGATE query with custom aggregations
+        agg_columns = []
+        for agg in query.aggregations:
+            if isinstance(agg, CountAggregation):
+                agg_columns.append(agg.to_expression(pivot_cte.c.entity_id))
+            else:
+                agg_columns.append(agg.to_expression(pivot_cte.c))
+        return agg_columns
+
+    # CountQuery without aggregations
+    count_agg = CountAggregation(type=AggregationType.COUNT, alias="count")
+    return [count_agg.to_expression(pivot_cte.c.entity_id)]
+
+
+def build_simple_count_query(base_query: Select) -> Select:
+    """Build a simple count query without grouping.
+
+    Args:
+        base_query: Base candidate query with filters applied
+
+    Returns:
+        Select statement that counts distinct entity IDs
+    """
+    return select(func.count(func.distinct(base_query.c.entity_id)).label("total_count")).select_from(
+        base_query.subquery()
+    )
+
+
+def build_aggregation_query(query: CountQuery | AggregateQuery, base_query: Select) -> tuple[Select, list[str]]:
+    """Build aggregation query with GROUP BY and aggregation functions.
+
+    Handles EAV storage by pivoting rows to columns, then applying SQL aggregations.
+    This function only handles grouped aggregations. Simple counts are handled directly
+    in the engine.
+
+    Args:
+        query: CountQuery or AggregateQuery with group_by and optional aggregations
+        base_query: Base candidate query with filters applied
+
+    Returns:
+        tuple: (query_stmt, group_column_names)
+            - query_stmt: SQLAlchemy Select statement for grouped aggregation
+            - group_column_names: List of column names that are grouping columns
+    """
+    pivot_cte = _build_pivot_cte(base_query, query.get_pivot_fields())
+    select_cols, group_cols, group_col_names = _build_grouping_columns(query, pivot_cte)
+    agg_cols = _build_aggregation_columns(query, pivot_cte)
+
+    stmt = select(*(select_cols + agg_cols)).select_from(pivot_cte)
+    if group_cols:
+        stmt = stmt.group_by(*group_cols)
+
+    return stmt, group_col_names