orchestrator-core 4.4.0rc2__py3-none-any.whl → 5.0.0a1__py3-none-any.whl

orchestrator/schemas/search.py

@@ -0,0 +1,101 @@
+from datetime import datetime
+from typing import Any, Generic, TypeVar
+from uuid import UUID
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from orchestrator.search.schemas.results import Highlight, PathInfo
+
+T = TypeVar("T")
+
+
+class PageInfoSchema(BaseModel):
+    total_items: int = Field(default=0, alias="totalItems")
+    start_cursor: int = Field(default=0, alias="startCursor")
+    has_previous_page: bool = Field(default=False, alias="hasPreviousPage")
+    has_next_page: bool = Field(default=False, alias="hasNextPage")
+    end_cursor: int = Field(default=0, alias="endCursor")
+    sort_fields: list[str] = Field(default_factory=list, alias="sortFields")
+    filter_fields: list[str] = Field(default_factory=list, alias="filterFields")
+
+
+class ProductSchema(BaseModel):
+    model_config = ConfigDict(from_attributes=True, populate_by_name=True)
+
+    name: str
+    tag: str
+    product_type: str = Field(alias="productType")
+
+
+class SubscriptionSearchResult(BaseModel):
+    score: float
+    highlight: Highlight | None = None
+
+    subscription: dict[str, Any]
+
+
+class ConnectionSchema(BaseModel, Generic[T]):
+    page: list[T]
+    page_info: PageInfoSchema = Field(alias="pageInfo")
+
+    model_config = ConfigDict(populate_by_name=True)
+
+
+class WorkflowProductSchema(BaseModel):
+    """Product associated with a workflow."""
+
+    model_config = ConfigDict(from_attributes=True, populate_by_name=True)
+
+    product_type: str = Field(alias="productType")
+    product_id: UUID = Field(alias="productId")
+    name: str
+
+
+class WorkflowSearchSchema(BaseModel):
+    """Schema for workflow search results."""
+
+    model_config = ConfigDict(from_attributes=True, populate_by_name=True)
+
+    name: str
+    products: list[WorkflowProductSchema]
+    description: str | None = None
+    created_at: datetime | None = Field(alias="createdAt", default=None)
+
+
+class ProductSearchSchema(BaseModel):
+    """Schema for product search results."""
+
+    model_config = ConfigDict(from_attributes=True, populate_by_name=True)
+
+    product_id: UUID = Field(alias="productId")
+    name: str
+    product_type: str = Field(alias="productType")
+    tag: str | None = None
+    description: str | None = None
+    status: str | None = None
+    created_at: datetime | None = Field(alias="createdAt", default=None)
+
+
+class ProcessSearchSchema(BaseModel):
+    """Schema for process search results."""
+
+    model_config = ConfigDict(from_attributes=True, populate_by_name=True)
+
+    process_id: UUID = Field(alias="processId")
+    workflow_name: str = Field(alias="workflowName")
+    workflow_id: UUID = Field(alias="workflowId")
+    status: str = Field(alias="last_status")
+    is_task: bool = Field(alias="isTask")
+    created_by: str | None = Field(alias="createdBy", default=None)
+    started_at: datetime = Field(alias="startedAt")
+    last_modified_at: datetime = Field(alias="lastModifiedAt")
+    last_step: str | None = Field(alias="lastStep", default=None)
+    failed_reason: str | None = Field(alias="failedReason", default=None)
+    subscription_ids: list[UUID] | None = Field(alias="subscriptionIds", default=None)
+
+
+class PathsResponse(BaseModel):
+    prefix: str
+    paths: list[PathInfo]
+
+    model_config = ConfigDict(extra="forbid", use_enum_values=True)

orchestrator/search/agent/__init__.py

@@ -0,0 +1 @@
+# This module requires: pydantic-ai==0.7.0, ag-ui-protocol>=0.1.8

orchestrator/search/agent/prompts.py

@@ -0,0 +1,62 @@
+import json
+from textwrap import dedent
+
+import structlog
+from pydantic_ai import RunContext
+from pydantic_ai.ag_ui import StateDeps
+
+from orchestrator.search.retrieval.validation import get_structured_filter_schema
+
+from .state import SearchState
+
+logger = structlog.get_logger(__name__)
+
+
+async def get_base_instructions() -> str:
+
+    try:
+        schema_dict = get_structured_filter_schema()
+        if schema_dict:
+            schema_info = "\n".join([f"    {path}: {field_type}" for path, field_type in schema_dict.items()])
+        else:
+            schema_info = "    No filterable fields available"
+    except Exception as e:
+        logger.warning(f"Failed to load schema for prompt: {e}")
+        schema_info = "    Schema temporarily unavailable"
+    logger.error(f"Generated schema for agent prompt:\n{schema_info}")
+
+    return dedent(
+        f"""
+    You are a helpful assistant for building and running database queries.
+
+    **Available Data Schema:**
+    Use the following schema to understand the available fields.
+    When you build filters, each `path` MUST be a valid path from this schema,
+    and the operator/value MUST match that path's type.
+    ```
+{schema_info}
+    ```
+    **Workflow (do in order):**
+    1) `set_search_parameters`  to define the main entity being searched.
+    2) Build a complete `FilterTree` (AND at root unless the user asks for OR).
+    3) `set_filter_tree(filters=<FilterTree or null>)`.
+    4) `execute_search()`.
+    5) Summarize the results for the user.
+    """
+    )
+
+
+async def get_dynamic_instructions(ctx: RunContext[StateDeps[SearchState]]) -> str:
+    """Dynamically generate the system prompt for the agent."""
+    param_state = json.dumps(ctx.deps.state.parameters, indent=2, default=str) if ctx.deps.state.parameters else "{}"
+
+    return dedent(
+        f"""
+        Current search parameters state:
+        {param_state}
+
+        Remember:
+        - If filters are missing or incomplete, construct a full FilterTree and call `set_filter_tree`.
+        - Then call `execute_search`.
+        """
+    )

orchestrator/search/agent/state.py

@@ -0,0 +1,8 @@
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class SearchState(BaseModel):
+    parameters: dict[str, Any] | None = None
+    results: list[dict[str, Any]] = Field(default_factory=list)

orchestrator/search/agent/tools.py

@@ -0,0 +1,122 @@
+import asyncio
+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 (
+    search_processes,
+    search_products,
+    search_subscriptions,
+    search_workflows,
+)
+from orchestrator.schemas.search import ConnectionSchema
+from orchestrator.search.core.types import ActionType, EntityType
+from orchestrator.search.filters import FilterTree
+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], ConnectionSchema[Any]] | Callable[[P], Awaitable[ConnectionSchema[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  # type: ignore[misc]
+async def set_search_parameters(
+    ctx: RunContext[StateDeps[SearchState]],
+    entity_type: EntityType,
+    action: str | ActionType = ActionType.SELECT,
+) -> StateSnapshotEvent:
+    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", "")
+
+    ctx.deps.state.parameters = {"action": action, "entity_type": entity_type, "filters": None, "query": final_query}
+    ctx.deps.state.results = []
+    logger.info(f"Set search parameters: entity_type={entity_type}, action={action}")
+
+    return StateSnapshotEvent(
+        type=EventType.STATE_SNAPSHOT,
+        snapshot=ctx.deps.state.model_dump(),
+    )
+
+
+@search_toolset.tool(retries=2)  # type: ignore[misc]
+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"])
+
+    try:
+        await validate_filter_tree(filters, entity_type)
+    except Exception as e:
+        raise ModelRetry(str(e))
+
+    ctx.deps.state.parameters["filters"] = None if filters is None else filters.model_dump(mode="json", by_alias=True)
+    logger.info(
+        "Set filter tree",
+        filters=None if filters is None else filters.model_dump(mode="json", by_alias=True),
+    )
+    return StateSnapshotEvent(type=EventType.STATE_SNAPSHOT, snapshot=ctx.deps.state.model_dump())
+
+
+@search_toolset.tool  # type: ignore[misc]
+async def execute_search(
+    ctx: RunContext[StateDeps[SearchState]],
+    limit: int = 5,
+) -> 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.info("Executing database search", **params.model_dump(mode="json"))
+
+    fn = SEARCH_FN_MAP[entity_type]
+    page_connection = await fn(params) if asyncio.iscoroutinefunction(fn) else fn(params)
+    ctx.deps.state.results = [item.model_dump(mode="json") for item in page_connection.page[:limit]]
+
+    return StateSnapshotEvent(type=EventType.STATE_SNAPSHOT, snapshot=ctx.deps.state.model_dump())

orchestrator/search/core/embedding.py

@@ -0,0 +1,64 @@
+import logging
+
+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.settings import app_settings
+
+logger = structlog.get_logger(__name__)
+
+# Its logging alot of noise such as embedding vectors.
+logging.getLogger("LiteLLM").setLevel(logging.WARNING)
+
+
+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=app_settings.EMBEDDING_MODEL,
+                input=[t.lower() for t in texts],
+                api_key=app_settings.OPENAI_API_KEY,
+                base_url=app_settings.OPENAI_BASE_URL,
+                timeout=app_settings.LLM_TIMEOUT,
+                max_retries=app_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=app_settings.EMBEDDING_MODEL,
+                input=[text.lower()],
+                api_key=app_settings.OPENAI_API_KEY,
+                base_url=app_settings.OPENAI_BASE_URL,
+                timeout=app_settings.LLM_TIMEOUT,
+                max_retries=app_settings.LLM_MAX_RETRIES,
+            )
+            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,16 @@
+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

orchestrator/search/core/types.py

@@ -0,0 +1,162 @@
+from dataclasses import dataclass
+from datetime import date, datetime
+from enum import Enum
+from typing import Any, NamedTuple, TypeAlias, TypedDict
+from uuid import UUID
+
+from sqlalchemy.orm.attributes import InstrumentedAttribute
+from sqlalchemy.sql.elements import ColumnElement
+from sqlalchemy_utils.types.ltree import Ltree
+
+from .validators import is_bool_string, is_iso_date, is_uuid
+
+SQLAColumn: TypeAlias = ColumnElement[Any] | InstrumentedAttribute[Any]
+
+
+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"
+
+
+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"
+
+    @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
+
+    def pg_cast(self) -> str:
+        return {
+            FieldType.STRING: "::text",
+            FieldType.INTEGER: "::integer",
+            FieldType.FLOAT: "::double precision",
+            FieldType.BOOLEAN: "::boolean",
+            FieldType.DATETIME: "::timestamptz",
+            FieldType.UUID: "::uuid",
+        }.get(self, "::text")
+
+    def is_embeddable(self) -> bool:
+        return self == 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

orchestrator/search/core/validators.py

@@ -0,0 +1,27 @@
+import uuid
+
+from dateutil.parser import isoparse
+
+
+def is_uuid(value: str) -> bool:
+    """Check if a string is a valid UUID."""
+    try:
+        uuid.UUID(value)
+        return True
+    except (ValueError, TypeError):
+        return False
+
+
+def is_iso_date(value: str) -> bool:
+    """Check if a string is a valid ISO 8601 date."""
+    try:
+        isoparse(value)
+        return True
+    except (ValueError, TypeError):
+        return False
+
+
+def is_bool_string(value: str) -> bool:
+    """Check if a string explicitly represents a boolean value with true/false."""
+
+    return value.strip().lower() in {"true", "false"}

orchestrator/search/docs/index.md

@@ -0,0 +1,37 @@
+# Search Indexing CLI
+
+Typer-based CLI for maintaining search indexes (subscriptions, products, processes, workflows).
+
+## Usage
+
+Run from project root:
+
+```
+dotenv run python main.py index [COMMAND] [OPTIONS]
+```
+
+### Commands
+
+- `subscriptions` – index `subscription_search_index`
+- `products` – index `product_search_index`
+- `processes` – index `process_search_index`
+- `workflows` – index `workflow_search_index`
+
+### Options
+
+- `--<id>` – UUID of a specific entity (default: all)
+- `--dry-run` – no DB writes
+- `--force-index` – re-index even if unchanged
+
+### Examples
+
+```
+# Index all subscriptions
+dotenv run python main.py index subscriptions
+
+# Re-index all subscriptions
+dotenv run python main.py index subscriptions --force-index
+
+# Index a single subscription
+dotenv run python main.py index subscriptions --subscription-id=<UUID>
+```

orchestrator/search/docs/running_local_text_embedding_inference.md

@@ -0,0 +1,45 @@
+# Running a local MiniLM embedding server with Hugging Face TEI
+
+Only **OpenAI-compatible endpoints** are supported locally.
+
+You can spin up a embedding API based on **sentence-transformers/all-MiniLM-L6-v2** using [Hugging Face TEI](https://github.com/huggingface/text-embeddings-inference):
+
+```bash
+docker run --rm -p 8080:80 ghcr.io/huggingface/text-embeddings-inference:cpu-1.8 \
+    --model-id sentence-transformers/all-MiniLM-L6-v2
+```
+
+---
+
+## Environment variables
+
+Point your backend to the local endpoint and declare the new vector size:
+
+```env
+OPENAI_BASE_URL=http://localhost:8080/v1
+EMBEDDING_DIMENSION=384
+```
+
+Depending on the model, you might want to change the `EMBEDDING_FALLBACK_MAX_TOKENS` and `EMBEDDING_MAX_BATCH_SIZE` settings, which are set conservatively and according to the requirements of the setup used in this example.
+
+---
+
+## Apply the schema change
+
+With these new settings run:
+
+```bash
+dotenv run python main.py embedding resize
+```
+
+**Note** that this will delete all records and you will have to re-index.
+
+---
+
+## Re-index embeddings
+
+```bash
+dotenv run python main.py index subscriptions
+```
+
+The search index now uses **384-dimension MiniLM vectors** served from your local Docker container. That’s it! 🚀

orchestrator/search/filters/__init__.py

@@ -0,0 +1,27 @@
+from .base import (
+    EqualityFilter,
+    FilterCondition,
+    FilterTree,
+    PathFilter,
+    StringFilter,
+)
+from .date_filters import DateFilter, DateRangeFilter, DateValueFilter
+from .ltree_filters import LtreeFilter
+from .numeric_filter import NumericFilter, NumericRangeFilter, NumericValueFilter
+
+__all__ = [
+    # Base filter classes
+    "PathFilter",
+    "FilterTree",
+    "FilterCondition",
+    "StringFilter",
+    "EqualityFilter",
+    # Filters for specific value types
+    "NumericValueFilter",
+    "NumericRangeFilter",
+    "DateValueFilter",
+    "DateRangeFilter",
+    "DateFilter",
+    "LtreeFilter",
+    "NumericFilter",
+]