krons 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

krons/work/engine.py

@@ -0,0 +1,333 @@
+# Copyright (c) 2025 - 2026, HaiyangLi <quantocean.li at gmail dot com>
+# SPDX-License-Identifier: Apache-2.0
+
+"""WorkerEngine - Execution driver for Worker workflows.
+
+The engine manages task execution, following worklinks to traverse the
+workflow graph defined by @work and @worklink decorators.
+
+Example:
+    worker = FileCoder()
+    engine = WorkerEngine(worker=worker, refresh_time=0.3)
+
+    # Add a task starting at a specific function
+    task = await engine.add_task(
+        form=my_form,
+        task_function="start_task",
+        task_max_steps=20,
+    )
+
+    # Run until all tasks complete
+    await engine.execute()
+
+    # Or run indefinitely
+    await engine.execute_lasting()
+"""
+
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+from uuid import UUID, uuid4
+
+from krons.utils import concurrency
+
+if TYPE_CHECKING:
+    from .worker import Worker
+
+__all__ = ("WorkerEngine", "WorkerTask")
+
+
+@dataclass
+class WorkerTask:
+    """A task being executed by the engine.
+
+    Attributes:
+        id: Unique task identifier
+        function: Current method name to execute
+        kwargs: Arguments for the method
+        status: PENDING, PROCESSING, COMPLETED, FAILED
+        result: Final result when completed
+        error: Exception if failed
+        max_steps: Max workflow steps before stopping
+        current_step: Current step count
+        history: List of (function, result) tuples for debugging
+    """
+
+    id: UUID = field(default_factory=uuid4)
+    function: str = ""
+    kwargs: dict[str, Any] = field(default_factory=dict)
+    status: str = "PENDING"
+    result: Any = None
+    error: Exception | None = None
+    max_steps: int = 100
+    current_step: int = 0
+    history: list[tuple[str, Any]] = field(default_factory=list)
+
+
+class WorkerEngine:
+    """Execution driver for Worker workflows.
+
+    Manages a queue of tasks, executing them through the workflow graph
+    defined by the worker's @work and @worklink decorators.
+
+    Attributes:
+        worker: The Worker instance to execute
+        refresh_time: Seconds between processing cycles
+        tasks: Dict of active tasks by ID
+        _task_queue: Async queue for pending work
+        _stopped: Stop flag
+
+    Example:
+        engine = WorkerEngine(worker=my_worker)
+        task = await engine.add_task(
+            form=my_form,
+            task_function="entry_point",
+        )
+        await engine.execute()
+    """
+
+    def __init__(
+        self,
+        worker: Worker,
+        refresh_time: float = 0.1,
+        max_concurrent: int = 10,
+    ) -> None:
+        """Initialize the engine.
+
+        Args:
+            worker: Worker instance with @work/@worklink methods
+            refresh_time: Seconds between processing cycles
+            max_concurrent: Max concurrent task executions
+        """
+        self.worker = worker
+        self.refresh_time = refresh_time
+        self.max_concurrent = max_concurrent
+
+        self.tasks: dict[UUID, WorkerTask] = {}
+        self._task_queue: asyncio.Queue[UUID] = asyncio.Queue()
+        self._stopped = False
+        self._semaphore = asyncio.Semaphore(max_concurrent)
+
+    async def add_task(
+        self,
+        task_function: str,
+        task_max_steps: int = 100,
+        **kwargs: Any,
+    ) -> WorkerTask:
+        """Add a new task to the execution queue.
+
+        Args:
+            task_function: Entry method name to start execution
+            task_max_steps: Max workflow steps before stopping
+            **kwargs: Arguments for the entry method
+
+        Returns:
+            WorkerTask instance (can be monitored for status)
+
+        Raises:
+            ValueError: If task_function not found in worker
+        """
+        if task_function not in self.worker._work_methods:
+            raise ValueError(
+                f"Method '{task_function}' not found. "
+                f"Available: {list(self.worker._work_methods.keys())}"
+            )
+
+        task = WorkerTask(
+            function=task_function,
+            kwargs=kwargs,
+            max_steps=task_max_steps,
+        )
+        self.tasks[task.id] = task
+        await self._task_queue.put(task.id)
+
+        return task
+
+    async def execute(self) -> None:
+        """Execute all queued tasks until queue is empty.
+
+        Processes tasks through their workflow graphs, following worklinks.
+        Returns when all tasks are completed or failed.
+        """
+        self._stopped = False
+        await self.worker.start()
+
+        while not self._stopped and not self._task_queue.empty():
+            await self._process_cycle()
+            await concurrency.sleep(self.refresh_time)
+
+    async def execute_lasting(self) -> None:
+        """Execute indefinitely until stop() is called.
+
+        Useful for long-running worker services that continuously
+        process incoming tasks.
+        """
+        self._stopped = False
+        await self.worker.start()
+
+        while not self._stopped:
+            await self._process_cycle()
+            await concurrency.sleep(self.refresh_time)
+
+    async def stop(self) -> None:
+        """Stop the execution loop."""
+        self._stopped = True
+        await self.worker.stop()
+
+    async def _process_cycle(self) -> None:
+        """Process one cycle of tasks."""
+        # Collect tasks to process this cycle
+        tasks_to_process: list[UUID] = []
+
+        while (
+            not self._task_queue.empty() and len(tasks_to_process) < self.max_concurrent
+        ):
+            try:
+                task_id = self._task_queue.get_nowait()
+                tasks_to_process.append(task_id)
+            except asyncio.QueueEmpty:
+                break
+
+        if not tasks_to_process:
+            return
+
+        # Process tasks concurrently
+        async with concurrency.create_task_group() as tg:
+            for task_id in tasks_to_process:
+                tg.start_soon(self._process_task, task_id)
+
+    async def _process_task(self, task_id: UUID) -> None:
+        """Process a single task through one workflow step."""
+        async with self._semaphore:
+            task = self.tasks.get(task_id)
+            if task is None or task.status in ("COMPLETED", "FAILED"):
+                return
+
+            # Check step limit
+            if task.current_step >= task.max_steps:
+                task.status = "COMPLETED"
+                return
+
+            task.status = "PROCESSING"
+            task.current_step += 1
+
+            try:
+                # Get the work method and config
+                method, config = self.worker._work_methods[task.function]
+
+                # Prepare kwargs with form binding
+                call_kwargs = dict(task.kwargs)
+                if config.form_param_key and config.assignment:
+                    form_id = call_kwargs.get(config.form_param_key)
+                    if form_id and form_id in self.worker.forms:
+                        form = self.worker.forms[form_id]
+                        # Bind input fields from form to kwargs
+                        for input_field in form.input_fields:
+                            if input_field in form.available_data:
+                                call_kwargs[input_field] = form.available_data[
+                                    input_field
+                                ]
+
+                # Execute with optional timeout
+                if config.timeout:
+                    result = await asyncio.wait_for(
+                        method(**call_kwargs),
+                        timeout=config.timeout,
+                    )
+                else:
+                    result = await method(**call_kwargs)
+
+                # Record history
+                task.history.append((task.function, result))
+                task.result = result
+
+                # Follow worklinks
+                next_tasks = await self._follow_links(task, result)
+
+                if next_tasks:
+                    # Continue with next step(s)
+                    for next_func, next_kwargs in next_tasks:
+                        task.function = next_func
+                        task.kwargs = next_kwargs
+                        task.status = "PENDING"
+                        await self._task_queue.put(task_id)
+                        break  # Only follow first matching link for now
+                else:
+                    # No more links - task complete
+                    task.status = "COMPLETED"
+
+            except Exception as e:
+                task.status = "FAILED"
+                task.error = e
+
+    async def _follow_links(
+        self, task: WorkerTask, result: Any
+    ) -> list[tuple[str, dict[str, Any]]]:
+        """Follow worklinks from current method.
+
+        Args:
+            task: Current task
+            result: Result from current method
+
+        Returns:
+            List of (next_function, kwargs) tuples for matching links
+        """
+        next_tasks: list[tuple[str, dict[str, Any]]] = []
+
+        for link in self.worker.get_links_from(task.function):
+            try:
+                # Get the handler method from worker and call it
+                handler = getattr(self.worker, link.handler_name)
+                next_kwargs = await handler(result)
+
+                # None means skip this edge
+                if next_kwargs is not None:
+                    next_tasks.append((link.to_, next_kwargs))
+
+            except Exception:
+                # Link handler failed - skip this edge
+                continue
+
+        return next_tasks
+
+    def get_task(self, task_id: UUID) -> WorkerTask | None:
+        """Get task by ID."""
+        return self.tasks.get(task_id)
+
+    def get_tasks_by_status(self, status: str) -> list[WorkerTask]:
+        """Get all tasks with given status."""
+        return [t for t in self.tasks.values() if t.status == status]
+
+    @property
+    def pending_tasks(self) -> list[WorkerTask]:
+        """Tasks waiting to be processed."""
+        return self.get_tasks_by_status("PENDING")
+
+    @property
+    def processing_tasks(self) -> list[WorkerTask]:
+        """Tasks currently being processed."""
+        return self.get_tasks_by_status("PROCESSING")
+
+    @property
+    def completed_tasks(self) -> list[WorkerTask]:
+        """Tasks that completed successfully."""
+        return self.get_tasks_by_status("COMPLETED")
+
+    @property
+    def failed_tasks(self) -> list[WorkerTask]:
+        """Tasks that failed with errors."""
+        return self.get_tasks_by_status("FAILED")
+
+    def status_counts(self) -> dict[str, int]:
+        """Count tasks by status."""
+        counts: dict[str, int] = {}
+        for task in self.tasks.values():
+            counts[task.status] = counts.get(task.status, 0) + 1
+        return counts
+
+    def __repr__(self) -> str:
+        counts = self.status_counts()
+        total = len(self.tasks)
+        return f"WorkerEngine(worker={self.worker.name}, tasks={total}, {counts})"

krons/work/form.py

@@ -0,0 +1,305 @@
+# Copyright (c) 2025 - 2026, HaiyangLi <quantocean.li at gmail dot com>
+# SPDX-License-Identifier: Apache-2.0
+
+"""Form - Data binding and scheduling for work units.
+
+A Form represents an instantiated work unit with:
+- Data binding (input values)
+- Execution state tracking (filled, workable)
+- Optional Phrase reference for typed I/O
+
+Forms are the stateful layer between Phrase (definition) and Operation (execution).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+from pydantic import Field
+
+from krons.core import Element
+
+if TYPE_CHECKING:
+    from .phrase import Phrase
+
+__all__ = ("Form", "ParsedAssignment", "parse_assignment", "parse_full_assignment")
+
+
+@dataclass
+class ParsedAssignment:
+    """Parsed form assignment with all components.
+
+    Attributes:
+        branch: Branch/worker name (e.g., "classifier1")
+        inputs: Input field names
+        outputs: Output field names
+        resource: Resource hint (e.g., "api:fast")
+        raw: Original assignment string
+    """
+
+    branch: str | None
+    inputs: list[str]
+    outputs: list[str]
+    resource: str | None
+    raw: str
+
+
+def parse_assignment(assignment: str) -> tuple[list[str], list[str]]:
+    """Parse 'inputs -> outputs' assignment DSL (simple form).
+
+    Args:
+        assignment: DSL string like "a, b -> c, d"
+
+    Returns:
+        Tuple of (input_fields, output_fields)
+
+    Raises:
+        ValueError: If assignment format is invalid
+    """
+    parsed = parse_full_assignment(assignment)
+    return parsed.inputs, parsed.outputs
+
+
+def parse_full_assignment(assignment: str) -> ParsedAssignment:
+    """Parse full assignment DSL with branch and resource hints.
+
+    Format: "branch: inputs -> outputs | resource"
+
+    Examples:
+        "a, b -> c"                           # Simple
+        "classifier: job -> role | api:fast"  # Full
+        "writer: context -> summary"          # Branch, no resource
+
+    Args:
+        assignment: DSL string
+
+    Returns:
+        ParsedAssignment with all components
+
+    Raises:
+        ValueError: If format is invalid
+    """
+    raw = assignment.strip()
+    branch = None
+    resource = None
+
+    # Extract resource hint (after |)
+    if "|" in raw:
+        main_part, resource_part = raw.rsplit("|", 1)
+        resource = resource_part.strip()
+        raw = main_part.strip()
+
+    # Extract branch name (before :)
+    if ":" in raw:
+        # Check it's not just inside the field list
+        colon_idx = raw.find(":")
+        arrow_idx = raw.find("->")
+        if arrow_idx == -1 or colon_idx < arrow_idx:
+            branch_part, raw = raw.split(":", 1)
+            branch = branch_part.strip()
+            raw = raw.strip()
+
+    # Parse inputs -> outputs
+    if "->" not in raw:
+        raise ValueError(f"Invalid assignment syntax (missing '->'): {assignment}")
+
+    parts = raw.split("->")
+    if len(parts) != 2:
+        raise ValueError(f"Invalid assignment syntax: {assignment}")
+
+    inputs = [f.strip() for f in parts[0].split(",") if f.strip()]
+    outputs = [f.strip() for f in parts[1].split(",") if f.strip()]
+
+    return ParsedAssignment(
+        branch=branch,
+        inputs=inputs,
+        outputs=outputs,
+        resource=resource,
+        raw=assignment,
+    )
+
+
+class Form(Element):
+    """Data binding container for work units.
+
+    A Form binds input data and tracks execution state. It can be created:
+    1. From a Phrase (typed I/O)
+    2. From an assignment string (dynamic fields)
+
+    Assignment DSL supports full format:
+        "branch: inputs -> outputs | resource"
+
+    Examples:
+        "a, b -> c"                           # Simple
+        "classifier: job -> role | api:fast"  # Full with branch and resource
+        "writer: context -> summary"          # Branch, no resource
+
+    Attributes:
+        assignment: DSL string 'branch: inputs -> outputs | resource'
+        branch: Worker/branch name for routing
+        resource: Resource hint for capability matching
+        input_fields: Fields required as inputs
+        output_fields: Fields produced as outputs
+        available_data: Current data values
+        output: Execution result
+        filled: Whether form has been executed
+        phrase: Optional Phrase reference for typed execution
+    """
+
+    assignment: str = Field(
+        default="",
+        description="Assignment DSL: 'branch: inputs -> outputs | resource'",
+    )
+    branch: str | None = Field(
+        default=None,
+        description="Worker/branch name for routing",
+    )
+    resource: str | None = Field(
+        default=None,
+        description="Resource hint (e.g., 'api:fast')",
+    )
+    input_fields: list[str] = Field(default_factory=list)
+    output_fields: list[str] = Field(default_factory=list)
+    available_data: dict[str, Any] = Field(default_factory=dict)
+    output: Any = Field(default=None)
+    filled: bool = Field(default=False)
+
+    # Optional phrase reference (set via from_phrase())
+    _phrase: "Phrase | None" = None
+
+    def model_post_init(self, _: Any) -> None:
+        """Parse assignment to derive fields if not already set."""
+        if self.assignment and not self.input_fields and not self.output_fields:
+            parsed = parse_full_assignment(self.assignment)
+            self.input_fields = parsed.inputs
+            self.output_fields = parsed.outputs
+            if parsed.branch and self.branch is None:
+                self.branch = parsed.branch
+            if parsed.resource and self.resource is None:
+                self.resource = parsed.resource
+
+    @classmethod
+    def from_phrase(
+        cls,
+        phrase: "Phrase",
+        **initial_data: Any,
+    ) -> "Form":
+        """Create Form from a Phrase with optional initial data.
+
+        Args:
+            phrase: Phrase defining typed I/O
+            **initial_data: Initial input values
+
+        Returns:
+            Form bound to the phrase
+        """
+        form = cls(
+            assignment=f"{', '.join(phrase.inputs)} -> {', '.join(phrase.outputs)}",
+            input_fields=list(phrase.inputs),
+            output_fields=list(phrase.outputs),
+            available_data=dict(initial_data),
+        )
+        form._phrase = phrase
+        return form
+
+    @property
+    def phrase(self) -> "Phrase | None":
+        """Get bound phrase if any."""
+        return self._phrase
+
+    def is_workable(self) -> bool:
+        """Check if form is ready for execution.
+
+        Returns:
+            True if all inputs available and not already filled
+        """
+        if self.filled:
+            return False
+
+        for field in self.input_fields:
+            if field not in self.available_data:
+                return False
+            if self.available_data[field] is None:
+                return False
+
+        return True
+
+    def get_inputs(self) -> dict[str, Any]:
+        """Extract input data for execution.
+
+        Returns:
+            Dict of input field values
+        """
+        return {
+            f: self.available_data[f]
+            for f in self.input_fields
+            if f in self.available_data
+        }
+
+    def fill(self, **data: Any) -> None:
+        """Add data to available_data.
+
+        Args:
+            **data: Field values to add
+        """
+        self.available_data.update(data)
+
+    def set_output(self, output: Any) -> None:
+        """Mark form as filled with output.
+
+        Args:
+            output: Execution result
+        """
+        self.output = output
+        self.filled = True
+
+        # Extract output field values from result
+        if output is not None:
+            for field in self.output_fields:
+                if hasattr(output, field):
+                    self.available_data[field] = getattr(output, field)
+                elif isinstance(output, dict) and field in output:
+                    self.available_data[field] = output[field]
+
+    def get_output_data(self) -> dict[str, Any]:
+        """Extract output field values.
+
+        Returns:
+            Dict mapping output field names to values
+        """
+        result = {}
+        for field in self.output_fields:
+            if field in self.available_data:
+                result[field] = self.available_data[field]
+        return result
+
+    async def execute(self, ctx: Any = None) -> Any:
+        """Execute the form if it has a bound phrase.
+
+        Args:
+            ctx: Execution context
+
+        Returns:
+            Execution result
+
+        Raises:
+            RuntimeError: If no phrase bound or form not workable
+        """
+        if self._phrase is None:
+            raise RuntimeError("Form has no bound phrase - cannot execute")
+
+        if not self.is_workable():
+            missing = [f for f in self.input_fields if f not in self.available_data]
+            raise RuntimeError(f"Form not workable - missing inputs: {missing}")
+
+        result = await self._phrase(self.get_inputs(), ctx)
+        self.set_output(result)
+        return result
+
+    def __repr__(self) -> str:
+        status = (
+            "filled" if self.filled else ("ready" if self.is_workable() else "pending")
+        )
+        phrase_info = f", phrase={self._phrase.name}" if self._phrase else ""
+        return f"Form('{self.assignment}', {status}{phrase_info})"

{kronos → krons/work}/operations/__init__.py

@@ -5,7 +5,7 @@
 
 Core types:
     Operation: Node + Event hybrid for graph-based execution.
-    OperationRegistry: Per-session factory mapping.
+    OperationRegistry: Per-session handler mapping.
     OperationGraphBuilder (Builder): Fluent DAG construction.
 
 Execution:
@@ -16,17 +16,20 @@ Execution:
 from __future__ import annotations
 
 from .builder import Builder, OperationGraphBuilder
+from .context import QueryFn, RequestContext
 from .flow import DependencyAwareExecutor, flow, flow_stream
-from .node import Operation, create_operation
-from .registry import OperationRegistry
+from .node import Operation
+from .registry import OperationHandler, OperationRegistry
 
 __all__ = (
     "Builder",
     "DependencyAwareExecutor",
     "Operation",
     "OperationGraphBuilder",
+    "OperationHandler",
     "OperationRegistry",
-    "create_operation",
+    "QueryFn",
+    "RequestContext",
     "flow",
     "flow_stream",
 )

{kronos → krons/work}/operations/builder.py

@@ -12,14 +12,14 @@ from __future__ import annotations
 from typing import TYPE_CHECKING, Any
 from uuid import UUID
 
-from kronos.core import Edge, Graph
-from kronos.types import Undefined, UndefinedType, is_sentinel, not_sentinel
-from kronos.utils._utils import to_uuid
+from krons.core import Edge, Graph
+from krons.core.types import Undefined, UndefinedType, is_sentinel, not_sentinel
+from krons.utils._utils import to_uuid
 
 from .node import Operation
 
 if TYPE_CHECKING:
-    from kronos.session import Branch
+    from krons.session import Branch
 
 __all__ = ("Builder", "OperationGraphBuilder")