agnt5 0.2.2__cp39-abi3-macosx_11_0_arm64.whl → 0.2.4__cp39-abi3-macosx_11_0_arm64.whl

agnt5/function.py

@@ -5,10 +5,11 @@ from __future__ import annotations
 import asyncio
 import functools
 import inspect
-import time
 import uuid
-from typing import Any, Awaitable, Callable, Dict, Optional, TypeVar, Union, cast, get_type_hints
+from typing import Any, Awaitable, Callable, Dict, Optional, TypeVar, Union, cast
 
+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 .exceptions import RetryError
 from .types import BackoffPolicy, BackoffType, FunctionConfig, HandlerFunc, RetryPolicy
@@ -18,13 +19,117 @@ T = TypeVar("T")
 # Global function registry
 _FUNCTION_REGISTRY: Dict[str, FunctionConfig] = {}
 
+class FunctionContext(Context):
+    """
+    Lightweight context for stateless functions.
+
+    AGNT5 Philosophy: Context is a convenience, not a requirement.
+    The best function is one that doesn't need context at all!
+
+    Provides only:
+    - Quick logging (ctx.log())
+    - Execution metadata (run_id, attempt)
+    - Smart retry helper (should_retry())
+    - Non-durable sleep
+
+    Does NOT provide:
+    - Orchestration (task, parallel, gather) - use workflows
+    - State management (get, set, delete) - functions are stateless
+    - Checkpointing (step) - functions are atomic
+    """
+
+    def __init__(
+        self,
+        run_id: str,
+        attempt: int = 0,
+        runtime_context: Optional[Any] = None,
+        retry_policy: Optional[Any] = None,
+    ) -> None:
+        """
+        Initialize function context.
+
+        Args:
+            run_id: Unique execution identifier
+            attempt: Retry attempt number (0-indexed)
+            runtime_context: RuntimeContext for trace correlation
+            retry_policy: RetryPolicy for should_retry() checks
+        """
+        super().__init__(run_id, attempt, runtime_context)
+        self._retry_policy = retry_policy
+
+    # === Quick Logging ===
+
+    def log(self, message: str, **extra) -> None:
+        """
+        Quick logging shorthand with structured data.
+
+        Example:
+            ctx.log("Processing payment", amount=100.50, user_id="123")
+        """
+        self._logger.info(message, extra=extra)
+
+    # === Smart Execution ===
+
+    def should_retry(self, error: Exception) -> bool:
+        """
+        Check if error is retryable based on configured policy.
+
+        Example:
+            try:
+                result = await external_api()
+            except Exception as e:
+                if not ctx.should_retry(e):
+                    raise  # Fail fast for non-retryable errors
+                # Otherwise let retry policy handle it
+                raise
+
+        Returns:
+            True if error is retryable, False otherwise
+        """
+        # TODO: Implement retry policy checks
+        # For now, all errors are retryable (let retry policy handle it)
+        return True
+
+    async def sleep(self, seconds: float) -> None:
+        """
+        Non-durable async sleep.
+
+        For durable sleep across failures, use workflows.
+
+        Args:
+            seconds: Number of seconds to sleep
+        """
+        import asyncio
+        await asyncio.sleep(seconds)
+
+
 
 class FunctionRegistry:
     """Registry for function handlers."""
 
     @staticmethod
     def register(config: FunctionConfig) -> None:
-        """Register a function handler."""
+        """Register a function handler.
+
+        Args:
+            config: Function configuration to register
+
+        Raises:
+            ValueError: If a function with the same name is already registered
+        """
+        # Check for name collision
+        if config.name in _FUNCTION_REGISTRY:
+            existing_config = _FUNCTION_REGISTRY[config.name]
+            existing_module = existing_config.handler.__module__
+            new_module = config.handler.__module__
+
+            raise ValueError(
+                f"Function name collision: '{config.name}' is already registered.\n"
+                f"  Existing: {existing_module}.{existing_config.handler.__name__}\n"
+                f"  New:      {new_module}.{config.handler.__name__}\n"
+                f"Please use a different function name or use name= parameter to specify a unique name."
+            )
+
         _FUNCTION_REGISTRY[config.name] = config
 
     @staticmethod
@@ -43,216 +148,72 @@ class FunctionRegistry:
         _FUNCTION_REGISTRY.clear()
 
 
-def _type_to_json_schema(python_type: Any) -> Dict[str, Any]:
-    """Convert Python type hint to JSON Schema."""
-    # Handle None type
-    if python_type is type(None):
-        return {"type": "null"}
-
-    # Handle basic types
-    if python_type is str:
-        return {"type": "string"}
-    if python_type is int:
-        return {"type": "integer"}
-    if python_type is float:
-        return {"type": "number"}
-    if python_type is bool:
-        return {"type": "boolean"}
-
-    # Handle typing module types
-    origin = getattr(python_type, "__origin__", None)
-
-    if origin is list:
-        args = getattr(python_type, "__args__", ())
-        if args:
-            return {"type": "array", "items": _type_to_json_schema(args[0])}
-        return {"type": "array"}
-
-    if origin is dict:
-        return {"type": "object"}
-
-    if origin is Union:
-        args = getattr(python_type, "__args__", ())
-        # Handle Optional[T] (Union[T, None])
-        if len(args) == 2 and type(None) in args:
-            non_none = args[0] if args[1] is type(None) else args[1]
-            schema = _type_to_json_schema(non_none)
-            # Mark as nullable in JSON Schema
-            return {**schema, "nullable": True}
-        # Handle other unions as anyOf
-        return {"anyOf": [_type_to_json_schema(arg) for arg in args]}
-
-    # Default to object for unknown types
-    return {"type": "object"}
-
-
-def _extract_function_schemas(func: Callable[..., Any]) -> tuple[Optional[Dict[str, Any]], Optional[Dict[str, Any]]]:
-    """Extract input and output schemas from function type hints.
-
-    Returns:
-        Tuple of (input_schema, output_schema) where each is a JSON Schema dict or None
-    """
-    try:
-        # Get type hints
-        hints = get_type_hints(func)
-        sig = inspect.signature(func)
-
-        # Build input schema from parameters (excluding 'ctx')
-        input_properties = {}
-        required_params = []
-
-        for param_name, param in sig.parameters.items():
-            if param_name == "ctx":
-                continue
-
-            # Get type hint for this parameter
-            if param_name in hints:
-                param_type = hints[param_name]
-                input_properties[param_name] = _type_to_json_schema(param_type)
-            else:
-                # No type hint, use generic object
-                input_properties[param_name] = {"type": "object"}
-
-            # Check if parameter is required (no default value)
-            if param.default is inspect.Parameter.empty:
-                required_params.append(param_name)
-
-        input_schema = None
-        if input_properties:
-            input_schema = {
-                "type": "object",
-                "properties": input_properties,
-            }
-            if required_params:
-                input_schema["required"] = required_params
-
-            # Add description from docstring if available
-            if func.__doc__:
-                docstring = inspect.cleandoc(func.__doc__)
-                first_line = docstring.split('\n')[0].strip()
-                if first_line:
-                    input_schema["description"] = first_line
-
-        # Build output schema from return type hint
-        output_schema = None
-        if "return" in hints:
-            return_type = hints["return"]
-            output_schema = _type_to_json_schema(return_type)
-
-        return input_schema, output_schema
-
-    except Exception:
-        # If schema extraction fails, return None schemas
-        return None, None
-
-
-def _extract_function_metadata(func: Callable[..., Any]) -> Dict[str, str]:
-    """Extract metadata from function including description from docstring.
-
-    Returns:
-        Dictionary with metadata fields like 'description'
-    """
-    metadata = {}
-
-    # Extract description from docstring
-    if func.__doc__:
-        # Get first line of docstring as description
-        docstring = inspect.cleandoc(func.__doc__)
-        first_line = docstring.split('\n')[0].strip()
-        if first_line:
-            metadata["description"] = first_line
-
-    return metadata
-
-
-def _calculate_backoff_delay(
-    attempt: int,
-    retry_policy: RetryPolicy,
-    backoff_policy: BackoffPolicy,
-) -> float:
-    """Calculate backoff delay in seconds based on attempt number."""
-    if backoff_policy.type == BackoffType.CONSTANT:
-        delay_ms = retry_policy.initial_interval_ms
-    elif backoff_policy.type == BackoffType.LINEAR:
-        delay_ms = retry_policy.initial_interval_ms * (attempt + 1)
-    else:  # EXPONENTIAL
-        delay_ms = retry_policy.initial_interval_ms * (backoff_policy.multiplier**attempt)
-
-    # Cap at max_interval_ms
-    delay_ms = min(delay_ms, retry_policy.max_interval_ms)
-    return delay_ms / 1000.0  # Convert to seconds
-
-
-async def _execute_with_retry(
-    handler: HandlerFunc,
-    ctx: Context,
-    retry_policy: RetryPolicy,
-    backoff_policy: BackoffPolicy,
-    *args: Any,
-    **kwargs: Any,
-) -> Any:
-    """Execute handler with retry logic."""
-    last_error: Optional[Exception] = None
-
-    for attempt in range(retry_policy.max_attempts):
-        try:
-            # Update context attempt number
-            ctx._attempt = attempt
-
-            # Execute handler
-            result = await handler(ctx, *args, **kwargs)
-            return result
-
-        except Exception as e:
-            last_error = e
-            ctx.logger.warning(
-                f"Function execution failed (attempt {attempt + 1}/{retry_policy.max_attempts}): {e}"
-            )
-
-            # If this was the last attempt, raise RetryError
-            if attempt == retry_policy.max_attempts - 1:
-                raise RetryError(
-                    f"Function failed after {retry_policy.max_attempts} attempts",
-                    attempts=retry_policy.max_attempts,
-                    last_error=e,
-                )
-
-            # Calculate backoff delay
-            delay = _calculate_backoff_delay(attempt, retry_policy, backoff_policy)
-            ctx.logger.info(f"Retrying in {delay:.2f} seconds...")
-            await asyncio.sleep(delay)
-
-    # Should never reach here, but for type safety
-    assert last_error is not None
-    raise RetryError(
-        f"Function failed after {retry_policy.max_attempts} attempts",
-        attempts=retry_policy.max_attempts,
-        last_error=last_error,
-    )
-
-
 def function(
     _func: Optional[Callable[..., Any]] = None,
     *,
     name: Optional[str] = None,
-    retries: Optional[RetryPolicy] = None,
-    backoff: Optional[BackoffPolicy] = None,
+    retries: Optional[Union[int, Dict[str, Any], RetryPolicy]] = None,
+    backoff: Optional[Union[str, Dict[str, Any], BackoffPolicy]] = None,
 ) -> Callable[..., Any]:
     """
     Decorator to mark a function as an AGNT5 durable function.
 
     Args:
         name: Custom function name (default: function's __name__)
-        retries: Retry policy configuration
-        backoff: Backoff policy for retries
+        retries: Retry policy configuration. Can be:
+            - int: max attempts (e.g., 5)
+            - dict: RetryPolicy params (e.g., {"max_attempts": 5, "initial_interval_ms": 1000})
+            - RetryPolicy: full policy object
+        backoff: Backoff policy for retries. Can be:
+            - str: backoff type ("constant", "linear", "exponential")
+            - dict: BackoffPolicy params (e.g., {"type": "exponential", "multiplier": 2.0})
+            - BackoffPolicy: full policy object
+
+    Note:
+        Sync Functions: Synchronous functions are automatically executed in a thread pool
+        to prevent blocking the event loop. This is ideal for I/O-bound operations
+        (requests.get(), file I/O, etc.). For CPU-bound operations or when you need
+        explicit control over concurrency, use async functions instead.
 
     Example:
+        # Basic function with context
         @function
-        async def greet(ctx: Context, name: str) -> str:
+        async def greet(ctx: FunctionContext, name: str) -> str:
+            ctx.log(f"Greeting {name}")  # AGNT5 shorthand!
             return f"Hello, {name}!"
 
-        @function(name="add_numbers", retries=RetryPolicy(max_attempts=5))
-        async def add(ctx: Context, a: int, b: int) -> int:
+        # Simple function without context (optional)
+        @function
+        async def add(a: int, b: int) -> int:
+            return a + b
+
+        # With Pydantic models (automatic validation + rich schemas)
+        from pydantic import BaseModel
+
+        class UserInput(BaseModel):
+            name: str
+            age: int
+
+        class UserOutput(BaseModel):
+            greeting: str
+            is_adult: bool
+
+        @function
+        async def process_user(ctx: FunctionContext, user: UserInput) -> UserOutput:
+            ctx.log(f"Processing user {user.name}")
+            return UserOutput(
+                greeting=f"Hello, {user.name}!",
+                is_adult=user.age >= 18
+            )
+
+        # Simple retry count
+        @function(retries=5)
+        async def with_retries(data: str) -> str:
+            return data.upper()
+
+        # Dict configuration
+        @function(retries={"max_attempts": 5}, backoff="exponential")
+        async def advanced(a: int, b: int) -> int:
             return a + b
     """
 
@@ -260,39 +221,43 @@ def function(
         # Get function name
         func_name = name or func.__name__
 
-        # Validate function signature
+        # Validate function signature and check if context is needed
         sig = inspect.signature(func)
         params = list(sig.parameters.values())
 
-        if not params or params[0].name != "ctx":
-            raise ValueError(
-                f"Function '{func_name}' must have 'ctx: Context' as first parameter"
-            )
+        # Check if function declares 'ctx' parameter
+        needs_context = params and params[0].name == "ctx"
 
         # Convert sync to async if needed
         # Note: Async generators should NOT be wrapped - they need to be returned as-is
         if inspect.iscoroutinefunction(func) or inspect.isasyncgenfunction(func):
             handler_func = cast(HandlerFunc, func)
         else:
-            # Wrap sync function in async
+            # Wrap sync function to run in thread pool (prevents blocking event loop)
             @functools.wraps(func)
             async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
-                return func(*args, **kwargs)
+                loop = asyncio.get_running_loop()
+                # Run sync function in thread pool executor to prevent blocking
+                return await loop.run_in_executor(None, lambda: func(*args, **kwargs))
 
             handler_func = cast(HandlerFunc, async_wrapper)
 
         # Extract schemas from type hints
-        input_schema, output_schema = _extract_function_schemas(func)
+        input_schema, output_schema = extract_function_schemas(func)
 
         # Extract metadata (description, etc.)
-        metadata = _extract_function_metadata(func)
+        metadata = extract_function_metadata(func)
+
+        # Parse retry and backoff policies from flexible formats
+        retry_policy = parse_retry_policy(retries)
+        backoff_policy = parse_backoff_policy(backoff)
 
         # Register function
         config = FunctionConfig(
             name=func_name,
             handler=handler_func,
-            retries=retries or RetryPolicy(),
-            backoff=backoff or BackoffPolicy(),
+            retries=retry_policy,
+            backoff=backoff_policy,
             input_schema=input_schema,
             output_schema=output_schema,
             metadata=metadata,
@@ -302,33 +267,41 @@ def function(
         # Create wrapper with retry logic
         @functools.wraps(func)
         async def wrapper(*args: Any, **kwargs: Any) -> Any:
-            # Create context if not provided
-            if not args or not isinstance(args[0], Context):
-                # Auto-create context for direct function calls
-                ctx = Context(
-                    run_id=f"local-{uuid.uuid4().hex[:8]}",
-                    component_type="function",
-                )
-                # Execute with retry
-                return await _execute_with_retry(
-                    handler_func,
-                    ctx,
-                    config.retries or RetryPolicy(),
-                    config.backoff or BackoffPolicy(),
-                    *args,
-                    **kwargs,
-                )
-            else:
-                # Context provided - use it
+            # Extract or create context based on function signature
+            if needs_context:
+                # Function declares ctx parameter - first argument must be FunctionContext
+                if not args or not isinstance(args[0], FunctionContext):
+                    raise TypeError(
+                        f"Function '{func_name}' requires FunctionContext as first argument. "
+                        f"Usage: await {func_name}(ctx, ...)"
+                    )
                 ctx = args[0]
-                return await _execute_with_retry(
-                    handler_func,
-                    ctx,
-                    config.retries or RetryPolicy(),
-                    config.backoff or BackoffPolicy(),
-                    *args[1:],
-                    **kwargs,
-                )
+                func_args = args[1:]
+            else:
+                # Function doesn't use context - create a minimal one for internal use
+                # But first check if a context was passed anyway (for Worker execution)
+                if args and isinstance(args[0], FunctionContext):
+                    # Context was provided by Worker - use it but don't pass to function
+                    ctx = args[0]
+                    func_args = args[1:]
+                else:
+                    # No context provided - create a default one
+                    ctx = FunctionContext(
+                        run_id=f"local-{uuid.uuid4().hex[:8]}",
+                        retry_policy=retry_policy
+                    )
+                    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,
+            )
 
         # Store config on wrapper for introspection
         wrapper._agnt5_config = config  # type: ignore

agnt5/lm.py

@@ -32,10 +32,13 @@ Supported Providers (via model prefix):
 
 from __future__ import annotations
 
+import json
 from dataclasses import dataclass, field
 from enum import Enum
 from typing import Any, AsyncIterator, Dict, List, Optional
 
+from ._schema_utils import detect_format_type
+
 try:
     from ._core import LanguageModel as RustLanguageModel
     from ._core import LanguageModelConfig as RustLanguageModelConfig
@@ -160,6 +163,39 @@ class GenerateResponse:
     usage: Optional[TokenUsage] = None
     finish_reason: Optional[str] = None
     tool_calls: Optional[List[Dict[str, Any]]] = None
+    _rust_response: Optional[Any] = field(default=None, repr=False)
+
+    @property
+    def structured_output(self) -> Optional[Any]:
+        """Parsed structured output (Pydantic model, dataclass, or dict).
+
+        Returns the parsed object when response_format is specified.
+        This is the recommended property name for accessing structured output.
+
+        Returns:
+            Parsed object according to the specified response_format, or None if not available
+        """
+        if self._rust_response and hasattr(self._rust_response, 'object'):
+            return self._rust_response.object
+        return None
+
+    @property
+    def parsed(self) -> Optional[Any]:
+        """Alias for structured_output (OpenAI SDK compatibility).
+
+        Returns:
+            Same as structured_output
+        """
+        return self.structured_output
+
+    @property
+    def object(self) -> Optional[Any]:
+        """Alias for structured_output.
+
+        Returns:
+            Same as structured_output
+        """
+        return self.structured_output
 
 
 @dataclass
@@ -172,6 +208,7 @@ class GenerateRequest:
     tools: List[ToolDefinition] = field(default_factory=list)
     tool_choice: Optional[ToolChoice] = None
     config: GenerationConfig = field(default_factory=GenerationConfig)
+    response_schema: Optional[str] = None  # JSON-encoded schema for structured output
 
 
 # Internal wrapper for the Rust-backed implementation
@@ -271,14 +308,19 @@ class _LanguageModel:
         if request.config.top_p is not None:
             kwargs["top_p"] = request.config.top_p
 
+        # Pass response schema for structured output if provided
+        if request.response_schema is not None:
+            kwargs["response_schema_kw"] = request.response_schema
+
         # TODO: Add tools and tool_choice support when needed
         # if request.tools:
         #     kwargs["tools"] = self._serialize_tools(request.tools)
         # if request.tool_choice:
         #     kwargs["tool_choice"] = request.tool_choice.value
 
-        # Call Rust implementation
-        rust_response = self._rust_lm.generate(prompt=prompt, **kwargs)
+        # 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
         return self._convert_response(rust_response)
@@ -326,8 +368,9 @@ class _LanguageModel:
         # if request.tool_choice:
         #     kwargs["tool_choice"] = request.tool_choice.value
 
-        # Call Rust implementation (it returns a list of chunks)
-        rust_chunks = self._rust_lm.stream(prompt=prompt, **kwargs)
+        # 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:
@@ -378,6 +421,7 @@ class _LanguageModel:
             usage=usage,
             finish_reason=None,  # TODO: Add finish_reason to Rust response
             tool_calls=None,  # TODO: Add tool calls support
+            _rust_response=rust_response,  # Store for .structured_output access
         )
 
 
@@ -394,6 +438,7 @@ async def generate(
     temperature: Optional[float] = None,
     max_tokens: Optional[int] = None,
     top_p: Optional[float] = None,
+    response_format: Optional[Any] = None,
 ) -> GenerateResponse:
     """Generate text using any LLM provider (simplified API).
 
@@ -408,9 +453,10 @@ async def generate(
         temperature: Sampling temperature (0.0-2.0)
         max_tokens: Maximum tokens to generate
         top_p: Nucleus sampling parameter
+        response_format: Pydantic model, dataclass, or JSON schema dict for structured output
 
     Returns:
-        GenerateResponse with text and usage information
+        GenerateResponse with text, usage, and optional structured output
 
     Examples:
         Simple prompt:
@@ -421,15 +467,21 @@ async def generate(
         ... )
         >>> print(response.text)
 
-        Multi-turn conversation:
+        Structured output with dataclass:
+        >>> from dataclasses import dataclass
+        >>>
+        >>> @dataclass
+        ... class CodeReview:
+        ...     issues: list[str]
+        ...     suggestions: list[str]
+        ...     overall_quality: int
+        >>>
         >>> response = await generate(
-        ...     model="anthropic/claude-3-5-haiku",
-        ...     messages=[
-        ...         {"role": "user", "content": "What is calculus?"},
-        ...         {"role": "assistant", "content": "Calculus is..."},
-        ...         {"role": "user", "content": "Give me an example"}
-        ...     ]
+        ...     model="openai/gpt-4o",
+        ...     prompt="Analyze this code...",
+        ...     response_format=CodeReview
         ... )
+        >>> review = response.structured_output  # Returns dict
     """
     # Validate input
     if not prompt and not messages:
@@ -446,6 +498,12 @@ async def generate(
 
     provider, model_name = model.split('/', 1)
 
+    # Convert response_format to JSON schema if provided
+    response_schema_json = None
+    if response_format is not None:
+        format_type, json_schema = detect_format_type(response_format)
+        response_schema_json = json.dumps(json_schema)
+
     # Create language model client
     lm = _LanguageModel(provider=provider.lower(), default_model=None)
 
@@ -478,6 +536,7 @@ async def generate(
         messages=message_objects,
         system_prompt=system_prompt,
         config=config,
+        response_schema=response_schema_json,
     )
 
     # Generate and return