orchestrator-core 4.4.1__py3-none-any.whl → 4.5.1a1__py3-none-any.whl

orchestrator/migrations/versions/schema/2025-08-12_52b37b5b2714_search_index_model_for_llm_integration.py

@@ -0,0 +1,95 @@
+"""Search index model for llm integration.
+
+Revision ID: 52b37b5b2714
+Revises: 850dccac3b02
+Create Date: 2025-08-12 22:34:26.694750
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+from pgvector.sqlalchemy import Vector
+from sqlalchemy.dialects import postgresql
+from sqlalchemy_utils import LtreeType
+
+from orchestrator.search.core.types import FieldType
+
+# revision identifiers, used by Alembic.
+revision = "52b37b5b2714"
+down_revision = "850dccac3b02"
+branch_labels = None
+depends_on = None
+
+TABLE = "ai_search_index"
+IDX_EMBED_HNSW = "ix_flat_embed_hnsw"
+IDX_PATH_GIST = "ix_flat_path_gist"
+IDX_PATH_BTREE = "ix_flat_path_btree"
+IDX_VALUE_TRGM = "ix_flat_value_trgm"
+IDX_CONTENT_HASH = "idx_ai_search_index_content_hash"
+
+TARGET_DIM = 1536
+
+
+def upgrade() -> None:
+    # Create PostgreSQL extensions
+    op.execute("CREATE EXTENSION IF NOT EXISTS ltree;")
+    op.execute("CREATE EXTENSION IF NOT EXISTS unaccent;")
+    op.execute("CREATE EXTENSION IF NOT EXISTS pg_trgm;")
+    op.execute("CREATE EXTENSION IF NOT EXISTS vector;")
+
+    # Create the ai_search_index table
+    op.create_table(
+        TABLE,
+        sa.Column("entity_type", sa.Text, nullable=False),
+        sa.Column("entity_id", postgresql.UUID, nullable=False),
+        sa.Column("path", LtreeType, nullable=False),
+        sa.Column("value", sa.Text, nullable=False),
+        sa.Column("embedding", Vector(TARGET_DIM), nullable=True),
+        sa.Column("content_hash", sa.String(64), nullable=False),
+        sa.PrimaryKeyConstraint("entity_id", "path", name="pk_ai_search_index"),
+    )
+
+    field_type_enum = sa.Enum(*[ft.value for ft in FieldType], name="field_type")
+    field_type_enum.create(op.get_bind(), checkfirst=True)
+    op.add_column(
+        TABLE,
+        sa.Column("value_type", field_type_enum, nullable=False, server_default=FieldType.STRING.value),
+    )
+    op.alter_column(TABLE, "value_type", server_default=None)
+
+    op.create_index(op.f("ix_ai_search_index_entity_id"), TABLE, ["entity_id"], unique=False)
+    op.create_index(IDX_CONTENT_HASH, TABLE, ["content_hash"])
+
+    op.create_index(
+        IDX_PATH_GIST,
+        TABLE,
+        ["path"],
+        postgresql_using="GIST",
+        postgresql_ops={"path": "gist_ltree_ops"},
+    )
+    op.create_index(IDX_PATH_BTREE, TABLE, ["path"])
+    op.create_index(IDX_VALUE_TRGM, TABLE, ["value"], postgresql_using="GIN", postgresql_ops={"value": "gin_trgm_ops"})
+
+    op.create_index(
+        IDX_EMBED_HNSW,
+        TABLE,
+        ["embedding"],
+        postgresql_using="HNSW",
+        postgresql_with={"m": 16, "ef_construction": 64},
+        postgresql_ops={"embedding": "vector_l2_ops"},
+    )
+
+
+def downgrade() -> None:
+    # Drop all indexes
+    op.drop_index(IDX_EMBED_HNSW, table_name=TABLE, if_exists=True)
+    op.drop_index(IDX_VALUE_TRGM, table_name=TABLE, if_exists=True)
+    op.drop_index(IDX_PATH_BTREE, table_name=TABLE, if_exists=True)
+    op.drop_index(IDX_PATH_GIST, table_name=TABLE, if_exists=True)
+    op.drop_index(IDX_CONTENT_HASH, table_name=TABLE, if_exists=True)
+    op.drop_index(op.f("ix_ai_search_index_entity_id"), table_name=TABLE, if_exists=True)
+
+    # Drop table and enum
+    op.drop_table(TABLE, if_exists=True)
+    field_type_enum = sa.Enum(name="field_type")
+    field_type_enum.drop(op.get_bind(), checkfirst=True)

orchestrator/schemas/search.py

@@ -0,0 +1,117 @@
+from datetime import datetime
+from typing import Any, Generic, TypeVar
+from uuid import UUID
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from orchestrator.search.core.types import SearchMetadata
+from orchestrator.search.schemas.results import ComponentInfo, LeafInfo, MatchingField
+
+T = TypeVar("T")
+
+
+class PageInfoSchema(BaseModel):
+    has_next_page: bool = False
+    next_page_cursor: str | None = None
+
+
+class ProductSchema(BaseModel):
+    model_config = ConfigDict(from_attributes=True)
+
+    name: str
+    tag: str
+    product_type: str
+
+
+class SubscriptionSearchResult(BaseModel):
+    score: float
+    perfect_match: int
+    matching_field: MatchingField | None = None
+    subscription: dict[str, Any]
+
+
+class SearchResultsSchema(BaseModel, Generic[T]):
+    data: list[T] = Field(default_factory=list)
+    page_info: PageInfoSchema = Field(default_factory=PageInfoSchema)
+    search_metadata: SearchMetadata | None = None
+
+
+class WorkflowProductSchema(BaseModel):
+    """Product associated with a workflow."""
+
+    model_config = ConfigDict(from_attributes=True)
+
+    product_type: str
+    product_id: UUID
+    name: str
+
+
+class WorkflowSearchSchema(BaseModel):
+    """Schema for workflow search results."""
+
+    model_config = ConfigDict(from_attributes=True)
+
+    name: str
+    products: list[WorkflowProductSchema]
+    description: str | None = None
+    created_at: datetime | None = None
+
+
+class ProductSearchSchema(BaseModel):
+    """Schema for product search results."""
+
+    model_config = ConfigDict(from_attributes=True)
+
+    product_id: UUID
+    name: str
+    product_type: str
+    tag: str | None = None
+    description: str | None = None
+    status: str | None = None
+    created_at: datetime | None = None
+
+
+class ProcessSearchSchema(BaseModel):
+    """Schema for process search results."""
+
+    model_config = ConfigDict(from_attributes=True)
+
+    process_id: UUID
+    workflow_name: str
+    workflow_id: UUID
+    last_status: str
+    is_task: bool
+    created_by: str | None = None
+    started_at: datetime
+    last_modified_at: datetime
+    last_step: str | None = None
+    failed_reason: str | None = None
+    subscription_ids: list[UUID] | None = None
+
+
+class WorkflowSearchResult(BaseModel):
+    score: float
+    perfect_match: int
+    matching_field: MatchingField | None = None
+    workflow: WorkflowSearchSchema
+
+
+class ProductSearchResult(BaseModel):
+    score: float
+    perfect_match: int
+    matching_field: MatchingField | None = None
+    product: ProductSearchSchema
+
+
+class ProcessSearchResult(BaseModel):
+    score: float
+    perfect_match: int
+    matching_field: MatchingField | None = None
+    process: ProcessSearchSchema
+
+
+class PathsResponse(BaseModel):
+    leaves: list[LeafInfo]
+    components: list[ComponentInfo]
+
+    model_config = ConfigDict(extra="forbid", use_enum_values=True)

orchestrator/search/__init__.py

@@ -0,0 +1,12 @@
+# Copyright 2019-2025 SURF.
+# 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/agent/__init__.py

@@ -0,0 +1,8 @@
+# This module requires: pydantic-ai==0.7.0, ag-ui-protocol>=0.1.8
+
+
+from orchestrator.search.agent.agent import build_agent_app
+
+__all__ = [
+    "build_agent_app",
+]

orchestrator/search/agent/agent.py

@@ -0,0 +1,47 @@
+from typing import Any
+
+import structlog
+from fastapi import FastAPI, HTTPException
+from pydantic_ai.ag_ui import StateDeps
+from pydantic_ai.agent import Agent
+from pydantic_ai.models.openai import OpenAIModel
+from pydantic_ai.settings import ModelSettings
+from pydantic_ai.toolsets import FunctionToolset
+from starlette.types import ASGIApp
+
+from orchestrator.search.agent.prompts import get_base_instructions, get_dynamic_instructions
+from orchestrator.search.agent.state import SearchState
+from orchestrator.search.agent.tools import search_toolset
+
+logger = structlog.get_logger(__name__)
+
+
+def _disabled_agent_app(reason: str) -> FastAPI:
+    app = FastAPI(title="Agent disabled")
+
+    @app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "PATCH", "DELETE", "OPTIONS", "HEAD"])
+    async def _disabled(path: str) -> None:
+        raise HTTPException(status_code=503, detail=f"Agent disabled: {reason}")
+
+    return app
+
+
+def build_agent_app(model: str | OpenAIModel, toolsets: list[FunctionToolset[Any]] | None = None) -> ASGIApp:
+    try:
+        toolsets = toolsets + [search_toolset] if toolsets else [search_toolset]
+
+        agent = Agent(
+            model=model,
+            deps_type=StateDeps[SearchState],
+            model_settings=ModelSettings(
+                parallel_tool_calls=False,
+            ),  # https://github.com/pydantic/pydantic-ai/issues/562
+            toolsets=toolsets,
+        )
+        agent.instructions(get_base_instructions)
+        agent.instructions(get_dynamic_instructions)
+
+        return agent.to_ag_ui(deps=StateDeps(SearchState()))
+    except Exception as e:
+        logger.error("Agent init failed; serving disabled stub.", error=str(e))
+        return _disabled_agent_app(str(e))

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.agent.state import SearchState
+from orchestrator.search.retrieval.validation import get_structured_filter_schema
+
+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,121 @@
+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 SearchResultsSchema
+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], 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:
+    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)
+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
+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]
+    search_results = await fn(params)
+    ctx.deps.state.results = search_results.data[: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.llm_settings import llm_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=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,22 @@
+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