flock-core 0.5.10__py3-none-any.whl → 0.5.20__py3-none-any.whl

flock/agent/component_lifecycle.py

@@ -0,0 +1,325 @@
+"""Agent component lifecycle management - hook execution and coordination.
+
+Phase 4: Extracted from agent.py to organize component hook execution logic.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
+
+from flock.logging.logging import get_logger
+
+
+if TYPE_CHECKING:
+    from flock.components import AgentComponent, EngineComponent
+    from flock.core import Agent
+    from flock.core.artifacts import Artifact
+    from flock.utils.runtime import Context, EvalInputs, EvalResult
+
+logger = get_logger(__name__)
+
+
+class ComponentLifecycle:
+    """Manages agent component lifecycle hook execution.
+
+    This module handles all component hook execution during agent lifecycle:
+    - on_initialize
+    - on_pre_consume
+    - on_pre_evaluate
+    - on_post_evaluate
+    - on_post_publish
+    - on_error
+    - on_terminate
+    """
+
+    def __init__(self, agent_name: str):
+        """Initialize ComponentLifecycle for a specific agent.
+
+        Args:
+            agent_name: Name of the agent (for logging)
+        """
+        self._agent_name = agent_name
+        self._logger = logging.getLogger(__name__)
+
+    def _component_display_name(
+        self, component: AgentComponent | EngineComponent
+    ) -> str:
+        """Get display name for component logging."""
+        return getattr(component, "name", None) or component.__class__.__name__
+
+    async def run_initialize(
+        self,
+        agent: Agent,
+        ctx: Context,
+        utilities: list[AgentComponent],
+        engines: list[EngineComponent],
+    ) -> None:
+        """Execute on_initialize hooks for all components.
+
+        Args:
+            agent: Agent instance
+            ctx: Execution context
+            utilities: List of utility components
+            engines: List of engine components
+
+        Raises:
+            Exception: If any component initialization fails
+        """
+        for component in utilities:
+            comp_name = self._component_display_name(component)
+            priority = getattr(component, "priority", 0)
+            logger.debug(
+                f"Agent initialize: agent={self._agent_name}, component={comp_name}, priority={priority}"
+            )
+            try:
+                await component.on_initialize(agent, ctx)
+            except Exception as exc:
+                logger.exception(
+                    f"Agent initialize failed: agent={self._agent_name}, component={comp_name}, "
+                    f"priority={priority}, error={exc!s}"
+                )
+                raise
+
+        for engine in engines:
+            await engine.on_initialize(agent, ctx)
+
+    async def run_pre_consume(
+        self,
+        agent: Agent,
+        ctx: Context,
+        inputs: list[Artifact],
+        utilities: list[AgentComponent],
+    ) -> list[Artifact]:
+        """Execute on_pre_consume hooks, allowing components to transform inputs.
+
+        Args:
+            agent: Agent instance
+            ctx: Execution context
+            inputs: Input artifacts to be consumed
+            utilities: List of utility components
+
+        Returns:
+            Transformed input artifacts after all components process them
+
+        Raises:
+            Exception: If any component pre_consume hook fails
+        """
+        current = inputs
+        for component in utilities:
+            comp_name = self._component_display_name(component)
+            priority = getattr(component, "priority", 0)
+            logger.debug(
+                f"Agent pre_consume: agent={self._agent_name}, component={comp_name}, "
+                f"priority={priority}, input_count={len(current)}"
+            )
+            try:
+                current = await component.on_pre_consume(agent, ctx, current)
+            except Exception as exc:
+                logger.exception(
+                    f"Agent pre_consume failed: agent={self._agent_name}, component={comp_name}, "
+                    f"priority={priority}, error={exc!s}"
+                )
+                raise
+        return current
+
+    async def run_pre_evaluate(
+        self,
+        agent: Agent,
+        ctx: Context,
+        inputs: EvalInputs,
+        utilities: list[AgentComponent],
+    ) -> EvalInputs:
+        """Execute on_pre_evaluate hooks, allowing components to transform evaluation inputs.
+
+        Args:
+            agent: Agent instance
+            ctx: Execution context
+            inputs: Evaluation inputs with artifacts and state
+            utilities: List of utility components
+
+        Returns:
+            Transformed evaluation inputs after all components process them
+
+        Raises:
+            Exception: If any component pre_evaluate hook fails
+        """
+        current = inputs
+        for component in utilities:
+            comp_name = self._component_display_name(component)
+            priority = getattr(component, "priority", 0)
+            logger.debug(
+                f"Agent pre_evaluate: agent={self._agent_name}, component={comp_name}, "
+                f"priority={priority}, artifact_count={len(current.artifacts)}"
+            )
+            try:
+                current = await component.on_pre_evaluate(agent, ctx, current)
+            except Exception as exc:
+                logger.exception(
+                    f"Agent pre_evaluate failed: agent={self._agent_name}, component={comp_name}, "
+                    f"priority={priority}, error={exc!s}"
+                )
+                raise
+        return current
+
+    async def run_post_evaluate(
+        self,
+        agent: Agent,
+        ctx: Context,
+        inputs: EvalInputs,
+        result: EvalResult,
+        utilities: list[AgentComponent],
+    ) -> EvalResult:
+        """Execute on_post_evaluate hooks, allowing components to transform results.
+
+        Args:
+            agent: Agent instance
+            ctx: Execution context
+            inputs: Original evaluation inputs
+            result: Evaluation result to be transformed
+            utilities: List of utility components
+
+        Returns:
+            Transformed evaluation result after all components process it
+
+        Raises:
+            Exception: If any component post_evaluate hook fails
+        """
+        current = result
+        for component in utilities:
+            comp_name = self._component_display_name(component)
+            priority = getattr(component, "priority", 0)
+            logger.debug(
+                f"Agent post_evaluate: agent={self._agent_name}, component={comp_name}, "
+                f"priority={priority}, artifact_count={len(current.artifacts)}"
+            )
+            try:
+                current = await component.on_post_evaluate(agent, ctx, inputs, current)
+            except Exception as exc:
+                logger.exception(
+                    f"Agent post_evaluate failed: agent={self._agent_name}, component={comp_name}, "
+                    f"priority={priority}, error={exc!s}"
+                )
+                raise
+        return current
+
+    async def run_post_publish(
+        self,
+        agent: Agent,
+        ctx: Context,
+        artifacts: Sequence[Artifact],
+        utilities: list[AgentComponent],
+    ) -> None:
+        """Execute on_post_publish hooks for each published artifact.
+
+        Args:
+            agent: Agent instance
+            ctx: Execution context
+            artifacts: Published artifacts
+            utilities: List of utility components
+
+        Raises:
+            Exception: If any component post_publish hook fails
+        """
+        for artifact in artifacts:
+            for component in utilities:
+                comp_name = self._component_display_name(component)
+                priority = getattr(component, "priority", 0)
+                logger.debug(
+                    f"Agent post_publish: agent={self._agent_name}, component={comp_name}, "
+                    f"priority={priority}, artifact_id={artifact.id}"
+                )
+                try:
+                    await component.on_post_publish(agent, ctx, artifact)
+                except Exception as exc:
+                    logger.exception(
+                        f"Agent post_publish failed: agent={self._agent_name}, component={comp_name}, "
+                        f"priority={priority}, artifact_id={artifact.id}, error={exc!s}"
+                    )
+                    raise
+
+    async def run_error(
+        self,
+        agent: Agent,
+        ctx: Context,
+        error: Exception,
+        utilities: list[AgentComponent],
+        engines: list[EngineComponent],
+    ) -> None:
+        """Execute on_error hooks for all components.
+
+        Args:
+            agent: Agent instance
+            ctx: Execution context
+            error: Exception that occurred
+            utilities: List of utility components
+            engines: List of engine components
+
+        Raises:
+            Exception: If any component error hook fails
+        """
+        for component in utilities:
+            comp_name = self._component_display_name(component)
+            priority = getattr(component, "priority", 0)
+
+            # Python 3.12+ TaskGroup raises BaseExceptionGroup - extract sub-exceptions
+            error_detail = str(error)
+            if isinstance(error, BaseExceptionGroup):
+                sub_exceptions = [f"{type(e).__name__}: {e}" for e in error.exceptions]
+                error_detail = f"{error!s} - Sub-exceptions: {sub_exceptions}"
+
+            logger.debug(
+                f"Agent error hook: agent={self._agent_name}, component={comp_name}, "
+                f"priority={priority}, error={error_detail}"
+            )
+            try:
+                await component.on_error(agent, ctx, error)
+            except Exception as exc:
+                logger.exception(
+                    f"Agent error hook failed: agent={self._agent_name}, component={comp_name}, "
+                    f"priority={priority}, original_error={error!s}, hook_error={exc!s}"
+                )
+                raise
+
+        for engine in engines:
+            await engine.on_error(agent, ctx, error)
+
+    async def run_terminate(
+        self,
+        agent: Agent,
+        ctx: Context,
+        utilities: list[AgentComponent],
+        engines: list[EngineComponent],
+    ) -> None:
+        """Execute on_terminate hooks for all components.
+
+        Args:
+            agent: Agent instance
+            ctx: Execution context
+            utilities: List of utility components
+            engines: List of engine components
+
+        Raises:
+            Exception: If any component termination fails
+        """
+        for component in utilities:
+            comp_name = self._component_display_name(component)
+            priority = getattr(component, "priority", 0)
+            logger.debug(
+                f"Agent terminate: agent={self._agent_name}, component={comp_name}, priority={priority}"
+            )
+            try:
+                await component.on_terminate(agent, ctx)
+            except Exception as exc:
+                logger.exception(
+                    f"Agent terminate failed: agent={self._agent_name}, component={comp_name}, "
+                    f"priority={priority}, error={exc!s}"
+                )
+                raise
+
+        for engine in engines:
+            await engine.on_terminate(agent, ctx)
+
+
+__all__ = ["ComponentLifecycle"]

flock/agent/context_resolver.py

@@ -0,0 +1,141 @@
+"""Agent context resolution - future context provider integration.
+
+Phase 4: Extracted from agent.py to organize context provider logic.
+
+NOTE: This module provides the foundation for context provider resolution.
+Currently minimal as context resolution happens at the orchestrator level.
+This extraction prepares for future refactoring where agents will have
+more control over their execution context.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any
+
+from flock.logging.logging import get_logger
+
+
+if TYPE_CHECKING:
+    from flock.core import Agent
+    from flock.core.artifacts import Artifact
+    from flock.core.subscription import Subscription
+    from flock.utils.runtime import Context
+
+logger = get_logger(__name__)
+
+
+class ContextResolver:
+    """Manages context provider resolution for agent execution.
+
+    This module handles determining which context provider to use
+    and potentially fetching context artifacts in future phases.
+
+    Currently minimal as context resolution is orchestrator-level.
+    This extraction prepares for Phase 6 refactoring where context
+    providers will be more agent-centric.
+    """
+
+    def __init__(self, agent_name: str):
+        """Initialize ContextResolver for a specific agent.
+
+        Args:
+            agent_name: Name of the agent (for logging)
+        """
+        self._agent_name = agent_name
+        self._logger = logging.getLogger(__name__)
+
+    def get_provider(
+        self, agent: Agent, default_provider: Any | None = None
+    ) -> Any | None:
+        """Determine which context provider to use for this agent.
+
+        Resolution order:
+        1. Agent-specific provider (agent.context_provider)
+        2. Orchestrator default provider (default_provider)
+        3. None (no provider configured)
+
+        Args:
+            agent: Agent instance
+            default_provider: Default provider from orchestrator
+
+        Returns:
+            Context provider to use, or None if not configured
+        """
+        # Check agent-specific provider first (Phase 3 security fix)
+        if agent.context_provider is not None:
+            logger.debug(
+                f"Agent context resolution: agent={self._agent_name}, "
+                f"using_agent_provider=True"
+            )
+            return agent.context_provider
+
+        # Fall back to orchestrator default
+        if default_provider is not None:
+            logger.debug(
+                f"Agent context resolution: agent={self._agent_name}, "
+                f"using_default_provider=True"
+            )
+            return default_provider
+
+        # No provider configured
+        logger.debug(
+            f"Agent context resolution: agent={self._agent_name}, "
+            f"no_provider_configured=True"
+        )
+        return None
+
+    async def resolve_context(
+        self,
+        agent: Agent,
+        subscription: Subscription,
+        trigger_artifacts: list[Artifact],
+        default_provider: Any | None = None,
+    ) -> Context:
+        """Resolve execution context for agent (future implementation).
+
+        NOTE: Currently returns a basic Context. Future phases will:
+        - Use provider to fetch additional context artifacts
+        - Build AgentContext with trigger + context artifacts
+        - Apply visibility filtering at security boundary
+
+        Args:
+            agent: Agent being executed
+            subscription: Subscription that triggered execution
+            trigger_artifacts: Artifacts that triggered agent
+            default_provider: Default context provider from orchestrator
+
+        Returns:
+            Resolved execution context (currently basic)
+        """
+        # Determine which provider to use
+        provider = self.get_provider(agent, default_provider)
+
+        # Future: Use provider to fetch context artifacts
+        # For now, we just log the resolution
+        if provider is None:
+            self._logger.debug(
+                f"Agent context: agent={self._agent_name}, "
+                f"no_context_provider, using_trigger_artifacts_only=True"
+            )
+
+        # Future implementation will:
+        # 1. Build context request from subscription
+        # 2. Call provider.get_artifacts(request)
+        # 3. Return AgentContext with both trigger + context artifacts
+
+        # For now, return a minimal context structure
+        # (actual Context building is orchestrator-level in current architecture)
+        from flock.utils.runtime import Context
+
+        return Context(
+            correlation_id=None,  # Will be filled by orchestrator
+            task_id="",  # Will be filled by orchestrator
+            state={},
+            is_batch=False,
+            artifacts=[],
+            agent_identity=None,
+        )
+
+
+__all__ = ["ContextResolver"]

flock/agent/mcp_integration.py

@@ -0,0 +1,212 @@
+"""Agent MCP integration - server configuration and tool loading.
+
+Phase 4: Extracted from agent.py to eliminate C-rated complexity in with_mcps() and _get_mcp_tools().
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable, Iterable
+from typing import TYPE_CHECKING, Any
+
+from flock.logging.logging import get_logger
+
+
+if TYPE_CHECKING:
+    from flock.agent import MCPServerConfig
+    from flock.core import Flock
+    from flock.utils.runtime import Context
+
+
+logger = get_logger(__name__)
+
+
+class MCPIntegration:
+    """Handles MCP server configuration and tool loading for an agent.
+
+    This module encapsulates all MCP-related logic including:
+    - Server configuration parsing (dict, list, mixed formats)
+    - Tool loading and whitelisting
+    - Graceful degradation on failures
+    """
+
+    def __init__(self, agent_name: str, orchestrator: Flock):
+        """Initialize MCPIntegration for a specific agent.
+
+        Args:
+            agent_name: Name of the agent (for error messages and logging)
+            orchestrator: Flock orchestrator instance (for MCP manager access)
+        """
+        self._agent_name = agent_name
+        self._orchestrator = orchestrator
+        self._logger = logging.getLogger(__name__)
+
+        # Agent MCP state
+        self.mcp_server_names: set[str] = set()
+        self.mcp_server_mounts: dict[str, list[str]] = {}
+        self.tool_whitelist: list[str] | None = None
+
+    async def get_mcp_tools(self, ctx: Context) -> list[Callable]:
+        """Lazy-load MCP tools from assigned servers.
+
+        Architecture Decision: AD001 - Two-Level Architecture
+        Agents fetch tools from servers registered at orchestrator level.
+
+        Architecture Decision: AD003 - Tool Namespacing
+        All tools are namespaced as {server}__{tool}.
+
+        Architecture Decision: AD007 - Graceful Degradation
+        If MCP loading fails, returns empty list so agent continues with native tools.
+
+        Args:
+            ctx: Current execution context with agent_id and run_id
+
+        Returns:
+            List of DSPy-compatible tool callables
+        """
+        if not self.mcp_server_names:
+            # No MCP servers assigned to this agent
+            return []
+
+        try:
+            # Get the MCP manager from orchestrator
+            manager = self._orchestrator.get_mcp_manager()
+
+            # Fetch tools from all assigned servers
+            tools_dict = await manager.get_tools_for_agent(
+                agent_id=self._agent_name,
+                run_id=ctx.task_id,
+                server_names=self.mcp_server_names,
+                server_mounts=self.mcp_server_mounts,  # Pass server-specific mounts
+            )
+
+            # Whitelisting logic
+            tool_whitelist = self.tool_whitelist
+            if (
+                tool_whitelist is not None
+                and isinstance(tool_whitelist, list)
+                and len(tool_whitelist) > 0
+            ):
+                filtered_tools: dict[str, Any] = {}
+                for tool_key, tool_entry in tools_dict.items():
+                    if isinstance(tool_entry, dict):
+                        original_name = tool_entry.get("original_name", None)
+                        if (
+                            original_name is not None
+                            and original_name in tool_whitelist
+                        ):
+                            filtered_tools[tool_key] = tool_entry
+
+                tools_dict = filtered_tools
+
+            # Convert to DSPy tool callables
+            dspy_tools = []
+            for namespaced_name, tool_info in tools_dict.items():
+                tool_info["server_name"]
+                flock_tool = tool_info["tool"]  # Already a FlockMCPTool
+                client = tool_info["client"]
+
+                # Convert to DSPy tool
+                dspy_tool = flock_tool.as_dspy_tool(server=client)
+
+                # Update name to include namespace
+                dspy_tool.name = namespaced_name
+
+                dspy_tools.append(dspy_tool)
+
+            return dspy_tools
+
+        except Exception as e:
+            # Architecture Decision: AD007 - Graceful Degradation
+            # Agent continues with native tools only
+            logger.error(
+                f"Failed to load MCP tools for agent {self._agent_name}: {e}",
+                exc_info=True,
+            )
+            return []
+
+    def configure_servers(
+        self,
+        servers: (Iterable[str] | dict[str, MCPServerConfig]),
+        registered_servers: set[str],
+    ) -> None:
+        """Configure MCP servers for this agent with optional server-specific mount points.
+
+        Architecture Decision: AD001 - Two-Level Architecture
+        Agents reference servers registered at orchestrator level.
+
+        Args:
+            servers: One of:
+                - List of server names (strings) - no specific mounts
+                - Dict mapping server names to MCPServerConfig
+            registered_servers: Set of server names registered with orchestrator (for validation)
+
+        Raises:
+            ValueError: If any server name is not registered with orchestrator
+            TypeError: If server specification format is invalid
+
+        Examples:
+            >>> # Simple: no mount restrictions
+            >>> integration.configure_servers(["filesystem", "github"], registered)
+
+            >>> # Server-specific config with roots and tool whitelist
+            >>> integration.configure_servers(
+            ...     {
+            ...         "filesystem": {
+            ...             "roots": ["/workspace/dir/data"],
+            ...             "tool_whitelist": ["read_file"],
+            ...         },
+            ...         "github": {},  # No restrictions for github
+            ...     },
+            ...     registered,
+            ... )
+        """
+        # Parse input into server_names and mounts
+        server_set: set[str] = set()
+        server_mounts: dict[str, list[str]] = {}
+        whitelist = None
+
+        if isinstance(servers, dict):
+            # Dict format: {"server": {"roots": ["/path1"], "tool_whitelist": ["tool1"]}}
+            for server_name, server_config in servers.items():
+                server_set.add(server_name)
+
+                if isinstance(server_config, dict):
+                    # MCPServerConfig dict with optional roots and tool_whitelist
+                    mounts = server_config.get("roots", None)
+                    if (
+                        mounts is not None
+                        and isinstance(mounts, list)
+                        and len(mounts) > 0
+                    ):
+                        server_mounts[server_name] = list(mounts)
+
+                    config_whitelist = server_config.get("tool_whitelist", None)
+                    if (
+                        config_whitelist is not None
+                        and isinstance(config_whitelist, list)
+                        and len(config_whitelist) > 0
+                    ):
+                        whitelist = config_whitelist
+        else:
+            # Assume it's an iterable of strings
+            server_set = set(servers)
+
+        # Validate all servers exist in orchestrator
+        invalid_servers = server_set - registered_servers
+
+        if invalid_servers:
+            available = list(registered_servers) if registered_servers else ["none"]
+            raise ValueError(
+                f"MCP servers not registered: {invalid_servers}. "
+                f"Available servers: {available}. "
+                f"Register servers using orchestrator.add_mcp() first."
+            )
+
+        # Store in integration
+        self.mcp_server_names = server_set
+        self.mcp_server_mounts = server_mounts
+        self.tool_whitelist = whitelist
+
+
+__all__ = ["MCPIntegration"]