openai-sdk-helpers 0.0.8__py3-none-any.whl → 0.1.0__py3-none-any.whl

openai_sdk_helpers/structure/plan/helpers.py

@@ -0,0 +1,172 @@
+"""Helper functions for creating and executing agent plans.
+
+This module provides convenience functions for working with PlanStructure
+and TaskStructure, simplifying common workflows like plan creation, task
+execution, and result aggregation.
+"""
+
+from __future__ import annotations
+
+from .enum import AgentEnum
+from .plan import PlanStructure
+from .task import TaskStructure
+from .types import AgentCallable, AgentRegistry
+
+
+def create_plan(*tasks: TaskStructure) -> PlanStructure:
+    """Create a PlanStructure from a sequence of tasks.
+
+    Convenience factory function that constructs a plan from individual
+    tasks. Tasks are executed in the order they are provided.
+
+    Parameters
+    ----------
+    *tasks : TaskStructure
+        Variable number of task definitions to include in the plan.
+
+    Returns
+    -------
+    PlanStructure
+        New plan containing the provided tasks in order.
+
+    Examples
+    --------
+    >>> task1 = TaskStructure(
+    ...     task_type=AgentEnum.WEB_SEARCH,
+    ...     prompt="Search for AI trends"
+    ... )
+    >>> task2 = TaskStructure(
+    ...     task_type=AgentEnum.SUMMARIZER,
+    ...     prompt="Summarize findings"
+    ... )
+    >>> plan = create_plan(task1, task2)
+    >>> len(plan)
+    2
+    """
+    return PlanStructure(tasks=list(tasks))
+
+
+def execute_task(
+    task: TaskStructure,
+    agent_callable: AgentCallable,
+) -> list[str]:
+    """Execute a single task with an agent callable.
+
+    Runs one task using the provided agent function. Updates task status,
+    timing, and results. Context from previous tasks is not supported in this
+    helper - use execute_plan() for multi-task execution with context passing.
+
+    Parameters
+    ----------
+    task : TaskStructure
+        Task definition containing prompt and metadata.
+    agent_callable : AgentCallable
+        Synchronous or asynchronous callable responsible for executing the task.
+        Should accept the task prompt and an optional context keyword argument.
+
+    Returns
+    -------
+    list[str]
+        Normalized string results from task execution.
+
+    Raises
+    ------
+    Exception
+        Any exception raised by the agent_callable is propagated after
+        task status is updated.
+
+    Examples
+    --------
+    >>> def agent_fn(prompt, context=None):
+    ...     return f"Result for {prompt}"
+    >>> task = TaskStructure(prompt="Test task")
+    >>> results = execute_task(task, agent_fn)
+    >>> task.status
+    'done'
+    """
+    from datetime import datetime, timezone
+
+    task.start_date = datetime.now(timezone.utc)
+    task.status = "running"
+
+    # Build plan with single task and execute
+    # Normalize task_type to string value for registry key to match PlanStructure.execute lookup
+    plan = PlanStructure(tasks=[task])
+    # Convert AgentEnum to its string value for registry key
+    registry_key = (
+        task.task_type.value
+        if isinstance(task.task_type, AgentEnum)
+        else task.task_type
+    )
+    registry: dict[str, AgentCallable] = {
+        registry_key: agent_callable,
+    }
+
+    # Execute the plan - it will update task status
+    aggregated = plan.execute(
+        agent_registry=registry,
+        halt_on_error=True,
+    )
+
+    # If task failed, raise the exception
+    if task.status == "error":
+        # Extract error message from results
+        error_msg = task.results[0] if task.results else "Task execution failed"
+        # Raise RuntimeError with the error message
+        # The original exception type information is lost but the message is preserved
+        raise RuntimeError(f"Task execution error: {error_msg}")
+
+    return aggregated
+
+
+def execute_plan(
+    plan: PlanStructure,
+    agent_registry: AgentRegistry,
+    halt_on_error: bool = True,
+) -> list[str]:
+    """Execute a plan using registered agent callables.
+
+    Convenience wrapper around PlanStructure.execute() for cleaner syntax.
+    Runs all tasks in sequence, passing results between tasks as context.
+
+    Parameters
+    ----------
+    plan : PlanStructure
+        Plan containing ordered tasks to execute.
+    agent_registry : AgentRegistry
+        Lookup of agent identifiers to callables. Keys may be AgentEnum
+        instances or their string values.
+    halt_on_error : bool, default True
+        Whether execution should stop when a task raises an exception.
+
+    Returns
+    -------
+    list[str]
+        Flattened list of normalized outputs from all executed tasks.
+
+    Raises
+    ------
+    KeyError
+        If a task references an agent not in the registry.
+
+    Examples
+    --------
+    >>> def search_agent(prompt, context=None):
+    ...     return ["search results"]
+    >>> def summary_agent(prompt, context=None):
+    ...     return ["summary"]
+    >>> registry = {
+    ...     AgentEnum.WEB_SEARCH: search_agent,
+    ...     AgentEnum.SUMMARIZER: summary_agent,
+    ... }
+    >>> plan = PlanStructure(tasks=[...])  # doctest: +SKIP
+    >>> results = execute_plan(plan, registry)  # doctest: +SKIP
+    """
+    return plan.execute(agent_registry, halt_on_error=halt_on_error)
+
+
+__all__ = [
+    "create_plan",
+    "execute_task",
+    "execute_plan",
+]

openai_sdk_helpers/structure/plan/plan.py

@@ -1,4 +1,8 @@
-"""Structured output model for agent plans."""
+"""Structured output model for agent plans.
+
+This module defines a Pydantic model for representing ordered lists of agent
+tasks, with support for sequential execution and result aggregation.
+"""
 
 from __future__ import annotations
 
@@ -6,16 +10,26 @@ import asyncio
 import inspect
 import threading
 from datetime import datetime, timezone
-from typing import Any, Callable, Dict, List, Mapping
+from typing import Any, Awaitable, Coroutine, cast
+from collections.abc import Mapping
 
 from .enum import AgentEnum
 from ..base import BaseStructure, spec_field
 from .task import TaskStructure
+from .types import AgentCallable, AgentRegistry
 
 
 class PlanStructure(BaseStructure):
     """Structured representation of an ordered list of agent tasks.
 
+    Represents a complete execution plan consisting of multiple agent tasks
+    to be run sequentially, with support for context passing between tasks.
+
+    Attributes
+    ----------
+    tasks : list[TaskStructure]
+        Ordered list of agent tasks to execute.
+
     Methods
     -------
     print()
@@ -23,12 +37,21 @@ class PlanStructure(BaseStructure):
     __len__()
         Return the count of tasks in the plan.
     append(task)
-        Append an ``TaskStructure`` to the plan.
+        Append a TaskStructure to the plan.
     execute(agent_registry, halt_on_error)
         Run tasks sequentially using the provided agent callables.
+
+    Examples
+    --------
+    >>> plan = PlanStructure(tasks=[
+    ...     TaskStructure(prompt="Task 1"),
+    ...     TaskStructure(prompt="Task 2")
+    ... ])
+    >>> len(plan)
+    2
     """
 
-    tasks: List[TaskStructure] = spec_field(
+    tasks: list[TaskStructure] = spec_field(
         "tasks",
         default_factory=list,
         description="Ordered list of agent tasks to execute.",
@@ -37,22 +60,15 @@ class PlanStructure(BaseStructure):
     def print(self) -> str:
         """Return a human-readable representation of the plan.
 
-        Parameters
-        ----------
-        None
-
         Returns
         -------
         str
-            Concatenated description of each plan step.
-
-        Raises
-        ------
-        None
+            Concatenated description of each task with task numbers.
 
         Examples
         --------
-        >>> PlanStructure().print()
+        >>> plan = PlanStructure()
+        >>> plan.print()
         'No tasks defined.'
         """
         if not self.tasks:
@@ -62,21 +78,13 @@ class PlanStructure(BaseStructure):
         )
 
     def __len__(self) -> int:
-        """Return the number of tasks contained in the plan.
-
-        Parameters
-        ----------
-        None
+        """Return the number of tasks in the plan.
 
         Returns
         -------
         int
             Count of stored agent tasks.
 
-        Raises
-        ------
-        None
-
         Examples
         --------
         >>> len(PlanStructure())
@@ -92,14 +100,6 @@ class PlanStructure(BaseStructure):
         task : TaskStructure
             Task to append to the plan.
 
-        Returns
-        -------
-        None
-
-        Raises
-        ------
-        None
-
         Examples
         --------
         >>> plan = PlanStructure()
@@ -109,18 +109,21 @@ class PlanStructure(BaseStructure):
 
     def execute(
         self,
-        agent_registry: Mapping[AgentEnum | str, Callable[..., Any]],
+        agent_registry: AgentRegistry,
         *,
         halt_on_error: bool = True,
     ) -> list[str]:
         """Execute tasks with registered agent callables and record outputs.
 
+        Runs each task in sequence, passing results as context to subsequent
+        tasks. Updates task status, timing, and results as execution proceeds.
+
         Parameters
         ----------
-        agent_registry : Mapping[AgentEnum | str, Callable[..., Any]]
-            Lookup of agent identifiers to callables. Keys may be ``AgentEnum``
+        agent_registry : AgentRegistry
+            Lookup of agent identifiers to callables. Keys may be AgentEnum
             instances or their string values. Each callable receives the task
-            prompt (augmented with prior context) and an optional ``context``
+            prompt (augmented with prior context) and an optional context
             keyword containing accumulated results.
         halt_on_error : bool, default=True
             Whether execution should stop when a task raises an exception.
@@ -133,16 +136,28 @@ class PlanStructure(BaseStructure):
         Raises
         ------
         KeyError
-            If a task does not have a corresponding callable in
-            ``agent_registry``.
+            If a task does not have a corresponding callable in agent_registry.
+
+        Examples
+        --------
+        >>> def agent_fn(prompt, context=None):
+        ...     return f"Result for {prompt}"
+        >>> registry = {AgentEnum.WEB_SEARCH: agent_fn}
+        >>> plan = PlanStructure(tasks=[TaskStructure(prompt="Test")])
+        >>> results = plan.execute(registry)  # doctest: +SKIP
         """
+        normalized_registry: dict[str, AgentCallable] = {
+            self._resolve_registry_key(key): value
+            for key, value in agent_registry.items()
+        }
+
         aggregated_results: list[str] = []
         for task in self.tasks:
             callable_key = self._resolve_registry_key(task.task_type)
-            if callable_key not in agent_registry:
+            if callable_key not in normalized_registry:
                 raise KeyError(f"No agent registered for '{callable_key}'.")
 
-            agent_callable = agent_registry[callable_key]
+            agent_callable = normalized_registry[callable_key]
             task.start_date = datetime.now(timezone.utc)
             task.status = "running"
 
@@ -171,7 +186,18 @@ class PlanStructure(BaseStructure):
 
     @staticmethod
     def _resolve_registry_key(task_type: AgentEnum | str) -> str:
-        """Return a normalized registry key for the given ``task_type``."""
+        """Return a normalized registry key for the given task_type.
+
+        Parameters
+        ----------
+        task_type : AgentEnum | str
+            Task type to normalize.
+
+        Returns
+        -------
+        str
+            Normalized key for agent registry lookup.
+        """
         if isinstance(task_type, AgentEnum):
             return task_type.value
         if task_type in AgentEnum.__members__:
@@ -185,16 +211,19 @@ class PlanStructure(BaseStructure):
     def _run_task(
         task: TaskStructure,
         *,
-        agent_callable: Callable[..., Any],
+        agent_callable: AgentCallable,
         aggregated_context: list[str],
-    ) -> Any:
+    ) -> object | Coroutine[Any, Any, object]:
         """Execute a single task using the supplied callable.
 
+        Combines task context with aggregated results from previous tasks,
+        then invokes the agent callable with the augmented prompt.
+
         Parameters
         ----------
         task : TaskStructure
             Task definition containing inputs and metadata.
-        agent_callable : Callable[..., Any]
+        agent_callable : AgentCallable
             Function responsible for performing the task.
         aggregated_context : list[str]
             Accumulated results from previously executed tasks.
@@ -218,26 +247,57 @@ class PlanStructure(BaseStructure):
             return agent_callable(prompt_with_context)
 
     @staticmethod
-    def _normalize_results(result: Any) -> list[str]:
-        """Convert callable outputs into a list of strings."""
+    def _normalize_results(result: object | Coroutine[Any, Any, object]) -> list[str]:
+        """Convert callable outputs into a list of strings.
+
+        Handles various result types including None, awaitables, lists,
+        and single values.
+
+        Parameters
+        ----------
+        result : Any
+            Raw result from agent callable.
+
+        Returns
+        -------
+        list[str]
+            Normalized list of string results.
+        """
         if result is None:
             return []
         if inspect.isawaitable(result):
-            return PlanStructure._normalize_results(PlanStructure._await_result(result))
+            awaited = PlanStructure._await_result(
+                cast(Coroutine[Any, Any, object], result)
+            )
+            return PlanStructure._normalize_results(awaited)
         if isinstance(result, list):
             return [str(item) for item in result]
         return [str(result)]
 
     @staticmethod
-    def _await_result(result: Any) -> Any:
-        """Await the provided result, handling running event loops."""
+    def _await_result(result: Coroutine[Any, Any, object]) -> object:
+        """Await the provided result, handling running event loops.
+
+        Properly handles awaiting results whether an event loop is running
+        or not, using a separate thread when necessary.
+
+        Parameters
+        ----------
+        result : Any
+            Awaitable result to resolve.
+
+        Returns
+        -------
+        Any
+            Resolved value from the awaitable.
+        """
         try:
             loop = asyncio.get_running_loop()
         except RuntimeError:
             return asyncio.run(result)
 
         if loop.is_running():
-            container: Dict[str, Any] = {"value": None}
+            container: dict[str, object | None] = {"value": None}
 
             def _runner() -> None:
                 container["value"] = asyncio.run(result)

openai_sdk_helpers/structure/plan/task.py

@@ -1,9 +1,13 @@
-"""Structured output model for agent tasks."""
+"""Structured output model for agent tasks.
+
+This module defines a Pydantic model for representing individual agent tasks
+within a plan, including task type, inputs, status tracking, and results.
+"""
 
 from __future__ import annotations
 
 from datetime import datetime
-from typing import List, Literal, Optional
+from typing import Literal
 
 from pydantic import field_validator
 
@@ -14,10 +18,38 @@ from ..base import BaseStructure, spec_field
 class TaskStructure(BaseStructure):
     """Structured representation of a single agent task.
 
+    Represents one task in an agent execution plan, including its type,
+    inputs, execution status, timing, and results.
+
+    Attributes
+    ----------
+    task_type : AgentEnum
+        Agent type responsible for executing the task.
+    prompt : str
+        Input passed to the agent.
+    context : list[str] or None
+        Additional context forwarded to the agent callable.
+    start_date : datetime or None
+        Timestamp marking when the task started (UTC).
+    end_date : datetime or None
+        Timestamp marking when the task completed (UTC).
+    status : Literal["waiting", "running", "done", "error"]
+        Current lifecycle state for the task.
+    results : list[str]
+        Normalized string outputs returned by the agent.
+
     Methods
     -------
     print()
         Return a formatted multi-line description of the task.
+
+    Examples
+    --------
+    >>> task = TaskStructure(
+    ...     task_type=AgentEnum.WEB_SEARCH,
+    ...     prompt="Research AI trends",
+    ...     status="waiting"
+    ... )
     """
 
     task_type: AgentEnum = spec_field(
@@ -30,17 +62,17 @@ class TaskStructure(BaseStructure):
         description="Input passed to the agent.",
         examples=["Research the latest trends in AI-assisted data analysis."],
     )
-    context: List[str] | None = spec_field(
+    context: list[str] | None = spec_field(
         "context",
         default_factory=list,
         description="Additional context forwarded to the agent callable.",
     )
-    start_date: Optional[datetime] = spec_field(
+    start_date: datetime | None = spec_field(
         "start_date",
         default=None,
         description="Timestamp marking when the task started (UTC).",
     )
-    end_date: Optional[datetime] = spec_field(
+    end_date: datetime | None = spec_field(
         "end_date",
         default=None,
         description="Timestamp marking when the task completed (UTC).",
@@ -50,7 +82,7 @@ class TaskStructure(BaseStructure):
         default="waiting",
         description="Current lifecycle state for the task.",
     )
-    results: List[str] = spec_field(
+    results: list[str] = spec_field(
         "results",
         default_factory=list,
         description="Normalized string outputs returned by the agent.",

openai_sdk_helpers/structure/plan/types.py

@@ -0,0 +1,15 @@
+"""Type aliases for plan execution helpers."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any, Callable, Coroutine, TypeAlias
+
+from .enum import AgentEnum
+
+AgentCallable = Callable[..., object | Coroutine[Any, Any, object]]
+AgentRegistry: TypeAlias = (
+    Mapping[str, AgentCallable] | Mapping[AgentEnum, AgentCallable]
+)
+
+__all__ = ["AgentCallable", "AgentRegistry"]

openai_sdk_helpers/structure/prompt.py

@@ -1,4 +1,8 @@
-"""Shared structured output model for prompts."""
+"""Structured output model for prompts.
+
+This module defines a simple Pydantic model for representing prompt text
+used in OpenAI API requests.
+"""
 
 from __future__ import annotations
 
@@ -6,12 +10,27 @@ from .base import BaseStructure, spec_field
 
 
 class PromptStructure(BaseStructure):
-    """The prompt text to use for the OpenAI API request.
+    """Structured representation of prompt text for OpenAI API requests.
+
+    Simple structure containing a single prompt string with examples.
+
+    Attributes
+    ----------
+    prompt : str
+        The prompt text to use for the OpenAI API request.
 
     Methods
     -------
     print()
         Return the formatted model fields.
+
+    Examples
+    --------
+    >>> prompt_struct = PromptStructure(
+    ...     prompt="What is the capital of France?"
+    ... )
+    >>> print(prompt_struct.prompt)
+    'What is the capital of France?'
     """
 
     prompt: str = spec_field(