diff --git a/engine/.env b/engine/.env index 1a90ad8c18..80958d80c9 100644 --- a/engine/.env +++ b/engine/.env @@ -48,6 +48,11 @@ STIRLING_RAG_TOP_K=20 # rather than chain more searches. STIRLING_RAG_MAX_SEARCHES=5 +# How many operations the edit planner sees, ranked against the request by embedding +# similarity. The full catalogue overruns a local model's context window. Set 0 to +# disable ranking and show every operation. +STIRLING_PLANNER_SHORTLIST_SIZE=20 + # Chunked reasoner settings: how big each per-worker slice is (in characters), # how many workers may run in parallel against the fast model, and how long # any single worker is allowed to wait for a response before being abandoned. diff --git a/engine/src/stirling/agents/operation_shortlist.py b/engine/src/stirling/agents/operation_shortlist.py new file mode 100644 index 0000000000..3254a67dd4 --- /dev/null +++ b/engine/src/stirling/agents/operation_shortlist.py @@ -0,0 +1,62 @@ +"""Narrows the operation catalogue to the candidates worth showing the planner.""" + +from __future__ import annotations + +import asyncio +import logging + +from stirling.documents import EmbeddingService +from stirling.models import OPERATIONS, ToolEndpoint + +logger = logging.getLogger(__name__) + + +def retrieval_text(operation: ToolEndpoint) -> str: + model = OPERATIONS[operation] + description = (model.model_json_schema().get("description") or "").strip() + parameters = [(field.description or "").strip() for field in model.model_fields.values() if field.description] + return f"{operation.name.replace('_', ' ').lower()}. {description} {' '.join(parameters)}".strip() + + +def _cosine(left: list[float], right: list[float]) -> float: + dot = sum(a * b for a, b in zip(left, right, strict=True)) + left_norm = sum(a * a for a in left) ** 0.5 + right_norm = sum(b * b for b in right) ** 0.5 + return dot / (left_norm * right_norm) if left_norm and right_norm else 0.0 + + +class OperationShortlist: + def __init__(self, embedder: EmbeddingService) -> None: + self._embedder = embedder + self._vectors: dict[ToolEndpoint, list[float]] | None = None + self._lock = asyncio.Lock() + + async def _catalogue_vectors(self) -> dict[ToolEndpoint, list[float]]: + async with self._lock: + if self._vectors is None: + operations = list(OPERATIONS) + embeddings = await self._embedder.embed_documents([retrieval_text(op) for op in operations]) + self._vectors = dict(zip(operations, embeddings, strict=True)) + return self._vectors + + async def select( + self, + message: str, + operations: list[ToolEndpoint], + limit: int, + ) -> list[ToolEndpoint]: + """The ``limit`` operations closest to ``message``, or all of them if ranking is unavailable. + + Falling back to the full list keeps a planner that would otherwise work on a deployment + with no reachable embedding provider, at the cost of a much larger prompt. + """ + if limit <= 0 or len(operations) <= limit: + return operations + try: + vectors = await self._catalogue_vectors() + query = await self._embedder.embed_query(message) + except Exception: # noqa: BLE001 - any embedding failure should fall back, never fail planning + logger.warning("[pdf-edit] operation ranking unavailable, showing the full catalogue", exc_info=True) + return operations + ranked = sorted(operations, key=lambda op: -_cosine(query, vectors[op])) + return ranked[:limit] diff --git a/engine/src/stirling/agents/pdf_edit.py b/engine/src/stirling/agents/pdf_edit.py index c46fdb9f24..05c792636e 100644 --- a/engine/src/stirling/agents/pdf_edit.py +++ b/engine/src/stirling/agents/pdf_edit.py @@ -9,6 +9,7 @@ from pydantic_ai import Agent from pydantic_ai.output import NativeOutput from stirling.agents._page_text import format_page_text, get_extracted_text_artifact, has_page_text +from stirling.agents.operation_shortlist import OperationShortlist from stirling.contracts import ( EditCannotDoResponse, EditClarificationRequest, @@ -175,6 +176,7 @@ class PdfEditAgent: def __init__(self, runtime: AppRuntime) -> None: self.runtime = runtime self.parameter_selector = PdfEditParameterSelector(runtime) + self.shortlist = OperationShortlist(runtime.documents.embedder) async def orchestrate(self, request: OrchestratorRequest) -> PdfEditResponse: """Entry point for the orchestrator delegate — adapts the orchestrator's @@ -277,13 +279,20 @@ class PdfEditAgent: repair_note: str = "", ) -> PdfEditPlanOutput: can_request_content = allow_need_content and not has_page_text(request.page_text) + available = list(supported_operations) + candidates = await self.shortlist.select( + request.user_message, + available, + self.runtime.settings.planner_shortlist_size, + ) + logger.info("[pdf-edit] showing %d of %d operations", len(candidates), len(available)) agent = self._build_selection_agent( - supported_operations, + candidates, unavailable_operations, allow_need_content=can_request_content, ) return await agent.select( - self._build_selection_prompt(request, supported_operations, unavailable_operations, repair_note) + self._build_selection_prompt(request, candidates, unavailable_operations, repair_note) ) def _build_selection_agent( @@ -359,19 +368,10 @@ class PdfEditAgent: f"Extracted page text:\n{format_page_text(request.page_text)}" ) - # Endpoints that exist on the server and are callable via the direct API or the manual UI, - # but are never offered to the AI agent as a routing option. - # - # Why: REDACT_EXECUTE is the preferred AI-driven redaction route. AUTO_REDACT and REDACT are - # legacy endpoints that remain fully functional for human callers (the manual redact UI, direct - # API consumers, pipelines) but would produce a worse experience if the AI routed to them — - # they accept a simpler, less expressive schema and pre-date the unified operation model. - # Hiding them here channels all AI redaction traffic through REDACT_EXECUTE without disabling - # the legacy endpoints for anyone else. - # - # How to reuse: add an endpoint here whenever a legacy endpoint has a preferred replacement - # that the AI should use exclusively. The endpoint remains live on the server; only the AI - # planner is prevented from selecting it. + # Hidden from the AI planner only; still live for the manual UI, direct API and pipelines. + # AUTO_REDACT and REDACT pre-date the unified operation model and take a less expressive + # schema, so AI redaction is channelled through REDACT_EXECUTE. Add an endpoint here when a + # legacy one has a preferred replacement the AI should use exclusively. _AGENT_HIDDEN_ENDPOINTS: frozenset[ToolEndpoint] = frozenset({ToolEndpoint.AUTO_REDACT, ToolEndpoint.REDACT}) def _classify_operations(self, request: PdfEditRequest) -> tuple[list[ToolEndpoint], list[ToolEndpoint]]: diff --git a/engine/src/stirling/config/settings.py b/engine/src/stirling/config/settings.py index d0cf6eba2f..d753884592 100644 --- a/engine/src/stirling/config/settings.py +++ b/engine/src/stirling/config/settings.py @@ -50,6 +50,10 @@ class AppSettings(BaseSettings): rag_chunk_overlap: int = Field(validation_alias="STIRLING_RAG_CHUNK_OVERLAP") rag_default_top_k: int = Field(validation_alias="STIRLING_RAG_TOP_K") rag_max_searches: int = Field(validation_alias="STIRLING_RAG_MAX_SEARCHES") + + # How many of the catalogue's operations the edit planner is shown. The full list + # overruns a local model's context; 0 disables ranking and shows every operation. + planner_shortlist_size: int = Field(default=20, validation_alias="STIRLING_PLANNER_SHORTLIST_SIZE") documents_reaper_interval_seconds: int = Field( default=900, validation_alias="STIRLING_DOCUMENTS_REAPER_INTERVAL_SECONDS", diff --git a/engine/tests/test_operation_shortlist.py b/engine/tests/test_operation_shortlist.py new file mode 100644 index 0000000000..69016dad77 --- /dev/null +++ b/engine/tests/test_operation_shortlist.py @@ -0,0 +1,76 @@ +from __future__ import annotations + +import pytest + +from stirling.agents.operation_shortlist import OperationShortlist, retrieval_text +from stirling.models import OPERATIONS, ToolEndpoint + + +class StubEmbedder: + """Embeds on a single axis: how often 'watermark' appears, so ranking is predictable.""" + + def __init__(self, fail: bool = False) -> None: + self.fail = fail + self.document_calls = 0 + + @staticmethod + def _vector(text: str) -> list[float]: + return [float(text.lower().count("watermark")), 1.0] + + async def embed_documents(self, texts: list[str]) -> list[list[float]]: + if self.fail: + raise RuntimeError("embedding provider unreachable") + self.document_calls += 1 + return [self._vector(text) for text in texts] + + async def embed_query(self, text: str) -> list[float]: + if self.fail: + raise RuntimeError("embedding provider unreachable") + return self._vector(text) + + +@pytest.mark.anyio +async def test_returns_every_operation_when_the_limit_is_not_binding() -> None: + shortlist = OperationShortlist(StubEmbedder()) + operations = list(OPERATIONS)[:5] + + assert await shortlist.select("watermark this", operations, len(operations)) == operations + assert await shortlist.select("watermark this", operations, 0) == operations + + +@pytest.mark.anyio +async def test_narrows_to_the_closest_operations() -> None: + shortlist = OperationShortlist(StubEmbedder()) + operations = list(OPERATIONS) + + selected = await shortlist.select("watermark watermark watermark", operations, 5) + + assert len(selected) == 5 + assert ToolEndpoint.ADD_WATERMARK in selected + + +@pytest.mark.anyio +async def test_falls_back_to_the_full_catalogue_when_embedding_fails() -> None: + shortlist = OperationShortlist(StubEmbedder(fail=True)) + operations = list(OPERATIONS) + + assert await shortlist.select("watermark this", operations, 5) == operations + + +@pytest.mark.anyio +async def test_embeds_the_catalogue_once() -> None: + embedder = StubEmbedder() + shortlist = OperationShortlist(embedder) + operations = list(OPERATIONS) + + await shortlist.select("watermark this", operations, 5) + await shortlist.select("rotate this", operations, 5) + + assert embedder.document_calls == 1 + + +def test_retrieval_text_carries_parameter_descriptions() -> None: + text = retrieval_text(ToolEndpoint.ADD_WATERMARK) + + assert "add watermark" in text + assert "opacity" in text.lower()