orchestrator-core 4.6.1__py3-none-any.whl → 4.6.3__py3-none-any.whl

orchestrator/search/agent/tools.py

@@ -12,10 +12,10 @@
 # limitations under the License.
 
 import json
-from typing import Any
+from typing import Any, cast
 
 import structlog
-from ag_ui.core import EventType, StateDeltaEvent, StateSnapshotEvent
+from ag_ui.core import EventType, StateSnapshotEvent
 from pydantic_ai import RunContext
 from pydantic_ai.ag_ui import StateDeps
 from pydantic_ai.exceptions import ModelRetry
@@ -26,17 +26,28 @@ from orchestrator.api.api_v1.endpoints.search import (
     get_definitions,
     list_paths,
 )
-from orchestrator.db import AgentRunTable, SearchQueryTable, db
-from orchestrator.search.agent.json_patch import JSONPatchOp
-from orchestrator.search.agent.state import ExportData, SearchResultsData, SearchState
+from orchestrator.db import db
+from orchestrator.search.agent.handlers import (
+    execute_aggregation_with_persistence,
+    execute_search_with_persistence,
+)
+from orchestrator.search.agent.state import SearchState
+from orchestrator.search.agent.validation import require_action
+from orchestrator.search.aggregations import Aggregation, FieldAggregation, TemporalGrouping
 from orchestrator.search.core.types import ActionType, EntityType, FilterOp
-from orchestrator.search.export import fetch_export_data
 from orchestrator.search.filters import FilterTree
-from orchestrator.search.retrieval.engine import execute_search
-from orchestrator.search.retrieval.exceptions import FilterValidationError, PathNotFoundError
-from orchestrator.search.retrieval.query_state import SearchQueryState
-from orchestrator.search.retrieval.validation import validate_filter_tree
-from orchestrator.search.schemas.parameters import BaseSearchParameters
+from orchestrator.search.query import engine
+from orchestrator.search.query.exceptions import PathNotFoundError, QueryValidationError
+from orchestrator.search.query.export import fetch_export_data
+from orchestrator.search.query.queries import AggregateQuery, CountQuery, Query, SelectQuery
+from orchestrator.search.query.results import AggregationResponse, AggregationResult, ExportData, VisualizationType
+from orchestrator.search.query.state import QueryState
+from orchestrator.search.query.validation import (
+    validate_aggregation_field,
+    validate_filter_path,
+    validate_filter_tree,
+    validate_temporal_grouping_field,
+)
 from orchestrator.settings import app_settings
 
 logger = structlog.get_logger(__name__)
@@ -53,27 +64,11 @@ def last_user_message(ctx: RunContext[StateDeps[SearchState]]) -> str | None:
     return None
 
 
-def _set_parameters(
-    ctx: RunContext[StateDeps[SearchState]],
-    entity_type: EntityType,
-    action: str | ActionType,
-    query: str,
-    filters: Any | None,
-) -> None:
-    """Internal helper to set parameters."""
-    ctx.deps.state.parameters = {
-        "action": action,
-        "entity_type": entity_type,
-        "filters": filters,
-        "query": query,
-    }
-
-
 @search_toolset.tool
 async def start_new_search(
     ctx: RunContext[StateDeps[SearchState]],
     entity_type: EntityType,
-    action: str | ActionType = ActionType.SELECT,
+    action: ActionType = ActionType.SELECT,
 ) -> StateSnapshotEvent:
     """Starts a completely new search, clearing all previous state.
 
@@ -90,13 +85,26 @@ async def start_new_search(
     )
 
     # Clear all state
-    ctx.deps.state.results_data = None
-    ctx.deps.state.export_data = None
-
-    # Set fresh parameters with no filters
-    _set_parameters(ctx, entity_type, action, final_query, None)
+    ctx.deps.state.results_count = None
+    ctx.deps.state.action = action
+
+    # Create the appropriate query object based on action
+    if action == ActionType.SELECT:
+        ctx.deps.state.query = SelectQuery(
+            entity_type=entity_type,
+            query_text=final_query,
+        )
+    elif action == ActionType.COUNT:
+        ctx.deps.state.query = CountQuery(
+            entity_type=entity_type,
+        )
+    else:  # ActionType.AGGREGATE
+        ctx.deps.state.query = AggregateQuery(
+            entity_type=entity_type,
+            aggregations=[],  # Will be set by set_aggregations tool
+        )
 
-    logger.debug("New search started", parameters=ctx.deps.state.parameters)
+    logger.debug("New search started", action=action.value, query_type=type(ctx.deps.state.query).__name__)
 
     return StateSnapshotEvent(
         type=EventType.STATE_SNAPSHOT,
@@ -108,18 +116,15 @@ async def start_new_search(
 async def set_filter_tree(
     ctx: RunContext[StateDeps[SearchState]],
     filters: FilterTree | None,
-) -> StateDeltaEvent:
+) -> StateSnapshotEvent:
     """Replace current filters atomically with a full FilterTree, or clear with None.
 
-    Requirements:
-    - Root/group operators must be 'AND' or 'OR' (uppercase).
-    - Provide either PathFilters or nested groups under `children`.
-    - See the FilterTree schema examples for the exact shape.
+    See FilterTree model for structure, operators, and examples.
     """
-    if ctx.deps.state.parameters is None:
-        raise ModelRetry("Search parameters are not initialized. Call start_new_search first.")
+    if ctx.deps.state.query is None:
+        raise ModelRetry("Search query is not initialized. Call start_new_search first.")
 
-    entity_type = EntityType(ctx.deps.state.parameters["entity_type"])
+    entity_type = ctx.deps.state.query.entity_type
 
     logger.debug(
         "Setting filter tree",
@@ -133,108 +138,112 @@ async def set_filter_tree(
     except PathNotFoundError as e:
         logger.debug(f"{PathNotFoundError.__name__}: {str(e)}")
         raise ModelRetry(f"{str(e)} Use discover_filter_paths tool to find valid paths.")
-    except FilterValidationError as e:
+    except QueryValidationError as e:
         # ModelRetry will trigger an agent retry, containing the specific validation error.
-        logger.debug(f"Filter validation failed: {str(e)}")
+        logger.debug(f"Query validation failed: {str(e)}")
         raise ModelRetry(str(e))
     except Exception as e:
         logger.error("Unexpected Filter validation exception", error=str(e))
         raise ModelRetry(f"Filter validation failed: {str(e)}. Please check your filter structure and try again.")
 
-    filter_data = None if filters is None else filters.model_dump(mode="json", by_alias=True)
-    filters_existed = "filters" in ctx.deps.state.parameters
-    ctx.deps.state.parameters["filters"] = filter_data
-    return StateDeltaEvent(
-        type=EventType.STATE_DELTA,
-        delta=[
-            JSONPatchOp.upsert(
-                path="/parameters/filters",
-                value=filter_data,
-                existed=filters_existed,
-            )
-        ],
+    ctx.deps.state.query = cast(Query, ctx.deps.state.query).model_copy(update={"filters": filters})
+
+    # Use snapshot to workaround state persistence issue
+    # TODO: Fix root cause; state tree may be empty on frontend when parameters are being set
+    return StateSnapshotEvent(
+        type=EventType.STATE_SNAPSHOT,
+        snapshot=ctx.deps.state.model_dump(),
     )
 
 
 @search_toolset.tool
+@require_action(ActionType.SELECT)
 async def run_search(
     ctx: RunContext[StateDeps[SearchState]],
     limit: int = 10,
-) -> StateDeltaEvent:
-    """Execute the search with the current parameters and save to database."""
-    if not ctx.deps.state.parameters:
-        raise ValueError("No search parameters set")
+) -> AggregationResponse:
+    """Execute a search to find and rank entities.
 
-    params = BaseSearchParameters.create(**ctx.deps.state.parameters)
-    logger.debug(
-        "Executing database search",
-        search_entity_type=params.entity_type.value,
-        limit=limit,
-        has_filters=params.filters is not None,
-        query=params.query,
-        action=params.action,
+    Use this tool for SELECT action to find entities matching your criteria.
+    For counting or computing statistics, use run_aggregation instead.
+    """
+    query = cast(SelectQuery, cast(Query, ctx.deps.state.query).model_copy(update={"limit": limit}))
+
+    search_response, run_id, query_id = await execute_search_with_persistence(query, db.session, ctx.deps.state.run_id)
+
+    ctx.deps.state.run_id = run_id
+    ctx.deps.state.query_id = query_id
+    ctx.deps.state.results_count = len(search_response.results)
+
+    # Convert SearchResults to AggregationResults for consistent rendering
+    aggregation_results = [
+        AggregationResult(
+            group_values={
+                "entity_id": result.entity_id,
+                "title": result.entity_title,
+                "entity_type": result.entity_type.value,
+            },
+            aggregations={"score": result.score},
+        )
+        for result in search_response.results
+    ]
+
+    # For now use the default table visualization for search results
+    aggregation_response = AggregationResponse(
+        results=aggregation_results,
+        total_groups=len(aggregation_results),
+        metadata=search_response.metadata,
+        visualization_type=VisualizationType(type="table"),
     )
 
-    if params.filters:
-        logger.debug("Search filters", filters=params.filters)
-
-    params.limit = limit
+    logger.debug(
+        "Search completed",
+        total_count=ctx.deps.state.results_count,
+        query_id=str(query_id),
+    )
 
-    changes: list[JSONPatchOp] = []
+    return aggregation_response
 
-    if not ctx.deps.state.run_id:
-        agent_run = AgentRunTable(agent_type="search")
 
-        db.session.add(agent_run)
-        db.session.commit()
-        db.session.expire_all()  # Release connection to prevent stacking while agent runs
+@search_toolset.tool
+@require_action(ActionType.COUNT, ActionType.AGGREGATE)
+async def run_aggregation(
+    ctx: RunContext[StateDeps[SearchState]],
+    visualization_type: VisualizationType,
+) -> AggregationResponse:
+    """Execute an aggregation to compute counts or statistics over entities.
 
-        ctx.deps.state.run_id = agent_run.run_id
-        logger.debug("Created new agent run", run_id=str(agent_run.run_id))
-        changes.append(JSONPatchOp(op="add", path="/run_id", value=str(ctx.deps.state.run_id)))
+    Use this tool for COUNT or AGGREGATE actions after setting up:
+    - Grouping fields with set_grouping or set_temporal_grouping
+    - Aggregation functions with set_aggregations (for AGGREGATE action)
+    """
+    query = cast(CountQuery | AggregateQuery, ctx.deps.state.query)
 
-    # Get query with embedding and save to DB
-    search_response = await execute_search(params, db.session)
-    query_embedding = search_response.query_embedding
-    query_state = SearchQueryState(parameters=params, query_embedding=query_embedding)
-    query_number = db.session.query(SearchQueryTable).filter_by(run_id=ctx.deps.state.run_id).count() + 1
-    search_query = SearchQueryTable.from_state(
-        state=query_state,
-        run_id=ctx.deps.state.run_id,
-        query_number=query_number,
+    logger.debug(
+        "Executing aggregation",
+        search_entity_type=query.entity_type.value,
+        has_filters=query.filters is not None,
+        action=query.action,
     )
-    db.session.add(search_query)
-    db.session.commit()
-    db.session.expire_all()
-
-    query_id_existed = ctx.deps.state.query_id is not None
-    ctx.deps.state.query_id = search_query.query_id
-    logger.debug("Saved search query", query_id=str(search_query.query_id), query_number=query_number)
-    changes.append(JSONPatchOp.upsert(path="/query_id", value=str(ctx.deps.state.query_id), existed=query_id_existed))
 
-    logger.debug(
-        "Search completed",
-        total_results=len(search_response.results),
+    aggregation_response, run_id, query_id = await execute_aggregation_with_persistence(
+        query, db.session, ctx.deps.state.run_id
     )
 
-    # Store results data for both frontend display and agent context
-    results_url = f"{app_settings.BASE_URL}/api/search/queries/{ctx.deps.state.query_id}"
+    ctx.deps.state.run_id = run_id
+    ctx.deps.state.query_id = query_id
+    ctx.deps.state.results_count = aggregation_response.total_groups
 
-    results_data_existed = ctx.deps.state.results_data is not None
-    ctx.deps.state.results_data = SearchResultsData(
-        query_id=str(ctx.deps.state.query_id),
-        results_url=results_url,
-        total_count=len(search_response.results),
-        message=f"Found {len(search_response.results)} results.",
-        results=search_response.results,  # Include actual results in state
-    )
-    changes.append(
-        JSONPatchOp.upsert(
-            path="/results_data", value=ctx.deps.state.results_data.model_dump(), existed=results_data_existed
-        )
+    aggregation_response.visualization_type = visualization_type
+
+    logger.debug(
+        "Aggregation completed",
+        total_groups=aggregation_response.total_groups,
+        visualization_type=visualization_type.type,
+        query_id=str(query_id),
     )
 
-    return StateDeltaEvent(type=EventType.STATE_DELTA, delta=changes)
+    return aggregation_response
 
 
 @search_toolset.tool
@@ -248,10 +257,11 @@ async def discover_filter_paths(
     Returns a dictionary where each key is a field_name from the input list and
     the value is its discovery result.
     """
-    if not entity_type and ctx.deps.state.parameters:
-        entity_type = EntityType(ctx.deps.state.parameters.get("entity_type"))
     if not entity_type:
-        entity_type = EntityType.SUBSCRIPTION
+        if ctx.deps.state.query:
+            entity_type = ctx.deps.state.query.entity_type
+        else:
+            raise ModelRetry("Entity type not specified and no query in state. Call start_new_search first.")
 
     all_results = {}
     for field_name in field_names:
@@ -332,17 +342,21 @@ async def fetch_entity_details(
         JSON string containing detailed entity information.
 
     Raises:
-        ValueError: If no search results are available.
+        ModelRetry: If no search has been executed.
     """
-    if not ctx.deps.state.results_data or not ctx.deps.state.results_data.results:
-        raise ValueError("No search results available. Run a search first before fetching entity details.")
+    if ctx.deps.state.query_id is None:
+        raise ModelRetry("No query_id found. Run a search first.")
 
-    if not ctx.deps.state.parameters:
-        raise ValueError("No search parameters found.")
+    # Load the saved query and re-execute it to get entity IDs
+    query_state = QueryState.load_from_id(ctx.deps.state.query_id, SelectQuery)
+    query = query_state.query.model_copy(update={"limit": limit})
+    search_response = await engine.execute_search(query, db.session)
+    entity_ids = [r.entity_id for r in search_response.results]
 
-    entity_type = EntityType(ctx.deps.state.parameters["entity_type"])
+    if not entity_ids:
+        return json.dumps({"message": "No entities found in search results."})
 
-    entity_ids = [r.entity_id for r in ctx.deps.state.results_data.results[:limit]]
+    entity_type = query.entity_type
 
     logger.debug(
         "Fetching detailed entity data",
@@ -356,23 +370,16 @@ async def fetch_entity_details(
 
 
 @search_toolset.tool
+@require_action(ActionType.SELECT)
 async def prepare_export(
     ctx: RunContext[StateDeps[SearchState]],
-) -> StateSnapshotEvent:
-    """Prepares export URL using the last executed search query."""
-    if not ctx.deps.state.query_id or not ctx.deps.state.run_id:
-        raise ValueError("No search has been executed yet. Run a search first before exporting.")
+) -> ExportData:
+    """Prepares export URL using the last executed search query.
 
-    if not ctx.deps.state.parameters:
-        raise ValueError("No search parameters found. Run a search first before exporting.")
-
-    # Validate that export is only available for SELECT actions
-    action = ctx.deps.state.parameters.get("action", ActionType.SELECT)
-    if action != ActionType.SELECT:
-        raise ValueError(
-            f"Export is only available for SELECT actions. Current action is '{action}'. "
-            "Please run a SELECT search first."
-        )
+    Returns export data which is displayed directly in the UI.
+    """
+    if not ctx.deps.state.query_id or not ctx.deps.state.run_id:
+        raise ModelRetry("No search has been executed yet. Run a search first before exporting.")
 
     logger.debug(
         "Prepared query for export",
@@ -381,16 +388,88 @@ async def prepare_export(
 
     download_url = f"{app_settings.BASE_URL}/api/search/queries/{ctx.deps.state.query_id}/export"
 
-    ctx.deps.state.export_data = ExportData(
+    export_data = ExportData(
         query_id=str(ctx.deps.state.query_id),
         download_url=download_url,
         message="Export ready for download.",
     )
 
-    logger.debug("Export data set in state", export_data=ctx.deps.state.export_data.model_dump())
+    logger.debug("Export prepared", query_id=export_data.query_id)
+
+    return export_data
+
+
+@search_toolset.tool(retries=2)
+@require_action(ActionType.COUNT, ActionType.AGGREGATE)
+async def set_grouping(
+    ctx: RunContext[StateDeps[SearchState]],
+    group_by_paths: list[str],
+) -> StateSnapshotEvent:
+    """Set which field paths to group results by for aggregation.
+
+    Only used with COUNT or AGGREGATE actions. Paths must exist in the schema; use discover_filter_paths to verify.
+    """
+    for path in group_by_paths:
+        field_type = validate_filter_path(path)
+        if field_type is None:
+            raise ModelRetry(
+                f"Path '{path}' not found in database schema. "
+                f"Use discover_filter_paths(['{path.split('.')[-1]}']) to find valid paths."
+            )
+
+    ctx.deps.state.query = cast(Query, ctx.deps.state.query).model_copy(update={"group_by": group_by_paths})
+
+    return StateSnapshotEvent(
+        type=EventType.STATE_SNAPSHOT,
+        snapshot=ctx.deps.state.model_dump(),
+    )
+
+
+@search_toolset.tool(retries=2)
+@require_action(ActionType.AGGREGATE)
+async def set_aggregations(
+    ctx: RunContext[StateDeps[SearchState]],
+    aggregations: list[Aggregation],
+) -> StateSnapshotEvent:
+    """Define what aggregations to compute over the matching records.
+
+    Only used with AGGREGATE action. See Aggregation model (CountAggregation, FieldAggregation) for structure and field requirements.
+    """
+    # Validate field paths for FieldAggregations
+    try:
+        for agg in aggregations:
+            if isinstance(agg, FieldAggregation):
+                validate_aggregation_field(agg.type, agg.field)
+    except ValueError as e:
+        raise ModelRetry(f"{str(e)} Use discover_filter_paths to find valid paths.")
+
+    ctx.deps.state.query = cast(Query, ctx.deps.state.query).model_copy(update={"aggregations": aggregations})
+
+    return StateSnapshotEvent(
+        type=EventType.STATE_SNAPSHOT,
+        snapshot=ctx.deps.state.model_dump(),
+    )
+
+
+@search_toolset.tool(retries=2)
+@require_action(ActionType.COUNT, ActionType.AGGREGATE)
+async def set_temporal_grouping(
+    ctx: RunContext[StateDeps[SearchState]],
+    temporal_groups: list[TemporalGrouping],
+) -> StateSnapshotEvent:
+    """Set temporal grouping to group datetime fields by time periods.
+
+    Only used with COUNT or AGGREGATE actions. See TemporalGrouping model for structure, periods, and examples.
+    """
+    # Validate that fields exist and are datetime types
+    try:
+        for tg in temporal_groups:
+            validate_temporal_grouping_field(tg.field)
+    except ValueError as e:
+        raise ModelRetry(f"{str(e)} Use discover_filter_paths to find datetime fields.")
+
+    ctx.deps.state.query = cast(Query, ctx.deps.state.query).model_copy(update={"temporal_group_by": temporal_groups})
 
-    # Should use StateDelta here? Use snapshot to workaround state persistence issue
-    # TODO: Fix root cause; state is empty on frontend when it should have data from run_search
     return StateSnapshotEvent(
         type=EventType.STATE_SNAPSHOT,
         snapshot=ctx.deps.state.model_dump(),

orchestrator/search/agent/validation.py

@@ -0,0 +1,80 @@
+# 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 functools import wraps
+from typing import TYPE_CHECKING, Any, Callable
+
+import structlog
+from pydantic_ai import RunContext
+from pydantic_ai.ag_ui import StateDeps
+from pydantic_ai.exceptions import ModelRetry
+
+from orchestrator.search.core.types import ActionType
+
+if TYPE_CHECKING:
+    from orchestrator.search.agent.state import SearchState
+
+logger = structlog.get_logger(__name__)
+
+
+def require_action(*allowed_actions: ActionType) -> Callable:
+    """Validate that the current search action is one of the allowed types.
+
+    This decorator is in preparation for a future finite state machine implementation
+    where we explicitly define which actions are valid in which states.
+
+    Example:
+        @search_toolset.tool
+        @require_action(ActionType.SELECT)
+        async def run_search(ctx, ...):
+            # Only callable when action is SELECT
+            ...
+
+    Args:
+        allowed_actions: One or more ActionType values that are valid for this tool.
+
+    Returns:
+        Decorated function that validates action before execution.
+
+    Raises:
+        ModelRetry: If current action is not in allowed_actions.
+    """
+
+    def decorator(func: Callable) -> Callable:
+        @wraps(func)
+        async def wrapper(ctx: "RunContext[StateDeps[SearchState]]", *args: Any, **kwargs: Any) -> Any:
+            if ctx.deps.state.action is None or ctx.deps.state.query is None:
+                logger.warning(f"Action validation failed for {func.__name__}: action or query is None")
+                raise ModelRetry("Search action and query are not initialized. Call start_new_search first.")
+
+            current_action = ctx.deps.state.action
+
+            if current_action not in allowed_actions:
+                allowed_names = ", ".join(a.value for a in allowed_actions)
+                logger.warning(
+                    "Invalid action for tool",
+                    tool=func.__name__,
+                    allowed_actions=allowed_names,
+                    current_action=current_action.value,
+                )
+                raise ModelRetry(
+                    f"{func.__name__} is only available for {allowed_names} action(s). "
+                    f"Current action is '{current_action.value}'."
+                )
+
+            return await func(ctx, *args, **kwargs)
+
+        return wrapper
+
+    return decorator

orchestrator/search/{schemas → aggregations}/__init__.py

@@ -10,3 +10,23 @@
 # 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 (
+    Aggregation,
+    AggregationType,
+    BaseAggregation,
+    CountAggregation,
+    FieldAggregation,
+    TemporalGrouping,
+    TemporalPeriod,
+)
+
+__all__ = [
+    "Aggregation",
+    "AggregationType",
+    "BaseAggregation",
+    "CountAggregation",
+    "FieldAggregation",
+    "TemporalGrouping",
+    "TemporalPeriod",
+]