fast-agent-mcp 0.1.13__py3-none-any.whl → 0.2.0__py3-none-any.whl

mcp_agent/core/direct_decorators.py

@@ -0,0 +1,443 @@
+"""
+Type-safe decorators for DirectFastAgent applications.
+These decorators provide type-safe function signatures and IDE support
+for creating agents in the DirectFastAgent framework.
+"""
+
+import inspect
+from functools import wraps
+from typing import (
+    Awaitable,
+    Callable,
+    List,
+    Literal,
+    Optional,
+    ParamSpec,
+    Protocol,
+    TypeVar,
+    Union,
+    cast,
+)
+
+from mcp_agent.agents.agent import AgentConfig
+from mcp_agent.core.agent_types import AgentType
+from mcp_agent.core.request_params import RequestParams
+
+# Type variables for the decorated function
+P = ParamSpec("P")  # Parameters
+R = TypeVar("R")  # Return type
+
+# Type for agent functions - can be either async or sync
+AgentCallable = Callable[P, Union[Awaitable[R], R]]
+
+
+# Protocol for decorated agent functions
+class DecoratedAgentProtocol(Protocol[P, R]):
+    """Protocol defining the interface of a decorated agent function."""
+
+    _agent_type: AgentType
+    _agent_config: AgentConfig
+
+    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> Union[Awaitable[R], R]: ...
+
+
+# Protocol for orchestrator functions
+class DecoratedOrchestratorProtocol(DecoratedAgentProtocol[P, R], Protocol):
+    """Protocol for decorated orchestrator functions with additional metadata."""
+
+    _child_agents: List[str]
+    _plan_type: Literal["full", "iterative"]
+
+
+# Protocol for router functions
+class DecoratedRouterProtocol(DecoratedAgentProtocol[P, R], Protocol):
+    """Protocol for decorated router functions with additional metadata."""
+
+    _router_agents: List[str]
+
+
+# Protocol for chain functions
+class DecoratedChainProtocol(DecoratedAgentProtocol[P, R], Protocol):
+    """Protocol for decorated chain functions with additional metadata."""
+
+    _chain_agents: List[str]
+
+
+# Protocol for parallel functions
+class DecoratedParallelProtocol(DecoratedAgentProtocol[P, R], Protocol):
+    """Protocol for decorated parallel functions with additional metadata."""
+
+    _fan_out: List[str]
+    _fan_in: str
+
+
+# Protocol for evaluator-optimizer functions
+class DecoratedEvaluatorOptimizerProtocol(DecoratedAgentProtocol[P, R], Protocol):
+    """Protocol for decorated evaluator-optimizer functions with additional metadata."""
+
+    _generator: str
+    _evaluator: str
+
+
+def _decorator_impl(
+    self,
+    agent_type: AgentType,
+    name: str,
+    instruction: str,
+    *,
+    servers: List[str] = [],
+    model: Optional[str] = None,
+    use_history: bool = True,
+    request_params: RequestParams | None = None,
+    human_input: bool = False,
+    **extra_kwargs,
+) -> Callable[[AgentCallable[P, R]], DecoratedAgentProtocol[P, R]]:
+    """
+    Core implementation for agent decorators with common behavior and type safety.
+
+    Args:
+        agent_type: Type of agent to create
+        name: Name of the agent
+        instruction: Base instruction for the agent
+        servers: List of server names the agent should connect to
+        model: Model specification string
+        use_history: Whether to maintain conversation history
+        request_params: Additional request parameters for the LLM
+        human_input: Whether to enable human input capabilities
+        **extra_kwargs: Additional agent/workflow-specific parameters
+    """
+
+    def decorator(func: AgentCallable[P, R]) -> DecoratedAgentProtocol[P, R]:
+        is_async = inspect.iscoroutinefunction(func)
+
+        # Handle both async and sync functions consistently
+        if is_async:
+
+            @wraps(func)
+            async def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+                # Call the original function
+                return await func(*args, **kwargs)
+        else:
+
+            @wraps(func)
+            def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+                # Call the original function
+                return func(*args, **kwargs)
+
+        # Create agent configuration
+        config = AgentConfig(
+            name=name,
+            instruction=instruction,
+            servers=servers,
+            model=model,
+            use_history=use_history,
+            human_input=human_input,
+        )
+
+        # Update request params if provided
+        if request_params:
+            config.default_request_params = request_params
+
+        # Store metadata on the wrapper function
+        agent_data = {
+            "config": config,
+            "type": agent_type.value,
+            "func": func,
+        }
+
+        # Add extra parameters specific to this agent type
+        for key, value in extra_kwargs.items():
+            agent_data[key] = value
+
+        # Store the configuration in the FastAgent instance
+        self.agents[name] = agent_data
+
+        # Store type information for IDE support
+        setattr(wrapper, "_agent_type", agent_type)
+        setattr(wrapper, "_agent_config", config)
+        for key, value in extra_kwargs.items():
+            setattr(wrapper, f"_{key}", value)
+
+        return cast("DecoratedAgentProtocol[P, R]", wrapper)
+
+    return decorator
+
+
+def agent(
+    self,
+    name: str = "default",
+    instruction_or_kwarg: Optional[str] = None,
+    *,
+    instruction: str = "You are a helpful agent.",
+    servers: List[str] = [],
+    model: Optional[str] = None,
+    use_history: bool = True,
+    request_params: RequestParams | None = None,
+    human_input: bool = False,
+) -> Callable[[AgentCallable[P, R]], DecoratedAgentProtocol[P, R]]:
+    """
+    Decorator to create and register a standard agent with type-safe signature.
+
+    Args:
+        name: Name of the agent
+        instruction_or_kwarg: Optional positional parameter for instruction
+        instruction: Base instruction for the agent (keyword arg)
+        servers: List of server names the agent should connect to
+        model: Model specification string
+        use_history: Whether to maintain conversation history
+        request_params: Additional request parameters for the LLM
+        human_input: Whether to enable human input capabilities
+
+    Returns:
+        A decorator that registers the agent with proper type annotations
+    """
+    final_instruction = instruction_or_kwarg if instruction_or_kwarg is not None else instruction
+
+    return _decorator_impl(
+        self,
+        AgentType.BASIC,
+        name=name,
+        instruction=final_instruction,
+        servers=servers,
+        model=model,
+        use_history=use_history,
+        request_params=request_params,
+        human_input=human_input,
+    )
+
+
+def orchestrator(
+    self,
+    name: str,
+    *,
+    agents: List[str],
+    instruction: Optional[str] = None,
+    model: Optional[str] = None,
+    use_history: bool = False,
+    request_params: RequestParams | None = None,
+    human_input: bool = False,
+    plan_type: Literal["full", "iterative"] = "full",
+    max_iterations: int = 30,
+) -> Callable[[AgentCallable[P, R]], DecoratedOrchestratorProtocol[P, R]]:
+    """
+    Decorator to create and register an orchestrator agent with type-safe signature.
+
+    Args:
+        name: Name of the orchestrator
+        agents: List of agent names this orchestrator can use
+        instruction: Base instruction for the orchestrator
+        model: Model specification string
+        use_history: Whether to maintain conversation history
+        request_params: Additional request parameters for the LLM
+        human_input: Whether to enable human input capabilities
+        plan_type: Planning approach - "full" or "iterative"
+        max_iterations: Maximum number of planning iterations
+
+    Returns:
+        A decorator that registers the orchestrator with proper type annotations
+    """
+    default_instruction = """
+    You are an expert planner. Given an objective task and a list of Agents 
+    (which are collections of capabilities), your job is to break down the objective 
+    into a series of steps, which can be performed by these agents.
+    """
+
+    # Create final request params with max_iterations
+
+    return cast(
+        "Callable[[AgentCallable[P, R]], DecoratedOrchestratorProtocol[P, R]]",
+        _decorator_impl(
+            self,
+            AgentType.ORCHESTRATOR,
+            name=name,
+            instruction=instruction or default_instruction,
+            servers=[],  # Orchestrators don't connect to servers directly
+            model=model,
+            use_history=use_history,
+            request_params=request_params,
+            human_input=human_input,
+            child_agents=agents,
+            plan_type=plan_type,
+            max_iterations=max_iterations,
+        ),
+    )
+
+
+def router(
+    self,
+    name: str,
+    *,
+    agents: List[str],
+    instruction: Optional[str] = None,
+    model: Optional[str] = None,
+    use_history: bool = False,
+    request_params: RequestParams | None = None,
+    human_input: bool = False,
+) -> Callable[[AgentCallable[P, R]], DecoratedRouterProtocol[P, R]]:
+    """
+    Decorator to create and register a router agent with type-safe signature.
+
+    Args:
+        name: Name of the router
+        agents: List of agent names this router can route to
+        instruction: Base instruction for the router
+        model: Model specification string
+        use_history: Whether to maintain conversation history
+        request_params: Additional request parameters for the LLM
+        human_input: Whether to enable human input capabilities
+
+    Returns:
+        A decorator that registers the router with proper type annotations
+    """
+    default_instruction = """
+    You are a router that determines which specialized agent should handle a given query.
+    Analyze the query and select the most appropriate agent to handle it.
+    """
+
+    return cast(
+        "Callable[[AgentCallable[P, R]], DecoratedRouterProtocol[P, R]]",
+        _decorator_impl(
+            self,
+            AgentType.ROUTER,
+            name=name,
+            instruction=instruction or default_instruction,
+            servers=[],  # Routers don't connect to servers directly
+            model=model,
+            use_history=use_history,
+            request_params=request_params,
+            human_input=human_input,
+            router_agents=agents,
+        ),
+    )
+
+
+def chain(
+    self,
+    name: str,
+    *,
+    sequence: List[str],
+    instruction: Optional[str] = None,
+    cumulative: bool = False,
+) -> Callable[[AgentCallable[P, R]], DecoratedChainProtocol[P, R]]:
+    """
+    Decorator to create and register a chain agent with type-safe signature.
+
+    Args:
+        name: Name of the chain
+        sequence: List of agent names in the chain, executed in sequence
+        instruction: Base instruction for the chain
+        cumulative: Whether to use cumulative mode (each agent sees all previous responses)
+
+    Returns:
+        A decorator that registers the chain with proper type annotations
+    """
+    # Validate sequence is not empty
+    if not sequence:
+        from mcp_agent.core.exceptions import AgentConfigError
+
+        raise AgentConfigError(f"Chain '{name}' requires at least one agent in the sequence")
+
+    default_instruction = """
+    You are a chain that processes requests through a series of specialized agents in sequence.
+    Pass the output of each agent to the next agent in the chain.
+    """
+
+    return cast(
+        "Callable[[AgentCallable[P, R]], DecoratedChainProtocol[P, R]]",
+        _decorator_impl(
+            self,
+            AgentType.CHAIN,
+            name=name,
+            instruction=instruction or default_instruction,
+            sequence=sequence,
+            cumulative=cumulative,
+        ),
+    )
+
+
+def parallel(
+    self,
+    name: str,
+    *,
+    fan_out: List[str],
+    fan_in: str | None = None,
+    instruction: Optional[str] = None,
+    include_request: bool = True,
+) -> Callable[[AgentCallable[P, R]], DecoratedParallelProtocol[P, R]]:
+    """
+    Decorator to create and register a parallel agent with type-safe signature.
+
+    Args:
+        name: Name of the parallel agent
+        fan_out: List of agents to execute in parallel
+        fan_in: Agent to aggregate results
+        instruction: Base instruction for the parallel agent
+        include_request: Whether to include the original request when aggregating
+
+    Returns:
+        A decorator that registers the parallel agent with proper type annotations
+    """
+    default_instruction = """
+    You are a parallel processor that executes multiple agents simultaneously 
+    and aggregates their results.
+    """
+
+    return cast(
+        "Callable[[AgentCallable[P, R]], DecoratedParallelProtocol[P, R]]",
+        _decorator_impl(
+            self,
+            AgentType.PARALLEL,
+            name=name,
+            instruction=instruction or default_instruction,
+            servers=[],  # Parallel agents don't connect to servers directly
+            fan_in=fan_in,
+            fan_out=fan_out,
+            include_request=include_request,
+        ),
+    )
+
+
+def evaluator_optimizer(
+    self,
+    name: str,
+    *,
+    generator: str,
+    evaluator: str,
+    instruction: Optional[str] = None,
+    min_rating: str = "GOOD",
+    max_refinements: int = 3,
+) -> Callable[[AgentCallable[P, R]], DecoratedEvaluatorOptimizerProtocol[P, R]]:
+    """
+    Decorator to create and register an evaluator-optimizer agent with type-safe signature.
+
+    Args:
+        name: Name of the evaluator-optimizer agent
+        generator: Name of the agent that generates responses
+        evaluator: Name of the agent that evaluates responses
+        instruction: Base instruction for the evaluator-optimizer
+        min_rating: Minimum acceptable quality rating (EXCELLENT, GOOD, FAIR, POOR)
+        max_refinements: Maximum number of refinement iterations
+
+    Returns:
+        A decorator that registers the evaluator-optimizer with proper type annotations
+    """
+    default_instruction = """
+    You implement an iterative refinement process where content is generated,
+    evaluated for quality, and then refined based on specific feedback until
+    it reaches an acceptable quality standard.
+    """
+
+    return cast(
+        "Callable[[AgentCallable[P, R]], DecoratedEvaluatorOptimizerProtocol[P, R]]",
+        _decorator_impl(
+            self,
+            AgentType.EVALUATOR_OPTIMIZER,
+            name=name,
+            instruction=instruction or default_instruction,
+            servers=[],  # Evaluator-optimizer doesn't connect to servers directly
+            generator=generator,
+            evaluator=evaluator,
+            min_rating=min_rating,
+            max_refinements=max_refinements,
+        ),
+    )