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

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}")

orchestrator/search/indexing/indexer.py

@@ -96,6 +96,7 @@ class Indexer:
         self.chunk_size = chunk_size
         self.embedding_model = llm_settings.EMBEDDING_MODEL
         self.logger = logger.bind(entity_kind=config.entity_kind.value)
+        self._entity_titles: dict[str, str] = {}
 
     def run(self, entities: Iterable[DatabaseEntity]) -> int:
         """Orchestrates the entire indexing process."""
@@ -138,6 +139,8 @@ class Indexer:
         if not entity_chunk:
             return 0, 0
 
+        self._entity_titles.clear()
+
         fields_to_upsert, paths_to_delete, identical_count = self._determine_changes(entity_chunk, session)
 
         if paths_to_delete and session is not None:
@@ -174,12 +177,15 @@ class Indexer:
                 entity, pk_name=self.config.pk_name, root_name=self.config.root_name
             )
 
+            entity_title = self.config.get_title_from_fields(current_fields)
+            self._entity_titles[entity_id] = entity_title
+
             entity_hashes = existing_hashes.get(entity_id, {})
             current_paths = set()
 
             for field in current_fields:
                 current_paths.add(field.path)
-                current_hash = self._compute_content_hash(field.path, field.value, field.value_type)
+                current_hash = self._compute_content_hash(field.path, field.value, field.value_type, entity_title)
                 if field.path not in entity_hashes or entity_hashes[field.path] != current_hash:
                     fields_to_upsert.append((entity_id, field))
                 else:
@@ -301,21 +307,23 @@ class Indexer:
         return f"{field.path}: {str(field.value)}"
 
     @staticmethod
-    def _compute_content_hash(path: str, value: Any, value_type: Any) -> str:
+    def _compute_content_hash(path: str, value: Any, value_type: Any, entity_title: str = "") -> str:
         v = "" if value is None else str(value)
-        content = f"{path}:{v}:{value_type}"
+        content = f"{path}:{v}:{value_type}:{entity_title}"
         return hashlib.sha256(content.encode("utf-8")).hexdigest()
 
     def _make_indexable_record(
         self, field: ExtractedField, entity_id: str, embedding: list[float] | None
     ) -> IndexableRecord:
+        entity_title = self._entity_titles[entity_id]
         return IndexableRecord(
             entity_id=entity_id,
             entity_type=self.config.entity_kind.value,
+            entity_title=entity_title,
             path=Ltree(field.path),
             value=field.value,
             value_type=field.value_type,
-            content_hash=self._compute_content_hash(field.path, field.value, field.value_type),
+            content_hash=self._compute_content_hash(field.path, field.value, field.value_type, entity_title),
             embedding=embedding if embedding else None,
         )
 
@@ -326,6 +334,7 @@ class Indexer:
         return stmt.on_conflict_do_update(
             index_elements=[AiSearchIndex.entity_id, AiSearchIndex.path],
             set_={
+                AiSearchIndex.entity_title: stmt.excluded.entity_title,
                 AiSearchIndex.value: stmt.excluded.value,
                 AiSearchIndex.value_type: stmt.excluded.value_type,
                 AiSearchIndex.content_hash: stmt.excluded.content_hash,

orchestrator/search/indexing/registry.py

@@ -25,7 +25,7 @@ from orchestrator.db import (
     WorkflowTable,
 )
 from orchestrator.db.database import BaseModel
-from orchestrator.search.core.types import EntityType
+from orchestrator.search.core.types import EntityType, ExtractedField
 
 from .traverse import (
     BaseTraverser,
@@ -48,6 +48,7 @@ class EntityConfig(Generic[ModelT]):
     traverser: "type[BaseTraverser]"
     pk_name: str
     root_name: str
+    title_paths: list[str]  # List of field paths to check for title (with fallback)
 
     def get_all_query(self, entity_id: str | None = None) -> Query | Select:
         query = self.table.query
@@ -56,6 +57,14 @@ class EntityConfig(Generic[ModelT]):
             query = query.filter(pk_column == UUID(entity_id))
         return query
 
+    def get_title_from_fields(self, fields: list[ExtractedField]) -> str:
+        """Extract title from fields using configured paths."""
+        for title_path in self.title_paths:
+            for field in fields:
+                if field.path == title_path and field.value:
+                    return str(field.value)
+        return "UNKNOWN"
+
 
 @dataclass(frozen=True)
 class WorkflowConfig(EntityConfig[WorkflowTable]):
@@ -76,6 +85,7 @@ ENTITY_CONFIG_REGISTRY: dict[EntityType, EntityConfig] = {
         traverser=SubscriptionTraverser,
         pk_name="subscription_id",
         root_name="subscription",
+        title_paths=["subscription.description"],
     ),
     EntityType.PRODUCT: EntityConfig(
         entity_kind=EntityType.PRODUCT,
@@ -83,6 +93,7 @@ ENTITY_CONFIG_REGISTRY: dict[EntityType, EntityConfig] = {
         traverser=ProductTraverser,
         pk_name="product_id",
         root_name="product",
+        title_paths=["product.description", "product.name"],
     ),
     EntityType.PROCESS: EntityConfig(
         entity_kind=EntityType.PROCESS,
@@ -90,6 +101,7 @@ ENTITY_CONFIG_REGISTRY: dict[EntityType, EntityConfig] = {
         traverser=ProcessTraverser,
         pk_name="process_id",
         root_name="process",
+        title_paths=["process.workflow_name"],
     ),
     EntityType.WORKFLOW: WorkflowConfig(
         entity_kind=EntityType.WORKFLOW,
@@ -97,5 +109,6 @@ ENTITY_CONFIG_REGISTRY: dict[EntityType, EntityConfig] = {
         traverser=WorkflowTraverser,
         pk_name="workflow_id",
         root_name="workflow",
+        title_paths=["workflow.description", "workflow.name"],
     ),
 }

orchestrator/search/llm_migration.py

@@ -37,6 +37,7 @@ def run_migration(connection: Connection) -> None:
         if llm_settings.LLM_FORCE_EXTENTION_MIGRATION or res.rowcount == 0:
             # Create PostgreSQL extensions
             logger.info("Attempting to run the extention creation;")
+            connection.execute(text('CREATE EXTENSION IF NOT EXISTS "uuid-ossp";'))
             connection.execute(text("CREATE EXTENSION IF NOT EXISTS ltree;"))
             connection.execute(text("CREATE EXTENSION IF NOT EXISTS unaccent;"))
             connection.execute(text("CREATE EXTENSION IF NOT EXISTS pg_trgm;"))
@@ -64,6 +65,7 @@ def run_migration(connection: Connection) -> None:
             CREATE TABLE IF NOT EXISTS {TABLE} (
                 entity_type TEXT NOT NULL,
                 entity_id UUID NOT NULL,
+                entity_title TEXT,
                 path LTREE NOT NULL,
                 value TEXT NOT NULL,
                 embedding VECTOR({TARGET_DIM}),
@@ -78,6 +80,23 @@ def run_migration(connection: Connection) -> None:
         # Drop default
         connection.execute(text(f"ALTER TABLE {TABLE} ALTER COLUMN value_type DROP DEFAULT;"))
 
+        # Add entity_title column if it doesn't exist (for existing installations)
+        connection.execute(
+            text(
+                f"""
+            DO $$
+            BEGIN
+                IF NOT EXISTS (
+                    SELECT 1 FROM information_schema.columns
+                    WHERE table_name = '{TABLE}' AND column_name = 'entity_title'
+                ) THEN
+                    ALTER TABLE {TABLE} ADD COLUMN entity_title TEXT;
+                END IF;
+            END $$;
+        """
+            )
+        )
+
         # Create indexes with IF NOT EXISTS
         connection.execute(text(f"CREATE INDEX IF NOT EXISTS ix_ai_search_index_entity_id ON {TABLE} (entity_id);"))
         connection.execute(
@@ -96,6 +115,42 @@ def run_migration(connection: Connection) -> None:
             )
         )
 
+        # Create agent_runs table
+        connection.execute(
+            text(
+                """
+            CREATE TABLE IF NOT EXISTS agent_runs (
+                run_id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
+                agent_type VARCHAR(50) NOT NULL,
+                created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP NOT NULL
+            );
+        """
+            )
+        )
+        connection.execute(text("CREATE INDEX IF NOT EXISTS ix_agent_runs_created_at ON agent_runs (created_at);"))
+
+        # Create search_queries table
+        connection.execute(
+            text(
+                f"""
+            CREATE TABLE IF NOT EXISTS search_queries (
+                query_id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
+                run_id UUID,
+                query_number INTEGER NOT NULL,
+                parameters JSONB NOT NULL,
+                query_embedding VECTOR({TARGET_DIM}),
+                executed_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP NOT NULL,
+                CONSTRAINT fk_search_queries_run_id FOREIGN KEY (run_id) REFERENCES agent_runs(run_id) ON DELETE CASCADE
+            );
+        """
+            )
+        )
+        connection.execute(text("CREATE INDEX IF NOT EXISTS ix_search_queries_run_id ON search_queries (run_id);"))
+        connection.execute(
+            text("CREATE INDEX IF NOT EXISTS ix_search_queries_executed_at ON search_queries (executed_at);")
+        )
+        connection.execute(text("CREATE INDEX IF NOT EXISTS ix_search_queries_query_id ON search_queries (query_id);"))
+
         connection.commit()
         logger.info("LLM migration completed successfully")
 

orchestrator/search/retrieval/__init__.py

@@ -11,6 +11,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .engine import execute_search
+from .engine import execute_search, execute_search_for_export
+from .query_state import SearchQueryState
 
-__all__ = ["execute_search"]
+__all__ = ["execute_search", "execute_search_for_export", "SearchQueryState"]

orchestrator/search/retrieval/builder.py

@@ -43,7 +43,11 @@ def build_candidate_query(params: BaseSearchParameters) -> Select:
         Select: The SQLAlchemy `Select` object representing the query.
     """
 
-    stmt = select(AiSearchIndex.entity_id).where(AiSearchIndex.entity_type == params.entity_type.value).distinct()
+    stmt = (
+        select(AiSearchIndex.entity_id, AiSearchIndex.entity_title)
+        .where(AiSearchIndex.entity_type == params.entity_type.value)
+        .distinct()
+    )
 
     if params.filters is not None:
         entity_id_col = AiSearchIndex.entity_id

orchestrator/search/retrieval/engine.py

@@ -17,13 +17,15 @@ import structlog
 from sqlalchemy.engine.row import RowMapping
 from sqlalchemy.orm import Session
 
+from orchestrator.search.core.embedding import QueryEmbedder
 from orchestrator.search.core.types import FilterOp, SearchMetadata
 from orchestrator.search.filters import FilterTree, LtreeFilter
 from orchestrator.search.schemas.parameters import BaseSearchParameters
 from orchestrator.search.schemas.results import MatchingField, SearchResponse, SearchResult
 
 from .builder import build_candidate_query
-from .pagination import PaginationParams
+from .pagination import PageCursor
+from .query_state import SearchQueryState
 from .retrievers import Retriever
 from .utils import generate_highlight_indices
 
@@ -74,9 +76,15 @@ def _format_response(
             # Structured search (filter-only)
             matching_field = _extract_matching_field_from_filters(search_params.filters)
 
+        entity_title = row.get("entity_title", "")
+        if not isinstance(entity_title, str):
+            entity_title = str(entity_title) if entity_title is not None else ""
+
         results.append(
             SearchResult(
                 entity_id=str(row.entity_id),
+                entity_type=search_params.entity_type,
+                entity_title=entity_title,
                 score=row.score,
                 perfect_match=row.get("perfect_match", 0),
                 matching_field=matching_field,
@@ -110,45 +118,80 @@ def _extract_matching_field_from_filters(filters: FilterTree) -> MatchingField |
     return MatchingField(text=text, path=pf.path, highlight_indices=[(0, len(text))])
 
 
-async def execute_search(
+async def _execute_search_internal(
     search_params: BaseSearchParameters,
     db_session: Session,
-    pagination_params: PaginationParams | None = None,
+    limit: int,
+    cursor: PageCursor | None = None,
+    query_embedding: list[float] | None = None,
 ) -> SearchResponse:
-    """Execute a hybrid search and return ranked results.
-
-    Builds a candidate entity query based on the given search parameters,
-    applies the appropriate ranking strategy, and executes the final ranked
-    query to retrieve results.
+    """Internal function to execute search with specified parameters.
 
     Args:
-        search_params (BaseSearchParameters): The search parameters specifying vector, fuzzy, or filter criteria.
-        db_session (Session): The active SQLAlchemy session for executing the query.
-        pagination_params (PaginationParams): Parameters controlling pagination of the search results.
-        limit (int, optional): The maximum number of search results to return, by default 5.
+        search_params: The search parameters specifying vector, fuzzy, or filter criteria.
+        db_session: The active SQLAlchemy session for executing the query.
+        limit: Maximum number of results to return.
+        cursor: Optional pagination cursor.
+        query_embedding: Optional pre-computed query embedding to use instead of generating a new one.
 
     Returns:
-        SearchResponse: A list of `SearchResult` objects containing entity IDs, scores,
-        and optional highlight metadata.
-
-    Notes:
-        If no vector query, filters, or fuzzy term are provided, a warning is logged
-        and an empty result set is returned.
+        SearchResponse with results and embedding (for internal use).
     """
-
     if not search_params.vector_query and not search_params.filters and not search_params.fuzzy_term:
         logger.warning("No search criteria provided (vector_query, fuzzy_term, or filters).")
         return SearchResponse(results=[], metadata=SearchMetadata.empty())
 
     candidate_query = build_candidate_query(search_params)
 
-    pagination_params = pagination_params or PaginationParams()
-    retriever = await Retriever.from_params(search_params, pagination_params)
+    if search_params.vector_query and not query_embedding:
+
+        query_embedding = await QueryEmbedder.generate_for_text_async(search_params.vector_query)
+
+    retriever = await Retriever.route(search_params, cursor, query_embedding)
     logger.debug("Using retriever", retriever_type=retriever.__class__.__name__)
 
     final_stmt = retriever.apply(candidate_query)
-    final_stmt = final_stmt.limit(search_params.limit)
+    final_stmt = final_stmt.limit(limit)
     logger.debug(final_stmt)
     result = db_session.execute(final_stmt).mappings().all()
 
-    return _format_response(result, search_params, retriever.metadata)
+    response = _format_response(result, search_params, retriever.metadata)
+    # Store embedding in response for agent to save to DB
+    response.query_embedding = query_embedding
+    return response
+
+
+async def execute_search(
+    search_params: BaseSearchParameters,
+    db_session: Session,
+    cursor: PageCursor | None = None,
+    query_embedding: list[float] | None = None,
+) -> SearchResponse:
+    """Execute a search and return ranked results."""
+    return await _execute_search_internal(search_params, db_session, search_params.limit, cursor, query_embedding)
+
+
+async def execute_search_for_export(
+    query_state: SearchQueryState,
+    db_session: Session,
+) -> list[dict]:
+    """Execute a search for export and fetch flattened entity data.
+
+    Args:
+        query_state: Query state containing parameters and query_embedding.
+        db_session: The active SQLAlchemy session for executing the query.
+
+    Returns:
+        List of flattened entity records suitable for export.
+    """
+    from orchestrator.search.export import fetch_export_data
+
+    search_response = await _execute_search_internal(
+        search_params=query_state.parameters,
+        db_session=db_session,
+        limit=query_state.parameters.export_limit,
+        query_embedding=query_state.query_embedding,
+    )
+
+    entity_ids = [res.entity_id for res in search_response.results]
+    return fetch_export_data(query_state.parameters.entity_type, entity_ids)