orchestrator-core 4.4.1__py3-none-any.whl → 4.5.0a2__py3-none-any.whl

orchestrator/cli/resize_embedding.py

@@ -0,0 +1,135 @@
+import structlog
+import typer
+from sqlalchemy import text
+from sqlalchemy.exc import SQLAlchemyError
+
+from orchestrator.db import db
+from orchestrator.db.models import AiSearchIndex
+from orchestrator.llm_settings import llm_settings
+
+logger = structlog.get_logger(__name__)
+
+app = typer.Typer(
+    name="embedding",
+    help="Resize vector dimensions of the embeddings.",
+)
+
+
+def get_current_embedding_dimension() -> int | None:
+    """Get the current dimension of the embedding column from ai_search_index table.
+
+    Returns:
+        Current dimension size or None if no records exist or column doesn't exist
+    """
+    try:
+        query = text(
+            """
+            SELECT vector_dims(embedding) as dimension
+            FROM ai_search_index
+            WHERE embedding IS NOT NULL
+            LIMIT 1
+        """
+        )
+        result = db.session.execute(query).fetchone()
+        if result and result[0]:
+            return result[0]
+        return None
+
+    except SQLAlchemyError as e:
+        logger.error("Failed to get current embedding dimension", error=str(e))
+        return None
+
+
+def drop_all_embeddings() -> int:
+    """Drop all records from the ai_search_index table.
+
+    Returns:
+        Number of records deleted
+    """
+    try:
+        result = db.session.query(AiSearchIndex).delete()
+        db.session.commit()
+        logger.info(f"Deleted {result} records from ai_search_index")
+        return result
+
+    except SQLAlchemyError as e:
+        db.session.rollback()
+        logger.error("Failed to drop embeddings records", error=str(e))
+        raise
+
+
+def alter_embedding_column_dimension(new_dimension: int) -> None:
+    """Alter the embedding column to use the new dimension size.
+
+    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)
+
+        add_query = text(f"ALTER TABLE ai_search_index ADD COLUMN embedding vector({new_dimension})")
+        db.session.execute(add_query)
+
+        db.session.commit()
+        logger.info(f"Altered embedding column to dimension {new_dimension}")
+
+    except SQLAlchemyError as e:
+        db.session.rollback()
+        logger.error("Failed to alter embedding column dimension", error=str(e))
+        raise
+
+
+@app.command("resize")
+def resize_embeddings_command() -> None:
+    """Resize vector dimensions of the ai_search_index embedding column.
+
+    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.
+    """
+    new_dimension = llm_settings.EMBEDDING_DIMENSION
+
+    logger.info("Starting embedding dimension resize", new_dimension=new_dimension)
+
+    current_dimension = get_current_embedding_dimension()
+
+    if current_dimension is None:
+        logger.warning("Could not determine current dimension for embedding column")
+
+    if current_dimension == new_dimension:
+        logger.info(
+            "Embedding dimensions match, no resize needed",
+            current_dimension=current_dimension,
+            new_dimension=new_dimension,
+        )
+        return
+
+    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?"):
+        logger.info("Operation cancelled by user")
+        return
+
+    try:
+        # Drop all records first.
+        logger.info("Dropping all embedding records...")
+        deleted_count = drop_all_embeddings()
+
+        # Then alter column dimension.
+        logger.info(f"Altering embedding column to dimension {new_dimension}...")
+        alter_embedding_column_dimension(new_dimension)
+
+        logger.info(
+            "Embedding dimension resize completed successfully",
+            records_deleted=deleted_count,
+            new_dimension=new_dimension,
+        )
+
+    except Exception as e:
+        logger.error("Embedding dimension resize failed", error=str(e))
+        raise typer.Exit(1)
+
+
+if __name__ == "__main__":
+    app()

orchestrator/cli/search_explore.py

@@ -0,0 +1,208 @@
+import asyncio
+
+import structlog
+import typer
+from pydantic import ValidationError
+
+from orchestrator.db import db
+from orchestrator.search.core.types import EntityType, FilterOp, UIType
+from orchestrator.search.filters import EqualityFilter, FilterTree, LtreeFilter, PathFilter
+from orchestrator.search.retrieval import execute_search
+from orchestrator.search.retrieval.utils import display_filtered_paths_only, display_results
+from orchestrator.search.retrieval.validation import get_structured_filter_schema
+from orchestrator.search.schemas.parameters import BaseSearchParameters
+
+app = typer.Typer(help="Experiment with the subscription search indexes.")
+
+logger = structlog.getLogger(__name__)
+
+
+@app.command()
+def structured(path: str, value: str, entity_type: EntityType = EntityType.SUBSCRIPTION, limit: int = 10) -> None:
+    """Finds subscriptions where a specific field path contains an exact value.
+
+    Example:
+        dotenv run python main.py search structured "subscription.status" "provisioning"
+        ...
+        {
+            "path": "subscription.status",
+            "value": "provisioning"
+        },
+        ...
+    """
+    path_filter = PathFilter(path=path, condition=EqualityFilter(op=FilterOp.EQ, value=value), value_kind=UIType.STRING)
+    search_params = BaseSearchParameters.create(
+        entity_type=entity_type, filters=FilterTree.from_flat_and([path_filter]), limit=limit
+    )
+    search_response = asyncio.run(execute_search(search_params=search_params, db_session=db.session))
+    display_filtered_paths_only(search_response.results, search_params, db.session)
+    display_results(search_response.results, db.session, "Match")
+
+
+@app.command()
+def semantic(query: str, entity_type: EntityType = EntityType.SUBSCRIPTION, limit: int = 10) -> None:
+    """Finds subscriptions that are conceptually most similar to the query text.
+
+    Example:
+        dotenv run python main.py search semantic "Shop for an alligator store"
+        ...
+        {
+            "path": "subscription.shop.shop_description",
+            "value": "Kingswood reptiles shop"
+        },
+        ...
+    """
+    search_params = BaseSearchParameters.create(entity_type=entity_type, query=query, limit=limit)
+    search_response = asyncio.run(execute_search(search_params=search_params, db_session=db.session))
+    display_results(search_response.results, db.session, "Distance")
+
+
+@app.command()
+def fuzzy(term: str, entity_type: EntityType = EntityType.SUBSCRIPTION, limit: int = 10) -> None:
+    """Finds subscriptions containing text similar to the query, tolerating typos.
+
+    Example:
+        dotenv run python main.py search fuzzy "Colonel"
+        ...
+        {
+          "path": "description",
+          "value": "X Follower WF for TimCoronel"
+        },
+        ...
+    """
+    search_params = BaseSearchParameters.create(entity_type=entity_type, query=term, limit=limit)
+    search_response = asyncio.run(execute_search(search_params=search_params, db_session=db.session))
+    display_results(search_response.results, db.session, "Similarity")
+
+
+@app.command()
+def hierarchical(
+    op: str = typer.Argument(..., help="The hierarchical operation to perform."),
+    path: str = typer.Argument(..., help="The ltree path or lquery pattern for the operation."),
+    query: str | None = typer.Option(None, "--query", "-f", help="An optional fuzzy term to rank the results."),
+    entity_type: EntityType = EntityType.SUBSCRIPTION,
+    limit: int = 10,
+) -> None:
+    """Performs a hierarchical search, optionally combined with fuzzy ranking.
+
+    Examples:
+        dotenv run python main.py search hierarchical is_descendant "subscription.shop" --query "Kingwood"
+        dotenv run python main.py search hierarchical matches_lquery "*.x_follower.x_follower_status*"
+    """
+    try:
+        condition = LtreeFilter(value=path, op=op)  # type: ignore[arg-type]
+    except (ValueError, ValidationError) as e:
+        raise typer.BadParameter(f"Invalid filter: {e}")
+
+    path_filter = PathFilter(path="ltree_hierarchical_filter", condition=condition, value_kind=UIType.STRING)
+
+    search_params = BaseSearchParameters.create(
+        entity_type=entity_type, filters=FilterTree.from_flat_and([path_filter]), query=query, limit=limit
+    )
+    search_response = asyncio.run(execute_search(search_params=search_params, db_session=db.session))
+    display_results(search_response.results, db.session, "Hierarchical Score")
+
+
+@app.command()
+def hybrid(query: str, term: str, entity_type: EntityType = EntityType.SUBSCRIPTION, limit: int = 10) -> None:
+    """Performs a hybrid search, combining semantic and fuzzy matching.
+
+    Example:
+        dotenv run python main.py search hybrid "reptile store" "Kingswood"
+    """
+    search_params = BaseSearchParameters.create(entity_type=entity_type, query=query, limit=limit)
+    logger.info("Executing Hybrid Search", query=query, term=term)
+    search_response = asyncio.run(execute_search(search_params=search_params, db_session=db.session))
+    display_results(search_response.results, db.session, "Hybrid Score")
+
+
+@app.command("generate-schema")
+def generate_schema() -> None:
+    """Generates and prints the dynamic filter schema from the live search index.
+
+    This queries the index for all distinct non-string paths to be used as
+    context for the LLM agent.
+
+    Example:
+        dotenv run python main.py search generate-schema
+    """
+
+    schema_map = get_structured_filter_schema()
+
+    if not schema_map:
+        logger.warning("No filterable paths found in the search index.")
+        return
+
+    logger.info("\nAvailable Structured Filters:\n")
+    for path, value_type in schema_map.items():
+        logger.info(f"- {path}: {value_type}")
+
+    logger.info("Successfully generated dynamic schema.", path_count=len(schema_map))
+
+
+@app.command("nested-demo")
+def nested_demo(entity_type: EntityType = EntityType.SUBSCRIPTION, limit: int = 10) -> None:
+    tree = FilterTree.model_validate(
+        {
+            "op": "AND",
+            "children": [
+                {
+                    "op": "OR",
+                    "children": [
+                        # First OR case: Active subscriptions from 2024
+                        {
+                            "op": "AND",
+                            "children": [
+                                {
+                                    "path": "subscription.status",
+                                    "condition": {"op": "eq", "value": "active"},
+                                    "value_kind": "string",
+                                },
+                                {
+                                    "path": "subscription.start_date",
+                                    "condition": {
+                                        "op": "between",
+                                        "value": {
+                                            "start": "2024-01-01T00:00:00Z",
+                                            "end": "2024-12-31T23:59:59Z",
+                                        },
+                                    },
+                                    "value_kind": "datetime",
+                                },
+                            ],
+                        },
+                        # Second OR case: Terminated subscriptions before 2026
+                        {
+                            "op": "AND",
+                            "children": [
+                                {
+                                    "path": "subscription.status",
+                                    "condition": {"op": "eq", "value": "terminated"},
+                                    "value_kind": "string",
+                                },
+                                {
+                                    "path": "subscription.end_date",
+                                    "condition": {"op": "lte", "value": "2025-12-31"},
+                                    "value_kind": "datetime",
+                                },
+                            ],
+                        },
+                    ],
+                },
+                {
+                    "path": "subscription.*.port_mode",
+                    "condition": {"op": "matches_lquery", "value": "*.port_mode"},
+                    "value_kind": "string",
+                },
+            ],
+        }
+    )
+
+    params = BaseSearchParameters.create(entity_type=entity_type, filters=tree, limit=limit)
+    search_response = asyncio.run(execute_search(params, db.session))
+
+    display_results(search_response.results, db.session, "Score")
+
+
+if __name__ == "__main__":
+    app()

orchestrator/cli/speedtest.py

@@ -0,0 +1,151 @@
+import asyncio
+import time
+from typing import Any
+
+import structlog
+import typer
+from rich.console import Console
+from rich.progress import Progress, SpinnerColumn, TextColumn, TimeElapsedColumn
+from rich.table import Table
+
+from orchestrator.db import db
+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__)
+console = Console()
+
+app = typer.Typer(name="speedtest", help="Search speed testing")
+
+DEFAULT_QUERIES = [
+    "network",
+    "fiber",
+    "port",
+    "network infrastructure",
+    "fiber connection",
+    "internet service",
+    "subscription",
+    "active",
+    "configuration",
+    "service provider",
+]
+
+
+async def generate_embeddings_for_queries(queries: list[str]) -> dict[str, list[float]]:
+    embedding_lookup = {}
+
+    for query in queries:
+        try:
+            embedding = await QueryEmbedder.generate_for_text_async(query)
+            if embedding:
+                embedding_lookup[query] = embedding
+            else:
+                logger.warning("Failed to generate embedding for query", query=query)
+        except Exception as e:
+            logger.error("Error generating embedding", query=query, error=str(e))
+
+    return embedding_lookup
+
+
+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)
+
+    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)
+
+    with db.session as session:
+        start_time = time.perf_counter()
+        response = await execute_search(search_params, session, pagination_params=pagination_params)
+        end_time = time.perf_counter()
+
+        return {
+            "query": query,
+            "time": end_time - start_time,
+            "results": len(response.results),
+            "search_type": response.metadata.search_type if hasattr(response, "metadata") else "unknown",
+        }
+
+
+@app.command()
+def quick(
+    queries: list[str] | None = typer.Option(None, "--query", "-q", help="Custom queries to test"),
+) -> None:
+    test_queries = queries if queries else DEFAULT_QUERIES
+
+    console.print(f"[bold blue]Quick Speed Test[/bold blue] - Testing {len(test_queries)} queries")
+
+    async def run_tests() -> list[dict[str, Any]]:
+        embedding_lookup = await generate_embeddings_for_queries(test_queries)
+
+        results = []
+
+        with Progress(
+            SpinnerColumn(),
+            TextColumn("[progress.description]{task.description}"),
+            TimeElapsedColumn(),
+            console=console,
+        ) as progress:
+            task = progress.add_task("Running queries...", total=len(test_queries))
+
+            for query in test_queries:
+                result = await run_single_query(query, embedding_lookup)
+                results.append(result)
+                progress.advance(task)
+
+        return results
+
+    results = asyncio.run(run_tests())
+
+    table = Table(show_header=True, header_style="bold magenta")
+    table.add_column("Query", style="dim", width=25)
+    table.add_column("Time", justify="right", style="cyan")
+    table.add_column("Type", justify="center", style="yellow")
+    table.add_column("Results", justify="right", style="green")
+
+    total_time = 0
+
+    for result in results:
+        time_ms = result["time"] * 1000
+        total_time += result["time"]
+
+        table.add_row(
+            result["query"][:24] + "..." if len(result["query"]) > 24 else result["query"],
+            f"{time_ms:.1f}ms",
+            result["search_type"],
+            str(result["results"]),
+        )
+
+    console.print(table)
+    console.print()
+
+    avg_time = total_time / len(results) * 1000
+    max_time = max(r["time"] for r in results) * 1000
+
+    console.print("[bold]Summary:[/bold]")
+    console.print(f"  Total time: {total_time * 1000:.1f}ms")
+    console.print(f"  Average: {avg_time:.1f}ms")
+    console.print(f"  Slowest: {max_time:.1f}ms")
+
+    by_type: dict[str, list[float]] = {}
+    for result in results:
+        search_type = result["search_type"]
+        if search_type not in by_type:
+            by_type[search_type] = []
+        by_type[search_type].append(result["time"] * 1000)
+
+    for search_type, times in by_type.items():
+        avg = sum(times) / len(times)
+        console.print(f"  {search_type.capitalize()}: {avg:.1f}ms avg ({len(times)} queries)")
+
+
+if __name__ == "__main__":
+    app()

orchestrator/db/models.py

@@ -20,7 +20,9 @@ from uuid import UUID
 import sqlalchemy
 import structlog
 from more_itertools import first_true
+from pgvector.sqlalchemy import Vector
 from sqlalchemy import (
+    TEXT,
     TIMESTAMP,
     Boolean,
     CheckConstraint,
@@ -29,6 +31,7 @@ from sqlalchemy import (
     ForeignKey,
     Index,
     Integer,
+    PrimaryKeyConstraint,
     Select,
     String,
     Table,
@@ -45,10 +48,12 @@ from sqlalchemy.ext.associationproxy import association_proxy
 from sqlalchemy.ext.orderinglist import ordering_list
 from sqlalchemy.orm import Mapped, deferred, mapped_column, object_session, relationship, undefer
 from sqlalchemy.sql.functions import GenericFunction
-from sqlalchemy_utils import TSVectorType, UUIDType
+from sqlalchemy_utils import LtreeType, TSVectorType, UUIDType
 
 from orchestrator.config.assignee import Assignee
 from orchestrator.db.database import BaseModel, SearchQuery
+from orchestrator.llm_settings import llm_settings
+from orchestrator.search.core.types import FieldType
 from orchestrator.targets import Target
 from orchestrator.utils.datetime import nowtz
 from orchestrator.version import GIT_COMMIT_HASH
@@ -685,3 +690,34 @@ class SubscriptionInstanceAsJsonFunction(GenericFunction):
 
     def __init__(self, sub_inst_id: UUID):
         super().__init__(sub_inst_id)
+
+
+class AiSearchIndex(BaseModel):
+
+    __tablename__ = "ai_search_index"
+
+    entity_type = mapped_column(
+        TEXT,
+        nullable=False,
+        index=True,
+    )
+    entity_id = mapped_column(
+        UUIDType,
+        nullable=False,
+    )
+
+    # Ltree path for hierarchical data
+    path = mapped_column(LtreeType, nullable=False, index=True)
+    value = mapped_column(TEXT, nullable=False)
+
+    value_type = mapped_column(
+        Enum(FieldType, name="field_type", values_callable=lambda obj: [e.value for e in obj]), nullable=False
+    )
+
+    # Embedding
+    embedding = mapped_column(Vector(llm_settings.EMBEDDING_DIMENSION), nullable=True)
+
+    # SHA-256
+    content_hash = mapped_column(String(64), nullable=False, index=True)
+
+    __table_args__ = (PrimaryKeyConstraint("entity_id", "path", name="pk_ai_search_index"),)

orchestrator/llm_settings.py

@@ -0,0 +1,51 @@
+# 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.
+from pydantic import Field, field_validator
+from pydantic_settings import BaseSettings
+from structlog import get_logger
+
+logger = get_logger(__name__)
+
+
+class LLMSettings(BaseSettings):
+    LLM_ENABLED: bool = False  # Default to false
+    # Pydantic-ai Agent settings
+    AGENT_MODEL: str = "gpt-4o-mini"  # See pydantic-ai docs for supported models.
+    AGENT_MODEL_VERSION: str = "2025-01-01-preview"
+    OPENAI_API_KEY: str = ""  # Change per provider (Azure, etc).
+    # Embedding settings
+    EMBEDDING_DIMENSION: int = 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
+    )
+
+    # The following settings are only needed for local models.
+    # By default, they are set conservative assuming a small model like All-MiniLM-L6-V2.
+    OPENAI_BASE_URL: str | None = None
+    EMBEDDING_FALLBACK_MAX_TOKENS: int | None = 512
+    EMBEDDING_MAX_BATCH_SIZE: int | None = 32
+
+    # General LiteLLM settings
+    LLM_MAX_RETRIES: int = 3
+    LLM_TIMEOUT: int = 30
+
+    @field_validator("EMBEDDING_MODEL")
+    def validate_embedding_model_format(cls, v: str) -> str:
+        """Validate that embedding model is in 'vendor/model' format."""
+        if "/" not in v:
+            raise ValueError("EMBEDDING_MODEL must be in format 'vendor/model'")
+        return v
+
+
+llm_settings = LLMSettings()