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

openai_sdk_helpers/agent/base.py

@@ -3,43 +3,126 @@
 from __future__ import annotations
 
 from pathlib import Path
-from typing import Any, Dict, Optional, Protocol
-
-from agents import Agent, RunResult, RunResultStreaming, Runner
+from typing import Any, Dict, Optional, Protocol, cast
+import uuid
+
+from agents import (
+    Agent,
+    Handoff,
+    InputGuardrail,
+    OutputGuardrail,
+    RunResultStreaming,
+    Session,
+)
+from agents.model_settings import ModelSettings
 from agents.run_context import RunContextWrapper
-from agents.tool import FunctionTool
+from agents.tool import Tool
 from jinja2 import Template
 
-from .runner import run_async, run_streamed, run_sync
+from ..utils.json.data_class import DataclassJSONSerializable
+from ..structure.base import StructureBase
 
+from ..utils import (
+    check_filepath,
+    log,
+)
 
-class AgentConfigLike(Protocol):
-    """Protocol describing the configuration attributes for AgentBase."""
+from .runner import run_async, run_streamed, run_sync
 
-    name: str
-    description: Optional[str]
-    model: Optional[str]
-    template_path: Optional[str]
-    input_type: Optional[Any]
-    output_type: Optional[Any]
-    tools: Optional[Any]
-    model_settings: Optional[Any]
 
+class AgentConfigurationLike(Protocol):
+    """Protocol describing the configuration attributes for AgentBase."""
 
-class AgentBase:
+    @property
+    def name(self) -> str:
+        """Agent name."""
+        ...
+
+    @property
+    def description(self) -> Optional[str]:
+        """Agent description."""
+        ...
+
+    @property
+    def model(self) -> Optional[str]:
+        """Model identifier."""
+        ...
+
+    @property
+    def template_path(self) -> Optional[str | Path]:
+        """Template path."""
+        ...
+
+    def resolve_prompt_path(self, prompt_dir: Path | None = None) -> Path | None:
+        """Resolve the prompt template path."""
+        ...
+
+    @property
+    def instructions(self) -> str | Path:
+        """Instructions."""
+        ...
+
+    @property
+    def instructions_text(self) -> str:
+        """Resolved instructions text."""
+        ...
+
+    @property
+    def input_structure(self) -> Optional[type[StructureBase]]:
+        """Input type."""
+        ...
+
+    @property
+    def output_structure(self) -> Optional[type[StructureBase]]:
+        """Output type."""
+        ...
+
+    @property
+    def tools(self) -> Optional[list]:
+        """Tools."""
+        ...
+
+    @property
+    def model_settings(self) -> Optional[ModelSettings]:
+        """Model settings."""
+        ...
+
+    @property
+    def handoffs(self) -> Optional[list[Agent | Handoff]]:
+        """Handoffs."""
+        ...
+
+    @property
+    def input_guardrails(self) -> Optional[list[InputGuardrail]]:
+        """Input guardrails."""
+        ...
+
+    @property
+    def output_guardrails(self) -> Optional[list[OutputGuardrail]]:
+        """Output guardrails."""
+        ...
+
+    @property
+    def session(self) -> Optional[Session]:
+        """Session."""
+        ...
+
+
+class AgentBase(DataclassJSONSerializable):
     """Factory for creating and configuring specialized agents.
 
     ``AgentBase`` provides the foundation for building OpenAI agents with support
-    for Jinja2 prompt templates, custom tools, and both synchronous and
-    asynchronous execution modes. All specialized agents in this package extend
-    this base class.
+    for Jinja2 prompt templates, custom tools, handoffs for agent delegation,
+    input and output guardrails for validation, session management for
+    conversation history, and both synchronous and asynchronous execution modes.
+    All specialized agents in this package extend this base class.
 
     Examples
     --------
     Create a basic agent from configuration:
 
-    >>> from openai_sdk_helpers.agent import AgentBase, AgentConfig
-    >>> config = AgentConfig(
+    >>> from openai_sdk_helpers.agent import AgentBase, AgentConfiguration
+    >>> config = AgentConfiguration(
     ...     name="my_agent",
     ...     description="A custom agent",
     ...     model="gpt-4o-mini"
@@ -49,7 +132,7 @@ class AgentBase:
 
     Use absolute path to template:
 
-    >>> config = AgentConfig(
+    >>> config = AgentConfiguration(
     ...     name="my_agent",
     ...     template_path="/absolute/path/to/template.jinja",
     ...     model="gpt-4o-mini"
@@ -66,30 +149,48 @@ class AgentBase:
 
     Methods
     -------
-    from_config(config, run_context_wrapper)
-        Instantiate a ``AgentBase`` from configuration.
     build_prompt_from_jinja(run_context_wrapper)
         Render the agent prompt using Jinja and optional context.
     get_prompt(run_context_wrapper, _)
         Render the agent prompt using the provided run context.
+    name
+        Return the name of this agent.
+    instructions_text
+        Return the resolved instructions for this agent.
+    tools
+        Return the tools configured for this agent.
+    output_structure
+        Return the output type configured for this agent.
+    model_settings
+        Return the model settings configured for this agent.
+    handoffs
+        Return the handoff configurations for this agent.
+    input_guardrails
+        Return the input guardrails configured for this agent.
+    output_guardrails
+        Return the output guardrails configured for this agent.
+    session
+        Return the session configured for this agent.
     get_agent()
         Construct the configured :class:`agents.Agent` instance.
-    run(input, context, output_type)
-        Execute the agent asynchronously (alias of ``run_async``).
-    run_async(input, context, output_type)
+    run_async(input, context, output_structure, session)
         Execute the agent asynchronously and optionally cast the result.
-    run_sync(input, context, output_type)
+    run_sync(input, context, output_structure, session)
         Execute the agent synchronously.
-    run_streamed(input, context, output_type)
+    run_streamed(input, context, output_structure, session)
         Return a streaming result for the agent execution.
     as_tool()
         Return the agent as a callable tool.
+    close()
+        Clean up agent resources (can be overridden by subclasses).
     """
 
     def __init__(
         self,
-        config: AgentConfigLike,
+        *,
+        config: AgentConfigurationLike,
         run_context_wrapper: Optional[RunContextWrapper[Dict[str, Any]]] = None,
+        data_path: Path | str | None = None,
         prompt_dir: Optional[Path] = None,
         default_model: Optional[str] = None,
     ) -> None:
@@ -97,7 +198,7 @@ class AgentBase:
 
         Parameters
         ----------
-        config : AgentConfigLike
+        config : AgentConfigurationLike
             Configuration describing this agent.
         run_context_wrapper : RunContextWrapper or None, default=None
             Optional wrapper providing runtime context for prompt rendering.
@@ -115,70 +216,49 @@ class AgentBase:
         if not model:
             raise ValueError("Model is required to construct the agent.")
 
-        prompt_path: Optional[Path]
-        if config.template_path:
-            prompt_path = Path(config.template_path)
-        elif prompt_dir is not None:
-            prompt_path = prompt_dir / f"{name}.jinja"
-        else:
-            prompt_path = None
+        prompt_path = config.resolve_prompt_path(prompt_dir)
 
+        # Build template from file or fall back to instructions
         if prompt_path is None:
-            self._template = Template("")
+            instructions_text = config.instructions_text
+            self._template = Template(instructions_text)
+            self._instructions = instructions_text
         elif prompt_path.exists():
-            self._template = Template(prompt_path.read_text())
+            self._template = Template(prompt_path.read_text(encoding="utf-8"))
+            self._instructions = None
         else:
             raise FileNotFoundError(
                 f"Prompt template for agent '{name}' not found at {prompt_path}."
             )
 
-        self.agent_name = name
+        self._name = name
+        self.uuid = uuid.uuid4()
         self.description = description
         self.model = model
 
-        self._input_type = config.input_type
-        self._output_type = config.output_type or config.input_type
+        # Resolve data_path with class name appended
+        class_name = self.__class__.__name__
+        if data_path is not None:
+            data_path_obj = Path(data_path)
+            if data_path_obj.name == class_name:
+                self._data_path = data_path_obj
+            else:
+                self._data_path = data_path_obj / class_name
+        else:
+            from ..environment import get_data_path
+
+            self._data_path = get_data_path(self.__class__.__name__)
+
+        self._input_structure = config.input_structure
+        self._output_structure = config.output_structure or config.input_structure
         self._tools = config.tools
         self._model_settings = config.model_settings
+        self._handoffs = config.handoffs
+        self._input_guardrails = config.input_guardrails
+        self._output_guardrails = config.output_guardrails
+        self._session = config.session
         self._run_context_wrapper = run_context_wrapper
 
-    @classmethod
-    def from_config(
-        cls,
-        config: AgentConfigLike,
-        *,
-        run_context_wrapper: Optional[RunContextWrapper[Dict[str, Any]]] = None,
-        prompt_dir: Optional[Path] = None,
-        default_model: Optional[str] = None,
-    ) -> AgentBase:
-        """Create an AgentBase instance from configuration.
-
-        Parameters
-        ----------
-        config : AgentConfigLike
-            Configuration describing the agent.
-        run_context_wrapper : RunContextWrapper or None, default=None
-            Optional wrapper providing runtime context.
-        prompt_dir : Path or None, default=None
-            Optional directory holding prompt templates. Used when
-            ``config.template_path`` is not provided or is relative. If
-            ``config.template_path`` is an absolute path, this parameter is
-            ignored.
-        default_model : str or None, default=None
-            Optional fallback model identifier.
-
-        Returns
-        -------
-        AgentBase
-            Instantiated agent.
-        """
-        return cls(
-            config=config,
-            run_context_wrapper=run_context_wrapper,
-            prompt_dir=prompt_dir,
-            default_model=default_model,
-        )
-
     def _build_prompt_from_jinja(self) -> str:
         """Render the instructions prompt for this agent.
 
@@ -233,6 +313,107 @@ class AgentBase:
         """
         return self.build_prompt_from_jinja(run_context_wrapper)
 
+    @property
+    def name(self) -> str:
+        """Return the name of this agent.
+
+        Returns
+        -------
+        str
+            Name used to identify the agent.
+        """
+        return self._name
+
+    @property
+    def instructions_text(self) -> str:
+        """Return the resolved instructions for this agent.
+
+        Returns
+        -------
+        str
+            Rendered instructions text using the current run context.
+        """
+        if self._instructions is not None:
+            return self._instructions
+        return self._build_prompt_from_jinja()
+
+    @property
+    def tools(self) -> Optional[list]:
+        """Return the tools configured for this agent.
+
+        Returns
+        -------
+        list or None
+            Tool definitions configured for the agent.
+        """
+        return self._tools
+
+    @property
+    def output_structure(self) -> Optional[type[StructureBase]]:
+        """Return the output type configured for this agent.
+
+        Returns
+        -------
+        type[StructureBase] or None
+            Output type used to cast responses.
+        """
+        return self._output_structure
+
+    @property
+    def model_settings(self) -> Optional[ModelSettings]:
+        """Return the model settings configured for this agent.
+
+        Returns
+        -------
+        ModelSettings or None
+            Model settings applied to the agent.
+        """
+        return self._model_settings
+
+    @property
+    def handoffs(self) -> Optional[list[Agent | Handoff]]:
+        """Return the handoff configurations for this agent.
+
+        Returns
+        -------
+        list[Agent or Handoff] or None
+            Handoff configurations available to the agent.
+        """
+        return self._handoffs
+
+    @property
+    def input_guardrails(self) -> Optional[list[InputGuardrail]]:
+        """Return the input guardrails configured for this agent.
+
+        Returns
+        -------
+        list[InputGuardrail] or None
+            Input guardrails applied to the agent.
+        """
+        return self._input_guardrails
+
+    @property
+    def output_guardrails(self) -> Optional[list[OutputGuardrail]]:
+        """Return the output guardrails configured for this agent.
+
+        Returns
+        -------
+        list[OutputGuardrail] or None
+            Output guardrails applied to the agent.
+        """
+        return self._output_guardrails
+
+    @property
+    def session(self) -> Optional[Session]:
+        """Return the session configured for this agent.
+
+        Returns
+        -------
+        Session or None
+            Session configuration used for maintaining conversation history.
+        """
+        return self._session
+
     def get_agent(self) -> Agent:
         """Construct and return the configured :class:`agents.Agent` instance.
 
@@ -242,16 +423,22 @@ class AgentBase:
             Initialized agent ready for execution.
         """
         agent_config: Dict[str, Any] = {
-            "name": self.agent_name,
+            "name": self._name,
             "instructions": self._build_prompt_from_jinja() or ".",
             "model": self.model,
         }
-        if self._output_type:
-            agent_config["output_type"] = self._output_type
+        if self._output_structure:
+            agent_config["output_structure"] = self._output_structure
         if self._tools:
             agent_config["tools"] = self._tools
         if self._model_settings:
             agent_config["model_settings"] = self._model_settings
+        if self._handoffs:
+            agent_config["handoffs"] = self._handoffs
+        if self._input_guardrails:
+            agent_config["input_guardrails"] = self._input_guardrails
+        if self._output_guardrails:
+            agent_config["output_guardrails"] = self._output_guardrails
 
         return Agent(**agent_config)
 
@@ -260,7 +447,8 @@ class AgentBase:
         input: str,
         *,
         context: Optional[Dict[str, Any]] = None,
-        output_type: Optional[Any] = None,
+        output_structure: Optional[type[StructureBase]] = None,
+        session: Optional[Any] = None,
     ) -> Any:
         """Execute the agent asynchronously.
 
@@ -270,21 +458,27 @@ class AgentBase:
             Prompt or query for the agent.
         context : dict or None, default=None
             Optional dictionary passed to the agent.
-        output_type : type or None, default=None
+        output_structure : type[StructureBase] or None, default=None
             Optional type used to cast the final output.
+        session : Session or None, default=None
+            Optional session for maintaining conversation history across runs.
+            If not provided, uses the session from config if available.
 
         Returns
         -------
         Any
-            Agent result, optionally converted to ``output_type``.
+            Agent result, optionally converted to ``output_structure``.
         """
-        if self._output_type is not None and output_type is None:
-            output_type = self._output_type
+        if self._output_structure is not None and output_structure is None:
+            output_structure = self._output_structure
+        # Use session from parameter, fall back to config session
+        session_to_use = session if session is not None else self._session
         return await run_async(
             agent=self.get_agent(),
             input=input,
             context=context,
-            output_type=output_type,
+            output_structure=output_structure,
+            session=session_to_use,
         )
 
     def run_sync(
@@ -292,7 +486,8 @@ class AgentBase:
         input: str,
         *,
         context: Optional[Dict[str, Any]] = None,
-        output_type: Optional[Any] = None,
+        output_structure: Optional[type[StructureBase]] = None,
+        session: Optional[Any] = None,
     ) -> Any:
         """Run the agent synchronously.
 
@@ -302,19 +497,27 @@ class AgentBase:
             Prompt or query for the agent.
         context : dict or None, default=None
             Optional dictionary passed to the agent.
-        output_type : type or None, default=None
+        output_structure : type[StructureBase] or None, default=None
             Optional type used to cast the final output.
+        session : Session or None, default=None
+            Optional session for maintaining conversation history across runs.
+            If not provided, uses the session from config if available.
 
         Returns
         -------
         Any
-            Agent result, optionally converted to ``output_type``.
+            Agent result, optionally converted to ``output_structure``.
         """
+        if self._output_structure is not None and output_structure is None:
+            output_structure = self._output_structure
+        # Use session from parameter, fall back to config session
+        session_to_use = session if session is not None else self._session
         return run_sync(
             agent=self.get_agent(),
             input=input,
             context=context,
-            output_type=output_type,
+            output_structure=output_structure,
+            session=session_to_use,
         )
 
     def run_streamed(
@@ -322,8 +525,9 @@ class AgentBase:
         input: str,
         *,
         context: Optional[Dict[str, Any]] = None,
-        output_type: Optional[Any] = None,
-    ) -> RunResultStreaming:
+        output_structure: Optional[type[StructureBase]] = None,
+        session: Optional[Any] = None,
+    ) -> RunResultStreaming | StructureBase:
         """Stream the agent execution results.
 
         Parameters
@@ -332,38 +536,119 @@ class AgentBase:
             Prompt or query for the agent.
         context : dict or None, default=None
             Optional dictionary passed to the agent.
-        output_type : type or None, default=None
+        output_structure : type[StructureBase] or None, default=None
             Optional type used to cast the final output.
+        session : Session or None, default=None
+            Optional session for maintaining conversation history across runs.
+            If not provided, uses the session from config if available.
 
         Returns
         -------
         RunResultStreaming
             Streaming output wrapper from the agent execution.
         """
+        # Use session from parameter, fall back to config session
+        session_to_use = session if session is not None else self._session
+        output_structure_to_use = output_structure or self._output_structure
         result = run_streamed(
             agent=self.get_agent(),
             input=input,
             context=context,
+            output_structure=output_structure_to_use,
+            session=session_to_use,
         )
-        if self._output_type and not output_type:
-            output_type = self._output_type
-        if output_type:
-            return result.final_output_as(output_type)
+        if output_structure_to_use and hasattr(result, "final_output_as"):
+            return cast(Any, result).final_output_as(output_structure_to_use)
         return result
 
-    def as_tool(self) -> FunctionTool:
+    def as_tool(self) -> Tool:
         """Return the agent as a callable tool.
 
         Returns
         -------
-        FunctionTool
+        Tool
             Tool instance wrapping this agent.
         """
         agent = self.get_agent()
-        tool_obj: FunctionTool = agent.as_tool(
-            tool_name=self.agent_name, tool_description=self.description
-        )  # type: ignore
+        tool_obj: Tool = agent.as_tool(
+            tool_name=self._name, tool_description=self.description
+        )
         return tool_obj
 
+    def __enter__(self) -> AgentBase:
+        """Enter the context manager for resource management.
+
+        Returns
+        -------
+        AgentBase
+            Self reference for use in with statements.
+        """
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        """Exit the context manager and clean up resources.
+
+        Parameters
+        ----------
+        exc_type : type or None
+            Exception type if an exception occurred, otherwise None.
+        exc_val : Exception or None
+            Exception instance if an exception occurred, otherwise None.
+        exc_tb : traceback or None
+            Traceback object if an exception occurred, otherwise None.
+        """
+        self.close()
+
+    def close(self) -> None:
+        """Clean up agent resources.
+
+        This method is called automatically when using the agent as a
+        context manager. Override in subclasses to implement custom
+        cleanup logic.
+
+        Examples
+        --------
+        >>> agent = AgentBase(config, default_model="gpt-4o-mini")
+        >>> try:
+        ...     result = agent.run_sync("query")
+        ... finally:
+        ...     agent.close()
+        """
+        log(f"Closing session {self.uuid} for {self.__class__.__name__}")
+        self.save()
+
+    def __repr__(self) -> str:
+        """Return a string representation of the AgentBase.
+
+        Returns
+        -------
+        str
+            String representation including agent name and model.
+        """
+        return f"<AgentBase name={self._name!r} model={self.model!r}>"
+
+    def save(self, filepath: str | Path | None = None) -> None:
+        """Serialize the message history to a JSON file.
+
+        Saves the current message history to a specified file path in JSON format.
+        If no file path is provided, it saves to a default location based on
+        the agent's UUID.
+
+        Parameters
+        ----------
+        filepath : str | Path | None, default=None
+            Optional file path to save the serialized history. If None,
+            uses a default filename based on the agent name.
+        """
+        if filepath is not None:
+            target = Path(filepath)
+        else:
+            filename = f"{str(self.uuid).lower()}.json"
+            target = self._data_path / self._name / filename
+
+        checked = check_filepath(filepath=target)
+        self.to_json_file(filepath=checked)
+        log(f"Saved messages to {target}")
+
 
-__all__ = ["AgentConfigLike", "AgentBase"]
+__all__ = ["AgentConfigurationLike", "AgentBase"]