orchestrator-core 4.4.1__py3-none-any.whl → 4.5.0__py3-none-any.whl

orchestrator/search/agent/tools.py

@@ -0,0 +1,258 @@
+# Copyright 2019-2025 SURF, GÉANT.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from collections.abc import Awaitable, Callable
+from typing import Any, TypeVar
+
+import structlog
+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
+from pydantic_ai.messages import ModelRequest, UserPromptPart
+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.search.core.types import ActionType, EntityType, FilterOp
+from orchestrator.search.filters import FilterTree
+from orchestrator.search.retrieval.exceptions import FilterValidationError, PathNotFoundError
+from orchestrator.search.retrieval.validation import validate_filter_tree
+from orchestrator.search.schemas.parameters import PARAMETER_REGISTRY, BaseSearchParameters
+
+from .state import SearchState
+
+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)
+
+
+def last_user_message(ctx: RunContext[StateDeps[SearchState]]) -> str | None:
+    for msg in reversed(ctx.messages):
+        if isinstance(msg, ModelRequest):
+            for part in msg.parts:
+                if isinstance(part, UserPromptPart) and isinstance(part.content, str):
+                    return part.content
+    return None
+
+
+@search_toolset.tool
+async def set_search_parameters(
+    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.
+
+    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.
+    """
+    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", "")
+
+    logger.debug(
+        "Setting search parameters",
+        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)
+
+    return StateSnapshotEvent(
+        type=EventType.STATE_SNAPSHOT,
+        snapshot=ctx.deps.state.model_dump(),
+    )
+
+
+@search_toolset.tool(retries=2)
+async def set_filter_tree(
+    ctx: RunContext[StateDeps[SearchState]],
+    filters: FilterTree | None,
+) -> 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.
+    """
+    if ctx.deps.state.parameters is None:
+        raise ModelRetry("Search parameters are not initialized. Call set_search_parameters first.")
+
+    entity_type = EntityType(ctx.deps.state.parameters["entity_type"])
+
+    logger.debug(
+        "Setting filter tree",
+        entity_type=entity_type.value,
+        has_filters=filters is not None,
+        filter_summary=f"{len(filters.get_all_leaves())} filters" if filters else "no filters",
+    )
+
+    try:
+        await validate_filter_tree(filters, entity_type)
+    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:
+        # ModelRetry will trigger an agent retry, containing the specific validation error.
+        logger.debug(f"Filter 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)
+    ctx.deps.state.parameters["filters"] = filter_data
+    return StateSnapshotEvent(type=EventType.STATE_SNAPSHOT, snapshot=ctx.deps.state.model_dump())
+
+
+@search_toolset.tool
+async def execute_search(
+    ctx: RunContext[StateDeps[SearchState]],
+    limit: int = 10,
+) -> StateSnapshotEvent:
+    """Execute the search with the current parameters."""
+    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)
+    logger.debug(
+        "Executing database search",
+        search_entity_type=entity_type.value,
+        limit=limit,
+        has_filters=params.filters is not None,
+        query=params.query,
+        action=params.action,
+    )
+
+    if params.filters:
+        logger.debug("Search filters", filters=params.filters)
+
+    params.limit = limit
+
+    fn = SEARCH_FN_MAP[entity_type]
+    search_results = await fn(params)
+
+    logger.debug(
+        "Search completed",
+        total_results=len(search_results.data) if search_results.data else 0,
+    )
+
+    ctx.deps.state.results = search_results.data
+
+    return StateSnapshotEvent(type=EventType.STATE_SNAPSHOT, snapshot=ctx.deps.state.model_dump())
+
+
+@search_toolset.tool
+async def discover_filter_paths(
+    ctx: RunContext[StateDeps[SearchState]],
+    field_names: list[str],
+    entity_type: EntityType | None = None,
+) -> dict[str, dict[str, Any]]:
+    """Discovers available filter paths for a list of field names.
+
+    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
+
+    all_results = {}
+    for field_name in field_names:
+        paths_response = await list_paths(prefix="", q=field_name, entity_type=entity_type, limit=100)
+
+        matching_leaves = []
+        for leaf in paths_response.leaves:
+            if field_name.lower() in leaf.name.lower():
+                matching_leaves.append(
+                    {
+                        "name": leaf.name,
+                        "value_kind": leaf.ui_types,
+                        "paths": leaf.paths,
+                    }
+                )
+
+        matching_components = []
+        for comp in paths_response.components:
+            if field_name.lower() in comp.name.lower():
+                matching_components.append(
+                    {
+                        "name": comp.name,
+                        "value_kind": comp.ui_types,
+                    }
+                )
+
+        result_for_field: dict[str, Any]
+        if not matching_leaves and not matching_components:
+            result_for_field = {
+                "status": "NOT_FOUND",
+                "guidance": f"No filterable paths found containing '{field_name}'. Do not create a filter for this.",
+                "leaves": [],
+                "components": [],
+            }
+        else:
+            result_for_field = {
+                "status": "OK",
+                "guidance": f"Found {len(matching_leaves)} field(s) and {len(matching_components)} component(s) for '{field_name}'.",
+                "leaves": matching_leaves,
+                "components": matching_components,
+            }
+
+        all_results[field_name] = result_for_field
+    logger.debug("Returning found fieldname - path mapping", all_results=all_results)
+    return all_results
+
+
+@search_toolset.tool
+async def get_valid_operators() -> dict[str, list[FilterOp]]:
+    """Gets the mapping of field types to their valid filter operators."""
+    definitions = await get_definitions()
+
+    operator_map = {}
+    for ui_type, type_def in definitions.items():
+        key = ui_type.value
+
+        if hasattr(type_def, "operators"):
+            operator_map[key] = type_def.operators
+    return operator_map

orchestrator/search/core/__init__.py

@@ -0,0 +1,12 @@
+# 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.

orchestrator/search/core/embedding.py

@@ -0,0 +1,73 @@
+# 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.
+
+
+import structlog
+from litellm import aembedding as llm_aembedding
+from litellm import embedding as llm_embedding
+from litellm import exceptions as llm_exc
+
+from orchestrator.llm_settings import llm_settings
+
+logger = structlog.get_logger(__name__)
+
+
+class EmbeddingIndexer:
+
+    @classmethod
+    def get_embeddings_from_api_batch(cls, texts: list[str], dry_run: bool) -> list[list[float]]:
+        if not texts:
+            return []
+        if dry_run:
+            logger.debug("Dry Run: returning empty embeddings")
+            return [[] for _ in texts]
+
+        try:
+            resp = llm_embedding(
+                model=llm_settings.EMBEDDING_MODEL,
+                input=[t.lower() for t in texts],
+                api_key=llm_settings.OPENAI_API_KEY,
+                api_base=llm_settings.OPENAI_BASE_URL,
+                timeout=llm_settings.LLM_TIMEOUT,
+                max_retries=llm_settings.LLM_MAX_RETRIES,
+            )
+            data = sorted(resp.data, key=lambda e: e["index"])
+            return [row["embedding"] for row in data]
+        except (llm_exc.APIError, llm_exc.APIConnectionError, llm_exc.RateLimitError, llm_exc.Timeout) as e:
+            logger.error("Embedding request failed", error=str(e))
+            return [[] for _ in texts]
+        except Exception as e:
+            logger.error("Unexpected embedding error", error=str(e))
+            return [[] for _ in texts]
+
+
+class QueryEmbedder:
+    """A stateless, async utility for embedding real-time user queries."""
+
+    @classmethod
+    async def generate_for_text_async(cls, text: str) -> list[float]:
+        if not text:
+            return []
+        try:
+            resp = await llm_aembedding(
+                model=llm_settings.EMBEDDING_MODEL,
+                input=[text.lower()],
+                api_key=llm_settings.OPENAI_API_KEY,
+                api_base=llm_settings.OPENAI_BASE_URL,
+                timeout=5.0,
+                max_retries=0,  # No retries, prioritize speed.
+            )
+            return resp.data[0]["embedding"]
+        except Exception as e:
+            logger.error("Async embedding generation failed", error=str(e))
+            return []

orchestrator/search/core/exceptions.py

@@ -0,0 +1,36 @@
+# 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.
+
+
+class SearchUtilsError(Exception):
+    """Base exception for this module."""
+
+    pass
+
+
+class ProductNotInRegistryError(SearchUtilsError):
+    """Raised when a product is not found in the model registry."""
+
+    pass
+
+
+class ModelLoadError(SearchUtilsError):
+    """Raised when a Pydantic model fails to load from a subscription."""
+
+    pass
+
+
+class InvalidCursorError(SearchUtilsError):
+    """Raised when cursor cannot be decoded."""
+
+    pass

orchestrator/search/core/types.py

@@ -0,0 +1,296 @@
+# 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 dataclasses import dataclass
+from datetime import date, datetime
+from enum import Enum, IntEnum
+from typing import Annotated, Any, Literal, NamedTuple, TypeAlias, TypedDict, get_args, get_origin
+from uuid import UUID
+
+from sqlalchemy.orm.attributes import InstrumentedAttribute
+from sqlalchemy.sql.elements import ColumnElement
+from sqlalchemy_utils.types.ltree import Ltree
+
+from orchestrator.types import filter_nonetype, get_origin_and_args, is_optional_type, is_union_type
+
+from .validators import is_bool_string, is_iso_date, is_uuid
+
+SQLAColumn: TypeAlias = ColumnElement[Any] | InstrumentedAttribute[Any]
+
+LTREE_SEPARATOR = "."
+
+
+@dataclass
+class SearchMetadata:
+    """Metadata about the search operation performed."""
+
+    search_type: str
+    description: str
+
+    @classmethod
+    def structured(cls) -> "SearchMetadata":
+        return cls(
+            search_type="structured", description="This search performs a filter-based search using structured queries."
+        )
+
+    @classmethod
+    def fuzzy(cls) -> "SearchMetadata":
+        return cls(
+            search_type="fuzzy",
+            description="This search performs a trigram similarity search.",
+        )
+
+    @classmethod
+    def semantic(cls) -> "SearchMetadata":
+        return cls(
+            search_type="semantic",
+            description="This search performs a vector similarity search, using L2 distance on embeddings with minimum distance scoring (normalized).",
+        )
+
+    @classmethod
+    def hybrid(cls) -> "SearchMetadata":
+        return cls(
+            search_type="hybrid",
+            description="This search performs reciprocal rank fusion combining trigram similarity, word_similarity, and L2 vector distance.",
+        )
+
+    @classmethod
+    def empty(cls) -> "SearchMetadata":
+        return cls(search_type="empty", description="Empty search - no criteria provided")
+
+
+class BooleanOperator(str, Enum):
+    AND = "AND"
+    OR = "OR"
+
+
+class FilterOp(str, Enum):
+    EQ = "eq"
+    NEQ = "neq"
+    LT = "lt"
+    LIKE = "like"
+    LTE = "lte"
+    GT = "gt"
+    GTE = "gte"
+    BETWEEN = "between"
+
+    MATCHES_LQUERY = "matches_lquery"  # The ~ operator for wildcard matching
+    IS_ANCESTOR = "is_ancestor"  # The @> operator
+    IS_DESCENDANT = "is_descendant"  # The <@ operator
+    PATH_MATCH = "path_match"
+
+    HAS_COMPONENT = "has_component"  # Path contains this segment
+    NOT_HAS_COMPONENT = "not_has_component"  # Path doesn't contain segment
+    ENDS_WITH = "ends_with"
+
+
+class EntityType(str, Enum):
+    SUBSCRIPTION = "SUBSCRIPTION"
+    PRODUCT = "PRODUCT"
+    WORKFLOW = "WORKFLOW"
+    PROCESS = "PROCESS"
+
+
+class ActionType(str, Enum):
+    """Defines the explicit, safe actions the agent can request."""
+
+    SELECT = "select"  # Retrieve a list of matching records.
+    # COUNT = "count"  # For phase1; the agent will not support this yet.
+
+
+class UIType(str, Enum):
+    STRING = "string"
+    NUMBER = "number"
+    BOOLEAN = "boolean"
+    DATETIME = "datetime"
+    COMPONENT = "component"
+
+    @classmethod
+    def from_field_type(cls, ft: "FieldType") -> "UIType":
+        """Create a UIType from a backend FieldType to indicate how a value must be rendered."""
+        if ft in (FieldType.INTEGER, FieldType.FLOAT):
+            return cls.NUMBER
+        if ft == FieldType.BOOLEAN:
+            return cls.BOOLEAN
+        if ft == FieldType.DATETIME:
+            return cls.DATETIME
+        return cls.STRING
+
+
+class FieldType(str, Enum):
+    STRING = "string"
+    INTEGER = "integer"
+    FLOAT = "float"
+    BOOLEAN = "boolean"
+    DATETIME = "datetime"
+    UUID = "uuid"
+    BLOCK = "block"
+    RESOURCE_TYPE = "resource_type"
+
+    @classmethod
+    def infer(cls, val: Any) -> "FieldType":
+        if isinstance(val, TypedValue):
+            return cls._infer_typed_value(val)
+
+        if isinstance(val, bool):
+            return cls.BOOLEAN
+        if isinstance(val, int):
+            return cls.INTEGER
+        if isinstance(val, float):
+            return cls.FLOAT
+        if isinstance(val, UUID):
+            return cls.UUID
+        if isinstance(val, (datetime, date)):
+            return cls.DATETIME
+        if isinstance(val, str):
+            return cls._infer_from_str(val)
+
+        return cls.STRING
+
+    @classmethod
+    def _infer_typed_value(cls, val: "TypedValue") -> "FieldType":
+        if val.type == cls.BLOCK:
+            return cls.BLOCK
+        if val.type == cls.RESOURCE_TYPE:
+            return cls.RESOURCE_TYPE
+        return cls.STRING
+
+    @classmethod
+    def _infer_from_str(cls, val: str) -> "FieldType":
+        if is_uuid(val):
+            return cls.UUID
+        if is_iso_date(val):
+            return cls.DATETIME
+        if is_bool_string(val):
+            return cls.BOOLEAN
+        if val.isdigit():
+            return cls.INTEGER
+        try:
+            float(val)
+            return cls.FLOAT
+        except ValueError:
+            return cls.STRING
+
+    @classmethod
+    def from_type_hint(cls, type_hint: object) -> "FieldType":
+        """Convert type hint to FieldType."""
+        _type_mapping = {
+            int: cls.INTEGER,
+            float: cls.FLOAT,
+            bool: cls.BOOLEAN,
+            str: cls.STRING,
+            datetime: cls.DATETIME,
+            UUID: cls.UUID,
+        }
+
+        if type_hint in _type_mapping:
+            return _type_mapping[type_hint]  # type: ignore[index]
+
+        if get_origin(type_hint) is Annotated:
+            inner_type = get_args(type_hint)[0]
+            return cls.from_type_hint(inner_type)
+
+        origin, args = get_origin_and_args(type_hint)
+
+        if origin is list:
+            return cls._handle_list_type(args)
+
+        if origin is Literal:
+            return cls._handle_literal_type(args)
+
+        if is_optional_type(type_hint) or is_union_type(type_hint):
+            return cls._handle_union_type(args)
+
+        if isinstance(type_hint, type):
+            return cls._handle_class_type(type_hint)
+
+        return cls.STRING
+
+    @classmethod
+    def _handle_list_type(cls, args: tuple) -> "FieldType":
+        if args:
+            element_type = args[0]
+            return cls.from_type_hint(element_type)
+        return cls.STRING
+
+    @classmethod
+    def _handle_literal_type(cls, args: tuple) -> "FieldType":
+        if not args:
+            return cls.STRING
+        first_value = args[0]
+        if isinstance(first_value, bool):
+            return cls.BOOLEAN
+        if isinstance(first_value, int):
+            return cls.INTEGER
+        if isinstance(first_value, str):
+            return cls.STRING
+        if isinstance(first_value, float):
+            return cls.FLOAT
+        return cls.STRING
+
+    @classmethod
+    def _handle_union_type(cls, args: tuple) -> "FieldType":
+        non_none_types = list(filter_nonetype(args))
+        if non_none_types:
+            return cls.from_type_hint(non_none_types[0])
+        return cls.STRING
+
+    @classmethod
+    def _handle_class_type(cls, type_hint: type) -> "FieldType":
+        if issubclass(type_hint, IntEnum):
+            return cls.INTEGER
+        if issubclass(type_hint, Enum):
+            return cls.STRING
+
+        from orchestrator.domain.base import ProductBlockModel
+
+        if issubclass(type_hint, ProductBlockModel):
+            return cls.BLOCK
+
+        return cls.STRING
+
+    def is_embeddable(self, value: str | None) -> bool:
+        """Check if a field should be embedded."""
+        if value is None:
+            return False
+
+        # If inference suggests it's not actually a string, don't embed it
+        return FieldType._infer_from_str(value) == FieldType.STRING
+
+
+@dataclass(frozen=True)
+class TypedValue:
+    value: Any
+    type: FieldType
+
+
+class ExtractedField(NamedTuple):
+    path: str
+    value: str
+    value_type: FieldType
+
+    @classmethod
+    def from_raw(cls, path: str, raw_value: Any) -> "ExtractedField":
+        value = str(raw_value.value if isinstance(raw_value, TypedValue) else raw_value)
+        value_type = FieldType.infer(raw_value)
+        return cls(path=path, value=value, value_type=value_type)
+
+
+class IndexableRecord(TypedDict):
+    entity_id: str
+    entity_type: str
+    path: Ltree
+    value: Any
+    value_type: Any
+    content_hash: str
+    embedding: list[float] | None