krons 0.1.1__py3-none-any.whl → 0.2.1__py3-none-any.whl

krons/{operations → 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",
 )

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

@@ -13,7 +13,7 @@ from typing import TYPE_CHECKING, Any
 from uuid import UUID
 
 from krons.core import Edge, Graph
-from krons.types import Undefined, UndefinedType, is_sentinel, not_sentinel
+from krons.core.types import Undefined, UndefinedType, is_sentinel, not_sentinel
 from krons.utils._utils import to_uuid
 
 from .node import Operation

krons/{enforcement → work/operations}/context.py

@@ -11,8 +11,8 @@ from datetime import datetime
 from typing import TYPE_CHECKING, Any, Protocol
 from uuid import UUID, uuid4
 
-from krons.types.base import DataClass
-from krons.types.identity import ID
+from krons.core.types.base import DataClass
+from krons.core.types.identity import ID
 
 if TYPE_CHECKING:
     from krons.session import Branch, Session
@@ -85,7 +85,7 @@ class RequestContext(DataClass):
     name: str
     id: UUID = field(default_factory=uuid4)
     session_id: ID[Session] | None = None
-    branch_id: ID[Branch] | None = None
+    branch: ID[Branch] | str | None = None
     metadata: dict[str, Any] = field(default_factory=dict)
     conn: Any | None = None
     query_fn: QueryFn | None = None
@@ -95,7 +95,7 @@ class RequestContext(DataClass):
         self,
         name: str,
         session_id: ID[Session] | None = None,
-        branch_id: ID[Branch] | None = None,
+        branch: ID[Branch] | str | None = None,
         id: UUID | None = None,
         conn: Any | None = None,
         query_fn: QueryFn | None = None,
@@ -105,7 +105,7 @@ class RequestContext(DataClass):
         self.name = name
         self.id = id or uuid4()
         self.session_id = session_id
-        self.branch_id = branch_id
+        self.branch = branch
         self.conn = conn
         self.query_fn = query_fn
         self.now = now
@@ -127,3 +127,34 @@ class RequestContext(DataClass):
             raise AttributeError(
                 f"'{type(self).__name__}' has no attribute '{name}'"
             ) from None
+
+    async def get_session(self) -> Session | None:
+        """Get the Session object.
+
+        Checks bound reference first (set by Operation.invoke),
+        then falls back to global registry lookup via session_id.
+
+        Returns None if no session available.
+        """
+        if "_bound_session" in self.metadata:
+            return self.metadata["_bound_session"]
+        if self.session_id is None:
+            return None
+        from krons.session.registry import get_session
+
+        return await get_session(self.session_id)
+
+    async def get_branch(self) -> Branch | None:
+        """Get the Branch object.
+
+        Checks bound reference first (set by Operation.invoke),
+        then falls back to session lookup via branch identifier.
+
+        Returns None if no branch available.
+        """
+        if "_bound_branch" in self.metadata:
+            return self.metadata["_bound_branch"]
+        session = await self.get_session()
+        if session is None or self.branch is None:
+            return None
+        return session.branches.get(self.branch)

krons/{operations → work/operations}/flow.py

@@ -16,7 +16,7 @@ from typing import TYPE_CHECKING, Any
 from uuid import UUID
 
 from krons.core import EventStatus, Graph
-from krons.types import Undefined, UndefinedType, is_sentinel
+from krons.core.types import Undefined, UndefinedType, is_sentinel
 from krons.utils import concurrency
 from krons.utils.concurrency import CapacityLimiter, CompletionStream
 
@@ -94,8 +94,12 @@ class DependencyAwareExecutor:
         """
         self.session = session
         self.graph = graph
-        resolved_max_concurrent = None if is_sentinel(max_concurrent) else max_concurrent
-        resolved_default_branch = None if is_sentinel(default_branch) else default_branch
+        resolved_max_concurrent = (
+            None if is_sentinel(max_concurrent) else max_concurrent
+        )
+        resolved_default_branch = (
+            None if is_sentinel(default_branch) else default_branch
+        )
         self.max_concurrent = resolved_max_concurrent
         self.stop_on_error = stop_on_error
         self.verbose = verbose
@@ -107,7 +111,9 @@ class DependencyAwareExecutor:
         self.operation_branches: dict[UUID, Branch | None] = {}
 
         self._limiter: CapacityLimiter | None = (
-            CapacityLimiter(resolved_max_concurrent) if resolved_max_concurrent else None
+            CapacityLimiter(resolved_max_concurrent)
+            if resolved_max_concurrent
+            else None
         )
 
         for node in graph.nodes:
@@ -258,7 +264,9 @@ class DependencyAwareExecutor:
                     self.operation_branches[node.id] = default_branch
 
         if self.verbose:
-            logger.debug("Pre-allocated branches for %d operations", len(self.operation_branches))
+            logger.debug(
+                "Pre-allocated branches for %d operations", len(self.operation_branches)
+            )
 
     async def _execute_operation(self, operation: Operation) -> Operation:
         """Execute single operation: wait deps -> acquire slot -> invoke -> signal."""

krons/{operations → work/operations}/node.py

@@ -1,5 +1,13 @@
 # Copyright (c) 2025 - 2026, HaiyangLi <quantocean.li at gmail dot com>
 # SPDX-License-Identifier: Apache-2.0
+
+"""Operation: executable graph node bridging session to handler.
+
+Operation.invoke() creates a RequestContext from the bound session/branch
+and calls the registered handler with (params, ctx). Handlers never need
+to know about the factory pattern — they receive a clean RequestContext.
+"""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Any
@@ -7,23 +15,36 @@ from typing import TYPE_CHECKING, Any
 from pydantic import Field, PrivateAttr
 
 from krons.core import Event, Node
-from krons.types import Undefined, UndefinedType, is_sentinel
+
+from .context import RequestContext
 
 if TYPE_CHECKING:
     from krons.session import Branch, Session
 
-__all__ = ("Operation", "create_operation")
+__all__ = ("Operation",)
 
 
 class Operation(Node, Event):
+    """Executable operation node.
+
+    Bridges session.conduct() to handler(params, ctx) by:
+    1. Storing bound session/branch references
+    2. Creating RequestContext with those references
+    3. Looking up the handler from session.operations registry
+    4. Calling handler(params, ctx)
+
+    The result is stored in execution.response (via Event.invoke).
+    """
+
     operation_type: str
     parameters: dict[str, Any] | Any = Field(
         default_factory=dict,
-        description="Operation parameters (dict or Pydantic model)",
+        description="Operation parameters (Params dataclass, dict, or model)",
     )
 
     _session: Any = PrivateAttr(default=None)
     _branch: Any = PrivateAttr(default=None)
+    _verbose: bool = PrivateAttr(default=False)
 
     def bind(self, session: Session, branch: Branch) -> Operation:
         """Bind session and branch for execution.
@@ -31,11 +52,11 @@ class Operation(Node, Event):
         Must be called before invoke() if not using Session.conduct().
 
         Args:
-            session: Session with operations registry and services
-            branch: Branch for message context
+            session: Session with operations registry and services.
+            branch: Branch for message context.
 
         Returns:
-            Self for chaining
+            Self for chaining.
         """
         self._session = session
         self._branch = branch
@@ -51,51 +72,32 @@ class Operation(Node, Event):
         return self._session, self._branch
 
     async def _invoke(self) -> Any:
-        """Execute via session's operation registry. Called by Event.invoke().
+        """Execute handler via session's operation registry.
+
+        Creates a RequestContext with bound session/branch references
+        and calls handler(params, ctx). Called by Event.invoke().
 
         Returns:
-            Factory result (stored in execution.response).
+            Handler result (stored in execution.response).
 
         Raises:
             RuntimeError: If not bound.
             KeyError: If operation_type not registered.
         """
         session, branch = self._require_binding()
-        factory = session.operations.get(self.operation_type)
-        return await factory(session, branch, self.parameters)
-
-    def __repr__(self) -> str:
-        bound = "bound" if self._session is not None else "unbound"
-        return (
-            f"Operation(type={self.operation_type}, status={self.execution.status.value}, {bound})"
+        handler = session.operations.get(self.operation_type)
+
+        ctx = RequestContext(
+            name=self.operation_type,
+            session_id=session.id,
+            branch=branch.name or str(branch.id),
+            _bound_session=session,
+            _bound_branch=branch,
+            _verbose=self._verbose,
         )
 
+        return await handler(self.parameters, ctx)
 
-def create_operation(
-    operation_type: str | UndefinedType = Undefined,
-    parameters: dict[str, Any] | UndefinedType = Undefined,
-    **kwargs,
-) -> Operation:
-    """Factory for Operation nodes.
-
-    Args:
-        operation_type: Registry key (required).
-        parameters: Factory arguments dict (default: {}).
-        **kwargs: Additional fields (metadata, timeout, etc.).
-
-    Returns:
-        Unbound Operation ready for bind() and invoke().
-
-    Raises:
-        ValueError: If operation_type not provided.
-    """
-    if is_sentinel(operation_type):
-        raise ValueError("operation_type is required")
-
-    resolved_params: dict[str, Any] = {} if is_sentinel(parameters) else parameters
-
-    return Operation(
-        operation_type=operation_type,
-        parameters=resolved_params,
-        **kwargs,
-    )
+    def __repr__(self) -> str:
+        bound = "bound" if self._session is not None else "unbound"
+        return f"Operation(type={self.operation_type}, status={self.execution.status.value}, {bound})"

krons/work/operations/registry.py

@@ -0,0 +1,103 @@
+# Copyright (c) 2025 - 2026, HaiyangLi <quantocean.li at gmail dot com>
+# SPDX-License-Identifier: Apache-2.0
+
+"""Per-session operation handler registry.
+
+Maps operation names to async handlers. Instantiated per-Session
+for isolation, testability, and per-session customization.
+
+Handler signature: async handler(params, ctx: RequestContext) -> result
+Operation._invoke() creates the RequestContext from bound session/branch.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable
+from typing import Any
+
+__all__ = ("OperationHandler", "OperationRegistry")
+
+OperationHandler = Callable[..., Awaitable[Any]]
+"""Handler signature: async (params, ctx: RequestContext) -> result"""
+
+
+class OperationRegistry:
+    """Map operation names to async handler functions.
+
+    Per-session registry (not global) for isolation and testability.
+
+    Example:
+        from krons.agent.operations import generate, structure, operate
+
+        registry = OperationRegistry()
+        registry.register("generate", generate)
+        registry.register("structure", structure)
+        registry.register("operate", operate)
+
+        # Called by Operation._invoke() — users call session.conduct()
+        handler = registry.get("generate")
+        result = await handler(params, ctx)
+    """
+
+    def __init__(self):
+        """Initialize empty registry."""
+        self._handlers: dict[str, OperationHandler] = {}
+
+    def register(
+        self,
+        operation_name: str,
+        handler: OperationHandler,
+        *,
+        override: bool = False,
+    ) -> None:
+        """Register handler for operation name.
+
+        Args:
+            operation_name: Lookup key (e.g. "generate", "operate").
+            handler: Async (params, ctx) -> result.
+            override: Allow replacing existing. Default False.
+
+        Raises:
+            ValueError: If name exists and override=False.
+        """
+        if operation_name in self._handlers and not override:
+            raise ValueError(
+                f"Operation '{operation_name}' already registered. "
+                "Use override=True to replace."
+            )
+        self._handlers[operation_name] = handler
+
+    def get(self, operation_name: str) -> OperationHandler:
+        """Get handler by name. Raises KeyError with available names if not found."""
+        if operation_name not in self._handlers:
+            raise KeyError(
+                f"Operation '{operation_name}' not registered. "
+                f"Available: {self.list_names()}"
+            )
+        return self._handlers[operation_name]
+
+    def has(self, operation_name: str) -> bool:
+        """Check if name is registered."""
+        return operation_name in self._handlers
+
+    def unregister(self, operation_name: str) -> bool:
+        """Remove registration. Returns True if existed."""
+        if operation_name in self._handlers:
+            del self._handlers[operation_name]
+            return True
+        return False
+
+    def list_names(self) -> list[str]:
+        """Return all registered operation names."""
+        return list(self._handlers.keys())
+
+    def __contains__(self, operation_name: str) -> bool:
+        """Support 'name in registry' syntax."""
+        return operation_name in self._handlers
+
+    def __len__(self) -> int:
+        """Count of registered operations."""
+        return len(self._handlers)
+
+    def __repr__(self) -> str:
+        return f"OperationRegistry(operations={self.list_names()})"

krons/work/report.py

@@ -0,0 +1,268 @@
+# Copyright (c) 2025 - 2026, HaiyangLi <quantocean.li at gmail dot com>
+# SPDX-License-Identifier: Apache-2.0
+
+"""Report - Multi-step workflow orchestration.
+
+A Report orchestrates multiple Forms based on data availability:
+- Schedules forms when their inputs become available
+- Groups forms by branch for sequential execution within branch
+- Propagates outputs between forms
+- Tracks overall workflow completion
+
+This is the scheduling layer that enables data-driven DAG execution.
+
+The Report pattern supports declarative workflow definition:
+    - Fields as class attributes (typed outputs)
+    - form_assignments DSL with branch/resource hints
+    - Implicit dependencies from data flow
+"""
+
+from __future__ import annotations
+
+from collections import defaultdict
+from typing import Any
+
+from pydantic import Field
+
+from krons.core import Element, Pile
+
+from .form import Form, parse_assignment
+
+__all__ = ("Report",)
+
+
+class Report(Element):
+    """Workflow orchestrator - schedules forms based on field availability.
+
+    A Report manages a collection of Forms, executing them as their
+    inputs become available. Forms are grouped by branch - forms on the
+    same branch execute sequentially, different branches can run in parallel.
+
+    Example (simple):
+        report = Report(
+            assignment="context -> final_score",
+            form_assignments=[
+                "context -> analysis",
+                "analysis -> score",
+                "score -> final_score",
+            ],
+        )
+        report.initialize(context="some input")
+
+        while not report.is_complete():
+            for form in report.next_forms():
+                await form.execute(ctx)
+                report.complete_form(form)
+
+    Example (with branches and resources):
+        class HiringBriefReport(Report):
+            role_classification: RoleClassification | None = None
+            strategic_context: StrategicContext | None = None
+
+            assignment: str = "job_input -> executive_summary"
+
+            form_assignments: list[str] = [
+                "classifier: job_input -> role_classification | api:fast",
+                "strategist: job_input, role_classification -> strategic_context | api:synthesis",
+                "writer: strategic_context -> executive_summary | api:reasoning",
+            ]
+
+    Attributes:
+        assignment: Overall workflow 'inputs -> final_outputs'
+        form_assignments: List of form assignments with optional branch/resource
+        input_fields: Workflow input fields
+        output_fields: Workflow output fields
+        forms: All forms in workflow
+        completed_forms: Forms that have finished
+        available_data: Current state of all field values
+    """
+
+    assignment: str = Field(
+        default="",
+        description="Overall workflow: 'inputs -> final_outputs'",
+    )
+    form_assignments: list[str] = Field(
+        default_factory=list,
+        description="List of form assignments: ['branch: a -> b | resource', ...]",
+    )
+
+    input_fields: list[str] = Field(default_factory=list)
+    output_fields: list[str] = Field(default_factory=list)
+
+    forms: Pile[Form] = Field(
+        default_factory=lambda: Pile(item_type=Form),
+        description="All forms in the workflow",
+    )
+    completed_forms: Pile[Form] = Field(
+        default_factory=lambda: Pile(item_type=Form),
+        description="Completed forms",
+    )
+    available_data: dict[str, Any] = Field(
+        default_factory=dict,
+        description="Current state of all field values",
+    )
+
+    # Branch tracking: branch_name -> list of form IDs in order
+    _branch_forms: dict[str, list[Form]] = {}
+    # Track last completed form per branch for sequential execution
+    _branch_progress: dict[str, int] = {}
+
+    def model_post_init(self, _: Any) -> None:
+        """Parse assignment and create forms."""
+        self._branch_forms = defaultdict(list)
+        self._branch_progress = defaultdict(int)
+
+        if not self.assignment:
+            return
+
+        # Parse overall assignment
+        self.input_fields, self.output_fields = parse_assignment(self.assignment)
+
+        # Create forms from form_assignments
+        for fa in self.form_assignments:
+            form = Form(assignment=fa)
+            self.forms.include(form)
+
+            # Track by branch
+            branch = form.branch or "_default"
+            self._branch_forms[branch].append(form)
+
+    def initialize(self, **inputs: Any) -> None:
+        """Provide initial input data.
+
+        Args:
+            **inputs: Initial field values
+
+        Raises:
+            ValueError: If required input is missing
+        """
+        for field in self.input_fields:
+            if field not in inputs:
+                raise ValueError(f"Missing required input: '{field}'")
+            self.available_data[field] = inputs[field]
+
+    def next_forms(self) -> list[Form]:
+        """Get forms that are ready to execute.
+
+        Forms with explicit branches execute sequentially within their branch.
+        Forms without branches (None) execute in parallel based on data availability.
+
+        Returns:
+            List of forms with all inputs available and not yet filled
+        """
+        ready = []
+
+        for branch, forms in self._branch_forms.items():
+            if branch == "_default":
+                # No explicit branch - parallel execution based on data
+                for form in forms:
+                    if form.filled:
+                        continue
+                    form.available_data = self.available_data.copy()
+                    if form.is_workable():
+                        ready.append(form)
+            else:
+                # Explicit branch - sequential execution
+                progress = self._branch_progress[branch]
+
+                # Only consider the next form in this branch
+                if progress < len(forms):
+                    form = forms[progress]
+                    if form.filled:
+                        # Already done, advance progress
+                        self._branch_progress[branch] += 1
+                        continue
+
+                    form.available_data = self.available_data.copy()
+                    if form.is_workable():
+                        ready.append(form)
+
+        return ready
+
+    def complete_form(self, form: Form) -> None:
+        """Mark a form as completed and update available data.
+
+        Args:
+            form: The completed form
+
+        Raises:
+            ValueError: If form is not filled
+        """
+        if not form.filled:
+            raise ValueError("Form is not filled")
+
+        self.completed_forms.include(form)
+
+        # Advance branch progress
+        branch = form.branch or "_default"
+        if branch in self._branch_forms:
+            forms = self._branch_forms[branch]
+            progress = self._branch_progress[branch]
+            if progress < len(forms) and forms[progress].id == form.id:
+                self._branch_progress[branch] += 1
+
+        # Update available data with form outputs
+        output_data = form.get_output_data()
+        self.available_data.update(output_data)
+
+    def is_complete(self) -> bool:
+        """Check if all output fields are available.
+
+        Returns:
+            True if workflow is complete
+        """
+        return all(field in self.available_data for field in self.output_fields)
+
+    def get_deliverable(self) -> dict[str, Any]:
+        """Get final output values.
+
+        Returns:
+            Dict of output field values
+        """
+        return {f: self.available_data.get(f) for f in self.output_fields}
+
+    @property
+    def progress(self) -> tuple[int, int]:
+        """Get progress as (completed, total).
+
+        Returns:
+            Tuple of (completed forms, total forms)
+        """
+        return len(self.completed_forms), len(self.forms)
+
+    def get_forms_by_branch(self, branch: str) -> list[Form]:
+        """Get all forms for a specific branch.
+
+        Args:
+            branch: Branch name
+
+        Returns:
+            List of forms on that branch (in order)
+        """
+        return list(self._branch_forms.get(branch, []))
+
+    def get_forms_by_resource(self, resource: str) -> list[Form]:
+        """Get all forms requiring a specific resource.
+
+        Args:
+            resource: Resource hint (e.g., 'api:fast')
+
+        Returns:
+            List of forms with that resource hint
+        """
+        return [f for f in self.forms if f.resource == resource]
+
+    @property
+    def branches(self) -> list[str]:
+        """Get all branch names in this report."""
+        return list(self._branch_forms.keys())
+
+    @property
+    def resources(self) -> set[str]:
+        """Get all resource hints used in this report."""
+        return {f.resource for f in self.forms if f.resource}
+
+    def __repr__(self) -> str:
+        completed, total = self.progress
+        branches = len(self._branch_forms)
+        return f"Report('{self.assignment}', {completed}/{total} forms, {branches} branches)"