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

orchestrator/__init__.py

@@ -13,7 +13,7 @@
 
 """This is the orchestrator workflow engine."""
 
-__version__ = "4.5.2"
+__version__ = "4.6.0rc1"
 
 
 from structlog import get_logger

orchestrator/agentic_app.py

@@ -27,7 +27,6 @@ from orchestrator.llm_settings import LLMSettings, llm_settings
 
 if TYPE_CHECKING:
     from pydantic_ai.models.openai import OpenAIModel
-    from pydantic_ai.toolsets import FunctionToolset
 
 logger = get_logger(__name__)
 
@@ -38,19 +37,17 @@ class LLMOrchestratorCore(OrchestratorCore):
         *args: Any,
         llm_settings: LLMSettings = llm_settings,
         agent_model: "OpenAIModel | str | None" = None,
-        agent_tools: "list[FunctionToolset] | None" = None,
         **kwargs: Any,
     ) -> None:
         """Initialize the `LLMOrchestratorCore` class.
 
         This class extends `OrchestratorCore` with LLM features (search and agent).
-        It runs the search migration and mounts the agent endpoint based on feature flags.
+        It runs the search migration based on feature flags.
 
         Args:
             *args: All the normal arguments passed to the `OrchestratorCore` class.
             llm_settings: A class of settings for the LLM
             agent_model: Override the agent model (defaults to llm_settings.AGENT_MODEL)
-            agent_tools: A list of tools that can be used by the agent
             **kwargs: Additional arguments passed to the `OrchestratorCore` class.
 
         Returns:
@@ -58,7 +55,6 @@ class LLMOrchestratorCore(OrchestratorCore):
         """
         self.llm_settings = llm_settings
         self.agent_model = agent_model or llm_settings.AGENT_MODEL
-        self.agent_tools = agent_tools
 
         super().__init__(*args, **kwargs)
 
@@ -79,22 +75,6 @@ class LLMOrchestratorCore(OrchestratorCore):
                 )
                 raise
 
-        # Mount agent endpoint if agent is enabled
-        if self.llm_settings.AGENT_ENABLED:
-            logger.info("Initializing agent features", model=self.agent_model)
-            try:
-                from orchestrator.search.agent import build_agent_router
-
-                agent_app = build_agent_router(self.agent_model, self.agent_tools)
-                self.mount("/agent", agent_app)
-            except ImportError as e:
-                logger.error(
-                    "Unable to initialize agent features. Please install agent dependencies: "
-                    "`pip install orchestrator-core[agent]`",
-                    error=str(e),
-                )
-                raise
-
 
 main_typer_app = typer.Typer()
 main_typer_app.add_typer(cli_app, name="orchestrator", help="The orchestrator CLI commands")

orchestrator/api/api_v1/api.py

@@ -95,3 +95,8 @@ if llm_settings.SEARCH_ENABLED:
     api_router.include_router(
         search.router, prefix="/search", tags=["Core", "Search"], dependencies=[Depends(authorize)]
     )
+
+if llm_settings.AGENT_ENABLED:
+    from orchestrator.api.api_v1.endpoints import agent
+
+    api_router.include_router(agent.router, prefix="/agent", tags=["Core", "Agent"], dependencies=[Depends(authorize)])

orchestrator/api/api_v1/endpoints/agent.py

@@ -0,0 +1,50 @@
+# 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 cache
+from typing import Annotated
+
+from fastapi import APIRouter, Depends, Request
+from pydantic_ai.ag_ui import StateDeps, handle_ag_ui_request
+from pydantic_ai.agent import Agent
+from starlette.responses import Response
+from structlog import get_logger
+
+from orchestrator.llm_settings import llm_settings
+from orchestrator.search.agent import build_agent_instance
+from orchestrator.search.agent.state import SearchState
+
+router = APIRouter()
+logger = get_logger(__name__)
+
+
+@cache
+def get_agent() -> Agent[StateDeps[SearchState], str]:
+    """Dependency to provide the agent instance.
+
+    The agent is built once and cached for the lifetime of the application.
+    """
+    return build_agent_instance(llm_settings.AGENT_MODEL)
+
+
+@router.post("/")
+async def agent_conversation(
+    request: Request,
+    agent: Annotated[Agent[StateDeps[SearchState], str], Depends(get_agent)],
+) -> Response:
+    """Agent conversation endpoint using pydantic-ai ag_ui protocol.
+
+    This endpoint handles the interactive agent conversation for search.
+    """
+    initial_state = SearchState()
+    return await handle_ag_ui_request(agent, request, deps=StateDeps(initial_state))

orchestrator/api/api_v1/endpoints/search.py

@@ -11,251 +11,121 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Literal, overload
-
+import structlog
 from fastapi import APIRouter, HTTPException, Query, status
-from sqlalchemy import case, select
-from sqlalchemy.orm import selectinload
-
-from orchestrator.db import (
-    ProcessTable,
-    ProductTable,
-    WorkflowTable,
-    db,
-)
-from orchestrator.domain.base import SubscriptionModel
-from orchestrator.domain.context_cache import cache_subscription_models
+
+from orchestrator.db import db
 from orchestrator.schemas.search import (
+    ExportResponse,
     PageInfoSchema,
     PathsResponse,
-    ProcessSearchResult,
-    ProcessSearchSchema,
-    ProductSearchResult,
-    ProductSearchSchema,
     SearchResultsSchema,
-    SubscriptionSearchResult,
-    WorkflowSearchResult,
-    WorkflowSearchSchema,
 )
-from orchestrator.search.core.exceptions import InvalidCursorError
+from orchestrator.search.core.exceptions import InvalidCursorError, QueryStateNotFoundError
 from orchestrator.search.core.types import EntityType, UIType
 from orchestrator.search.filters.definitions import generate_definitions
-from orchestrator.search.indexing.registry import ENTITY_CONFIG_REGISTRY
-from orchestrator.search.retrieval import execute_search
+from orchestrator.search.retrieval import SearchQueryState, execute_search, execute_search_for_export
 from orchestrator.search.retrieval.builder import build_paths_query, create_path_autocomplete_lquery, process_path_rows
-from orchestrator.search.retrieval.pagination import (
-    create_next_page_cursor,
-    process_pagination_cursor,
-)
+from orchestrator.search.retrieval.pagination import PageCursor, encode_next_page_cursor
 from orchestrator.search.retrieval.validation import is_lquery_syntactically_valid
 from orchestrator.search.schemas.parameters import (
-    BaseSearchParameters,
     ProcessSearchParameters,
     ProductSearchParameters,
+    SearchParameters,
     SubscriptionSearchParameters,
     WorkflowSearchParameters,
 )
 from orchestrator.search.schemas.results import SearchResult, TypeDefinition
-from orchestrator.services.subscriptions import format_special_types
 
 router = APIRouter()
+logger = structlog.get_logger(__name__)
 
 
-def _create_search_result_item(
-    entity: WorkflowTable | ProductTable | ProcessTable, entity_type: EntityType, search_info: SearchResult
-) -> WorkflowSearchResult | ProductSearchResult | ProcessSearchResult | None:
-    match entity_type:
-        case EntityType.WORKFLOW:
-            workflow_data = WorkflowSearchSchema.model_validate(entity)
-            return WorkflowSearchResult(
-                workflow=workflow_data,
-                score=search_info.score,
-                perfect_match=search_info.perfect_match,
-                matching_field=search_info.matching_field,
-            )
-        case EntityType.PRODUCT:
-            product_data = ProductSearchSchema.model_validate(entity)
-            return ProductSearchResult(
-                product=product_data,
-                score=search_info.score,
-                perfect_match=search_info.perfect_match,
-                matching_field=search_info.matching_field,
-            )
-        case EntityType.PROCESS:
-            process_data = ProcessSearchSchema.model_validate(entity)
-            return ProcessSearchResult(
-                process=process_data,
-                score=search_info.score,
-                perfect_match=search_info.perfect_match,
-                matching_field=search_info.matching_field,
-            )
-        case _:
-            return None
-
-
-@overload
 async def _perform_search_and_fetch(
-    search_params: BaseSearchParameters,
-    entity_type: Literal[EntityType.WORKFLOW],
-    eager_loads: list[Any],
+    search_params: SearchParameters | None = None,
     cursor: str | None = None,
-) -> SearchResultsSchema[WorkflowSearchResult]: ...
-
+    query_id: str | None = None,
+) -> SearchResultsSchema[SearchResult]:
+    """Execute search with optional pagination.
+
+    Args:
+        search_params: Search parameters for new search
+        cursor: Pagination cursor (loads saved query state)
+        query_id: Saved query ID to retrieve and execute
+
+    Returns:
+        Search results with entity_id, score, and matching_field.
+    """
+    try:
+        page_cursor: PageCursor | None = None
+
+        if cursor:
+            page_cursor = PageCursor.decode(cursor)
+            query_state = SearchQueryState.load_from_id(page_cursor.query_id)
+        elif query_id:
+            query_state = SearchQueryState.load_from_id(query_id)
+        elif search_params:
+            query_state = SearchQueryState(parameters=search_params, query_embedding=None)
+        else:
+            raise HTTPException(
+                status_code=status.HTTP_400_BAD_REQUEST,
+                detail="Either search_params, cursor, or query_id must be provided",
+            )
 
-@overload
-async def _perform_search_and_fetch(
-    search_params: BaseSearchParameters,
-    entity_type: Literal[EntityType.PRODUCT],
-    eager_loads: list[Any],
-    cursor: str | None = None,
-) -> SearchResultsSchema[ProductSearchResult]: ...
+        search_response = await execute_search(
+            query_state.parameters, db.session, page_cursor, query_state.query_embedding
+        )
+        if not search_response.results:
+            return SearchResultsSchema(search_metadata=search_response.metadata)
 
+        next_page_cursor = encode_next_page_cursor(search_response, page_cursor, query_state.parameters)
+        has_next_page = next_page_cursor is not None
+        page_info = PageInfoSchema(has_next_page=has_next_page, next_page_cursor=next_page_cursor)
 
-@overload
-async def _perform_search_and_fetch(
-    search_params: BaseSearchParameters,
-    entity_type: Literal[EntityType.PROCESS],
-    eager_loads: list[Any],
-    cursor: str | None = None,
-) -> SearchResultsSchema[ProcessSearchResult]: ...
+        return SearchResultsSchema(
+            data=search_response.results, page_info=page_info, search_metadata=search_response.metadata
+        )
+    except (InvalidCursorError, ValueError) as e:
+        raise HTTPException(status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail=str(e))
+    except QueryStateNotFoundError as e:
+        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(e))
+    except Exception as e:
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail=f"Search failed: {str(e)}",
+        )
 
 
-async def _perform_search_and_fetch(
-    search_params: BaseSearchParameters,
-    entity_type: EntityType,
-    eager_loads: list[Any],
-    cursor: str | None = None,
-) -> SearchResultsSchema[Any]:
-    try:
-        pagination_params = await process_pagination_cursor(cursor, search_params)
-    except InvalidCursorError:
-        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid pagination cursor")
-
-    search_response = await execute_search(
-        search_params=search_params,
-        db_session=db.session,
-        pagination_params=pagination_params,
-    )
-    if not search_response.results:
-        return SearchResultsSchema(search_metadata=search_response.metadata)
-
-    next_page_cursor = create_next_page_cursor(search_response.results, pagination_params, search_params.limit)
-    has_next_page = next_page_cursor is not None
-    page_info = PageInfoSchema(has_next_page=has_next_page, next_page_cursor=next_page_cursor)
-
-    config = ENTITY_CONFIG_REGISTRY[entity_type]
-    entity_ids = [res.entity_id for res in search_response.results]
-    pk_column = getattr(config.table, config.pk_name)
-    ordering_case = case({entity_id: i for i, entity_id in enumerate(entity_ids)}, value=pk_column)
-
-    stmt = select(config.table).options(*eager_loads).filter(pk_column.in_(entity_ids)).order_by(ordering_case)
-    entities = db.session.scalars(stmt).all()
-
-    search_info_map = {res.entity_id: res for res in search_response.results}
-    data = []
-    for entity in entities:
-        entity_id = getattr(entity, config.pk_name)
-        search_info = search_info_map.get(str(entity_id))
-        if not search_info:
-            continue
-
-        search_result_item = _create_search_result_item(entity, entity_type, search_info)
-        if search_result_item:
-            data.append(search_result_item)
-
-    return SearchResultsSchema(data=data, page_info=page_info, search_metadata=search_response.metadata)
-
-
-@router.post(
-    "/subscriptions",
-    response_model=SearchResultsSchema[SubscriptionSearchResult],
-)
+@router.post("/subscriptions", response_model=SearchResultsSchema[SearchResult])
 async def search_subscriptions(
     search_params: SubscriptionSearchParameters,
     cursor: str | None = None,
-) -> SearchResultsSchema[SubscriptionSearchResult]:
-    try:
-        pagination_params = await process_pagination_cursor(cursor, search_params)
-    except InvalidCursorError:
-        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid pagination cursor")
-
-    search_response = await execute_search(
-        search_params=search_params,
-        db_session=db.session,
-        pagination_params=pagination_params,
-    )
-
-    if not search_response.results:
-        return SearchResultsSchema(search_metadata=search_response.metadata)
-
-    next_page_cursor = create_next_page_cursor(search_response.results, pagination_params, search_params.limit)
-    has_next_page = next_page_cursor is not None
-    page_info = PageInfoSchema(has_next_page=has_next_page, next_page_cursor=next_page_cursor)
-
-    search_info_map = {res.entity_id: res for res in search_response.results}
-
-    with cache_subscription_models():
-        subscriptions_data = {
-            sub_id: SubscriptionModel.from_subscription(sub_id).model_dump(exclude_unset=False)
-            for sub_id in search_info_map
-        }
-
-    results_data = [
-        SubscriptionSearchResult(
-            subscription=format_special_types(subscriptions_data[sub_id]),
-            score=search_info.score,
-            perfect_match=search_info.perfect_match,
-            matching_field=search_info.matching_field,
-        )
-        for sub_id, search_info in search_info_map.items()
-    ]
-
-    return SearchResultsSchema(data=results_data, page_info=page_info, search_metadata=search_response.metadata)
+) -> SearchResultsSchema[SearchResult]:
+    return await _perform_search_and_fetch(search_params, cursor)
 
 
-@router.post("/workflows", response_model=SearchResultsSchema[WorkflowSearchResult])
+@router.post("/workflows", response_model=SearchResultsSchema[SearchResult])
 async def search_workflows(
     search_params: WorkflowSearchParameters,
     cursor: str | None = None,
-) -> SearchResultsSchema[WorkflowSearchResult]:
-    return await _perform_search_and_fetch(
-        search_params=search_params,
-        entity_type=EntityType.WORKFLOW,
-        eager_loads=[selectinload(WorkflowTable.products)],
-        cursor=cursor,
-    )
+) -> SearchResultsSchema[SearchResult]:
+    return await _perform_search_and_fetch(search_params, cursor)
 
 
-@router.post("/products", response_model=SearchResultsSchema[ProductSearchResult])
+@router.post("/products", response_model=SearchResultsSchema[SearchResult])
 async def search_products(
     search_params: ProductSearchParameters,
     cursor: str | None = None,
-) -> SearchResultsSchema[ProductSearchResult]:
-    return await _perform_search_and_fetch(
-        search_params=search_params,
-        entity_type=EntityType.PRODUCT,
-        eager_loads=[
-            selectinload(ProductTable.workflows),
-            selectinload(ProductTable.fixed_inputs),
-            selectinload(ProductTable.product_blocks),
-        ],
-        cursor=cursor,
-    )
-
-
-@router.post("/processes", response_model=SearchResultsSchema[ProcessSearchResult])
+) -> SearchResultsSchema[SearchResult]:
+    return await _perform_search_and_fetch(search_params, cursor)
+
+
+@router.post("/processes", response_model=SearchResultsSchema[SearchResult])
 async def search_processes(
     search_params: ProcessSearchParameters,
     cursor: str | None = None,
-) -> SearchResultsSchema[ProcessSearchResult]:
-    return await _perform_search_and_fetch(
-        search_params=search_params,
-        entity_type=EntityType.PROCESS,
-        eager_loads=[
-            selectinload(ProcessTable.workflow),
-        ],
-        cursor=cursor,
-    )
+) -> SearchResultsSchema[SearchResult]:
+    return await _perform_search_and_fetch(search_params, cursor)
 
 
 @router.get(
@@ -294,3 +164,52 @@ async def list_paths(
 async def get_definitions() -> dict[UIType, TypeDefinition]:
     """Provide a static definition of operators and schemas for each UI type."""
     return generate_definitions()
+
+
+@router.get(
+    "/queries/{query_id}",
+    response_model=SearchResultsSchema[SearchResult],
+    summary="Retrieve saved search results by query_id",
+)
+async def get_by_query_id(
+    query_id: str,
+    cursor: str | None = None,
+) -> SearchResultsSchema[SearchResult]:
+    """Retrieve and execute a saved search by query_id."""
+    return await _perform_search_and_fetch(query_id=query_id, cursor=cursor)
+
+
+@router.get(
+    "/queries/{query_id}/export",
+    summary="Export query results by query_id",
+    response_model=ExportResponse,
+)
+async def export_by_query_id(query_id: str) -> ExportResponse:
+    """Export search results using query_id.
+
+    The query is retrieved from the database, re-executed, and results are returned
+    as flattened records suitable for CSV download.
+
+    Args:
+        query_id: Query UUID
+
+    Returns:
+        ExportResponse containing 'page' with an array of flattened entity records.
+
+    Raises:
+        HTTPException: 404 if query not found, 400 if invalid data
+    """
+    try:
+        query_state = SearchQueryState.load_from_id(query_id)
+        export_records = await execute_search_for_export(query_state, db.session)
+        return ExportResponse(page=export_records)
+    except ValueError as e:
+        raise HTTPException(status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail=str(e))
+    except QueryStateNotFoundError as e:
+        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(e))
+    except Exception as e:
+        logger.error(e)
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail=f"Error executing export: {str(e)}",
+        )

orchestrator/cli/database.py

@@ -23,6 +23,8 @@ from structlog import get_logger
 
 import orchestrator.workflows
 from orchestrator.cli.domain_gen_helpers.types import ModelUpdates
+from orchestrator.cli.generate import create_writer, get_template_environment
+from orchestrator.cli.generator.generator.migration import create_data_head_if_not_exists
 from orchestrator.cli.helpers.print_helpers import COLOR, str_fmt
 from orchestrator.cli.migrate_domain_models import create_domain_models_migration_sql
 from orchestrator.cli.migrate_tasks import create_tasks_migration_wizard
@@ -223,6 +225,7 @@ def revision(
             --head TEXT                     Determine the head you need to add your migration to.
         ```
     """
+    create_data_head_if_not_exists({"writer": create_writer(), "environment": get_template_environment()})
     command.revision(alembic_cfg(), message, version_path=version_path, autogenerate=autogenerate, head=head)
 
 

orchestrator/cli/generate.py

@@ -14,6 +14,7 @@
 # ruff: noqa: S603
 import subprocess
 from pathlib import Path
+from typing import Callable
 
 import structlog
 import typer
@@ -78,7 +79,7 @@ def ruff(content: str) -> str:
     return content
 
 
-def create_context(config_file: Path, dryrun: bool, force: bool, python_version: str, tdd: bool | None = False) -> dict:
+def create_writer(dryrun: bool = False, force: bool = False) -> Callable[..., None]:
     def writer(path: Path, content: str, append: bool = False) -> None:
         content = ruff(content) if path.suffix == ".py" else content
         if dryrun:
@@ -88,9 +89,15 @@ def create_context(config_file: Path, dryrun: bool, force: bool, python_version:
         else:
             write_file(path, content, append=append, force=force)
 
+    return writer
+
+
+def get_template_environment() -> Environment:
     search_path = (settings.CUSTOM_TEMPLATES, Path(__file__).parent / "generator" / "templates")
-    environment = Environment(loader=FileSystemLoader(search_path), autoescape=True, keep_trailing_newline=True)
+    return Environment(loader=FileSystemLoader(search_path), autoescape=True, keep_trailing_newline=True)
 
+
+def create_context(config_file: Path, dryrun: bool, force: bool, python_version: str, tdd: bool | None = False) -> dict:
     config = read_config(config_file)
     config["variable"] = get_variable(config)
     for pb in config["product_blocks"]:
@@ -98,10 +105,10 @@ def create_context(config_file: Path, dryrun: bool, force: bool, python_version:
 
     return {
         "config": config,
-        "environment": environment,
+        "environment": get_template_environment(),
         "python_version": python_version,
         "tdd": tdd,
-        "writer": writer,
+        "writer": create_writer(dryrun=dryrun, force=force),
     }
 
 

orchestrator/cli/generator/generator/migration.py

@@ -79,6 +79,12 @@ def create_data_head(context: dict, depends_on: str) -> None:
     writer(path, content)
 
 
+def create_data_head_if_not_exists(context: dict) -> None:
+    heads = get_heads()
+    if "data" not in heads:
+        create_data_head(context=context, depends_on=heads["schema"])
+
+
 def extract_revision_info(content: list[str]) -> dict:
     def process() -> Generator:
         for line in content:
@@ -136,9 +142,7 @@ def generate_product_migration(context: dict) -> None:
     environment = context["environment"]
     writer = context["writer"]
 
-    heads = get_heads()
-    if "data" not in heads:
-        create_data_head(context=context, depends_on=heads["schema"])
+    create_data_head_if_not_exists(context=context)
 
     if not (migration_file := create_migration_file(message=f"add {config['name']}", head="data")):
         logger.error("Could not create migration file")