openai-sdk-helpers 0.3.0__py3-none-any.whl → 0.4.0__py3-none-any.whl

openai_sdk_helpers/agent/search/base.py

@@ -12,29 +12,54 @@ from abc import ABC, abstractmethod
 from pathlib import Path
 from typing import Generic, List, Optional, TypeVar, Union
 
-from ..base import BaseAgent
+from ..base import AgentBase
 from ..config import AgentConfiguration
+from ...structure.base import StructureBase
 
 # Type variables for search workflow components
-ItemType = TypeVar("ItemType")  # Search item structure (e.g., WebSearchItemStructure)
+ItemType = TypeVar(
+    "ItemType", bound=StructureBase
+)  # Search item structure (e.g., WebSearchItemStructure)
 ResultType = TypeVar("ResultType")  # Individual search result
-PlanType = TypeVar("PlanType")  # Complete search plan structure
-ReportType = TypeVar("ReportType")  # Final report structure
-OutputType = TypeVar("OutputType")  # Generic output type
+PlanType = TypeVar("PlanType", bound=StructureBase)  # Complete search plan structure
+ReportType = TypeVar("ReportType", bound=StructureBase)  # Final report structure
 
 
-class SearchPlanner(BaseAgent, Generic[PlanType]):
+class SearchPlanner(AgentBase, Generic[PlanType]):
     """Generic planner agent for search workflows.
 
     Subclasses implement specific planner logic by overriding the
     `_configure_agent` method and specifying the output type.
 
+    Parameters
+    ----------
+    prompt_dir : Path, optional
+        Directory containing prompt templates.
+    default_model : str, optional
+        Default model identifier to use when not defined in config.
+
     Methods
     -------
     run_agent(query)
         Generate a search plan for the provided query.
     _configure_agent()
         Return AgentConfiguration for this planner instance.
+
+    Raises
+    ------
+    ValueError
+        If the default model is not provided.
+
+    Examples
+    --------
+    >>> class MyPlanner(SearchPlanner):
+    ...     def _configure_agent(self):
+    ...         return AgentConfiguration(
+    ...             name="my_planner",
+    ...             description="Plans searches",
+    ...             output_structure=WebSearchPlanStructure,
+    ...         )
+    >>> planner = MyPlanner(default_model="gpt-4o-mini")
     """
 
     def __init__(
@@ -42,15 +67,7 @@ class SearchPlanner(BaseAgent, Generic[PlanType]):
         prompt_dir: Optional[Path] = None,
         default_model: Optional[str] = None,
     ) -> None:
-        """Initialize the planner agent.
-
-        Parameters
-        ----------
-        prompt_dir : Path, optional
-            Directory containing prompt templates.
-        default_model : str, optional
-            Default model identifier to use when not defined in config.
-        """
+        """Initialize the planner agent."""
         config = self._configure_agent()
         super().__init__(
             config=config,
@@ -65,14 +82,14 @@ class SearchPlanner(BaseAgent, Generic[PlanType]):
         Returns
         -------
         AgentConfiguration
-            Configuration with name, description, and output_type set.
+            Configuration with name, description, and output_structure set.
 
         Examples
         --------
         >>> config = AgentConfiguration(
         ...     name="web_planner",
         ...     description="Plan web searches",
-        ...     output_type=WebSearchPlanStructure,
+        ...     output_structure=WebSearchPlanStructure,
         ... )
         >>> return config
         """
@@ -93,18 +110,27 @@ class SearchPlanner(BaseAgent, Generic[PlanType]):
         """
         result: PlanType = await self.run_async(
             input=query,
-            output_type=self._output_type,
+            output_structure=self._output_structure,
         )
         return result
 
 
-class SearchToolAgent(BaseAgent, Generic[ItemType, ResultType, PlanType]):
+class SearchToolAgent(AgentBase, Generic[ItemType, ResultType, PlanType]):
     """Generic tool agent for executing search workflows.
 
     Executes individual searches in a plan with concurrency control.
     Subclasses implement search execution logic by overriding the
     `_configure_agent` and `run_search` methods.
 
+    Parameters
+    ----------
+    prompt_dir : Path, optional
+        Directory containing prompt templates.
+    default_model : str, optional
+        Default model identifier to use when not defined in config.
+    max_concurrent_searches : int, default=10
+        Maximum number of concurrent search operations.
+
     Methods
     -------
     run_agent(search_plan)
@@ -113,6 +139,24 @@ class SearchToolAgent(BaseAgent, Generic[ItemType, ResultType, PlanType]):
         Execute a single search item.
     _configure_agent()
         Return AgentConfiguration for this tool agent.
+
+    Raises
+    ------
+    ValueError
+        If the default model is not provided.
+
+    Examples
+    --------
+    >>> class MyTool(SearchToolAgent):
+    ...     def _configure_agent(self):
+    ...         return AgentConfiguration(
+    ...             name="my_tool",
+    ...             description="Executes searches",
+    ...             input_structure=WebSearchPlanStructure,
+    ...         )
+    ...     async def run_search(self, item):
+    ...         return "result"
+    >>> tool = MyTool(default_model="gpt-4o-mini")
     """
 
     def __init__(
@@ -122,17 +166,7 @@ class SearchToolAgent(BaseAgent, Generic[ItemType, ResultType, PlanType]):
         default_model: Optional[str] = None,
         max_concurrent_searches: int = 10,
     ) -> None:
-        """Initialize the search tool agent.
-
-        Parameters
-        ----------
-        prompt_dir : Path, optional
-            Directory containing prompt templates.
-        default_model : str, optional
-            Default model identifier to use when not defined in config.
-        max_concurrent_searches : int, default=10
-            Maximum number of concurrent search operations.
-        """
+        """Initialize the search tool agent."""
         self._max_concurrent_searches = max_concurrent_searches
         config = self._configure_agent()
         super().__init__(
@@ -148,14 +182,14 @@ class SearchToolAgent(BaseAgent, Generic[ItemType, ResultType, PlanType]):
         Returns
         -------
         AgentConfiguration
-            Configuration with name, description, input_type, and tools set.
+            Configuration with name, description, input_structure, and tools set.
 
         Examples
         --------
         >>> config = AgentConfiguration(
         ...     name="web_search",
         ...     description="Perform web searches",
-        ...     input_type=WebSearchPlanStructure,
+        ...     input_structure=WebSearchPlanStructure,
         ...     tools=[WebSearchTool()],
         ... )
         >>> return config
@@ -205,18 +239,41 @@ class SearchToolAgent(BaseAgent, Generic[ItemType, ResultType, PlanType]):
         return [result for result in results if result is not None]
 
 
-class SearchWriter(BaseAgent, Generic[ReportType]):
+class SearchWriter(AgentBase, Generic[ReportType]):
     """Generic writer agent for search workflow reports.
 
     Synthesizes search results into a final report. Subclasses implement
     specific report generation logic by overriding the `_configure_agent` method.
 
+    Parameters
+    ----------
+    prompt_dir : Path, optional
+        Directory containing prompt templates.
+    default_model : str, optional
+        Default model identifier to use when not defined in config.
+
     Methods
     -------
     run_agent(query, search_results)
         Generate a report from search results.
     _configure_agent()
         Return AgentConfiguration for this writer instance.
+
+    Raises
+    ------
+    ValueError
+        If the default model is not provided.
+
+    Examples
+    --------
+    >>> class MyWriter(SearchWriter):
+    ...     def _configure_agent(self):
+    ...         return AgentConfiguration(
+    ...             name="my_writer",
+    ...             description="Writes reports",
+    ...             output_structure=WebSearchReportStructure,
+    ...         )
+    >>> writer = MyWriter(default_model="gpt-4o-mini")
     """
 
     def __init__(
@@ -224,15 +281,7 @@ class SearchWriter(BaseAgent, Generic[ReportType]):
         prompt_dir: Optional[Path] = None,
         default_model: Optional[str] = None,
     ) -> None:
-        """Initialize the writer agent.
-
-        Parameters
-        ----------
-        prompt_dir : Path, optional
-            Directory containing prompt templates.
-        default_model : str, optional
-            Default model identifier to use when not defined in config.
-        """
+        """Initialize the writer agent."""
         config = self._configure_agent()
         super().__init__(
             config=config,
@@ -247,14 +296,14 @@ class SearchWriter(BaseAgent, Generic[ReportType]):
         Returns
         -------
         AgentConfiguration
-            Configuration with name, description, and output_type set.
+            Configuration with name, description, and output_structure set.
 
         Examples
         --------
         >>> config = AgentConfiguration(
         ...     name="web_writer",
         ...     description="Write web search report",
-        ...     output_type=WebSearchReportStructure,
+        ...     output_structure=WebSearchReportStructure,
         ... )
         >>> return config
         """
@@ -286,7 +335,7 @@ class SearchWriter(BaseAgent, Generic[ReportType]):
         result: ReportType = await self.run_async(
             input=query,
             context=template_context,
-            output_type=self._output_type,
+            output_structure=self._output_structure,
         )
         return result
 

openai_sdk_helpers/agent/search/vector.py

@@ -26,24 +26,32 @@ MAX_CONCURRENT_SEARCHES = 10
 class VectorSearchPlanner(SearchPlanner[VectorSearchPlanStructure]):
     """Plan vector searches to satisfy a user query.
 
+    Parameters
+    ----------
+    prompt_dir : Path or None, default=None
+        Directory containing prompt templates.
+    default_model : str or None, default=None
+        Default model identifier to use when not defined in config.
+
     Methods
     -------
     run_agent(query)
         Generate a vector search plan for the provided query.
+
+    Raises
+    ------
+    ValueError
+        If the default model is not provided.
+
+    Examples
+    --------
+    >>> planner = VectorSearchPlanner(default_model="gpt-4o-mini")
     """
 
     def __init__(
         self, prompt_dir: Optional[Path] = None, default_model: Optional[str] = None
     ) -> None:
-        """Initialize the planner agent.
-
-        Parameters
-        ----------
-        prompt_dir : Path or None, default=None
-            Directory containing prompt templates.
-        default_model : str or None, default=None
-            Default model identifier to use when not defined in config.
-        """
+        """Initialize the planner agent."""
         super().__init__(prompt_dir=prompt_dir, default_model=default_model)
 
     def _configure_agent(self) -> AgentConfiguration:
@@ -58,7 +66,7 @@ class VectorSearchPlanner(SearchPlanner[VectorSearchPlanStructure]):
             name="vector_planner",
             instructions="Agent instructions",
             description="Plan vector searches based on a user query.",
-            output_type=VectorSearchPlanStructure,
+            output_structure=VectorSearchPlanStructure,
         )
 
 
@@ -71,12 +79,37 @@ class VectorSearchTool(
 ):
     """Execute vector searches defined in a search plan.
 
+    Parameters
+    ----------
+    prompt_dir : Path or None, default=None
+        Directory containing prompt templates.
+    default_model : str or None, default=None
+        Default model identifier to use when not defined in config.
+    store_name : str or None, default=None
+        Name of the vector store to query.
+    max_concurrent_searches : int, default=MAX_CONCURRENT_SEARCHES
+        Maximum number of concurrent vector search tasks to run.
+    vector_storage : VectorStorage or None, default=None
+        Optional preconfigured vector storage instance to reuse.
+    vector_storage_factory : Callable or None, default=None
+        Factory for constructing a VectorStorage when one is not provided.
+        Receives ``store_name`` as an argument.
+
     Methods
     -------
     run_agent(search_plan)
         Execute searches described by the plan.
     run_search(item)
         Perform a single vector search and summarise the result.
+
+    Raises
+    ------
+    ValueError
+        If the default model is not provided.
+
+    Examples
+    --------
+    >>> tool = VectorSearchTool(default_model="gpt-4o-mini", store_name="my_store")
     """
 
     def __init__(
@@ -89,24 +122,7 @@ class VectorSearchTool(
         vector_storage: Optional[VectorStorage] = None,
         vector_storage_factory: Optional[Callable[[str], VectorStorage]] = None,
     ) -> None:
-        """Initialize the search tool agent.
-
-        Parameters
-        ----------
-        prompt_dir : Path or None, default=None
-            Directory containing prompt templates.
-        default_model : str or None, default=None
-            Default model identifier to use when not defined in config.
-        store_name : str or None, default=None
-            Name of the vector store to query.
-        max_concurrent_searches : int, default=MAX_CONCURRENT_SEARCHES
-            Maximum number of concurrent vector search tasks to run.
-        vector_storage : VectorStorage or None, default=None
-            Optional preconfigured vector storage instance to reuse.
-        vector_storage_factory : Callable or None, default=None
-            Factory for constructing a VectorStorage when one is not provided.
-            Receives ``store_name`` as an argument.
-        """
+        """Initialize the search tool agent."""
         self._vector_storage: Optional[VectorStorage] = None
         self._store_name = store_name or "editorial"
         self._vector_storage_factory = vector_storage_factory
@@ -130,8 +146,8 @@ class VectorSearchTool(
             name="vector_search",
             instructions="Agent instructions",
             description="Perform vector searches based on a search plan.",
-            input_type=VectorSearchPlanStructure,
-            output_type=VectorSearchItemResultsStructure,
+            input_structure=VectorSearchPlanStructure,
+            output_structure=VectorSearchItemResultsStructure,
         )
 
     def _get_vector_storage(self) -> VectorStorage:
@@ -181,24 +197,32 @@ class VectorSearchTool(
 class VectorSearchWriter(SearchWriter[VectorSearchReportStructure]):
     """Generate reports summarizing vector search results.
 
+    Parameters
+    ----------
+    prompt_dir : Path or None, default=None
+        Directory containing prompt templates.
+    default_model : str or None, default=None
+        Default model identifier to use when not defined in config.
+
     Methods
     -------
     run_agent(query, search_results)
         Compile a final report from search results.
+
+    Raises
+    ------
+    ValueError
+        If the default model is not provided.
+
+    Examples
+    --------
+    >>> writer = VectorSearchWriter(default_model="gpt-4o-mini")
     """
 
     def __init__(
         self, prompt_dir: Optional[Path] = None, default_model: Optional[str] = None
     ) -> None:
-        """Initialize the writer agent.
-
-        Parameters
-        ----------
-        prompt_dir : Path or None, default=None
-            Directory containing prompt templates.
-        default_model : str or None, default=None
-            Default model identifier to use when not defined in config.
-        """
+        """Initialize the writer agent."""
         super().__init__(prompt_dir=prompt_dir, default_model=default_model)
 
     def _configure_agent(self) -> AgentConfiguration:
@@ -213,7 +237,7 @@ class VectorSearchWriter(SearchWriter[VectorSearchReportStructure]):
             name="vector_writer",
             instructions="Agent instructions",
             description="Write a report based on search results.",
-            output_type=VectorSearchReportStructure,
+            output_structure=VectorSearchReportStructure,
         )
 
 
@@ -225,6 +249,22 @@ class VectorSearch:
     comprehensive reports. It combines ``VectorSearchPlanner``,
     ``VectorSearchTool``, and ``VectorSearchWriter`` into a single workflow.
 
+    Parameters
+    ----------
+    prompt_dir : Path or None, default=None
+        Directory containing prompt templates.
+    default_model : str or None, default=None
+        Default model identifier to use when not defined in config.
+    vector_store_name : str or None, default=None
+        Name of the vector store to query.
+    max_concurrent_searches : int, default=MAX_CONCURRENT_SEARCHES
+        Maximum number of concurrent search tasks to run.
+    vector_storage : VectorStorage or None, default=None
+        Optional preconfigured vector storage instance to reuse.
+    vector_storage_factory : callable, default=None
+        Factory used to construct a VectorStorage when one is not provided.
+        Receives ``vector_store_name`` as an argument.
+
     Examples
     --------
     Basic vector search:
@@ -256,6 +296,11 @@ class VectorSearch:
         Convenience asynchronous entry point for the workflow.
     run_vector_agent_sync(search_query)
         Convenience synchronous entry point for the workflow.
+
+    Raises
+    ------
+    ValueError
+        If the default model is not provided.
     """
 
     def __init__(
@@ -268,24 +313,7 @@ class VectorSearch:
         vector_storage: Optional[VectorStorage] = None,
         vector_storage_factory: Optional[Callable[[str], VectorStorage]] = None,
     ) -> None:
-        """Create the main VectorSearch agent.
-
-        Parameters
-        ----------
-        prompt_dir : Path or None, default=None
-            Directory containing prompt templates.
-        default_model : str or None, default=None
-            Default model identifier to use when not defined in config.
-        vector_store_name : str or None, default=None
-            Name of the vector store to query.
-        max_concurrent_searches : int, default=MAX_CONCURRENT_SEARCHES
-            Maximum number of concurrent search tasks to run.
-        vector_storage : VectorStorage or None, default=None
-            Optional preconfigured vector storage instance to reuse.
-        vector_storage_factory : callable, default=None
-            Factory used to construct a VectorStorage when one is not provided.
-            Receives ``vector_store_name`` as an argument.
-        """
+        """Create the main VectorSearch agent."""
         self._prompt_dir = prompt_dir
         self._default_model = default_model
         self._vector_store_name = vector_store_name