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

orchestrator/cli/search/resize_embedding.py

@@ -4,7 +4,7 @@ from sqlalchemy import text
 from sqlalchemy.exc import SQLAlchemyError
 
 from orchestrator.db import db
-from orchestrator.db.models import AiSearchIndex
+from orchestrator.db.models import AiSearchIndex, SearchQueryTable
 from orchestrator.llm_settings import llm_settings
 
 logger = structlog.get_logger(__name__)
@@ -40,17 +40,20 @@ def get_current_embedding_dimension() -> int | None:
         return None
 
 
-def drop_all_embeddings() -> int:
-    """Drop all records from the ai_search_index table.
+def drop_all_embeddings() -> tuple[int, int]:
+    """Drop all records from ai_search_index and search_queries tables.
 
     Returns:
-        Number of records deleted
+        Tuple of (ai_search_index records deleted, search_queries records deleted)
     """
     try:
-        result = db.session.query(AiSearchIndex).delete()
+        index_deleted = db.session.query(AiSearchIndex).delete()
+        query_deleted = db.session.query(SearchQueryTable).delete()
         db.session.commit()
-        logger.info(f"Deleted {result} records from ai_search_index")
-        return result
+        logger.info(
+            f"Deleted {index_deleted} records from ai_search_index and {query_deleted} records from search_queries"
+        )
+        return index_deleted, query_deleted
 
     except SQLAlchemyError as e:
         db.session.rollback()
@@ -59,34 +62,34 @@ def drop_all_embeddings() -> int:
 
 
 def alter_embedding_column_dimension(new_dimension: int) -> None:
-    """Alter the embedding column to use the new dimension size.
+    """Alter the embedding columns in both ai_search_index and search_queries tables.
 
     Args:
         new_dimension: New vector dimension size
     """
     try:
-        drop_query = text("ALTER TABLE ai_search_index DROP COLUMN IF EXISTS embedding")
-        db.session.execute(drop_query)
+        db.session.execute(text("ALTER TABLE ai_search_index DROP COLUMN IF EXISTS embedding"))
+        db.session.execute(text(f"ALTER TABLE ai_search_index ADD COLUMN embedding vector({new_dimension})"))
 
-        add_query = text(f"ALTER TABLE ai_search_index ADD COLUMN embedding vector({new_dimension})")
-        db.session.execute(add_query)
+        db.session.execute(text("ALTER TABLE search_queries DROP COLUMN IF EXISTS query_embedding"))
+        db.session.execute(text(f"ALTER TABLE search_queries ADD COLUMN query_embedding vector({new_dimension})"))
 
         db.session.commit()
-        logger.info(f"Altered embedding column to dimension {new_dimension}")
+        logger.info(f"Altered embedding columns to dimension {new_dimension} in ai_search_index and search_queries")
 
     except SQLAlchemyError as e:
         db.session.rollback()
-        logger.error("Failed to alter embedding column dimension", error=str(e))
+        logger.error("Failed to alter embedding column dimensions", error=str(e))
         raise
 
 
 @app.command("resize")
 def resize_embeddings_command() -> None:
-    """Resize vector dimensions of the ai_search_index embedding column.
+    """Resize vector dimensions of embedding columns in ai_search_index and search_queries tables.
 
     Compares the current embedding dimension in the database with the configured
-    dimension in llm_settings. If they differ, drops all records and alters the
-    column to match the new dimension.
+    dimension in llm_settings. If they differ, drops all records and alters both
+    embedding columns to match the new dimension.
     """
     new_dimension = llm_settings.EMBEDDING_DIMENSION
 
@@ -107,22 +110,25 @@ def resize_embeddings_command() -> None:
 
     logger.info("Dimension mismatch detected", current_dimension=current_dimension, new_dimension=new_dimension)
 
-    if not typer.confirm("This will DELETE ALL RECORDS from ai_search_index and alter the embedding column. Continue?"):
+    if not typer.confirm(
+        "This will DELETE ALL RECORDS from ai_search_index and search_queries tables and alter embedding columns. Continue?"
+    ):
         logger.info("Operation cancelled by user")
         return
 
     try:
         # Drop all records first.
         logger.info("Dropping all embedding records...")
-        deleted_count = drop_all_embeddings()
+        index_deleted, query_deleted = drop_all_embeddings()
 
-        # Then alter column dimension.
-        logger.info(f"Altering embedding column to dimension {new_dimension}...")
+        # Then alter column dimensions.
+        logger.info(f"Altering embedding columns to dimension {new_dimension}...")
         alter_embedding_column_dimension(new_dimension)
 
         logger.info(
             "Embedding dimension resize completed successfully",
-            records_deleted=deleted_count,
+            index_records_deleted=index_deleted,
+            query_records_deleted=query_deleted,
             new_dimension=new_dimension,
         )
 

orchestrator/cli/search/speedtest.py

@@ -13,7 +13,6 @@ from orchestrator.search.core.embedding import QueryEmbedder
 from orchestrator.search.core.types import EntityType
 from orchestrator.search.core.validators import is_uuid
 from orchestrator.search.retrieval.engine import execute_search
-from orchestrator.search.retrieval.pagination import PaginationParams
 from orchestrator.search.schemas.parameters import BaseSearchParameters
 
 logger = structlog.get_logger(__name__)
@@ -54,17 +53,16 @@ async def generate_embeddings_for_queries(queries: list[str]) -> dict[str, list[
 async def run_single_query(query: str, embedding_lookup: dict[str, list[float]]) -> dict[str, Any]:
     search_params = BaseSearchParameters(entity_type=EntityType.SUBSCRIPTION, query=query, limit=30)
 
+    query_embedding = None
+
     if is_uuid(query):
-        pagination_params = PaginationParams()
         logger.debug("Using fuzzy-only ranking for full UUID", query=query)
     else:
-
-        cached_embedding = embedding_lookup[query]
-        pagination_params = PaginationParams(q_vec_override=cached_embedding)
+        query_embedding = embedding_lookup[query]
 
     with db.session as session:
         start_time = time.perf_counter()
-        response = await execute_search(search_params, session, pagination_params=pagination_params)
+        response = await execute_search(search_params, session, cursor=None, query_embedding=query_embedding)
         end_time = time.perf_counter()
 
         return {

orchestrator/db/__init__.py

@@ -17,6 +17,7 @@ from structlog import get_logger
 from orchestrator.db.database import BaseModel as DbBaseModel
 from orchestrator.db.database import Database, transactional
 from orchestrator.db.models import (  # noqa: F401
+    AgentRunTable,
     EngineSettingsTable,
     FixedInputTable,
     InputStateTable,
@@ -26,6 +27,7 @@ from orchestrator.db.models import (  # noqa: F401
     ProductBlockTable,
     ProductTable,
     ResourceTypeTable,
+    SearchQueryTable,
     SubscriptionCustomerDescriptionTable,
     SubscriptionInstanceRelationTable,
     SubscriptionInstanceTable,
@@ -74,6 +76,8 @@ def init_database(settings: AppSettings) -> Database:
 
 __all__ = [
     "transactional",
+    "SearchQueryTable",
+    "AgentRunTable",
     "SubscriptionTable",
     "ProcessSubscriptionTable",
     "ProcessTable",
@@ -97,6 +101,8 @@ __all__ = [
 ]
 
 ALL_DB_MODELS: list[type[DbBaseModel]] = [
+    SearchQueryTable,
+    AgentRunTable,
     FixedInputTable,
     ProcessStepTable,
     ProcessSubscriptionTable,

orchestrator/db/models.py

@@ -15,6 +15,7 @@ from __future__ import annotations
 
 import enum
 from datetime import datetime, timezone
+from typing import TYPE_CHECKING
 from uuid import UUID
 
 import sqlalchemy
@@ -58,6 +59,9 @@ from orchestrator.targets import Target
 from orchestrator.utils.datetime import nowtz
 from orchestrator.version import GIT_COMMIT_HASH
 
+if TYPE_CHECKING:
+    from orchestrator.search.retrieval.query_state import SearchQueryState
+
 logger = structlog.get_logger(__name__)
 
 TAG_LENGTH = 20
@@ -674,6 +678,76 @@ class SubscriptionSearchView(BaseModel):
     subscription = relationship("SubscriptionTable", foreign_keys=[subscription_id])
 
 
+class AgentRunTable(BaseModel):
+    """Agent conversation/session tracking."""
+
+    __tablename__ = "agent_runs"
+
+    run_id = mapped_column("run_id", UUIDType, server_default=text("uuid_generate_v4()"), primary_key=True)
+    agent_type = mapped_column(String(50), nullable=False)
+    created_at = mapped_column(UtcTimestamp, server_default=text("current_timestamp()"), nullable=False)
+
+    queries = relationship("SearchQueryTable", back_populates="run", cascade="delete", passive_deletes=True)
+
+    __table_args__ = (Index("ix_agent_runs_created_at", "created_at"),)
+
+
+class SearchQueryTable(BaseModel):
+    """Search query execution - used by both agent runs and regular API searches.
+
+    When run_id is NULL: standalone API search query
+    When run_id is NOT NULL: query belongs to an agent conversation run
+    """
+
+    __tablename__ = "search_queries"
+
+    query_id = mapped_column("query_id", UUIDType, server_default=text("uuid_generate_v4()"), primary_key=True)
+    run_id = mapped_column(
+        "run_id", UUIDType, ForeignKey("agent_runs.run_id", ondelete="CASCADE"), nullable=True, index=True
+    )
+    query_number = mapped_column(Integer, nullable=False)
+
+    # Search parameters as JSONB (maps to BaseSearchParameters subclasses)
+    parameters = mapped_column(pg.JSONB, nullable=False)
+
+    # Query embedding for semantic search (pgvector)
+    query_embedding = mapped_column(Vector(1536), nullable=True)
+
+    executed_at = mapped_column(UtcTimestamp, server_default=text("current_timestamp()"), nullable=False)
+
+    run = relationship("AgentRunTable", back_populates="queries")
+
+    __table_args__ = (
+        Index("ix_search_queries_run_id", "run_id"),
+        Index("ix_search_queries_executed_at", "executed_at"),
+        Index("ix_search_queries_query_id", "query_id"),
+    )
+
+    @classmethod
+    def from_state(
+        cls,
+        state: "SearchQueryState",
+        run_id: "UUID | None" = None,
+        query_number: int = 1,
+    ) -> "SearchQueryTable":
+        """Create a SearchQueryTable instance from a SearchQueryState.
+
+        Args:
+            state: The search query state with parameters and embedding
+            run_id: Optional agent run ID (NULL for regular API searches)
+            query_number: Query number within the run (default=1)
+
+        Returns:
+            SearchQueryTable instance ready to be added to the database.
+        """
+        return cls(
+            run_id=run_id,
+            query_number=query_number,
+            parameters=state.parameters.model_dump(),
+            query_embedding=state.query_embedding,
+        )
+
+
 class EngineSettingsTable(BaseModel):
     __tablename__ = "engine_settings"
     global_lock = mapped_column(Boolean(), default=False, nullable=False, primary_key=True)
@@ -705,6 +779,7 @@ class AiSearchIndex(BaseModel):
         UUIDType,
         nullable=False,
     )
+    entity_title = mapped_column(TEXT, nullable=True)
 
     # Ltree path for hierarchical data
     path = mapped_column(LtreeType, nullable=False, index=True)

orchestrator/migrations/helpers.py

@@ -170,49 +170,57 @@ def create_workflow(conn: sa.engine.Connection, workflow: dict) -> None:
         }
         >>> create_workflow(conn, workflow)
     """
-    if not workflow.get("is_task", False):
-        workflow["is_task"] = False
+    params = workflow.copy()
+    params.setdefault("is_task", False)
+    params.setdefault("product_tag", None)
 
-    if not workflow.get("product_tag"):
-        workflow["product_tag"] = None
+    query_parts = []
 
     if has_table_column(table_name="workflows", column_name="is_task", conn=conn):
-        query = """
-                WITH new_workflow AS (
-                    INSERT INTO workflows (name, target, is_task, description)
-                        VALUES (:name, :target, :is_task, :description)
-                        ON CONFLICT DO NOTHING
-                        RETURNING workflow_id)
-                INSERT
-                INTO products_workflows (product_id, workflow_id)
-                SELECT p.product_id,
-                       nw.workflow_id
-                FROM products AS p
-                         CROSS JOIN new_workflow AS nw
-                WHERE p.product_type = :product_type
-                AND (:product_tag IS NULL OR p.tag = :product_tag)
-                ON CONFLICT DO NOTHING
-                """
+        query_parts.append(
+            """
+            WITH new_workflow AS (
+                INSERT INTO workflows (name, target, is_task, description)
+                    VALUES (:name, :target, :is_task, :description)
+                    ON CONFLICT DO NOTHING
+                    RETURNING workflow_id
+            )
+            """
+        )
     else:
-        # Remove is_task from workflow dict and insert SQL
-        workflow = {k: v for k, v in workflow.items() if k != "is_task"}
-        query = """
-                WITH new_workflow AS (
-                    INSERT INTO workflows (name, target, description)
-                        VALUES (:name, :target, :description)
-                        ON CONFLICT DO NOTHING
-                        RETURNING workflow_id)
-                INSERT
-                INTO products_workflows (product_id, workflow_id)
-                SELECT p.product_id, nw.workflow_id
-                FROM products AS p
-                         CROSS JOIN new_workflow AS nw
-                WHERE p.product_type = :product_type
-                AND (:product_tag IS NULL OR p.tag = :product_tag)
-                ON CONFLICT DO NOTHING
-                """
+        params.pop("is_task", None)
+        query_parts.append(
+            """
+            WITH new_workflow AS (
+                INSERT INTO workflows (name, target, description)
+                    VALUES (:name, :target, :description)
+                    ON CONFLICT DO NOTHING
+                    RETURNING workflow_id
+            )
+            """
+        )
+
+    query_parts.append(
+        """
+        INSERT INTO products_workflows (product_id, workflow_id)
+        SELECT p.product_id, nw.workflow_id
+        FROM products AS p
+                 CROSS JOIN new_workflow AS nw
+        """
+    )
+
+    query_parts.append("WHERE p.product_type = :product_type")
+
+    if params.get("product_tag") is not None:
+        query_parts.append("AND p.tag = :product_tag")
+    else:
+        params.pop("product_tag", None)
+
+    query_parts.append("ON CONFLICT DO NOTHING")
+
+    query = "\n".join(query_parts)
 
-    conn.execute(sa.text(query), workflow)
+    conn.execute(sa.text(query), params)
 
 
 def create_task(conn: sa.engine.Connection, task: dict) -> None:

orchestrator/schedules/scheduler.py

@@ -17,16 +17,16 @@ from datetime import datetime
 from typing import Any, Generator
 
 from apscheduler.executors.pool import ThreadPoolExecutor
-from apscheduler.jobstores.sqlalchemy import SQLAlchemyJobStore
+from apscheduler.jobstores.sqlalchemy import Job, SQLAlchemyJobStore
 from apscheduler.schedulers.background import BackgroundScheduler
 from more_itertools import partition
 from pydantic import BaseModel
 
+from orchestrator.db import db
 from orchestrator.db.filters import Filter
 from orchestrator.db.filters.filters import CallableErrorHandler
 from orchestrator.db.sorting import Sort
 from orchestrator.db.sorting.sorting import SortOrder
-from orchestrator.settings import app_settings
 from orchestrator.utils.helpers import camel_to_snake, to_camel
 
 executors = {
@@ -40,18 +40,37 @@ scheduler = BackgroundScheduler(executors=executors, job_defaults=job_defaults)
 
 
 @contextmanager
-def get_paused_scheduler() -> Generator[BackgroundScheduler, Any, None]:
+def get_scheduler_store() -> Generator[SQLAlchemyJobStore, Any, None]:
+    store = SQLAlchemyJobStore(engine=db.engine)
     try:
-        scheduler.add_jobstore(SQLAlchemyJobStore(url=str(app_settings.DATABASE_URI)))
-    except ValueError:
-        pass
-    scheduler.start(paused=True)
-
-    try:
-        yield scheduler
+        yield store
     finally:
-        scheduler.shutdown()
-        scheduler._jobstores["default"].engine.dispose()
+        store.shutdown()
+
+
+def get_all_scheduler_tasks() -> list[Job]:
+    with get_scheduler_store() as scheduler_store:
+        return scheduler_store.get_all_jobs()
+
+
+def get_scheduler_task(job_id: str) -> Job | None:
+    with get_scheduler_store() as scheduler_store:
+        return scheduler_store.lookup_job(job_id)
+
+
+@contextmanager
+def get_scheduler(paused: bool = False) -> Generator[BackgroundScheduler, Any, None]:
+    with get_scheduler_store() as store:
+        try:
+            scheduler.add_jobstore(store)
+        except ValueError:
+            pass
+        scheduler.start(paused=paused)
+
+        try:
+            yield scheduler
+        finally:
+            scheduler.shutdown()
 
 
 class ScheduledTask(BaseModel):
@@ -149,9 +168,7 @@ def get_scheduler_tasks(
     sort_by: list[Sort] | None = None,
     error_handler: CallableErrorHandler = default_error_handler,
 ) -> tuple[list[ScheduledTask], int]:
-    with get_paused_scheduler() as pauzed_scheduler:
-        scheduled_tasks = pauzed_scheduler.get_jobs()
-
+    scheduled_tasks = get_all_scheduler_tasks()
     scheduled_tasks = filter_scheduled_tasks(scheduled_tasks, error_handler, filter_by)
     scheduled_tasks = sort_scheduled_tasks(scheduled_tasks, error_handler, sort_by)
 

orchestrator/schedules/validate_products.py

@@ -29,7 +29,7 @@ def validate_products() -> None:
     uncompleted_products = db.session.scalar(
         select(func.count())
         .select_from(ProcessTable)
-        .filter(ProcessTable.workflow.name == "validate_products", ProcessTable.last_status != "completed")
+        .filter(ProcessTable.workflow.has(name="validate_products"), ProcessTable.last_status != "completed")
     )
     if not uncompleted_products:
         start_process("task_validate_products")

orchestrator/schemas/search.py

@@ -11,14 +11,12 @@
 # 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 typing import Generic, TypeVar
 
 from pydantic import BaseModel, ConfigDict, Field
 
 from orchestrator.search.core.types import SearchMetadata
-from orchestrator.search.schemas.results import ComponentInfo, LeafInfo, MatchingField
+from orchestrator.search.schemas.results import ComponentInfo, LeafInfo
 
 T = TypeVar("T")
 
@@ -36,95 +34,20 @@ class ProductSchema(BaseModel):
     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)
+
+
+class ExportResponse(BaseModel):
+    page: list[dict]
+
+    model_config = ConfigDict(extra="forbid")

orchestrator/search/agent/__init__.py

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

orchestrator/search/agent/agent.py

@@ -14,13 +14,11 @@
 from typing import Any
 
 import structlog
-from fastapi import APIRouter, HTTPException, Request
-from pydantic_ai.ag_ui import StateDeps, handle_ag_ui_request
+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.responses import Response
 
 from orchestrator.search.agent.prompts import get_base_instructions, get_dynamic_instructions
 from orchestrator.search.agent.state import SearchState
@@ -29,34 +27,32 @@ from orchestrator.search.agent.tools import search_toolset
 logger = structlog.get_logger(__name__)
 
 
-def build_agent_router(model: str | OpenAIModel, toolsets: list[FunctionToolset[Any]] | None = None) -> APIRouter:
-    router = APIRouter()
+def build_agent_instance(
+    model: str | OpenAIModel, agent_tools: list[FunctionToolset[Any]] | None = None
+) -> Agent[StateDeps[SearchState], str]:
+    """Build and configure the search agent instance.
 
-    try:
-        toolsets = toolsets + [search_toolset] if toolsets else [search_toolset]
+    Args:
+        model: The LLM model to use (string or OpenAIModel instance)
+        agent_tools: Optional list of additional toolsets to include
 
-        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)
+    Returns:
+        Configured Agent instance with StateDeps[SearchState] dependencies
 
-        @router.post("/")
-        async def agent_endpoint(request: Request) -> Response:
-            return await handle_ag_ui_request(agent, request, deps=StateDeps(SearchState()))
+    Raises:
+        Exception: If agent initialization fails
+    """
+    toolsets = agent_tools + [search_toolset] if agent_tools else [search_toolset]
 
-        return router
-    except Exception as e:
-        logger.error("Agent init failed; serving disabled stub.", error=str(e))
-        error_msg = f"Agent disabled: {str(e)}"
+    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)
 
-        @router.api_route("/{path:path}", methods=["GET", "POST", "PUT", "PATCH", "DELETE", "OPTIONS", "HEAD"])
-        async def _disabled(path: str) -> None:
-            raise HTTPException(status_code=503, detail=error_msg)
-
-        return router
+    return agent