orchestrator-core 4.4.1__py3-none-any.whl → 4.5.0a3__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,130 @@
+# 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 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/schemas/workflow.py

@@ -60,6 +60,7 @@ class SubscriptionWorkflowListsSchema(OrchestratorBaseModel):
     modify: list[WorkflowListItemSchema]
     terminate: list[WorkflowListItemSchema]
     system: list[WorkflowListItemSchema]
+    reconcile: list[WorkflowListItemSchema]
     validate_: list[WorkflowListItemSchema] = Field(default_factory=list, alias="validate")
 
 

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,21 @@
+# 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.
+
+# 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,60 @@
+# 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 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,100 @@
+# 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 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
+
+logger = structlog.get_logger(__name__)
+
+
+async def get_base_instructions() -> str:
+    return dedent(
+        """
+        You are an expert assistant designed to find relevant information by building and running database queries.
+
+        ---
+        ### 1. Your Goal and Method
+
+        Your ultimate goal is to **find information** that answers the user's request.
+
+        To do this, you will perform either a broad search or a filtered search.
+        For **filtered searches**, your primary method is to **construct a valid `FilterTree` object**.
+        To do this correctly, you must infer the exact structure, operators, and nesting rules from the Pydantic schema of the `set_filter_tree` tool itself.
+
+        ---
+        ### 2. Information-Gathering Tools
+
+        **If you determine that a `FilterTree` is needed**, use these tools to gather information first:
+
+        - **discover_filter_paths(field_names: list[str])**: Use this to discover all valid filter paths for a list of field names in a single call.
+        - **get_valid_operators()**: Use this to get the JSON map of all valid operators for each field type.
+
+        ---
+        ### 3. Execution Workflow
+
+        Follow these steps in strict order:
+
+        1.  **Set Context**: Always begin by calling `set_search_parameters`.
+        2.  **Analyze for Filters**: Based on the user's request, decide if specific filters are necessary.
+            - **If filters ARE required**, follow these sub-steps:
+                a. **Gather Intel**: Identify all needed field names, then call `discover_filter_paths` and `get_valid_operators` **once each** to get all required information.
+                b. **Construct FilterTree**: Build the `FilterTree` object.
+                c. **Set Filters**: Call `set_filter_tree`.
+        3.  **Execute**: Call `execute_search`. This is done for both filtered and non-filtered searches.
+        4.  **Report**: Answer the users' question directly and summarize when appropiate.
+
+        ---
+        ### 4. Critical Rules
+
+        - **NEVER GUESS PATHS**: You *must* verify every filter path by calling `discover_filter_paths` first. If a path does not exist, you must inform the user and not include it in the `FilterTree`.
+        - **USE FULL PATHS**: Always use the full, unambiguous path returned by the discovery tool.
+        - **MATCH OPERATORS**: Only use operators that are compatible with the field type as confirmed by `get_filter_operators`.
+        """
+    )
+
+
+async def get_dynamic_instructions(ctx: RunContext[StateDeps[SearchState]]) -> str:
+    """Dynamically provides 'next step' coaching based on the current state."""
+    state = ctx.deps.state
+    param_state_str = json.dumps(state.parameters, indent=2, default=str) if state.parameters else "Not set."
+
+    next_step_guidance = ""
+    if not state.parameters or not state.parameters.get("entity_type"):
+        next_step_guidance = (
+            "INSTRUCTION: The search context is not set. Your next action is to call `set_search_parameters`."
+        )
+    else:
+        next_step_guidance = (
+            "INSTRUCTION: Context is set. Now, analyze the user's request. "
+            "If specific filters ARE required, use the information-gathering tools to build a `FilterTree` and call `set_filter_tree`. "
+            "If no specific filters are needed, you can proceed directly to `execute_search`."
+        )
+    return dedent(
+        f"""
+        ---
+        ### Current State & Next Action
+
+        **Current Search Parameters:**
+        ```json
+        {param_state_str}
+        ```
+
+        **{next_step_guidance}**
+        """
+    )

orchestrator/search/agent/state.py

@@ -0,0 +1,21 @@
+# 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 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)