flock-core 0.5.0b2__py3-none-any.whl → 0.5.0b5__py3-none-any.whl

flock/core/agent/flock_agent_integration.py

@@ -1,17 +1,80 @@
 # src/flock/core/agent/flock_agent_integration.py
 """Tool and server integration functionality for FlockAgent."""
 
-from typing import TYPE_CHECKING, Any
+from collections.abc import Callable
+from functools import wraps
+from inspect import Parameter, signature
+from typing import TYPE_CHECKING, Any, TypeVar, cast
 
 from flock.core.context.context import FlockContext
 from flock.core.logging.logging import get_logger
-from flock.core.mcp.flock_mcp_server import FlockMCPServerBase
+from flock.core.mcp.flock_mcp_server import FlockMCPServer
 
 if TYPE_CHECKING:
     from flock.core.flock_agent import FlockAgent
 
 logger = get_logger("agent.integration")
 
+R = TypeVar("R", bound=str)
+
+
+def adapt(prop_name: str, fn: Callable[..., R]) -> Callable[[FlockContext], R]:
+    """Coerce *fn* into the canonical ``(ctx: FlockContext) -> str`` form.
+
+    Acceptable signatures
+    ---------------------
+    1. ``() -> str``                           (no parameters)
+    2. ``(ctx: FlockContext) -> str``          (exactly one positional parameter)
+
+    Anything else raises ``TypeError``.
+
+    The wrapper also enforces at runtime that the result is ``str``.
+    """
+    if not callable(fn):
+        raise TypeError(f"{prop_name} must be a callable, got {type(fn).__name__}")
+
+    sig = signature(fn)
+    params = list(sig.parameters.values())
+
+    def _validate_result(res: object) -> R:
+        if not isinstance(res, str):
+            raise TypeError(
+                f"{prop_name} callable must return str, got {type(res).__name__}"
+            )
+        return cast(R, res)
+
+    # ── Case 1: () -> str ────────────────────────────────────────────────────
+    if len(params) == 0:
+
+        @wraps(fn)
+        def _wrapped(ctx: FlockContext) -> R:
+            return _validate_result(fn())
+
+        return _wrapped
+
+    # ── Case 2: (ctx) -> str ────────────────────────────────────────────────
+    if len(params) == 1:
+        p: Parameter = params[0]
+        valid_kind = p.kind in (
+            Parameter.POSITIONAL_ONLY,
+            Parameter.POSITIONAL_OR_KEYWORD,
+        )
+        valid_annotation = p.annotation in (Parameter.empty, FlockContext)
+        has_no_default = p.default is Parameter.empty
+
+        if valid_kind and valid_annotation and has_no_default:
+
+            @wraps(fn)
+            def _wrapped(ctx: FlockContext) -> R:
+                return _validate_result(fn(ctx))  # type: ignore[arg-type]
+
+            return _wrapped
+
+    # ── Anything else: reject ───────────────────────────────────────────────
+    raise TypeError(
+        f"{prop_name} callable must be () -> str or (ctx: FlockContext) -> str; "
+        f"got signature {sig}"
+    )
 
 class FlockAgentIntegration:
     """Handles tool and server integration for FlockAgent including MCP servers and callable tools."""
@@ -19,30 +82,26 @@ class FlockAgentIntegration:
     def __init__(self, agent: "FlockAgent"):
         self.agent = agent
 
-    def resolve_callables(self, context: FlockContext | None = None) -> None:
-        """Resolves callable fields (description, input, output) using context."""
-        if callable(self.agent.description):
-            self.agent.description = self.agent.description(
-                context
-            )  # Pass context if needed by callable
-        if callable(self.agent.input):
-            self.agent.input = self.agent.input(context)
-        if callable(self.agent.output):
-            self.agent.output = self.agent.output(context)
-
-    def resolve_description(self, context: FlockContext | None = None) -> str:
-        if callable(self.agent.description):
-            try:
-                # Attempt to call without context first.
-                return self.agent.description(context)
-            except Exception as e:
-                logger.error(
-                    f"Error resolving callable description for agent '{self.agent.name}': {e}"
-                )
-                return None
-        elif isinstance(self.agent.description, str):
-            return self.agent.description
-        return None
+    def _resolve(self, raw: str | Callable[..., str], name: str, ctx: FlockContext | None) -> str | None:
+        if callable(raw):
+            raw = adapt(name, raw)(ctx or FlockContext())
+        return raw
+
+    def resolve_description(self, context: FlockContext | None = None) -> str | None:
+        """Resolve the agent's description, handling callable descriptions."""
+        return self._resolve(self.agent.description_spec, "description", context)
+
+    def resolve_input(self, context: FlockContext | None = None) -> str | None:
+        """Resolve the agent's input, handling callable inputs."""
+        return self._resolve(self.agent.input_spec, "input", context)
+
+    def resolve_output(self, context: FlockContext | None = None) -> str | None:
+        """Resolve the agent's output, handling callable outputs."""
+        return self._resolve(self.agent.output_spec, "output", context)
+
+    def resolve_next_agent(self, context: FlockContext | None = None) -> str | None:
+        """Resolve the next agent, handling callable next agents."""
+        return self._resolve(self.agent.next_agent_spec, "next_agent", context)
 
     async def get_mcp_tools(self) -> list[Any]:
         """Get tools from registered MCP servers."""
@@ -52,9 +111,9 @@ class FlockAgentIntegration:
 
             registry = get_registry()  # Get the registry
             for server in self.agent.servers:
-                registered_server: FlockMCPServerBase | None = None
+                registered_server: FlockMCPServer | None = None
                 server_tools = []
-                if isinstance(server, FlockMCPServerBase):
+                if isinstance(server, FlockMCPServer):
                     # check if registered
                     server_name = server.config.name
                     registered_server = registry.get_server(
@@ -72,7 +131,7 @@ class FlockAgentIntegration:
                     )
                 else:
                     logger.warning(
-                        f"No Server with name '{server.config.name if isinstance(server, FlockMCPServerBase) else server}' registered! Skipping."
+                        f"No Server with name '{server.config.name if isinstance(server, FlockMCPServer) else server}' registered! Skipping."
                     )
                 mcp_tools = mcp_tools + server_tools
         return mcp_tools

flock/core/agent/flock_agent_serialization.py

@@ -8,7 +8,7 @@ from typing import TYPE_CHECKING, Any, TypeVar
 
 # Legacy component imports removed
 from flock.core.logging.logging import get_logger
-from flock.core.mcp.flock_mcp_server import FlockMCPServerBase
+from flock.core.mcp.flock_mcp_server import FlockMCPServer
 from flock.core.serialization.json_encoder import FlockJSONEncoder
 from flock.core.serialization.serialization_utils import (
     deserialize_component,
@@ -68,20 +68,20 @@ class FlockAgentSerialization:
         is_output_callable = False
         is_next_agent_callable = False
         # if self.agent.description is a callable, exclude it
-        if callable(self.agent.description):
+        if callable(self.agent.description_spec):
             is_description_callable = True
-            exclude.append("description")
+            exclude.append("description_spec")
         # if self.agent.input is a callable, exclude it
-        if callable(self.agent.input):
+        if callable(self.agent.input_spec):
             is_input_callable = True
-            exclude.append("input")
+            exclude.append("input_spec")
         # if self.agent.output is a callable, exclude it
-        if callable(self.agent.output):
+        if callable(self.agent.output_spec):
             is_output_callable = True
-            exclude.append("output")
-        if callable(self.agent.next_agent):
+            exclude.append("output_spec")
+        if callable(self.agent.next_agent_spec):
             is_next_agent_callable = True
-            exclude.append("next_agent")
+            exclude.append("next_agent_spec")
 
         logger.debug(f"Serializing agent '{self.agent.name}' to dict.")
         # Use Pydantic's dump, exclude manually handled fields and runtime context
@@ -122,7 +122,7 @@ class FlockAgentSerialization:
             )
             serialized_servers = []
             for server in self.agent.servers:
-                if isinstance(server, FlockMCPServerBase):
+                if isinstance(server, FlockMCPServer):
                     serialized_servers.append(server.config.name)
                 else:
                     # Write it down as a list of server names.
@@ -167,7 +167,7 @@ class FlockAgentSerialization:
                 )
 
         if is_description_callable:
-            path_str = registry.get_callable_path_string(self.agent.description)
+            path_str = registry.get_callable_path_string(self.agent.description_spec)
             if path_str:
                 func_name = path_str.split(".")[-1]
                 data["description_callable"] = func_name
@@ -176,11 +176,11 @@ class FlockAgentSerialization:
                 )
             else:
                 logger.warning(
-                    f"Could not get path string for description {self.agent.description} in agent '{self.agent.name}'. Skipping."
+                    f"Could not get path string for description {self.agent.description_spec} in agent '{self.agent.name}'. Skipping."
                 )
 
         if is_input_callable:
-            path_str = registry.get_callable_path_string(self.agent.input)
+            path_str = registry.get_callable_path_string(self.agent.input_spec)
             if path_str:
                 func_name = path_str.split(".")[-1]
                 data["input_callable"] = func_name
@@ -189,11 +189,11 @@ class FlockAgentSerialization:
                 )
             else:
                 logger.warning(
-                    f"Could not get path string for input {self.agent.input} in agent '{self.agent.name}'. Skipping."
+                    f"Could not get path string for input {self.agent.input_spec} in agent '{self.agent.name}'. Skipping."
                 )
 
         if is_output_callable:
-            path_str = registry.get_callable_path_string(self.agent.output)
+            path_str = registry.get_callable_path_string(self.agent.output_spec)
             if path_str:
                 func_name = path_str.split(".")[-1]
                 data["output_callable"] = func_name
@@ -202,20 +202,20 @@ class FlockAgentSerialization:
                 )
             else:
                 logger.warning(
-                    f"Could not get path string for output {self.agent.output} in agent '{self.agent.name}'. Skipping."
+                    f"Could not get path string for output {self.agent.output_spec} in agent '{self.agent.name}'. Skipping."
                 )
 
         if is_next_agent_callable:
-            path_str = registry.get_callable_path_string(self.agent.next_agent)
+            path_str = registry.get_callable_path_string(self.agent.next_agent_spec)
             if path_str:
                 func_name = path_str.split(".")[-1]
-                data["next_agent"] = func_name
+                data["next_agent_callable"] = func_name
                 logger.debug(
                     f"Added next_agent '{func_name}' (from path '{path_str}') to agent '{self.agent.name}'"
                 )
             else:
                 logger.warning(
-                    f"Could not get path string for next_agent {self.agent.next_agent} in agent '{self.agent.name}'. Skipping."
+                    f"Could not get path string for next_agent {self.agent.next_agent_spec} in agent '{self.agent.name}'. Skipping."
                 )
 
 
@@ -245,6 +245,7 @@ class FlockAgentSerialization:
             "description_callable",
             "input_callable",
             "output_callable",
+            "next_agent_callable",
         ]
 
         for key in callable_keys:
@@ -322,7 +323,7 @@ class FlockAgentSerialization:
                 if isinstance(server_name, str):
                     # Case 1 (default behavior): A server name is passed.
                     agent.servers.append(server_name)
-                elif isinstance(server_name, FlockMCPServerBase):
+                elif isinstance(server_name, FlockMCPServer):
                     # Case 2 (highly unlikely): If someone somehow manages to pass
                     # an instance of a server during the deserialization step (however that might be achieved)
                     # check the registry, if the server is already registered, if not, register it
@@ -373,6 +374,8 @@ class FlockAgentSerialization:
         resolve_and_assign("description", "description_callable")
         resolve_and_assign("input", "input_callable")
         resolve_and_assign("output", "output_callable")
+        resolve_and_assign("next_agent", "next_agent_callable")
+        # --- Finalize ---
 
         logger.info(f"Successfully deserialized agent '{agent.name}'.")
         return agent

flock/core/api/endpoints.py

@@ -89,7 +89,7 @@ def create_api_router() -> APIRouter:
                     typed_inputs = inputs_task # Simplified for now
 
                     result = await flock_instance.run_async(
-                        start_agent=agent_name_task, input=typed_inputs
+                        agent=agent_name_task, input=typed_inputs
                     )
                     run_store.update_run_result(run_id_task, result)
                 except Exception as e_task:
@@ -246,7 +246,7 @@ def create_api_router() -> APIRouter:
         """List all available agents in the currently loaded Flock."""
         logger.debug("API request: list agents")
         agents_list = [
-            {"name": agent.name, "description": agent.resolved_description or agent.name}
+            agent.to_dict()
             for agent in flock_instance.agents.values()
         ]
         return {"agents": agents_list}

flock/core/api/service.py

@@ -33,7 +33,7 @@ class FlockApiService:
             )
             # Flock.run_async now handles context creation and execution
             result = await self.flock.run_async(
-                start_agent=agent_name, input=typed_inputs
+                agent=agent_name, input=typed_inputs
             )
             self.run_store.update_run_result(run_id, result)
 
@@ -100,7 +100,7 @@ class FlockApiService:
                         try:
                             # Call Flock's run_async for a single item
                             item_result = await self.flock.run_async(
-                                start_agent=request.agent_name,
+                                agent=request.agent_name,
                                 input=item_inputs,
                                 box_result=request.box_results,
                             )

flock/core/component/__init__.py

@@ -2,14 +2,14 @@
 """Unified component system for Flock agents."""
 
 from .agent_component_base import AgentComponent, AgentComponentConfig
-from .evaluation_component_base import EvaluationComponentBase
-from .routing_component_base import RoutingComponentBase
-from .utility_component_base import UtilityComponentBase
+from .evaluation_component import EvaluationComponent
+from .routing_component import RoutingComponent
+from .utility_component import UtilityComponent
 
 __all__ = [
     "AgentComponent",
-    "AgentComponentConfig", 
-    "EvaluationComponentBase",
-    "RoutingComponentBase",
-    "UtilityComponentBase",
+    "AgentComponentConfig",
+    "EvaluationComponent",
+    "RoutingComponent",
+    "UtilityComponent",
 ]

flock/core/component/{evaluation_component_base.py → evaluation_component.py}

@@ -9,7 +9,7 @@ from flock.core.context.context import FlockContext
 from .agent_component_base import AgentComponent
 
 
-class EvaluationComponentBase(AgentComponent):
+class EvaluationComponent(AgentComponent):
     """Base class for evaluation components.
     
     Evaluation components implement the core intelligence/logic of an agent.
@@ -23,7 +23,7 @@ class EvaluationComponentBase(AgentComponent):
     - ScriptEvaluationComponent (Python script-based)
     - LLMEvaluationComponent (direct LLM API)
     """
-    
+
     @abstractmethod
     async def evaluate_core(
         self,

flock/core/component/{routing_component_base.py → routing_component.py}

@@ -6,11 +6,10 @@ from typing import Any
 
 from flock.core.context.context import FlockContext
 
-
 from .agent_component_base import AgentComponent
 
 
-class RoutingComponentBase(AgentComponent):
+class RoutingComponent(AgentComponent):
     """Base class for routing components.
     
     Routing components determine the next step in a workflow based on the
@@ -26,7 +25,7 @@ class RoutingComponentBase(AgentComponent):
     - DefaultRoutingModule (simple next-agent routing)
     - ListGeneratorRoutingModule (dynamic agent creation)
     """
-    
+
     @abstractmethod
     async def determine_next_step(
         self,
@@ -54,7 +53,7 @@ class RoutingComponentBase(AgentComponent):
         raise NotImplementedError(
             f"{self.__class__.__name__} must implement determine_next_step()"
         )
-        
+
     async def evaluate_core(
         self,
         agent: Any,

flock/core/component/{utility_component_base.py → utility_component.py}

@@ -4,12 +4,12 @@
 from typing import Any
 
 from flock.core.context.context import FlockContext
-# HandOffRequest removed - using agent.next_agent directly
 
+# HandOffRequest removed - using agent.next_agent directly
 from .agent_component_base import AgentComponent
 
 
-class UtilityComponentBase(AgentComponent):
+class UtilityComponent(AgentComponent):
     """Base class for utility/enhancement components.
     
     Utility components add cross-cutting concerns to agents without being
@@ -32,7 +32,7 @@ class UtilityComponentBase(AgentComponent):
     - MetricsUtilityModule (performance tracking)
     - AssertionUtilityModule (result validation)
     """
-    
+
     async def evaluate_core(
         self,
         agent: Any,

flock/core/evaluation/utils.py

@@ -17,6 +17,7 @@ from opik.evaluation import evaluate
 
 from flock.core.flock import Flock
 from flock.core.flock_agent import FlockAgent
+
 # Legacy FlockEvaluator import removed
 from flock.core.logging.logging import get_logger
 
@@ -58,7 +59,7 @@ def evaluate_with_opik(
 
         # Use the shared Flock instance instead of creating a new one
         result_flock = shared_flock.run(
-            start_agent=start_agent, input=agent_input, box_result=False
+            agent=start_agent, input=agent_input, box_result=False
         )
 
         # agent_output = result_flock.get(answer_mapping[key], "No answer found")

flock/core/execution/batch_executor.py

@@ -228,7 +228,7 @@ class BatchProcessor:
             try:
                 # Use the synchronous wrapper to avoid nested event-loop issues inside threads
                 result = self.flock.run(
-                    start_agent=start_agent,
+                    agent=start_agent,
                     input=full_input,
                     box_result=box_results,
                 )

flock/core/execution/evaluation_executor.py

@@ -188,7 +188,7 @@ class EvaluationExecutor:
                 try:
                     # Run the agent/flock for this item
                     agent_output = await self.flock.run_async(
-                        start_agent=start_agent,  # Name or instance
+                        agent=start_agent,  # Name or instance
                         input=agent_inputs_with_static,
                         box_result=True,  # Use Box for easier access via dot notation
                         # context=... # Assuming isolated context for now

flock/core/execution/opik_executor.py

@@ -64,7 +64,7 @@ class OpikExecutor:
 
         # Evaluation task
         def evaluation_task(dataset_item):
-          flock_result = self.flock.run(start_agent=start_agent, input=dataset_item, box_result=False)
+          flock_result = self.flock.run(agent=start_agent, input=dataset_item, box_result=False)
 
           result = {
               "input": dataset_item.get("test"),

flock/core/flock.py

@@ -17,7 +17,7 @@ from typing import (
 from box import Box
 from temporalio import workflow
 
-from flock.core.mcp.flock_mcp_server import FlockMCPServerBase
+from flock.core.mcp.flock_mcp_server import FlockMCPServer
 
 with workflow.unsafe.imports_passed_through():
     from datasets import Dataset  # type: ignore
@@ -125,7 +125,7 @@ class Flock(BaseModel, Serializable):
     _start_input: dict = {}  # For potential pre-configuration
 
     # Internal server storage - not part of the Pydantic model for direct serialization
-    _servers: dict[str, FlockMCPServerBase]
+    _servers: dict[str, FlockMCPServer]
 
     # Note: _mgr is now handled by the server manager helper
 
@@ -207,7 +207,7 @@ class Flock(BaseModel, Serializable):
         enable_temporal: bool = False,
         enable_opik: bool = False,
         agents: list[FlockAgent] | None = None,
-        servers: list[FlockMCPServerBase] | None = None,
+        servers: list[FlockMCPServer] | None = None,
         temporal_config: TemporalWorkflowConfig | None = None,
         temporal_start_in_process_worker: bool = True,
         **kwargs,
@@ -279,7 +279,7 @@ class Flock(BaseModel, Serializable):
             agent_input = {self.benchmark_input_field: msg_content}
 
             result = await self.run_async(
-                start_agent=self.benchmark_agent_name,
+                agent=self.benchmark_agent_name,
                 input=agent_input,
                 box_result=False,
             )
@@ -296,7 +296,7 @@ class Flock(BaseModel, Serializable):
 
 
 
-    def add_server(self, server: FlockMCPServerBase) -> FlockMCPServerBase:
+    def add_server(self, server: FlockMCPServer) -> FlockMCPServer:
         """Adds a server instance to this Flock configuration and registry."""
         return self._server_manager.add_server(server)
 
@@ -349,24 +349,24 @@ class Flock(BaseModel, Serializable):
         return self._agents
 
     @property
-    def servers(self) -> dict[str, FlockMCPServerBase]:
+    def servers(self) -> dict[str, FlockMCPServer]:
         """Returns the dictionary of servers managed by this Flock instance."""
         return self._server_manager.servers
 
     def run(
         self,
-        start_agent: FlockAgent | str | None = None,
+        agent: FlockAgent | str | None = None,
         input: dict | None = None,
         context: FlockContext | None = None,
         run_id: str = "",
         box_result: bool = True,
         agents: list[FlockAgent] | None = None,
-        servers: list[FlockMCPServerBase] | None = None,
+        servers: list[FlockMCPServer] | None = None,
         memo: dict[str, Any] | None = None,
     ) -> Box | dict:
         """Synchronous execution wrapper."""
         return self._execution.run(
-            start_agent=start_agent,
+            agent=agent,
             input=input,
             context=context,
             run_id=run_id,
@@ -378,18 +378,18 @@ class Flock(BaseModel, Serializable):
 
     async def run_async(
         self,
-        start_agent: FlockAgent | str | None = None,
+        agent: FlockAgent | str | None = None,
         input: dict | None = None,
         context: FlockContext | None = None,
         run_id: str = "",
         box_result: bool = True,
         agents: list[FlockAgent] | None = None,
-        servers: list[FlockMCPServerBase] | None = None,
+        servers: list[FlockMCPServer] | None = None,
         memo: dict[str, Any] | None = None,
     ) -> Box | dict:
         """Entry point for running an agent system asynchronously."""
         return await self._execution.run_async(
-            start_agent=start_agent,
+            agent=agent,
             input=input,
             context=context,
             run_id=run_id,