knowledge2 0.4.0__py3-none-any.whl

sdk/resources/pipeline_builder.py

@@ -0,0 +1,257 @@
+"""Fluent builder for Pipeline Spec topologies."""
+
+from __future__ import annotations
+
+import re
+from typing import Any, Literal
+
+from sdk._preview import preview_resource
+
+try:
+    from typing import Self
+except ImportError:
+    from typing_extensions import Self
+
+_UUID_RE = re.compile(
+    r"^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$", re.IGNORECASE
+)
+
+
+@preview_resource
+class PipelineBuilder:
+    """Build a Pipeline Spec topology using a fluent API.
+
+    Example::
+
+        topology = (
+            PipelineBuilder()
+            .add_corpus("my-corpus", description="Main corpus")
+            .add_agent("my-agent", "summarize", corpus_ref="my-corpus")
+            .add_feed("my-feed", query="latest news", agent_ref="my-agent")
+            .subscribe(agent_ref="my-agent", feed_ref="my-feed", role="input")
+            .build()
+        )
+    """
+
+    def __init__(self) -> None:
+        self._corpora: list[dict[str, Any]] = []
+        self._agents: list[dict[str, Any]] = []
+        self._feeds: list[dict[str, Any]] = []
+        self._subscriptions: list[dict[str, Any]] = []
+        self._metadata: dict[str, Any] | None = None
+
+    def add_corpus(
+        self, name: str, *, id: str | None = None, description: str | None = None
+    ) -> Self:
+        """Add a corpus node to the topology.
+
+        Args:
+            name: Name for the corpus.
+            id: Optional existing corpus ID to reference.
+            description: Optional description.
+
+        Returns:
+            Self for method chaining.
+        """
+        node: dict[str, Any] = {"name": name}
+        if id is not None:
+            node["id"] = id
+        if description is not None:
+            node["description"] = description
+        self._corpora.append(node)
+        return self
+
+    def add_agent(
+        self,
+        name: str,
+        task_type: str,
+        *,
+        corpus_ref: str | None = None,
+        id: str | None = None,
+        description: str | None = None,
+        instructions: str | None = None,
+        model: str | None = None,
+        tool_config: dict[str, Any] | None = None,
+    ) -> Self:
+        """Add an agent node to the topology.
+
+        Args:
+            name: Name for the agent.
+            task_type: Task type for the agent (required).
+            corpus_ref: Optional reference to a corpus (name or existing ID).
+            id: Optional existing agent ID to reference.
+            description: Optional description.
+            instructions: Optional instructions for the agent.
+            model: Optional model identifier.
+            tool_config: Optional tool configuration dict.
+
+        Returns:
+            Self for method chaining.
+        """
+        node: dict[str, Any] = {"name": name, "task_type": task_type}
+        if corpus_ref is not None:
+            node["corpus_ref"] = corpus_ref
+        if id is not None:
+            node["id"] = id
+        if description is not None:
+            node["description"] = description
+        if instructions is not None:
+            node["instructions"] = instructions
+        if model is not None:
+            node["model"] = model
+        if tool_config is not None:
+            node["tool_config"] = tool_config
+        self._agents.append(node)
+        return self
+
+    def add_feed(
+        self,
+        name: str,
+        *,
+        query: str | None = None,
+        agent_ref: str | None = None,
+        id: str | None = None,
+        persistent: bool = False,
+        target_corpus_ref: str | None = None,
+        reactive: bool = False,
+        schedule_interval: str | None = None,
+    ) -> Self:
+        """Add a feed node to the topology.
+
+        Args:
+            name: Name for the feed.
+            query: Optional query string for the feed.
+            agent_ref: Optional reference to an agent (name or existing ID).
+            id: Optional existing feed ID to reference.
+            persistent: Whether the feed persists results.
+            target_corpus_ref: Optional reference to a target corpus.
+            reactive: Whether the feed is reactive.
+            schedule_interval: Optional schedule interval.
+
+        Returns:
+            Self for method chaining.
+        """
+        node: dict[str, Any] = {"name": name}
+        if query is not None:
+            node["query"] = query
+        if agent_ref is not None:
+            node["agent_ref"] = agent_ref
+        node["persistent"] = persistent
+        node["reactive"] = reactive
+        if id is not None:
+            node["id"] = id
+        if target_corpus_ref is not None:
+            node["target_corpus_ref"] = target_corpus_ref
+        if schedule_interval is not None:
+            node["schedule_interval"] = schedule_interval
+        self._feeds.append(node)
+        return self
+
+    def subscribe(self, agent_ref: str, feed_ref: str, role: Literal["input", "output"]) -> Self:
+        """Add a subscription linking an agent and a feed.
+
+        Args:
+            agent_ref: Reference to an agent (name or existing ID).
+            feed_ref: Reference to a feed (name or existing ID).
+            role: The role of the subscription (``"input"`` or ``"output"``).
+
+        Returns:
+            Self for method chaining.
+        """
+        self._subscriptions.append({"agent_ref": agent_ref, "feed_ref": feed_ref, "role": role})
+        return self
+
+    def set_metadata(self, metadata: dict[str, Any]) -> Self:
+        """Set arbitrary metadata on the topology.
+
+        Args:
+            metadata: Metadata dictionary.
+
+        Returns:
+            Self for method chaining.
+        """
+        self._metadata = metadata
+        return self
+
+    def _is_existing_id(self, ref: str) -> bool:
+        return bool(_UUID_RE.match(ref))
+
+    def _validate_ref(self, ref: str, field_name: str, allowed: set[str]) -> None:
+        if ref not in allowed and not self._is_existing_id(ref):
+            raise ValueError(
+                f"{field_name} '{ref}' does not match any defined node name "
+                f"and is not a valid existing entity ID (UUID)"
+            )
+
+    def build(self) -> dict[str, Any]:
+        """Build and return a TopologyDocument dict.
+
+        Validates that all refs resolve to defined node names of the
+        correct entity type or existing entity IDs (UUIDs).
+
+        Raises:
+            ValueError: If topology is empty or has unresolved refs.
+        """
+        if not self._corpora and not self._agents and not self._feeds:
+            raise ValueError("Topology must contain at least one entity")
+
+        corpus_names: set[str] = {c["name"] for c in self._corpora}
+        agent_names: set[str] = {a["name"] for a in self._agents}
+        feed_names: set[str] = {f["name"] for f in self._feeds}
+
+        # Validate agent corpus_ref (must reference a corpus)
+        for agent in self._agents:
+            if "corpus_ref" in agent:
+                self._validate_ref(agent["corpus_ref"], "corpus_ref", corpus_names)
+
+        # Validate feed agent_ref (must reference an agent) and target_corpus_ref (must reference a corpus)
+        for feed in self._feeds:
+            if "agent_ref" in feed:
+                self._validate_ref(feed["agent_ref"], "agent_ref", agent_names)
+            if "target_corpus_ref" in feed:
+                self._validate_ref(feed["target_corpus_ref"], "target_corpus_ref", corpus_names)
+
+        # Validate subscription refs (agent_ref → agent, feed_ref → feed)
+        for sub in self._subscriptions:
+            self._validate_ref(sub["agent_ref"], "agent_ref", agent_names)
+            self._validate_ref(sub["feed_ref"], "feed_ref", feed_names)
+
+        doc: dict[str, Any] = {
+            "corpora": list(self._corpora),
+            "agents": list(self._agents),
+            "feeds": list(self._feeds),
+            "subscriptions": list(self._subscriptions),
+        }
+        if self._metadata is not None:
+            doc["metadata"] = self._metadata
+        return doc
+
+    def create_and_apply(
+        self,
+        client: Any,
+        project_id: str,
+        name: str,
+        *,
+        description: str | None = None,
+        activate_entities: bool = True,
+    ) -> dict[str, Any]:
+        """Build topology, create a pipeline spec, and apply it.
+
+        Args:
+            client: A Knowledge2 client instance (sync).
+            project_id: Project to create the pipeline spec in.
+            name: Name for the pipeline spec.
+            description: Optional description.
+            activate_entities: Whether to activate created entities on apply.
+
+        Returns:
+            The ApplyResult dict from the apply call.
+        """
+        topology = self.build()
+        spec = client.create_pipeline_spec(
+            project_id=project_id,
+            name=name,
+            topology=topology,
+            description=description,
+        )
+        return client.apply_pipeline_spec(spec["id"], activate_entities=activate_entities)