openai-sdk-helpers 0.0.7__py3-none-any.whl → 0.0.9__py3-none-any.whl

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,7 +10,8 @@ import asyncio
 import inspect
 import threading
 from datetime import datetime, timezone
-from typing import Any, Callable, Dict, List, Mapping
+from typing import Any, Awaitable, Callable, Coroutine, cast
+from collections.abc import Mapping
 
 from .enum import AgentEnum
 from ..base import BaseStructure, spec_field
@@ -16,6 +21,14 @@ from .task import TaskStructure
 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 +36,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 +59,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 +77,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 +99,6 @@ class PlanStructure(BaseStructure):
         task : TaskStructure
             Task to append to the plan.
 
-        Returns
-        -------
-        None
-
-        Raises
-        ------
-        None
-
         Examples
         --------
         >>> plan = PlanStructure()
@@ -109,18 +108,23 @@ class PlanStructure(BaseStructure):
 
     def execute(
         self,
-        agent_registry: Mapping[AgentEnum | str, Callable[..., Any]],
+        agent_registry: Mapping[
+            AgentEnum | str, Callable[..., object | Coroutine[Any, Any, object]]
+        ],
         *,
         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``
+            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,8 +137,15 @@ 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
         """
         aggregated_results: list[str] = []
         for task in self.tasks:
@@ -171,7 +182,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,11 +207,14 @@ class PlanStructure(BaseStructure):
     def _run_task(
         task: TaskStructure,
         *,
-        agent_callable: Callable[..., Any],
+        agent_callable: Callable[..., object | Coroutine[Any, Any, object]],
         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
@@ -218,26 +243,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/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(

openai_sdk_helpers/structure/responses.py

@@ -1,8 +1,12 @@
-"""OpenAI response and tool helpers for structured outputs."""
+"""OpenAI response and tool helpers for structured outputs.
 
-from __future__ import annotations
+This module provides helper functions for creating OpenAI API response formats
+and tool definitions from BaseStructure classes. These helpers simplify the
+process of configuring structured outputs for both Assistant and chat
+completion APIs.
+"""
 
-from typing import Type
+from __future__ import annotations
 
 from openai.types.responses.response_format_text_json_schema_config_param import (
     ResponseFormatTextJSONSchemaConfigParam,
@@ -14,10 +18,13 @@ from ..utils import log
 
 
 def assistant_tool_definition(
-    structure: Type[BaseStructure], name: str, description: str
+    structure: type[BaseStructure], name: str, description: str
 ) -> dict:
     """Build a function tool definition for OpenAI Assistants.
 
+    Creates a tool definition compatible with the Assistant API that uses
+    the structure's schema as function parameters.
+
     Parameters
     ----------
     structure : type[BaseStructure]
@@ -30,7 +37,16 @@ def assistant_tool_definition(
     Returns
     -------
     dict
-        Assistant tool definition payload.
+        Assistant tool definition payload in OpenAI format.
+
+    Examples
+    --------
+    >>> from openai_sdk_helpers.structure import BaseStructure
+    >>> tool = assistant_tool_definition(
+    ...     BaseStructure,
+    ...     "process_data",
+    ...     "Process input data"
+    ... )
     """
     log(f"{structure.__name__}::assistant_tool_definition")
     return {
@@ -43,9 +59,12 @@ def assistant_tool_definition(
     }
 
 
-def assistant_format(structure: Type[BaseStructure]) -> dict:
+def assistant_format(structure: type[BaseStructure]) -> dict:
     """Build a response format definition for OpenAI Assistants.
 
+    Creates a response format specification that instructs the Assistant API
+    to return structured output matching the provided schema.
+
     Parameters
     ----------
     structure : type[BaseStructure]
@@ -54,7 +73,11 @@ def assistant_format(structure: Type[BaseStructure]) -> dict:
     Returns
     -------
     dict
-        Assistant response format definition.
+        Assistant response format definition in OpenAI format.
+
+    Examples
+    --------
+    >>> format_def = assistant_format(BaseStructure)
     """
     log(f"{structure.__name__}::assistant_format")
     return {
@@ -67,12 +90,15 @@ def assistant_format(structure: Type[BaseStructure]) -> dict:
 
 
 def response_tool_definition(
-    structure: Type[BaseStructure],
+    structure: type[BaseStructure],
     tool_name: str,
     tool_description: str,
 ) -> dict:
     """Build a tool definition for OpenAI chat completions.
 
+    Creates a function tool definition compatible with the chat completions
+    API, using the structure's schema as parameters.
+
     Parameters
     ----------
     structure : type[BaseStructure]
@@ -85,7 +111,15 @@ def response_tool_definition(
     Returns
     -------
     dict
-        Tool definition payload for chat completions.
+        Tool definition payload for chat completions API.
+
+    Examples
+    --------
+    >>> tool = response_tool_definition(
+    ...     BaseStructure,
+    ...     "analyze",
+    ...     "Analyze data"
+    ... )
     """
     log(f"{structure.__name__}::response_tool_definition")
     return {
@@ -98,9 +132,12 @@ def response_tool_definition(
     }
 
 
-def response_format(structure: Type[BaseStructure]) -> ResponseTextConfigParam:
+def response_format(structure: type[BaseStructure]) -> ResponseTextConfigParam:
     """Build a response format for OpenAI chat completions.
 
+    Creates a response format specification that instructs the chat
+    completions API to return structured output matching the schema.
+
     Parameters
     ----------
     structure : type[BaseStructure]
@@ -109,7 +146,11 @@ def response_format(structure: Type[BaseStructure]) -> ResponseTextConfigParam:
     Returns
     -------
     ResponseTextConfigParam
-        Response format definition.
+        Response format definition for chat completions API.
+
+    Examples
+    --------
+    >>> format_spec = response_format(BaseStructure)
     """
     log(f"{structure.__name__}::response_format")
     response_format_text_JSONSchema_config_param = (

openai_sdk_helpers/structure/summary.py

@@ -1,8 +1,10 @@
-"""Shared structured output models for summaries."""
+"""Structured output models for summaries.
 
-from __future__ import annotations
+This module defines Pydantic models for representing summarization results,
+including topic-level summaries with citations and consolidated text summaries.
+"""
 
-from typing import List
+from __future__ import annotations
 
 from .base import BaseStructure, spec_field
 
@@ -10,10 +12,30 @@ from .base import BaseStructure, spec_field
 class SummaryTopic(BaseStructure):
     """Capture a topic-level summary with supporting citations.
 
+    Represents a single topic or micro-trend identified in source excerpts,
+    along with a summary and supporting citations.
+
+    Attributes
+    ----------
+    topic : str
+        Topic or micro-trend identified in the provided excerpts.
+    summary : str
+        Concise explanation of what the excerpts convey about the topic.
+    citations : list[str]
+        Indices or short quotes that justify the topic summary.
+
     Methods
     -------
     print()
         Return a formatted string representation of the stored fields.
+
+    Examples
+    --------
+    >>> topic = SummaryTopic(
+    ...     topic="AI Trends",
+    ...     summary="Growing adoption of AI",
+    ...     citations=["Source 1", "Source 2"]
+    ... )
     """
 
     topic: str = spec_field(
@@ -26,7 +48,7 @@ class SummaryTopic(BaseStructure):
         default=...,
         description="Concise explanation of what the excerpts convey about the topic.",
     )
-    citations: List[str] = spec_field(
+    citations: list[str] = spec_field(
         "citations",
         default_factory=list,
         description="Indices or short quotes that justify the topic summary.",
@@ -34,12 +56,23 @@ class SummaryTopic(BaseStructure):
 
 
 class SummaryStructure(BaseStructure):
-    """Defines the consolidated summary returned by the summarizer agent.
+    """Consolidated summary returned by the summarizer agent.
+
+    Represents a synthesized summary text derived from multiple source excerpts.
+
+    Attributes
+    ----------
+    text : str
+        Combined summary synthesized from the supplied excerpts.
 
     Methods
     -------
     print()
         Return a formatted string representation of the stored fields.
+
+    Examples
+    --------
+    >>> summary = SummaryStructure(text="This is a summary")
     """
 
     text: str = spec_field(
@@ -50,15 +83,30 @@ class SummaryStructure(BaseStructure):
 
 
 class ExtendedSummaryStructure(SummaryStructure):
-    """Extend ``SummaryStructure`` with optional topic breakdown metadata.
+    """Extended summary with optional topic breakdown metadata.
+
+    Extends SummaryStructure to include topic-level summaries with citations,
+    providing more granular insight into the summarization.
+
+    Attributes
+    ----------
+    metadata : list[SummaryTopic]
+        Optional topic-level summaries with supporting citations.
 
     Methods
     -------
     print()
         Return a formatted string representation of the stored fields.
+
+    Examples
+    --------
+    >>> extended = ExtendedSummaryStructure(
+    ...     text="Overall summary",
+    ...     metadata=[SummaryTopic(topic="T1", summary="S1", citations=[])]
+    ... )
     """
 
-    metadata: List[SummaryTopic] = spec_field(
+    metadata: list[SummaryTopic] = spec_field(
         "metadata",
         default_factory=list,
         description="Optional topic-level summaries with supporting citations.",