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

orchestrator/cli/scheduler.py

@@ -12,26 +12,23 @@
 # limitations under the License.
 
 
-import logging
 import time
 
 import typer
 
 from orchestrator.schedules.scheduler import (
-    get_paused_scheduler,
+    get_all_scheduler_tasks,
+    get_scheduler,
+    get_scheduler_task,
 )
 
-log = logging.getLogger(__name__)
-
 app: typer.Typer = typer.Typer()
 
 
 @app.command()
 def run() -> None:
     """Start scheduler and loop eternally to keep thread alive."""
-    with get_paused_scheduler() as scheduler:
-        scheduler.resume()
-
+    with get_scheduler():
         while True:
             time.sleep(1)
 
@@ -42,27 +39,23 @@ def show_schedule() -> None:
 
     in cli underscore is replaced by a dash `show-schedule`
     """
-    with get_paused_scheduler() as scheduler:
-        jobs = scheduler.get_jobs()
-
-    for job in jobs:
-        typer.echo(f"[{job.id}] Next run: {job.next_run_time} | Trigger: {job.trigger}")
+    for task in get_all_scheduler_tasks():
+        typer.echo(f"[{task.id}] Next run: {task.next_run_time} | Trigger: {task.trigger}")
 
 
 @app.command()
-def force(job_id: str) -> None:
-    """Force the execution of (a) scheduler(s) based on a job_id."""
-    with get_paused_scheduler() as scheduler:
-        job = scheduler.get_job(job_id)
+def force(task_id: str) -> None:
+    """Force the execution of (a) scheduler(s) based on a task_id."""
+    task = get_scheduler_task(task_id)
 
-    if not job:
-        typer.echo(f"Job '{job_id}' not found.")
+    if not task:
+        typer.echo(f"Task '{task_id}' not found.")
         raise typer.Exit(code=1)
 
-    typer.echo(f"Running job [{job.id}] now...")
+    typer.echo(f"Running Task [{task.id}] now...")
     try:
-        job.func(*job.args or (), **job.kwargs or {})
-        typer.echo("Job executed successfully.")
+        task.func(*task.args or (), **task.kwargs or {})
+        typer.echo("Task executed successfully.")
     except Exception as e:
-        typer.echo(f"Job execution failed: {e}")
+        typer.echo(f"Task execution failed: {e}")
         raise typer.Exit(code=1)

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(llm_settings.EMBEDDING_DIMENSION), 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/llm_settings.py

@@ -10,14 +10,31 @@
 # 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 Annotated
+
 from pydantic import Field, field_validator
 from pydantic_settings import BaseSettings
 from structlog import get_logger
 
 logger = get_logger(__name__)
 
+EMBEDDING_DIMENSION_MIN = 100
+EMBEDDING_DIMENSION_MAX = 2000
+EMBEDDING_DIMENSION_DEFAULT = 1536
+
+EMBEDDING_DIMENSION_FIELD = Annotated[
+    int,
+    Field(
+        ge=EMBEDDING_DIMENSION_MIN,
+        le=EMBEDDING_DIMENSION_MAX,
+        default=EMBEDDING_DIMENSION_DEFAULT,
+        description="Embedding dimension: when embeddings are generated at a higher resolution than this setting, the least significant numbers will be truncated",
+    ),
+]
+
 
 class LLMSettings(BaseSettings):
+
     # Feature flags for LLM functionality
     SEARCH_ENABLED: bool = False  # Enable search/indexing with embeddings
     AGENT_ENABLED: bool = False  # Enable agentic functionality
@@ -27,7 +44,7 @@ class LLMSettings(BaseSettings):
     AGENT_MODEL_VERSION: str = "2025-01-01-preview"
     OPENAI_API_KEY: str = ""  # Change per provider (Azure, etc).
     # Embedding settings
-    EMBEDDING_DIMENSION: int = 1536
+    EMBEDDING_DIMENSION: EMBEDDING_DIMENSION_FIELD = 1536
     EMBEDDING_MODEL: str = "openai/text-embedding-3-small"  # See litellm docs for supported models.
     EMBEDDING_SAFE_MARGIN_PERCENT: float = Field(
         0.1, description="Safety margin as a percentage (e.g., 0.1 for 10%) for token budgeting.", ge=0, le=1

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:
@@ -229,7 +237,7 @@ def create_task(conn: sa.engine.Connection, task: dict) -> None:
             "name": "task_name",
             "description": "task description",
         }
-        >>> create_workflow(conn, task)
+        >>> create_task(conn, task)
     """
     if has_table_column(table_name="workflows", column_name="is_task", conn=conn):
         query = """

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")