agnt5 0.2.8a7__cp310-abi3-macosx_11_0_arm64.whl → 0.2.8a9__cp310-abi3-macosx_11_0_arm64.whl

agnt5/client.py

@@ -43,6 +43,8 @@ class Client:
         component: str,
         input_data: Optional[Dict[str, Any]] = None,
         component_type: str = "function",
+        session_id: Optional[str] = None,
+        user_id: Optional[str] = None,
     ) -> Dict[str, Any]:
         """Execute a component synchronously and wait for the result.
 
@@ -52,6 +54,8 @@ class Client:
             component: Name of the component to execute
             input_data: Input data for the component (will be sent as JSON body)
             component_type: Type of component - "function", "workflow", "agent", "tool" (default: "function")
+            session_id: Session identifier for multi-turn conversations (optional)
+            user_id: User identifier for user-scoped memory (optional)
 
         Returns:
             Dictionary containing the component's output
@@ -68,6 +72,12 @@ class Client:
             # Workflow execution (explicit)
             result = client.run("order_fulfillment", {"order_id": "123"}, component_type="workflow")
 
+            # Multi-turn conversation with session
+            result = client.run("chat", {"message": "Hello"}, session_id="session-123")
+
+            # User-scoped memory
+            result = client.run("assistant", {"message": "Help me"}, user_id="user-456")
+
             # No input data
             result = client.run("get_status")
             ```
@@ -78,11 +88,18 @@ class Client:
         # Build URL with component type
         url = urljoin(self.gateway_url + "/", f"v1/run/{component_type}/{component}")
 
+        # Build headers with memory scoping identifiers
+        headers = {"Content-Type": "application/json"}
+        if session_id:
+            headers["X-Session-ID"] = session_id
+        if user_id:
+            headers["X-User-ID"] = user_id
+
         # Make request
         response = self._client.post(
             url,
             json=input_data,
-            headers={"Content-Type": "application/json"},
+            headers=headers,
         )
 
         # Handle errors

agnt5/context.py

@@ -2,11 +2,19 @@
 
 from __future__ import annotations
 
+import contextvars
 import logging
 from typing import Any, Awaitable, Callable, Dict, List, Optional, TypeVar, Union
 
 T = TypeVar("T")
 
+# Task-local context variable for automatic context propagation
+# This is NOT a global variable - contextvars provide task-isolated storage
+# Each asyncio task gets its own independent copy, preventing cross-contamination
+_current_context: contextvars.ContextVar[Optional["Context"]] = contextvars.ContextVar(
+    "_current_context", default=None
+)
+
 
 class _CorrelationFilter(logging.Filter):
     """Inject correlation IDs (run_id, trace_id, span_id) into every log record."""
@@ -81,4 +89,90 @@ class Context:
         return self._logger
 
 
+def get_current_context() -> Optional[Context]:
+    """
+    Get the current execution context from task-local storage.
+
+    This function retrieves the context that was set by the nearest enclosing
+    decorator (@function, @workflow) or Agent.run() call in the current asyncio task.
+
+    Returns:
+        Current Context if available (WorkflowContext, FunctionContext, AgentContext),
+        None if no context is set (e.g., running outside AGNT5 execution)
+
+    Example:
+        >>> ctx = get_current_context()
+        >>> if ctx:
+        ...     ctx.logger.info("Logging from anywhere in the call stack!")
+        ...     runtime = ctx._runtime_context  # Access tracing context
+
+    Note:
+        This uses Python's contextvars which provide task-local (NOT global) storage.
+        Each asyncio task has its own isolated context, preventing cross-contamination
+        between concurrent executions.
+    """
+    return _current_context.get()
+
+
+def set_current_context(ctx: Context) -> contextvars.Token:
+    """
+    Set the current execution context in task-local storage.
+
+    This is typically called by decorators and framework code, not by user code.
+    Returns a token that can be used to reset the context to its previous value.
+
+    Args:
+        ctx: Context to set as current
+
+    Returns:
+        Token for resetting the context later (use with contextvars.Token.reset())
+
+    Example:
+        >>> token = set_current_context(my_context)
+        >>> try:
+        ...     # Context is available via get_current_context()
+        ...     do_work()
+        >>> finally:
+        ...     _current_context.reset(token)  # Restore previous context
+
+    Note:
+        Always use try/finally to ensure context is properly reset, even if
+        an exception occurs. This prevents context leakage between executions.
+    """
+    return _current_context.set(ctx)
+
+
+def get_workflow_context() -> Optional["WorkflowContext"]:
+    """
+    Get the WorkflowContext from the current context or its parent chain.
+
+    This function traverses the context hierarchy to find a WorkflowContext,
+    which is needed for emitting workflow checkpoints from nested contexts
+    like AgentContext or FunctionContext.
+
+    Returns:
+        WorkflowContext if found in the context chain, None otherwise
+
+    Example:
+        >>> # Inside an agent called from a workflow
+        >>> workflow_ctx = get_workflow_context()
+        >>> if workflow_ctx:
+        ...     workflow_ctx._send_checkpoint("workflow.lm.started", {...})
+    """
+    from .workflow import WorkflowContext
+
+    ctx = get_current_context()
+
+    # Traverse up the context chain looking for WorkflowContext
+    while ctx is not None:
+        if isinstance(ctx, WorkflowContext):
+            return ctx
+        # Check if this context has a parent_context attribute
+        if hasattr(ctx, 'parent_context'):
+            ctx = ctx.parent_context
+        else:
+            break
+
+    return None
+
 

agnt5/exceptions.py

@@ -65,6 +65,8 @@ class WaitingForUserInputException(AGNT5Error):
         input_type: Type of input ("text", "approval", or "choice")
         options: List of options for approval/choice inputs
         checkpoint_state: Current workflow state for resume
+        agent_context: Optional agent execution state for agent-level HITL
+            Contains: agent_name, iteration, messages, tool_results, pending_tool_call, etc.
     """
 
     def __init__(
@@ -73,6 +75,7 @@ class WaitingForUserInputException(AGNT5Error):
         input_type: str,
         options: Optional[List[Dict]],
         checkpoint_state: Dict,
+        agent_context: Optional[Dict] = None,
     ) -> None:
         """Initialize WaitingForUserInputException.
 
@@ -81,9 +84,19 @@ class WaitingForUserInputException(AGNT5Error):
             input_type: Type of input - "text", "approval", or "choice"
             options: List of option dicts (for approval/choice)
             checkpoint_state: Workflow state snapshot for resume
+            agent_context: Optional agent execution state for resuming agents
+                Required fields when provided:
+                - agent_name: Name of the agent that paused
+                - iteration: Current iteration number (0-indexed)
+                - messages: LLM conversation history as list of dicts
+                - tool_results: Partial tool results for current iteration
+                - pending_tool_call: The HITL tool call awaiting response
+                - all_tool_calls: All tool calls made so far
+                - model_config: Model settings for resume
         """
         super().__init__(f"Waiting for user input: {question}")
         self.question = question
         self.input_type = input_type
         self.options = options or []
         self.checkpoint_state = checkpoint_state
+        self.agent_context = agent_context

agnt5/function.py

@@ -10,7 +10,7 @@ from typing import Any, Awaitable, Callable, Dict, Optional, TypeVar, Union, cas
 
 from ._retry_utils import execute_with_retry, parse_backoff_policy, parse_retry_policy
 from ._schema_utils import extract_function_metadata, extract_function_schemas
-from .context import Context
+from .context import Context, set_current_context
 from .exceptions import RetryError
 from .types import BackoffPolicy, BackoffType, FunctionConfig, HandlerFunc, RetryPolicy
 
@@ -292,16 +292,23 @@ def function(
                     )
                     func_args = args
 
-            # Execute with retry
-            return await execute_with_retry(
-                handler_func,
-                ctx,
-                config.retries or RetryPolicy(),
-                config.backoff or BackoffPolicy(),
-                needs_context,
-                *func_args,
-                **kwargs,
-            )
+            # Set context in task-local storage for automatic propagation
+            token = set_current_context(ctx)
+            try:
+                # Execute with retry
+                return await execute_with_retry(
+                    handler_func,
+                    ctx,
+                    config.retries or RetryPolicy(),
+                    config.backoff or BackoffPolicy(),
+                    needs_context,
+                    *func_args,
+                    **kwargs,
+                )
+            finally:
+                # Always reset context to prevent leakage
+                from .context import _current_context
+                _current_context.reset(token)
 
         # Store config on wrapper for introspection
         wrapper._agnt5_config = config  # type: ignore

agnt5/lm.py

@@ -39,6 +39,7 @@ from enum import Enum
 from typing import Any, AsyncIterator, Dict, List, Optional
 
 from ._schema_utils import detect_format_type
+from .context import get_current_context
 
 try:
     from ._core import LanguageModel as RustLanguageModel
@@ -366,12 +367,55 @@ class _LanguageModel(LanguageModel):
             # Serialize tool_choice to JSON for Rust
             kwargs["tool_choice"] = json.dumps(request.tool_choice.value)
 
-        # Call Rust implementation - it returns a proper Python coroutine now
-        # Using pyo3-async-runtimes for truly async HTTP calls without blocking
-        rust_response = await self._rust_lm.generate(prompt=prompt, **kwargs)
+        # Pass runtime_context for proper trace linking
+        # Try to get from current context if available
+        current_ctx = get_current_context()
+        if current_ctx and hasattr(current_ctx, '_runtime_context') and current_ctx._runtime_context:
+            kwargs["runtime_context"] = current_ctx._runtime_context
+
+        # Emit checkpoint if called within a workflow context
+        from .context import get_workflow_context
+        workflow_ctx = get_workflow_context()
+        if workflow_ctx:
+            workflow_ctx._send_checkpoint("workflow.lm.started", {
+                "model": model,
+                "provider": self._provider,
+                "temperature": kwargs.get("temperature"),
+                "max_tokens": kwargs.get("max_tokens"),
+            })
 
-        # Convert Rust response to Python
-        return self._convert_response(rust_response)
+        try:
+            # Call Rust implementation - it returns a proper Python coroutine now
+            # Using pyo3-async-runtimes for truly async HTTP calls without blocking
+            rust_response = await self._rust_lm.generate(prompt=prompt, **kwargs)
+
+            # Convert Rust response to Python
+            response = self._convert_response(rust_response)
+
+            # Emit completion checkpoint with usage stats
+            if workflow_ctx:
+                usage_dict = None
+                if response.usage:
+                    usage_dict = {
+                        "prompt_tokens": response.usage.prompt_tokens,
+                        "completion_tokens": response.usage.completion_tokens,
+                        "total_tokens": response.usage.total_tokens,
+                    }
+                workflow_ctx._send_checkpoint("workflow.lm.completed", {
+                    "model": model,
+                    "usage": usage_dict,
+                })
+
+            return response
+        except Exception as e:
+            # Emit error checkpoint for observability
+            if workflow_ctx:
+                workflow_ctx._send_checkpoint("workflow.lm.error", {
+                    "model": model,
+                    "error": str(e),
+                    "error_type": type(e).__name__,
+                })
+            raise
 
     async def stream(self, request: GenerateRequest) -> AsyncIterator[str]:
         """Stream completion from LLM.
@@ -427,14 +471,43 @@ class _LanguageModel(LanguageModel):
             # Serialize tool_choice to JSON for Rust
             kwargs["tool_choice"] = json.dumps(request.tool_choice.value)
 
-        # Call Rust implementation - it returns a proper Python coroutine now
-        # Using pyo3-async-runtimes for truly async streaming without blocking
-        rust_chunks = await self._rust_lm.stream(prompt=prompt, **kwargs)
+        # Emit checkpoint if called within a workflow context
+        from .context import get_workflow_context
+        workflow_ctx = get_workflow_context()
+        if workflow_ctx:
+            workflow_ctx._send_checkpoint("workflow.lm.started", {
+                "model": model,
+                "provider": self._provider,
+                "temperature": kwargs.get("temperature"),
+                "max_tokens": kwargs.get("max_tokens"),
+                "streaming": True,
+            })
 
-        # Yield each chunk
-        for chunk in rust_chunks:
-            if chunk.text:
-                yield chunk.text
+        try:
+            # Call Rust implementation - it returns a proper Python coroutine now
+            # Using pyo3-async-runtimes for truly async streaming without blocking
+            rust_chunks = await self._rust_lm.stream(prompt=prompt, **kwargs)
+
+            # Yield each chunk
+            for chunk in rust_chunks:
+                if chunk.text:
+                    yield chunk.text
+
+            # Emit completion checkpoint after streaming finishes
+            if workflow_ctx:
+                workflow_ctx._send_checkpoint("workflow.lm.completed", {
+                    "model": model,
+                    "streaming": True,
+                })
+        except Exception as e:
+            # Emit error checkpoint for observability
+            if workflow_ctx:
+                workflow_ctx._send_checkpoint("workflow.lm.error", {
+                    "model": model,
+                    "error": str(e),
+                    "error_type": type(e).__name__,
+                })
+            raise
 
     def _build_prompt_messages(self, request: GenerateRequest) -> List[Dict[str, str]]:
         """Build structured message list for Rust.
@@ -603,8 +676,12 @@ async def generate(
         response_schema=response_schema_json,
     )
 
+    # Checkpoints are emitted by _LanguageModel.generate() internally
+    # to avoid duplication. No need to emit them here.
+
     # Generate and return
-    return await lm.generate(request)
+    result = await lm.generate(request)
+    return result
 
 
 async def stream(
@@ -700,6 +777,37 @@ async def stream(
         config=config,
     )
 
-    # Stream and yield chunks
-    async for chunk in lm.stream(request):
-        yield chunk
+    # Emit checkpoint if called within a workflow context
+    from .context import get_workflow_context
+
+    workflow_ctx = get_workflow_context()
+    if workflow_ctx:
+        workflow_ctx._send_checkpoint("workflow.lm.started", {
+            "model": model,
+            "provider": provider,
+            "max_tokens": max_tokens,
+            "temperature": temperature,
+            "streaming": True,
+        })
+
+    try:
+        # Stream and yield chunks
+        async for chunk in lm.stream(request):
+            yield chunk
+
+        # Emit completion checkpoint (note: no usage stats for streaming)
+        if workflow_ctx:
+            workflow_ctx._send_checkpoint("workflow.lm.completed", {
+                "model": model,
+                "streaming": True,
+            })
+    except Exception as e:
+        # Emit error checkpoint for observability
+        if workflow_ctx:
+            workflow_ctx._send_checkpoint("workflow.lm.error", {
+                "model": model,
+                "error": str(e),
+                "error_type": type(e).__name__,
+                "streaming": True,
+            })
+        raise

agnt5/tool.py

@@ -13,7 +13,7 @@ from typing import Any, Awaitable, Callable, Dict, List, Optional, TypeVar, get_
 
 from docstring_parser import parse as parse_docstring
 
-from .context import Context
+from .context import Context, set_current_context
 from .exceptions import ConfigurationError
 from ._telemetry import setup_module_logger
 
@@ -229,26 +229,61 @@ class Tool:
                 f"Tool '{self.name}' requires confirmation but confirmation is not yet implemented"
             )
 
-        # Create span for tool execution with trace linking
-        from ._core import create_span
+        # Emit checkpoint if called within a workflow context
+        from .context import get_workflow_context
 
-        logger.debug(f"Invoking tool '{self.name}' with args: {list(kwargs.keys())}")
-
-        # Create span with runtime_context for parent-child span linking
-        with create_span(
-            self.name,
-            "tool",
-            ctx._runtime_context if hasattr(ctx, "_runtime_context") else None,
-            {
+        workflow_ctx = get_workflow_context()
+        if workflow_ctx:
+            workflow_ctx._send_checkpoint("workflow.tool.started", {
                 "tool.name": self.name,
-                "tool.args": ",".join(kwargs.keys()),
-            },
-        ) as span:
-            # Handler is already async (validated in tool() decorator)
-            result = await self.handler(ctx, **kwargs)
-
-            logger.debug(f"Tool '{self.name}' completed successfully")
-            return result
+                "tool.args": list(kwargs.keys()),
+            })
+
+        # Set context in task-local storage for automatic propagation to nested calls
+        token = set_current_context(ctx)
+        try:
+            try:
+                # Create span for tool execution with trace linking
+                from ._core import create_span
+
+                logger.debug(f"Invoking tool '{self.name}' with args: {list(kwargs.keys())}")
+
+                # Create span with runtime_context for parent-child span linking
+                with create_span(
+                    self.name,
+                    "tool",
+                    ctx._runtime_context if hasattr(ctx, "_runtime_context") else None,
+                    {
+                        "tool.name": self.name,
+                        "tool.args": ",".join(kwargs.keys()),
+                    },
+                ) as span:
+                    # Handler is already async (validated in tool() decorator)
+                    result = await self.handler(ctx, **kwargs)
+
+                    logger.debug(f"Tool '{self.name}' completed successfully")
+
+                    # Emit completion checkpoint
+                    if workflow_ctx:
+                        workflow_ctx._send_checkpoint("workflow.tool.completed", {
+                            "tool.name": self.name,
+                            "tool.success": True,
+                        })
+
+                    return result
+            except Exception as e:
+                # Emit error checkpoint for observability
+                if workflow_ctx:
+                    workflow_ctx._send_checkpoint("workflow.tool.error", {
+                        "tool.name": self.name,
+                        "error": str(e),
+                        "error_type": type(e).__name__,
+                    })
+                raise
+        finally:
+            # Always reset context to prevent leakage
+            from .context import _current_context
+            _current_context.reset(token)
 
     def get_schema(self) -> Dict[str, Any]:
         """
@@ -448,17 +483,18 @@ class AskUserTool(Tool):
         ```
     """
 
-    def __init__(self, context: "WorkflowContext"):  # type: ignore
+    def __init__(self, context: Optional["WorkflowContext"] = None):  # type: ignore
         """
         Initialize AskUserTool.
 
         Args:
-            context: Workflow context with wait_for_user capability
+            context: Optional workflow context with wait_for_user capability.
+                     If not provided, will attempt to get from task-local contextvar.
         """
         # Import here to avoid circular dependency
         from .workflow import WorkflowContext
 
-        if not isinstance(context, WorkflowContext):
+        if context is not None and not isinstance(context, WorkflowContext):
             raise ConfigurationError(
                 "AskUserTool requires a WorkflowContext. "
                 "This tool can only be used within workflows."
@@ -477,13 +513,35 @@ class AskUserTool(Tool):
         Ask user a question and wait for their response.
 
         Args:
-            ctx: Execution context (unused, required by Tool signature)
+            ctx: Execution context (may contain WorkflowContext via contextvar)
             question: Question to ask the user
 
         Returns:
             User's text response
         """
-        return await self.context.wait_for_user(question, input_type="text")
+        # Import here to avoid circular dependency
+        from .workflow import WorkflowContext
+        from .context import get_current_context
+
+        # Use explicit context if provided during __init__
+        workflow_ctx = self.context
+
+        # If not provided, try to get from task-local contextvar
+        if workflow_ctx is None:
+            current = get_current_context()
+            if isinstance(current, WorkflowContext):
+                workflow_ctx = current
+            elif hasattr(current, '_workflow_entity'):
+                # Current context has workflow entity (is WorkflowContext)
+                workflow_ctx = current  # type: ignore
+
+        if workflow_ctx is None:
+            raise ConfigurationError(
+                "AskUserTool requires WorkflowContext. "
+                "Either pass context to __init__ or ensure tool is used within a workflow."
+            )
+
+        return await workflow_ctx.wait_for_user(question, input_type="text")
 
 
 class RequestApprovalTool(Tool):
@@ -516,17 +574,18 @@ class RequestApprovalTool(Tool):
         ```
     """
 
-    def __init__(self, context: "WorkflowContext"):  # type: ignore
+    def __init__(self, context: Optional["WorkflowContext"] = None):  # type: ignore
         """
         Initialize RequestApprovalTool.
 
         Args:
-            context: Workflow context with wait_for_user capability
+            context: Optional workflow context with wait_for_user capability.
+                     If not provided, will attempt to get from task-local contextvar.
         """
         # Import here to avoid circular dependency
         from .workflow import WorkflowContext
 
-        if not isinstance(context, WorkflowContext):
+        if context is not None and not isinstance(context, WorkflowContext):
             raise ConfigurationError(
                 "RequestApprovalTool requires a WorkflowContext. "
                 "This tool can only be used within workflows."
@@ -545,19 +604,41 @@ class RequestApprovalTool(Tool):
         Request approval from user for an action.
 
         Args:
-            ctx: Execution context (unused, required by Tool signature)
+            ctx: Execution context (may contain WorkflowContext via contextvar)
             action: The action requiring approval
             details: Additional details about the action
 
         Returns:
             "approve" or "reject" based on user's decision
         """
+        # Import here to avoid circular dependency
+        from .workflow import WorkflowContext
+        from .context import get_current_context
+
+        # Use explicit context if provided during __init__
+        workflow_ctx = self.context
+
+        # If not provided, try to get from task-local contextvar
+        if workflow_ctx is None:
+            current = get_current_context()
+            if isinstance(current, WorkflowContext):
+                workflow_ctx = current
+            elif hasattr(current, '_workflow_entity'):
+                # Current context has workflow entity (is WorkflowContext)
+                workflow_ctx = current  # type: ignore
+
+        if workflow_ctx is None:
+            raise ConfigurationError(
+                "RequestApprovalTool requires WorkflowContext. "
+                "Either pass context to __init__ or ensure tool is used within a workflow."
+            )
+
         question = f"Action: {action}"
         if details:
             question += f"\n\nDetails:\n{details}"
         question += "\n\nDo you approve?"
 
-        return await self.context.wait_for_user(
+        return await workflow_ctx.wait_for_user(
             question,
             input_type="approval",
             options=[