hexdag 0.5.0.dev1__py3-none-any.whl

hexdag/core/ports_builder.py

@@ -0,0 +1,623 @@
+"""Fluent builder for organizing orchestrator ports into logical categories.
+
+This builder provides a clean, type-safe interface for configuring orchestrator
+dependencies while maintaining backward compatibility with the flat dictionary format.
+
+Enhanced with per-node and per-type port configuration support for fine-grained
+control over port assignment across different node types and specific nodes.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Self
+
+if TYPE_CHECKING:
+    from hexdag.core.ports import (
+        LLM,
+        APICall,
+        DatabasePort,
+        Memory,
+        ObserverManagerPort,
+        ToolRouter,
+    )
+
+    # Type alias for all valid port types
+    PortType = (
+        LLM
+        | APICall
+        | DatabasePort
+        | Memory
+        | ObserverManagerPort
+        | ToolRouter
+        | Any  # Allow Any for backward compatibility with custom ports
+    )
+else:
+    # At runtime, just use Any (backward compatible)
+    PortType = Any
+
+from hexdag.core.orchestration.models import PortConfig, PortsConfiguration
+
+
+class PortsBuilder:
+    """Fluent builder for organizing ports into logical categories.
+
+    Provides a type-safe, discoverable API for configuring orchestrator
+    dependencies while maintaining backward compatibility.
+
+    Example
+    -------
+        ```python
+        # Traditional flat dictionary approach
+        ports = {
+            "llm": OpenAIAdapter(),
+            "database": PostgresAdapter(),
+            "observer_manager": ObserverManager(),
+            # ... many more mixed together
+        }
+
+        # New builder approach - organized and type-safe
+        ports = (
+            PortsBuilder()
+            .with_llm(OpenAIAdapter())
+            .with_database(PostgresAdapter())
+            .with_defaults()
+            .build()
+        )
+
+        # Or use with_defaults() for automatic setup
+        ports = PortsBuilder().with_defaults().build()
+        ```
+    """
+
+    def __init__(self) -> None:
+        """Initialize an empty ports builder."""
+        self._ports: dict[str, PortType] = {}
+        self._type_ports: dict[str, dict[str, PortType]] = {}
+        self._node_ports: dict[str, dict[str, PortType]] = {}
+
+    def _add_port(self, key: str, port: PortType) -> Self:
+        """Add a port to the internal registry.
+
+        Args
+        ----
+            key: Port identifier
+            port: Port implementation
+
+        Returns
+        -------
+            Self for method chaining
+        """
+        self._ports[key] = port
+        return self
+
+    # Core AI Capabilities
+    # --------------------
+
+    def with_llm(self, llm: LLM) -> Self:
+        """Add a Language Model adapter.
+
+        Args
+        ----
+            llm: LLM adapter instance (OpenAI, Anthropic, etc.)
+
+        Returns
+        -------
+            Self for method chaining
+        """
+        return self._add_port("llm", llm)
+
+    def with_tool_router(self, router: ToolRouter) -> Self:
+        """Add a tool router for function calling.
+
+        Args
+        ----
+            router: Tool router instance for managing tool execution
+
+        Returns
+        -------
+            Self for method chaining
+        """
+        return self._add_port("tool_router", router)
+
+    # Storage & Persistence
+    # ---------------------
+
+    def with_database(self, database: DatabasePort) -> Self:
+        """Add a database adapter.
+
+        Args
+        ----
+            database: Database adapter instance
+
+        Returns
+        -------
+            Self for method chaining
+        """
+        return self._add_port("database", database)
+
+    def with_memory(self, memory: Memory) -> Self:
+        """Add a memory system for agents.
+
+        Args
+        ----
+            memory: Memory adapter for conversation history
+
+        Returns
+        -------
+            Self for method chaining
+        """
+        return self._add_port("memory", memory)
+
+    # Event & Control Systems
+    # -----------------------
+
+    def with_observer_manager(self, manager: ObserverManagerPort) -> Self:
+        """Add an observer manager for event monitoring.
+
+        Args
+        ----
+            manager: Observer manager for read-only event handling
+
+        Returns
+        -------
+            Self for method chaining
+        """
+        return self._add_port("observer_manager", manager)
+
+    # External Integrations
+    # ---------------------
+
+    def with_api_call(self, api_call: APICall) -> Self:
+        """Add an API call adapter.
+
+        Args
+        ----
+            api_call: API call adapter for external services
+
+        Returns
+        -------
+            Self for method chaining
+        """
+        return self._add_port("api_call", api_call)
+
+    # Convenience Methods
+    # ------------------
+
+    def with_defaults(self) -> Self:
+        """Add default implementations for common ports.
+
+        This method provides sensible defaults:
+        - LocalObserverManager for event observation
+        - MockLLM for testing (should be overridden in production)
+
+        Returns
+        -------
+            Self for method chaining
+        """
+        # Only add defaults if not already configured
+        if "observer_manager" not in self._ports:
+            try:
+                from hexdag.builtin.adapters.local import LocalObserverManager
+
+                self._ports["observer_manager"] = LocalObserverManager()
+            except ImportError:
+                pass  # Optional dependency
+
+        if "llm" not in self._ports:
+            try:
+                from hexdag.builtin.adapters.mock import MockLLM
+
+                self._ports["llm"] = MockLLM()
+            except ImportError:
+                pass  # Optional dependency
+
+        return self
+
+    def with_custom(self, key: str, port: Any) -> Self:
+        """Add a custom port with any key.
+
+        This provides backward compatibility and flexibility for
+        custom port implementations.
+
+        Args
+        ----
+            key: The port key name
+            port: The port implementation
+
+        Returns
+        -------
+            Self for method chaining
+        """
+        return self._add_port(key, port)
+
+    def update(self, ports: dict[str, Any]) -> Self:
+        """Update multiple ports at once from a dictionary.
+
+        Useful for migrating from existing dictionary-based configs.
+
+        Args
+        ----
+            ports: Dictionary of port implementations
+
+        Returns
+        -------
+            Self for method chaining
+        """
+        self._ports.update(ports)
+        return self
+
+    # Builder Methods
+    # --------------
+
+    def build(self) -> dict[str, Any]:
+        """Build and return the final ports dictionary.
+
+        Returns
+        -------
+            Dictionary of configured ports for orchestrator use
+        """
+        return self._ports.copy()
+
+    def clear(self) -> Self:
+        """Clear all configured ports.
+
+        Returns
+        -------
+            Self for method chaining
+        """
+        self._ports.clear()
+        return self
+
+    # Inspection Methods
+    # -----------------
+
+    def has(self, key: str) -> bool:
+        """Check if a port is configured.
+
+        Args
+        ----
+            key: The port key to check
+
+        Returns
+        -------
+            True if the port is configured
+        """
+        return key in self._ports
+
+    def __contains__(self, key: str) -> bool:
+        """Check if a port is configured using 'in' operator.
+
+        Parameters
+        ----------
+        key : str
+            Port key to check
+
+        Returns
+        -------
+        bool
+            True if the port is configured at global level
+
+        Examples
+        --------
+        >>> builder = PortsBuilder()
+        >>> class MockPort: pass
+        >>> builder["llm"] = MockPort()
+        >>> "llm" in builder
+        True
+        >>> "database" in builder
+        False
+        """
+        return key in self._ports
+
+    def get(self, key: str, default: Any = None) -> Any:
+        """Get a configured port by key.
+
+        Args
+        ----
+            key: The port key
+            default: Default value if not found
+
+        Returns
+        -------
+            The port instance or default value
+        """
+        return self._ports.get(key, default)
+
+    def keys(self) -> list[str]:
+        """Get list of configured port keys.
+
+        Returns
+        -------
+            List of port keys
+        """
+        return list(self._ports.keys())
+
+    def __len__(self) -> int:
+        """Get number of configured ports.
+
+        Returns
+        -------
+            Number of configured ports
+        """
+        return len(self._ports)
+
+    def __repr__(self) -> str:
+        """Get string representation of builder state.
+
+        Returns
+        -------
+            String showing configured ports
+        """
+        configured = ", ".join(self._ports.keys()) if self._ports else "none"
+        return f"PortsBuilder(configured: {configured})"
+
+    def __getitem__(self, key: str | tuple) -> Any:
+        """Get port(s) using dictionary-style access with tuple key support.
+
+        Supports simple and hierarchical access patterns:
+        - Simple: builder["llm"] - Get global port
+        - Tuple (2 elements): builder["agent", "llm"] - Get type-level port
+        - Tuple (3 elements): builder["researcher", "agent", "llm"] - Get node-level port
+
+        Parameters
+        ----------
+        key : str | tuple
+            Port key or tuple of (node, type, port) for hierarchical lookup
+
+        Returns
+        -------
+        Any
+            The port instance
+
+        Raises
+        ------
+        KeyError
+            If the port is not configured
+
+        Examples
+        --------
+        >>> builder = PortsBuilder()
+        >>> class MockLLM: pass
+        >>> class AgentLLM: pass
+        >>> builder["llm"] = MockLLM()
+        >>> port = builder["llm"]  # Simple access
+        >>> isinstance(port, MockLLM)
+        True
+        >>> result = builder.for_type("agent", llm=AgentLLM())  # Returns self
+        >>> port = builder["agent", "llm"]  # Type-level access
+        >>> isinstance(port, AgentLLM)
+        True
+        """
+        if isinstance(key, tuple):
+            if len(key) == 2:
+                # (type, port_name) - Get type-level port
+                node_type, port_name = key
+                if node_type in self._type_ports and port_name in self._type_ports[node_type]:
+                    return self._type_ports[node_type][port_name]
+                # Fall back to global
+                return self._ports.get(port_name)
+            if len(key) == 3:
+                # (node_name, type, port_name) - Get node-level port with resolution
+                node_name, node_type, port_name = key
+                # Try node-level first
+                if node_name in self._node_ports and port_name in self._node_ports[node_name]:
+                    return self._node_ports[node_name][port_name]
+                # Fall back to type-level
+                if node_type in self._type_ports and port_name in self._type_ports[node_type]:
+                    return self._type_ports[node_type][port_name]
+                # Fall back to global
+                return self._ports.get(port_name)
+            raise KeyError(f"Tuple key must have 2 or 3 elements, got {len(key)}")
+        return self._ports[key]
+
+    def __setitem__(self, key: str | tuple, value: Any) -> None:
+        """Set port(s) using dictionary-style access with tuple key support.
+
+        Supports simple and hierarchical configuration:
+        - Simple: builder["llm"] = adapter - Set global port
+        - Tuple (2 elements): builder["agent", "llm"] = adapter - Set type-level port
+        - Tuple (3 elements): builder["node", "agent", "llm"] = adapter - Set node-level port
+
+        Parameters
+        ----------
+        key : str | tuple
+            Port key or tuple of (node, type, port) for hierarchical configuration
+        value : Any
+            Port implementation to configure
+
+        Examples
+        --------
+        >>> builder = PortsBuilder()
+        >>> class MockLLM: pass
+        >>> class OpenAIAdapter: pass
+        >>> class ClaudeAdapter: pass
+        >>> builder["llm"] = MockLLM()  # Global
+        >>> builder["agent", "llm"] = OpenAIAdapter()  # Type-level
+        >>> builder["researcher", "agent", "llm"] = ClaudeAdapter()  # Node-level
+        >>> isinstance(builder["llm"], MockLLM)
+        True
+        >>> isinstance(builder["agent", "llm"], OpenAIAdapter)
+        True
+        >>> isinstance(builder["researcher", "agent", "llm"], ClaudeAdapter)
+        True
+        """
+        if isinstance(key, tuple):
+            if len(key) == 2:
+                # (type, port_name) - Set type-level port
+                node_type, port_name = key
+                self.for_type(node_type, **{port_name: value})
+            elif len(key) == 3:
+                # (node_name, type, port_name) - Set node-level port
+                node_name, _node_type, port_name = key
+                self.for_node(node_name, **{port_name: value})
+            else:
+                raise KeyError(f"Tuple key must have 2 or 3 elements, got {len(key)}")
+        else:
+            self._add_port(key, value)
+
+    def __delitem__(self, key: str) -> None:
+        """Remove a port from the builder.
+
+        Parameters
+        ----------
+        key : str
+            Port key to remove
+
+        Examples
+        --------
+        .. code-block:: python
+
+            builder = PortsBuilder()
+            builder["llm"] = MockLLM()
+            del builder["llm"]
+            assert "llm" not in builder
+        """
+        del self._ports[key]
+
+    # Enhanced Configuration Methods
+    # ------------------------------
+
+    def for_type(self, node_type: str, **ports: Any) -> Self:
+        """Configure ports for all nodes of a specific type.
+
+        This method allows setting default ports for all nodes of a given type
+        (e.g., "agent", "llm", "function"). These type-level defaults override
+        global defaults but are overridden by per-node configurations.
+
+        Args
+        ----
+            node_type: The node type to configure (e.g., "agent", "llm", "function")
+            **ports: Port implementations as keyword arguments
+
+        Returns
+        -------
+            Self for method chaining
+
+        Examples
+        --------
+        Example usage::
+
+            from hexdag.builtin.adapters.openai import OpenAIAdapter
+            from hexdag.builtin.adapters.mock import MockLLM
+
+            builder = (
+                PortsBuilder()
+                .with_llm(MockLLM())  # Global default
+                .for_type("agent", llm=OpenAIAdapter(model="gpt-4"))  # Agent nodes
+                .build_configuration()
+            )
+
+            # All agent nodes will use OpenAI, other nodes use MockLLM
+            config = builder.build_configuration()
+            agent_ports = config.resolve_ports("my_agent", "agent")
+            assert isinstance(agent_ports["llm"].port, OpenAIAdapter)
+        """
+        if node_type not in self._type_ports:
+            self._type_ports[node_type] = {}
+        self._type_ports[node_type].update(ports)
+        return self
+
+    def for_node(self, node_name: str, **ports: Any) -> Self:
+        """Configure ports for a specific node by name.
+
+        This method allows overriding ports for individual nodes, providing
+        the highest level of configuration precedence. Perfect for nodes that
+        require special adapters or configurations.
+
+        Args
+        ----
+            node_name: The node name to configure
+            **ports: Port implementations as keyword arguments
+
+        Returns
+        -------
+            Self for method chaining
+
+        Examples
+        --------
+        Example usage::
+
+            from hexdag.builtin.adapters.anthropic import AnthropicAdapter
+            from hexdag.builtin.adapters.openai import OpenAIAdapter
+
+            builder = (
+                PortsBuilder()
+                .for_type("agent", llm=OpenAIAdapter(model="gpt-4"))  # Agent default
+                .for_node("researcher", llm=AnthropicAdapter(model="claude-3"))  # Override
+                .build_configuration()
+            )
+
+            # Researcher node gets Claude, other agents get GPT-4
+            config = builder.build_configuration()
+            researcher_ports = config.resolve_ports("researcher", "agent")
+            assert isinstance(researcher_ports["llm"].port, AnthropicAdapter)
+        """
+        if node_name not in self._node_ports:
+            self._node_ports[node_name] = {}
+        self._node_ports[node_name].update(ports)
+        return self
+
+    def build_configuration(self) -> PortsConfiguration:
+        """Build a PortsConfiguration with full inheritance support.
+
+        Creates a PortsConfiguration object that encapsulates global, per-type,
+        and per-node port configurations. This provides more flexibility than
+        the flat dictionary returned by build().
+
+        Returns
+        -------
+        PortsConfiguration
+            Configuration with port inheritance and resolution support
+
+        Examples
+        --------
+        Example usage::
+
+            from hexdag.builtin.adapters.mock import MockLLM
+            from hexdag.builtin.adapters.openai import OpenAIAdapter
+            from hexdag.builtin.adapters.anthropic import AnthropicAdapter
+
+            config = (
+                PortsBuilder()
+                .with_llm(MockLLM())  # Global default
+                .for_type("agent", llm=OpenAIAdapter(model="gpt-4"))  # Agent default
+                .for_node("researcher", llm=AnthropicAdapter(model="claude-3"))  # Override
+                .build_configuration()
+            )
+
+            # Resolve ports for different nodes
+            researcher = config.to_flat_dict("researcher", "agent")  # AnthropicAdapter
+            analyzer = config.to_flat_dict("analyzer", "agent")  # OpenAIAdapter
+            transformer = config.to_flat_dict("transformer", "function")  # MockLLM
+
+        Notes
+        -----
+        Resolution order: per-node > per-type > global defaults
+
+        See Also
+        --------
+        build : For backward-compatible flat dictionary output
+        """
+
+        def _wrap_ports(ports_dict: dict[str, PortType] | None) -> dict[str, PortConfig] | None:
+            """Helper to wrap ports in PortConfig."""
+            return {k: PortConfig(port=v) for k, v in ports_dict.items()} if ports_dict else None
+
+        def _wrap_nested_ports(
+            nested: dict[str, dict[str, PortType]] | None,
+        ) -> dict[str, dict[str, PortConfig]] | None:
+            """Helper to wrap nested ports in PortConfig."""
+            return (
+                {
+                    name: {k: PortConfig(port=v) for k, v in ports.items()}
+                    for name, ports in nested.items()
+                }
+                if nested
+                else None
+            )
+
+        return PortsConfiguration(
+            global_ports=_wrap_ports(self._ports),
+            type_ports=_wrap_nested_ports(self._type_ports),
+            node_ports=_wrap_nested_ports(self._node_ports),
+        )