orchestrator-core 4.5.3__py3-none-any.whl → 4.6.0rc1__py3-none-any.whl

orchestrator/search/agent/prompts.py

@@ -50,14 +50,15 @@ async def get_base_instructions() -> str:
 
         Follow these steps in strict order:
 
-        1.  **Set Context**: Always begin by calling `set_search_parameters`.
+        1.  **Set Context**: If the user is asking for a NEW search, call `start_new_search`.
         2.  **Analyze for Filters**: Based on the user's request, decide if specific filters are necessary.
             - **If filters ARE required**, follow these sub-steps:
                 a. **Gather Intel**: Identify all needed field names, then call `discover_filter_paths` and `get_valid_operators` **once each** to get all required information.
                 b. **Construct FilterTree**: Build the `FilterTree` object.
                 c. **Set Filters**: Call `set_filter_tree`.
-        3.  **Execute**: Call `execute_search`. This is done for both filtered and non-filtered searches.
-        4.  **Report**: Answer the users' question directly and summarize when appropiate.
+        3.  **Execute**: Call `run_search`. This is done for both filtered and non-filtered searches.
+
+        After search execution, follow the dynamic instructions based on the current state.
 
         ---
         ### 4. Critical Rules
@@ -73,28 +74,53 @@ async def get_dynamic_instructions(ctx: RunContext[StateDeps[SearchState]]) -> s
     """Dynamically provides 'next step' coaching based on the current state."""
     state = ctx.deps.state
     param_state_str = json.dumps(state.parameters, indent=2, default=str) if state.parameters else "Not set."
+    results_count = state.results_data.total_count if state.results_data else 0
 
-    next_step_guidance = ""
-    if not state.parameters or not state.parameters.get("entity_type"):
+    if state.export_data:
+        next_step_guidance = (
+            "INSTRUCTION: Export has been prepared successfully. "
+            "Simply confirm to the user that the export is ready for download. "
+            "DO NOT include or mention the download URL - the UI will display it automatically."
+        )
+    elif not state.parameters or not state.parameters.get("entity_type"):
         next_step_guidance = (
-            "INSTRUCTION: The search context is not set. Your next action is to call `set_search_parameters`."
+            "INSTRUCTION: The search context is not set. Your next action is to call `start_new_search`."
+        )
+    elif results_count > 0:
+        next_step_guidance = dedent(
+            f"""
+            INSTRUCTION: Search completed successfully.
+            Found {results_count} results containing only: entity_id, title, score.
+
+            Choose your next action based on what the user requested:
+            1. **Broad/generic search** (e.g., 'show me subscriptions'): Confirm search completed and report count. Do nothing else.
+            2. **Question answerable with entity_id/title/score**: Answer directly using the current results.
+            3. **Question requiring other details**: Call `fetch_entity_details` first, then answer with the detailed data.
+            4. **Export request** (phrases like 'export', 'download', 'save as CSV'): Call `prepare_export` directly.
+            """
         )
     else:
         next_step_guidance = (
             "INSTRUCTION: Context is set. Now, analyze the user's request. "
             "If specific filters ARE required, use the information-gathering tools to build a `FilterTree` and call `set_filter_tree`. "
-            "If no specific filters are needed, you can proceed directly to `execute_search`."
+            "If no specific filters are needed, you can proceed directly to `run_search`."
         )
+
     return dedent(
         f"""
         ---
-        ### Current State & Next Action
+        ## CURRENT STATE
 
         **Current Search Parameters:**
         ```json
         {param_state_str}
         ```
 
-        **{next_step_guidance}**
+        **Current Results Count:** {results_count}
+
+        ---
+        ## NEXT ACTION REQUIRED
+
+        {next_step_guidance}
         """
     )

orchestrator/search/agent/state.py

@@ -12,10 +12,36 @@
 # limitations under the License.
 
 from typing import Any
+from uuid import UUID
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel
+
+from orchestrator.search.schemas.results import SearchResult
+
+
+class ExportData(BaseModel):
+    """Export metadata for download."""
+
+    action: str = "export"
+    query_id: str
+    download_url: str
+    message: str
+
+
+class SearchResultsData(BaseModel):
+    """Search results data for frontend display and agent context."""
+
+    action: str = "view_results"
+    query_id: str
+    results_url: str
+    total_count: int
+    message: str
+    results: list[SearchResult] = []
 
 
 class SearchState(BaseModel):
+    run_id: UUID | None = None
+    query_id: UUID | None = None
     parameters: dict[str, Any] | None = None
-    results: list[dict[str, Any]] = Field(default_factory=list)
+    results_data: SearchResultsData | None = None
+    export_data: ExportData | None = None

orchestrator/search/agent/tools.py

@@ -11,11 +11,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from collections.abc import Awaitable, Callable
-from typing import Any, TypeVar
+import json
+from typing import Any
 
 import structlog
-from ag_ui.core import EventType, StateSnapshotEvent
+from ag_ui.core import EventType, StateDeltaEvent, StateSnapshotEvent
 from pydantic_ai import RunContext
 from pydantic_ai.ag_ui import StateDeps
 from pydantic_ai.exceptions import ModelRetry
@@ -25,34 +25,22 @@ from pydantic_ai.toolsets import FunctionToolset
 from orchestrator.api.api_v1.endpoints.search import (
     get_definitions,
     list_paths,
-    search_processes,
-    search_products,
-    search_subscriptions,
-    search_workflows,
 )
-from orchestrator.schemas.search import SearchResultsSchema
+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.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 PARAMETER_REGISTRY, BaseSearchParameters
-
-from .state import SearchState
+from orchestrator.search.schemas.parameters import BaseSearchParameters
+from orchestrator.settings import app_settings
 
 logger = structlog.get_logger(__name__)
 
-
-P = TypeVar("P", bound=BaseSearchParameters)
-
-SearchFn = Callable[[P], Awaitable[SearchResultsSchema[Any]]]
-
-SEARCH_FN_MAP: dict[EntityType, SearchFn] = {
-    EntityType.SUBSCRIPTION: search_subscriptions,
-    EntityType.WORKFLOW: search_workflows,
-    EntityType.PRODUCT: search_products,
-    EntityType.PROCESS: search_processes,
-}
-
 search_toolset: FunctionToolset[StateDeps[SearchState]] = FunctionToolset(max_retries=1)
 
 
@@ -65,32 +53,50 @@ 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 set_search_parameters(
+async def start_new_search(
     ctx: RunContext[StateDeps[SearchState]],
     entity_type: EntityType,
     action: str | ActionType = ActionType.SELECT,
 ) -> StateSnapshotEvent:
-    """Sets the initial search context, like the entity type and the user's query.
+    """Starts a completely new search, clearing all previous state.
 
-    This MUST be the first tool called to start any new search.
-    Warning: Calling this tool will erase any existing filters and search results from the state.
+    This MUST be the first tool called when the user asks for a NEW search.
+    Warning: This will erase any existing filters, results, and search state.
     """
-    params = ctx.deps.state.parameters or {}
-    is_new_search = params.get("entity_type") != entity_type.value
-    final_query = (last_user_message(ctx) or "") if is_new_search else params.get("query", "")
+    final_query = last_user_message(ctx) or ""
 
     logger.debug(
-        "Setting search parameters",
+        "Starting new search",
         entity_type=entity_type.value,
         action=action,
-        is_new_search=is_new_search,
         query=final_query,
     )
 
-    ctx.deps.state.parameters = {"action": action, "entity_type": entity_type, "filters": None, "query": final_query}
-    ctx.deps.state.results = []
-    logger.debug("Search parameters set", parameters=ctx.deps.state.parameters)
+    # 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)
+
+    logger.debug("New search started", parameters=ctx.deps.state.parameters)
 
     return StateSnapshotEvent(
         type=EventType.STATE_SNAPSHOT,
@@ -102,7 +108,7 @@ async def set_search_parameters(
 async def set_filter_tree(
     ctx: RunContext[StateDeps[SearchState]],
     filters: FilterTree | None,
-) -> StateSnapshotEvent:
+) -> StateDeltaEvent:
     """Replace current filters atomically with a full FilterTree, or clear with None.
 
     Requirements:
@@ -111,7 +117,7 @@ async def set_filter_tree(
     - See the FilterTree schema examples for the exact shape.
     """
     if ctx.deps.state.parameters is None:
-        raise ModelRetry("Search parameters are not initialized. Call set_search_parameters first.")
+        raise ModelRetry("Search parameters are not initialized. Call start_new_search first.")
 
     entity_type = EntityType(ctx.deps.state.parameters["entity_type"])
 
@@ -136,28 +142,33 @@ async def set_filter_tree(
         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 StateSnapshotEvent(type=EventType.STATE_SNAPSHOT, snapshot=ctx.deps.state.model_dump())
+    return StateDeltaEvent(
+        type=EventType.STATE_DELTA,
+        delta=[
+            JSONPatchOp.upsert(
+                path="/parameters/filters",
+                value=filter_data,
+                existed=filters_existed,
+            )
+        ],
+    )
 
 
 @search_toolset.tool
-async def execute_search(
+async def run_search(
     ctx: RunContext[StateDeps[SearchState]],
     limit: int = 10,
-) -> StateSnapshotEvent:
-    """Execute the search with the current parameters."""
+) -> StateDeltaEvent:
+    """Execute the search with the current parameters and save to database."""
     if not ctx.deps.state.parameters:
         raise ValueError("No search parameters set")
 
-    entity_type = EntityType(ctx.deps.state.parameters["entity_type"])
-    param_class = PARAMETER_REGISTRY.get(entity_type)
-    if not param_class:
-        raise ValueError(f"Unknown entity type: {entity_type}")
-
-    params = param_class(**ctx.deps.state.parameters)
+    params = BaseSearchParameters.create(**ctx.deps.state.parameters)
     logger.debug(
         "Executing database search",
-        search_entity_type=entity_type.value,
+        search_entity_type=params.entity_type.value,
         limit=limit,
         has_filters=params.filters is not None,
         query=params.query,
@@ -169,17 +180,61 @@ async def execute_search(
 
     params.limit = limit
 
-    fn = SEARCH_FN_MAP[entity_type]
-    search_results = await fn(params)
+    changes: list[JSONPatchOp] = []
+
+    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
+
+        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)))
+
+    # 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,
+    )
+    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_results.data) if search_results.data else 0,
+        total_results=len(search_response.results),
     )
 
-    ctx.deps.state.results = search_results.data
+    # 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}"
+
+    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
+        )
+    )
 
-    return StateSnapshotEvent(type=EventType.STATE_SNAPSHOT, snapshot=ctx.deps.state.model_dump())
+    return StateDeltaEvent(type=EventType.STATE_DELTA, delta=changes)
 
 
 @search_toolset.tool
@@ -256,3 +311,87 @@ async def get_valid_operators() -> dict[str, list[FilterOp]]:
         if hasattr(type_def, "operators"):
             operator_map[key] = type_def.operators
     return operator_map
+
+
+@search_toolset.tool
+async def fetch_entity_details(
+    ctx: RunContext[StateDeps[SearchState]],
+    limit: int = 10,
+) -> str:
+    """Fetch detailed entity information to answer user questions.
+
+    Use this tool when you need detailed information about entities from the search results
+    to answer the user's question. This provides the same detailed data that would be
+    included in an export (e.g., subscription status, product details, workflow info, etc.).
+
+    Args:
+        ctx: Runtime context for agent (injected).
+        limit: Maximum number of entities to fetch details for (default 10).
+
+    Returns:
+        JSON string containing detailed entity information.
+
+    Raises:
+        ValueError: If no search results are available.
+    """
+    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 not ctx.deps.state.parameters:
+        raise ValueError("No search parameters found.")
+
+    entity_type = EntityType(ctx.deps.state.parameters["entity_type"])
+
+    entity_ids = [r.entity_id for r in ctx.deps.state.results_data.results[:limit]]
+
+    logger.debug(
+        "Fetching detailed entity data",
+        entity_type=entity_type.value,
+        entity_count=len(entity_ids),
+    )
+
+    detailed_data = fetch_export_data(entity_type, entity_ids)
+
+    return json.dumps(detailed_data, indent=2)
+
+
+@search_toolset.tool
+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.")
+
+    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."
+        )
+
+    logger.debug(
+        "Prepared query for export",
+        query_id=str(ctx.deps.state.query_id),
+    )
+
+    download_url = f"{app_settings.BASE_URL}/api/search/queries/{ctx.deps.state.query_id}/export"
+
+    ctx.deps.state.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())
+
+    # 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/core/exceptions.py

@@ -34,3 +34,9 @@ class InvalidCursorError(SearchUtilsError):
     """Raised when cursor cannot be decoded."""
 
     pass
+
+
+class QueryStateNotFoundError(SearchUtilsError):
+    """Raised when a query state cannot be found in the database."""
+
+    pass

orchestrator/search/core/types.py

@@ -289,6 +289,7 @@ class ExtractedField(NamedTuple):
 class IndexableRecord(TypedDict):
     entity_id: str
     entity_type: str
+    entity_title: str
     path: Ltree
     value: Any
     value_type: Any

orchestrator/search/export.py

@@ -0,0 +1,199 @@
+# 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 uuid import UUID
+
+from sqlalchemy import select
+from sqlalchemy.orm import selectinload
+
+from orchestrator.db import (
+    ProcessTable,
+    ProductTable,
+    SubscriptionTable,
+    WorkflowTable,
+    db,
+)
+from orchestrator.search.core.types import EntityType
+
+
+def fetch_subscription_export_data(entity_ids: list[str]) -> list[dict]:
+    """Fetch subscription data for export.
+
+    Args:
+        entity_ids: List of subscription IDs as strings
+
+    Returns:
+        List of flattened subscription dictionaries with fields:
+        subscription_id, description, status, insync, start_date, end_date,
+        note, product_name, tag, product_type, customer_id
+    """
+    stmt = (
+        select(
+            SubscriptionTable.subscription_id,
+            SubscriptionTable.description,
+            SubscriptionTable.status,
+            SubscriptionTable.insync,
+            SubscriptionTable.start_date,
+            SubscriptionTable.end_date,
+            SubscriptionTable.note,
+            SubscriptionTable.customer_id,
+            ProductTable.name.label("product_name"),
+            ProductTable.tag,
+            ProductTable.product_type,
+        )
+        .join(ProductTable, SubscriptionTable.product_id == ProductTable.product_id)
+        .filter(SubscriptionTable.subscription_id.in_([UUID(sid) for sid in entity_ids]))
+    )
+
+    rows = db.session.execute(stmt).all()
+
+    return [
+        {
+            "subscription_id": str(row.subscription_id),
+            "description": row.description,
+            "status": row.status,
+            "insync": row.insync,
+            "start_date": row.start_date.isoformat() if row.start_date else None,
+            "end_date": row.end_date.isoformat() if row.end_date else None,
+            "note": row.note,
+            "product_name": row.product_name,
+            "tag": row.tag,
+            "product_type": row.product_type,
+            "customer_id": row.customer_id,
+        }
+        for row in rows
+    ]
+
+
+def fetch_workflow_export_data(entity_ids: list[str]) -> list[dict]:
+    """Fetch workflow data for export.
+
+    Args:
+        entity_ids: List of workflow names as strings
+
+    Returns:
+        List of flattened workflow dictionaries with fields:
+        name, description, created_at, product_names (comma-separated),
+        product_ids (comma-separated), product_types (comma-separated)
+    """
+    stmt = (
+        select(WorkflowTable).options(selectinload(WorkflowTable.products)).filter(WorkflowTable.name.in_(entity_ids))
+    )
+    workflows = db.session.scalars(stmt).all()
+
+    return [
+        {
+            "name": w.name,
+            "description": w.description,
+            "created_at": w.created_at.isoformat() if w.created_at else None,
+            "product_names": ", ".join(p.name for p in w.products),
+            "product_ids": ", ".join(str(p.product_id) for p in w.products),
+            "product_types": ", ".join(p.product_type for p in w.products),
+        }
+        for w in workflows
+    ]
+
+
+def fetch_product_export_data(entity_ids: list[str]) -> list[dict]:
+    """Fetch product data for export.
+
+    Args:
+        entity_ids: List of product IDs as strings
+
+    Returns:
+        List of flattened product dictionaries with fields:
+        product_id, name, product_type, tag, description, status, created_at
+    """
+    stmt = (
+        select(ProductTable)
+        .options(
+            selectinload(ProductTable.workflows),
+            selectinload(ProductTable.fixed_inputs),
+            selectinload(ProductTable.product_blocks),
+        )
+        .filter(ProductTable.product_id.in_([UUID(pid) for pid in entity_ids]))
+    )
+    products = db.session.scalars(stmt).all()
+
+    return [
+        {
+            "product_id": str(p.product_id),
+            "name": p.name,
+            "product_type": p.product_type,
+            "tag": p.tag,
+            "description": p.description,
+            "status": p.status,
+            "created_at": p.created_at.isoformat() if p.created_at else None,
+        }
+        for p in products
+    ]
+
+
+def fetch_process_export_data(entity_ids: list[str]) -> list[dict]:
+    """Fetch process data for export.
+
+    Args:
+        entity_ids: List of process IDs as strings
+
+    Returns:
+        List of flattened process dictionaries with fields:
+        process_id, workflow_name, workflow_id, last_status, is_task,
+        created_by, started_at, last_modified_at, last_step
+    """
+    stmt = (
+        select(ProcessTable)
+        .options(selectinload(ProcessTable.workflow))
+        .filter(ProcessTable.process_id.in_([UUID(pid) for pid in entity_ids]))
+    )
+    processes = db.session.scalars(stmt).all()
+
+    return [
+        {
+            "process_id": str(p.process_id),
+            "workflow_name": p.workflow.name if p.workflow else None,
+            "workflow_id": str(p.workflow_id),
+            "last_status": p.last_status,
+            "is_task": p.is_task,
+            "created_by": p.created_by,
+            "started_at": p.started_at.isoformat() if p.started_at else None,
+            "last_modified_at": p.last_modified_at.isoformat() if p.last_modified_at else None,
+            "last_step": p.last_step,
+        }
+        for p in processes
+    ]
+
+
+def fetch_export_data(entity_type: EntityType, entity_ids: list[str]) -> list[dict]:
+    """Fetch export data for any entity type.
+
+    Args:
+        entity_type: The type of entities to fetch
+        entity_ids: List of entity IDs/names as strings
+
+    Returns:
+        List of flattened entity dictionaries ready for CSV export
+
+    Raises:
+        ValueError: If entity_type is not supported
+    """
+    match entity_type:
+        case EntityType.SUBSCRIPTION:
+            return fetch_subscription_export_data(entity_ids)
+        case EntityType.WORKFLOW:
+            return fetch_workflow_export_data(entity_ids)
+        case EntityType.PRODUCT:
+            return fetch_product_export_data(entity_ids)
+        case EntityType.PROCESS:
+            return fetch_process_export_data(entity_ids)
+        case _:
+            raise ValueError(f"Unsupported entity type: {entity_type}")