fast-agent-mcp 0.2.12__py3-none-any.whl → 0.2.14__py3-none-any.whl

mcp_agent/executor/temporal.py

@@ -1,383 +0,0 @@
-"""
-Temporal based orchestrator for the MCP Agent.
-Temporal provides durable execution and robust workflow orchestration,
-as well as dynamic control flow, making it a good choice for an AI agent orchestrator.
-Read more: https://docs.temporal.io/develop/python/core-application
-"""
-
-import asyncio
-import functools
-import uuid
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    AsyncIterator,
-    Callable,
-    Coroutine,
-    Dict,
-    List,
-    Optional,
-)
-
-from pydantic import ConfigDict
-from temporalio import activity, exceptions, workflow
-from temporalio.client import Client as TemporalClient
-from temporalio.worker import Worker
-
-from mcp_agent.config import TemporalSettings
-from mcp_agent.executor.executor import Executor, ExecutorConfig, R
-from mcp_agent.executor.workflow_signal import (
-    BaseSignalHandler,
-    Signal,
-    SignalHandler,
-    SignalRegistration,
-    SignalValueT,
-)
-
-if TYPE_CHECKING:
-    from mcp_agent.context import Context
-
-
-class TemporalSignalHandler(BaseSignalHandler[SignalValueT]):
-    """Temporal-based signal handling using workflow signals"""
-
-    async def wait_for_signal(self, signal, timeout_seconds=None) -> SignalValueT:
-        if not workflow._Runtime.current():
-            raise RuntimeError(
-                "TemporalSignalHandler.wait_for_signal must be called from within a workflow"
-            )
-
-        unique_signal_name = f"{signal.name}_{uuid.uuid4()}"
-        registration = SignalRegistration(
-            signal_name=signal.name,
-            unique_name=unique_signal_name,
-            workflow_id=workflow.info().workflow_id,
-        )
-
-        # Container for signal value
-        container = {"value": None, "completed": False}
-
-        # Define the signal handler for this specific registration
-        @workflow.signal(name=unique_signal_name)
-        def signal_handler(value: SignalValueT) -> None:
-            container["value"] = value
-            container["completed"] = True
-
-        async with self._lock:
-            # Register both the signal registration and handler atomically
-            self._pending_signals.setdefault(signal.name, []).append(registration)
-            self._handlers.setdefault(signal.name, []).append((unique_signal_name, signal_handler))
-
-        try:
-            # Wait for signal with optional timeout
-            await workflow.wait_condition(lambda: container["completed"], timeout=timeout_seconds)
-
-            return container["value"]
-        except asyncio.TimeoutError as exc:
-            raise TimeoutError(f"Timeout waiting for signal {signal.name}") from exc
-        finally:
-            async with self._lock:
-                # Remove ourselves from _pending_signals
-                if signal.name in self._pending_signals:
-                    self._pending_signals[signal.name] = [
-                        sr
-                        for sr in self._pending_signals[signal.name]
-                        if sr.unique_name != unique_signal_name
-                    ]
-                    if not self._pending_signals[signal.name]:
-                        del self._pending_signals[signal.name]
-
-                # Remove ourselves from _handlers
-                if signal.name in self._handlers:
-                    self._handlers[signal.name] = [
-                        h for h in self._handlers[signal.name] if h[0] != unique_signal_name
-                    ]
-                    if not self._handlers[signal.name]:
-                        del self._handlers[signal.name]
-
-    def on_signal(self, signal_name):
-        """Decorator to register a signal handler."""
-
-        def decorator(func: Callable) -> Callable:
-            # Create unique signal name for this handler
-            unique_signal_name = f"{signal_name}_{uuid.uuid4()}"
-
-            # Create the actual handler that will be registered with Temporal
-            @workflow.signal(name=unique_signal_name)
-            async def wrapped(signal_value: SignalValueT) -> None:
-                # Create a signal object to pass to the handler
-                signal = Signal(
-                    name=signal_name,
-                    payload=signal_value,
-                    workflow_id=workflow.info().workflow_id,
-                )
-                if asyncio.iscoroutinefunction(func):
-                    await func(signal)
-                else:
-                    func(signal)
-
-            # Register the handler under the original signal name
-            self._handlers.setdefault(signal_name, []).append((unique_signal_name, wrapped))
-            return func
-
-        return decorator
-
-    async def signal(self, signal) -> None:
-        self.validate_signal(signal)
-
-        workflow_handle = workflow.get_external_workflow_handle(workflow_id=signal.workflow_id)
-
-        # Send the signal to all registrations of this signal
-        async with self._lock:
-            signal_tasks = []
-
-            if signal.name in self._pending_signals:
-                for pending_signal in self._pending_signals[signal.name]:
-                    registration = pending_signal.registration
-                    if registration.workflow_id == signal.workflow_id:
-                        # Only signal for registrations of that workflow
-                        signal_tasks.append(
-                            workflow_handle.signal(registration.unique_name, signal.payload)
-                        )
-                    else:
-                        continue
-
-            # Notify any registered handler functions
-            if signal.name in self._handlers:
-                for unique_name, _ in self._handlers[signal.name]:
-                    signal_tasks.append(workflow_handle.signal(unique_name, signal.payload))
-
-        await asyncio.gather(*signal_tasks, return_exceptions=True)
-
-    def validate_signal(self, signal) -> None:
-        super().validate_signal(signal)
-        # Add TemporalSignalHandler-specific validation
-        if signal.workflow_id is None:
-            raise ValueError(
-                "No workflow_id provided on Signal. That is required for Temporal signals"
-            )
-
-
-class TemporalExecutorConfig(ExecutorConfig, TemporalSettings):
-    """Configuration for Temporal executors."""
-
-    model_config = ConfigDict(extra="allow", arbitrary_types_allowed=True)
-
-
-class TemporalExecutor(Executor):
-    """Executor that runs @workflows as Temporal workflows, with @workflow_tasks as Temporal activities"""
-
-    def __init__(
-        self,
-        config: TemporalExecutorConfig | None = None,
-        signal_bus: SignalHandler | None = None,
-        client: TemporalClient | None = None,
-        context: Optional["Context"] = None,
-        **kwargs,
-    ) -> None:
-        signal_bus = signal_bus or TemporalSignalHandler()
-        super().__init__(
-            engine="temporal",
-            config=config,
-            signal_bus=signal_bus,
-            context=context,
-            **kwargs,
-        )
-        self.config: TemporalExecutorConfig = (
-            config or self.context.config.temporal or TemporalExecutorConfig()
-        )
-        self.client = client
-        self._worker = None
-        self._activity_semaphore = None
-
-        if config.max_concurrent_activities is not None:
-            self._activity_semaphore = asyncio.Semaphore(self.config.max_concurrent_activities)
-
-    @staticmethod
-    def wrap_as_activity(
-        activity_name: str,
-        func: Callable[..., R] | Coroutine[Any, Any, R],
-        **kwargs: Any,
-    ) -> Coroutine[Any, Any, R]:
-        """
-        Convert a function into a Temporal activity and return its info.
-        """
-
-        @activity.defn(name=activity_name)
-        async def wrapped_activity(*args, **local_kwargs):
-            try:
-                if asyncio.iscoroutinefunction(func):
-                    return await func(*args, **local_kwargs)
-                elif asyncio.iscoroutine(func):
-                    return await func
-                else:
-                    return func(*args, **local_kwargs)
-            except Exception as e:
-                # Handle exceptions gracefully
-                raise e
-
-        return wrapped_activity
-
-    async def _execute_task_as_async(
-        self, task: Callable[..., R] | Coroutine[Any, Any, R], **kwargs: Any
-    ) -> R | BaseException:
-        async def run_task(task: Callable[..., R] | Coroutine[Any, Any, R]) -> R:
-            try:
-                if asyncio.iscoroutine(task):
-                    return await task
-                elif asyncio.iscoroutinefunction(task):
-                    return await task(**kwargs)
-                else:
-                    # Execute the callable and await if it returns a coroutine
-                    loop = asyncio.get_running_loop()
-
-                    # If kwargs are provided, wrap the function with partial
-                    if kwargs:
-                        wrapped_task = functools.partial(task, **kwargs)
-                        result = await loop.run_in_executor(None, wrapped_task)
-                    else:
-                        result = await loop.run_in_executor(None, task)
-
-                    # Handle case where the sync function returns a coroutine
-                    if asyncio.iscoroutine(result):
-                        return await result
-
-                    return result
-            except Exception as e:
-                # TODO: saqadri - adding logging or other error handling here
-                return e
-
-        if self._activity_semaphore:
-            async with self._activity_semaphore:
-                return await run_task(task)
-        else:
-            return await run_task(task)
-
-    async def _execute_task(
-        self, task: Callable[..., R] | Coroutine[Any, Any, R], **kwargs: Any
-    ) -> R | BaseException:
-        func = task.func if isinstance(task, functools.partial) else task
-        is_workflow_task = getattr(func, "is_workflow_task", False)
-        if not is_workflow_task:
-            return await asyncio.create_task(self._execute_task_as_async(task, **kwargs))
-
-        execution_metadata: Dict[str, Any] = getattr(func, "execution_metadata", {})
-
-        # Derive stable activity name, e.g. module + qualname
-        activity_name = execution_metadata.get("activity_name")
-        if not activity_name:
-            activity_name = f"{func.__module__}.{func.__qualname__}"
-
-        schedule_to_close = execution_metadata.get(
-            "schedule_to_close_timeout", self.config.timeout_seconds
-        )
-
-        retry_policy = execution_metadata.get("retry_policy", None)
-
-        _task_activity = self.wrap_as_activity(activity_name=activity_name, func=task)
-
-        # # For partials, we pass the partial's arguments
-        # args = task.args if isinstance(task, functools.partial) else ()
-        try:
-            result = await workflow.execute_activity(
-                activity_name,
-                args=kwargs.get("args", ()),
-                task_queue=self.config.task_queue,
-                schedule_to_close_timeout=schedule_to_close,
-                retry_policy=retry_policy,
-                **kwargs,
-            )
-            return result
-        except Exception as e:
-            # Properly propagate activity errors
-            if isinstance(e, exceptions.ActivityError):
-                raise e.cause if e.cause else e
-            raise
-
-    async def execute(
-        self,
-        *tasks: Callable[..., R] | Coroutine[Any, Any, R],
-        **kwargs: Any,
-    ) -> List[R | BaseException]:
-        # Must be called from within a workflow
-        if not workflow._Runtime.current():
-            raise RuntimeError("TemporalExecutor.execute must be called from within a workflow")
-
-        # TODO: saqadri - validate if async with self.execution_context() is needed here
-        async with self.execution_context():
-            return await asyncio.gather(
-                *(self._execute_task(task, **kwargs) for task in tasks),
-                return_exceptions=True,
-            )
-
-    async def execute_streaming(
-        self,
-        *tasks: Callable[..., R] | Coroutine[Any, Any, R],
-        **kwargs: Any,
-    ) -> AsyncIterator[R | BaseException]:
-        if not workflow._Runtime.current():
-            raise RuntimeError(
-                "TemporalExecutor.execute_streaming must be called from within a workflow"
-            )
-
-        # TODO: saqadri - validate if async with self.execution_context() is needed here
-        async with self.execution_context():
-            # Create futures for all tasks
-            futures = [self._execute_task(task, **kwargs) for task in tasks]
-            pending = set(futures)
-
-            while pending:
-                done, pending = await workflow.wait(pending, return_when=asyncio.FIRST_COMPLETED)
-                for future in done:
-                    try:
-                        result = await future
-                        yield result
-                    except Exception as e:
-                        yield e
-
-    async def ensure_client(self):
-        """Ensure we have a connected Temporal client."""
-        if self.client is None:
-            self.client = await TemporalClient.connect(
-                target_host=self.config.host,
-                namespace=self.config.namespace,
-                api_key=self.config.api_key,
-            )
-
-        return self.client
-
-    async def start_worker(self) -> None:
-        """
-        Start a worker in this process, auto-registering all tasks
-        from the global registry. Also picks up any classes decorated
-        with @workflow_defn as recognized workflows.
-        """
-        await self.ensure_client()
-
-        if self._worker is None:
-            # We'll collect the activities from the global registry
-            # and optionally wrap them with `activity.defn` if we want
-            # (Not strictly required if your code calls `execute_activity("name")` by name)
-            activity_registry = self.context.task_registry
-            activities = []
-            for name in activity_registry.list_activities():
-                activities.append(activity_registry.get_activity(name))
-
-            # Now we attempt to discover any classes that are recognized as workflows
-            # But in this simple example, we rely on the user specifying them or
-            # we might do a dynamic scan.
-            # For demonstration, we'll just assume the user is only using
-            # the workflow classes they decorate with `@workflow_defn`.
-            # We'll rely on them passing the classes or scanning your code.
-
-            self._worker = Worker(
-                client=self.client,
-                task_queue=self.config.task_queue,
-                activities=activities,
-                workflows=[],  # We'll auto-load by Python scanning or let the user specify
-            )
-            print(
-                f"Starting Temporal Worker on task queue '{self.config.task_queue}' with {len(activities)} activities."
-            )
-
-        await self._worker.run()

mcp_agent/executor/workflow.py

@@ -1,195 +0,0 @@
-from abc import ABC, abstractmethod
-from datetime import datetime
-from typing import (
-    Any,
-    Dict,
-    Generic,
-    TypeVar,
-    Union,
-)
-
-from pydantic import BaseModel, ConfigDict, Field
-
-from mcp_agent.executor.executor import Executor
-
-T = TypeVar("T")
-
-
-class WorkflowState(BaseModel):
-    """
-    Simple container for persistent workflow state.
-    This can hold fields that should persist across tasks.
-    """
-
-    status: str = "initialized"
-    metadata: Dict[str, Any] = Field(default_factory=dict)
-    updated_at: float | None = None
-    error: Dict[str, Any] | None = None
-
-    model_config = ConfigDict(extra="allow", arbitrary_types_allowed=True)
-
-    def record_error(self, error: Exception) -> None:
-        self.error = {
-            "type": type(error).__name__,
-            "message": str(error),
-            "timestamp": datetime.utcnow().timestamp(),
-        }
-
-
-class WorkflowResult(BaseModel, Generic[T]):
-    value: Union[T, None] = None
-    metadata: Dict[str, Any] = Field(default_factory=dict)
-    start_time: float | None = None
-    end_time: float | None = None
-
-
-class Workflow(ABC, Generic[T]):
-    """
-    Base class for user-defined workflows.
-    Handles execution and state management.
-    Some key notes:
-        - To enable the executor engine to recognize and orchestrate the workflow,
-            - the class MUST be decorated with @workflow.
-            - the main entrypoint method MUST be decorated with @workflow_run.
-            - any task methods MUST be decorated with @workflow_task.
-
-        - Persistent state: Provides a simple `state` object for storing data across tasks.
-    """
-
-    def __init__(
-        self,
-        executor: Executor,
-        name: str | None = None,
-        metadata: Dict[str, Any] | None = None,
-        **kwargs: Any,
-    ) -> None:
-        self.executor = executor
-        self.name = name or self.__class__.__name__
-        self.init_kwargs = kwargs
-        # TODO: handle logging
-        # self._logger = logging.getLogger(self.name)
-
-        # A simple workflow state object
-        # If under Temporal, storing it as a field on this class
-        # means it can be replayed automatically
-        self.state = WorkflowState(name=name, metadata=metadata or {})
-
-    @abstractmethod
-    async def run(self, *args: Any, **kwargs: Any) -> "WorkflowResult[T]":
-        """
-        Main workflow implementation. Must be overridden by subclasses.
-        """
-
-    async def update_state(self, **kwargs) -> None:
-        """Syntactic sugar to update workflow state."""
-        for key, value in kwargs.items():
-            self.state[key] = value
-            setattr(self.state, key, value)
-
-        self.state.updated_at = datetime.utcnow().timestamp()
-
-    async def wait_for_input(self, description: str = "Provide input") -> str:
-        """
-        Convenience method for human input. Uses `human_input` signal
-        so we can unify local (console input) and Temporal signals.
-        """
-        return await self.executor.wait_for_signal("human_input", description=description)
-
-
-# ############################
-# # Example: DocumentWorkflow
-# ############################
-
-
-# @workflow_defn  # <-- This becomes @temporal_workflow.defn if in Temporal mode, else no-op
-# class DocumentWorkflow(Workflow[List[Dict[str, Any]]]):
-#     """
-#     Example workflow with persistent state.
-#     If run locally, `self.state` is ephemeral.
-#     If run in Temporal mode, `self.state` is replayed automatically.
-#     """
-
-#     @workflow_task(
-#         schedule_to_close_timeout=timedelta(minutes=10),
-#         retry_policy={"initial_interval": 1, "max_attempts": 3},
-#     )
-#     async def process_document(self, doc_id: str) -> Dict[str, Any]:
-#         """Activity that simulates document processing."""
-#         await asyncio.sleep(1)
-#         # Optionally mutate workflow state
-#         self.state.metadata.setdefault("processed_docs", []).append(doc_id)
-#         return {
-#             "doc_id": doc_id,
-#             "status": "processed",
-#             "timestamp": datetime.utcnow().isoformat(),
-#         }
-
-#     @workflow_run  # <-- This becomes @temporal_workflow.run(...) if Temporal is used
-#     async def _run_impl(
-#         self, documents: List[str], batch_size: int = 2
-#     ) -> List[Dict[str, Any]]:
-#         """Main workflow logic, which becomes the official 'run' in Temporal mode."""
-#         self._logger.info("Workflow starting, state=%s", self.state)
-#         self.state.update_status("running")
-
-#         all_results = []
-#         for i in range(0, len(documents), batch_size):
-#             batch = documents[i : i + batch_size]
-#             tasks = [self.process_document(doc) for doc in batch]
-#             results = await self.executor.execute(*tasks)
-
-#             for res in results:
-#                 if isinstance(res.value, Exception):
-#                     self._logger.error(
-#                         f"Error processing document: {res.metadata.get('error')}"
-#                     )
-#                 else:
-#                     all_results.append(res.value)
-
-#         self.state.update_status("completed")
-#         return all_results
-
-
-# ########################
-# # 12. Example Local Usage
-# ########################
-
-
-# async def run_example_local():
-#     from . import AsyncIOExecutor, DocumentWorkflow  # if in a package
-
-#     executor = AsyncIOExecutor()
-#     wf = DocumentWorkflow(executor)
-
-#     documents = ["doc1", "doc2", "doc3", "doc4"]
-#     result = await wf.run(documents, batch_size=2)
-
-#     print("Local results:", result.value)
-#     print("Local workflow final state:", wf.state)
-#     # Notice `wf.state.metadata['processed_docs']` has the processed doc IDs.
-
-
-# ########################
-# # Example Temporal Usage
-# ########################
-
-
-# async def run_example_temporal():
-#     from . import TemporalExecutor, DocumentWorkflow  # if in a package
-
-#     # 1) Create a TemporalExecutor (client side)
-#     executor = TemporalExecutor(task_queue="my_task_queue")
-#     await executor.ensure_client()
-
-#     # 2) Start a worker in the same process (or do so in a separate process)
-#     asyncio.create_task(executor.start_worker())
-#     await asyncio.sleep(2)  # Wait for worker to be up
-
-#     # 3) Now we can run the workflow by normal means if we like,
-#     #    or rely on the Worker picking it up. Typically, you'd do:
-#     #    handle = await executor._client.start_workflow(...)
-#     #    but let's keep it simple and show conceptually
-#     #    that 'DocumentWorkflow' is now recognized as a real Temporal workflow
-#     print(
-#         "Temporal environment is running. Use the Worker logs or CLI to start 'DocumentWorkflow'."
-#     )