openai-sdk-helpers 0.0.5__py3-none-any.whl → 0.0.6__py3-none-any.whl

openai_sdk_helpers/agent/web_search.py

@@ -0,0 +1,404 @@
+"""Core workflow management for ``web search``."""
+
+from __future__ import annotations
+
+import asyncio
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Union
+
+from agents import custom_span, gen_trace_id, trace
+from agents.model_settings import ModelSettings
+from agents.tool import WebSearchTool
+
+from ..structure.web_search import (
+    WebSearchItemStructure,
+    WebSearchItemResultStructure,
+    WebSearchStructure,
+    WebSearchPlanStructure,
+    WebSearchReportStructure,
+)
+from .base import AgentBase
+from .config import AgentConfig
+from .utils import run_coroutine_agent_sync
+
+MAX_CONCURRENT_SEARCHES = 10
+
+
+class WebAgentPlanner(AgentBase):
+    """Plan web searches to satisfy a user query.
+
+    Methods
+    -------
+    run_agent(query)
+        Generate a search plan for the provided query.
+    """
+
+    def __init__(
+        self, prompt_dir: Optional[Path] = None, default_model: Optional[str] = None
+    ) -> None:
+        """Initialize the planner agent.
+
+        Parameters
+        ----------
+        prompt_dir : pathlib.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.
+
+        Returns
+        -------
+        None
+        """
+        config = AgentConfig(
+            name="web_planner",
+            description="Agent that plans web searches based on a user query.",
+            output_type=WebSearchPlanStructure,
+        )
+        super().__init__(
+            config=config, prompt_dir=prompt_dir, default_model=default_model
+        )
+
+    async def run_agent(self, query: str) -> WebSearchPlanStructure:
+        """Plan searches for ``query``.
+
+        Parameters
+        ----------
+        query : str
+            User search query.
+
+        Returns
+        -------
+        WebSearchPlanStructure
+            Plan describing searches to perform.
+        """
+        result: WebSearchPlanStructure = await self.run_async(
+            input=query,
+            output_type=self._output_type,
+        )
+
+        return result
+
+
+class WebSearchToolAgent(AgentBase):
+    """Execute web searches defined in a plan.
+
+    Methods
+    -------
+    run_agent(search_plan)
+        Execute searches described by the plan.
+    run_search(item)
+        Perform a single web search and summarise the result.
+    """
+
+    def __init__(
+        self, prompt_dir: Optional[Path] = None, default_model: Optional[str] = None
+    ) -> None:
+        """Initialize the search tool agent.
+
+        Parameters
+        ----------
+        prompt_dir : pathlib.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.
+
+        Returns
+        -------
+        None
+        """
+        config = AgentConfig(
+            name="web_search",
+            description="Agent that performs web searches and summarizes results.",
+            input_type=WebSearchPlanStructure,
+            tools=[WebSearchTool()],
+            model_settings=ModelSettings(tool_choice="required"),
+        )
+        super().__init__(
+            config=config, prompt_dir=prompt_dir, default_model=default_model
+        )
+
+    async def run_agent(
+        self, search_plan: WebSearchPlanStructure
+    ) -> List[WebSearchItemResultStructure]:
+        """Execute all searches in the plan with a progress bar.
+
+        Parameters
+        ----------
+        search_plan : WebSearchPlanStructure
+            Plan describing each search to perform.
+
+        Returns
+        -------
+        list[WebSearchItemResultStructure]
+            Completed search results.
+        """
+        with custom_span("Search the web"):
+            semaphore = asyncio.Semaphore(MAX_CONCURRENT_SEARCHES)
+
+            async def _bounded_search(
+                item: WebSearchItemStructure,
+            ) -> WebSearchItemResultStructure:
+                """Execute a single search within the concurrency limit.
+
+                Parameters
+                ----------
+                item : WebSearchItemStructure
+                    Search item to process.
+
+                Returns
+                -------
+                WebSearchItemResultStructure
+                    Search result for ``item``.
+                """
+                async with semaphore:
+                    return await self.run_search(item)
+
+            tasks = [
+                asyncio.create_task(_bounded_search(item))
+                for item in search_plan.searches
+            ]
+            results = await asyncio.gather(*tasks)
+            return [result for result in results if result is not None]
+
+    async def run_search(
+        self, item: WebSearchItemStructure
+    ) -> WebSearchItemResultStructure:
+        """Perform a single web search using the search agent.
+
+        Parameters
+        ----------
+        item : WebSearchItemStructure
+            Search item containing the query and reason.
+
+        Returns
+        -------
+        WebSearchItemResultStructure
+            Search result summarising the page.
+        """
+        template_context: Dict[str, Any] = {
+            "search_term": item.query,
+            "reason": item.reason,
+        }
+
+        result = await super().run_async(
+            input=item.query,
+            context=template_context,
+            output_type=str,
+        )
+        return self._coerce_item_result(result)
+
+    @staticmethod
+    def _coerce_item_result(
+        result: Union[str, WebSearchItemResultStructure, Any],
+    ) -> WebSearchItemResultStructure:
+        """Return a WebSearchItemResultStructure from varied agent outputs."""
+        if isinstance(result, WebSearchItemResultStructure):
+            return result
+        try:
+            return WebSearchItemResultStructure(text=str(result))
+        except Exception:
+            return WebSearchItemResultStructure(text="")
+
+
+class WebAgentWriter(AgentBase):
+    """Summarize search results into a human-readable report.
+
+    Methods
+    -------
+    run_agent(query, search_results)
+        Compile a report from search results.
+    """
+
+    def __init__(
+        self, prompt_dir: Optional[Path] = None, default_model: Optional[str] = None
+    ) -> None:
+        """Initialize the writer agent.
+
+        Parameters
+        ----------
+        prompt_dir : pathlib.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.
+
+        Returns
+        -------
+        None
+        """
+        config = AgentConfig(
+            name="web_writer",
+            description="Agent that writes a report based on web search results.",
+            output_type=WebSearchReportStructure,
+        )
+        super().__init__(
+            config=config, prompt_dir=prompt_dir, default_model=default_model
+        )
+
+    async def run_agent(
+        self, query: str, search_results: List[WebSearchItemResultStructure]
+    ) -> WebSearchReportStructure:
+        """Compile a report from search results.
+
+        Parameters
+        ----------
+        query : str
+            Original search query.
+        search_results : list[WebSearchItemResultStructure]
+            Results produced by the search step.
+
+        Returns
+        -------
+        WebSearchReportStructure
+            Generated report for the query.
+        """
+        template_context: Dict[str, Any] = {
+            "original_query": query,
+            "search_results": search_results,
+        }
+        result: WebSearchReportStructure = await self.run_async(
+            input=query,
+            context=template_context,
+            output_type=self._output_type,
+        )
+
+        return result
+
+
+class WebAgentSearch(AgentBase):
+    """Manage the complete web search workflow.
+
+    Methods
+    -------
+    run_agent(search_query)
+        Execute the research workflow asynchronously.
+    run_agent_sync(search_query)
+        Execute the research workflow synchronously.
+    run_web_agent(search_query)
+        Convenience asynchronous entry point for the workflow.
+    run_web_agent_sync(search_query)
+        Convenience synchronous entry point for the workflow.
+    """
+
+    def __init__(
+        self,
+        config: Optional[AgentConfig] = None,
+        prompt_dir: Optional[Path] = None,
+        default_model: Optional[str] = None,
+    ) -> None:
+        """Create the main web search agent.
+
+        Parameters
+        ----------
+        config : AgentConfig or None, default=None
+            Optional configuration for the agent.
+        prompt_dir : pathlib.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.
+
+        Returns
+        -------
+        None
+        """
+        if config is None:
+            config = AgentConfig(
+                name="web_agent",
+                description="Agent that coordinates web searches and report writing.",
+                output_type=WebSearchStructure,
+            )
+        super().__init__(
+            config=config, prompt_dir=prompt_dir, default_model=default_model
+        )
+        self._prompt_dir = prompt_dir
+
+    async def run_agent(self, search_query: str) -> WebSearchStructure:
+        """Execute the entire research workflow for ``search_query``.
+
+        Parameters
+        ----------
+        search_query : str
+            User's research query.
+
+        Returns
+        -------
+        WebSearchStructure
+            Completed research output.
+        """
+        trace_id = gen_trace_id()
+        with trace("WebAgentSearch trace", trace_id=trace_id):
+            planner = WebAgentPlanner(
+                prompt_dir=self._prompt_dir, default_model=self.model
+            )
+            tool = WebSearchToolAgent(
+                prompt_dir=self._prompt_dir, default_model=self.model
+            )
+            writer = WebAgentWriter(
+                prompt_dir=self._prompt_dir, default_model=self.model
+            )
+            search_plan = await planner.run_agent(query=search_query)
+            search_results = await tool.run_agent(search_plan=search_plan)
+            search_report = await writer.run_agent(search_query, search_results)
+        return WebSearchStructure(
+            query=search_query,
+            web_search_plan=search_plan,
+            web_search_results=search_results,
+            web_search_report=search_report,
+        )
+
+    def run_agent_sync(self, search_query: str) -> WebSearchStructure:
+        """Run :meth:`run_agent` synchronously for ``search_query``.
+
+        Parameters
+        ----------
+        search_query : str
+            User's research query.
+
+        Returns
+        -------
+        WebSearchStructure
+            Completed research output.
+        """
+        return run_coroutine_agent_sync(self.run_agent(search_query))
+
+    @staticmethod
+    async def run_web_agent(search_query: str) -> WebSearchStructure:
+        """Return a research report for the given query using ``WebAgentSearch``.
+
+        Parameters
+        ----------
+        search_query : str
+            User's research query.
+
+        Returns
+        -------
+        WebSearchStructure
+            Completed research output.
+        """
+        return await WebAgentSearch().run_agent(search_query=search_query)
+
+    @staticmethod
+    def run_web_agent_sync(search_query: str) -> WebSearchStructure:
+        """Run :meth:`run_web_agent` synchronously for ``search_query``.
+
+        Parameters
+        ----------
+        search_query : str
+            User's research query.
+
+        Returns
+        -------
+        WebSearchStructure
+            Completed research output.
+        """
+        return run_coroutine_agent_sync(
+            WebAgentSearch.run_web_agent(search_query=search_query)
+        )
+
+
+__all__ = [
+    "MAX_CONCURRENT_SEARCHES",
+    "WebAgentPlanner",
+    "WebSearchToolAgent",
+    "WebAgentWriter",
+    "WebAgentSearch",
+]

openai_sdk_helpers/config.py

@@ -0,0 +1,153 @@
+"""Shared configuration for OpenAI SDK usage."""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Any, Dict, Mapping, Optional
+
+from dotenv import dotenv_values
+from openai import OpenAI
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class OpenAISettings(BaseModel):
+    """Configuration helpers for constructing OpenAI clients.
+
+    Methods
+    -------
+    from_env(dotenv_path, **overrides)
+        Build settings from environment variables and optional overrides.
+    client_kwargs()
+        Return keyword arguments for ``OpenAI`` initialization.
+    create_client()
+        Instantiate an ``OpenAI`` client using the stored configuration.
+    """
+
+    model_config = ConfigDict(extra="ignore")
+
+    api_key: Optional[str] = Field(
+        default=None,
+        description=(
+            "API key used to authenticate requests. Defaults to ``OPENAI_API_KEY``"
+            " from the environment."
+        ),
+    )
+    org_id: Optional[str] = Field(
+        default=None,
+        description=(
+            "Organization identifier applied to outgoing requests. Defaults to"
+            " ``OPENAI_ORG_ID``."
+        ),
+    )
+    project_id: Optional[str] = Field(
+        default=None,
+        description=(
+            "Project identifier used for billing and resource scoping. Defaults to"
+            " ``OPENAI_PROJECT_ID``."
+        ),
+    )
+    base_url: Optional[str] = Field(
+        default=None,
+        description=(
+            "Custom base URL for self-hosted or proxied deployments. Defaults to"
+            " ``OPENAI_BASE_URL``."
+        ),
+    )
+    default_model: Optional[str] = Field(
+        default=None,
+        description=(
+            "Model name used when constructing agents if no model is explicitly"
+            " provided. Defaults to ``OPENAI_MODEL``."
+        ),
+    )
+
+    @classmethod
+    def from_env(
+        cls, dotenv_path: Optional[Path] = None, **overrides: Any
+    ) -> "OpenAISettings":
+        """Load settings from the environment and optional overrides.
+
+        Parameters
+        ----------
+        dotenv_path : Path | None
+            Path to a ``.env`` file to load before reading environment
+            variables. Default ``None``.
+        overrides : Any
+            Keyword overrides applied on top of environment values.
+
+        Returns
+        -------
+        OpenAISettings
+        Settings instance populated from environment variables and overrides.
+        """
+        env_file_values: Mapping[str, Optional[str]]
+        if dotenv_path is not None:
+            env_file_values = dotenv_values(dotenv_path)
+        else:
+            env_file_values = dotenv_values()
+
+        values: Dict[str, Optional[str]] = {
+            "api_key": overrides.get("api_key")
+            or env_file_values.get("OPENAI_API_KEY")
+            or os.getenv("OPENAI_API_KEY"),
+            "org_id": overrides.get("org_id")
+            or env_file_values.get("OPENAI_ORG_ID")
+            or os.getenv("OPENAI_ORG_ID"),
+            "project_id": overrides.get("project_id")
+            or env_file_values.get("OPENAI_PROJECT_ID")
+            or os.getenv("OPENAI_PROJECT_ID"),
+            "base_url": overrides.get("base_url")
+            or env_file_values.get("OPENAI_BASE_URL")
+            or os.getenv("OPENAI_BASE_URL"),
+            "default_model": overrides.get("default_model")
+            or env_file_values.get("OPENAI_MODEL")
+            or os.getenv("OPENAI_MODEL"),
+        }
+
+        settings = cls(**values)
+        if not settings.api_key:
+            source_hint = (
+                f" from {dotenv_path}"
+                if dotenv_path is not None
+                else " from environment"
+            )
+            raise ValueError(
+                "OPENAI_API_KEY is required to configure the OpenAI client"
+                f" and was not found{source_hint}."
+            )
+
+        return settings
+
+    def client_kwargs(self) -> Dict[str, Any]:
+        """Return keyword arguments for constructing an ``OpenAI`` client.
+
+        Returns
+        -------
+        dict
+        Keyword arguments populated with available authentication and routing
+        values.
+        """
+        kwargs: Dict[str, Any] = {}
+        if self.api_key:
+            kwargs["api_key"] = self.api_key
+        if self.org_id:
+            kwargs["organization"] = self.org_id
+        if self.project_id:
+            kwargs["project"] = self.project_id
+        if self.base_url:
+            kwargs["base_url"] = self.base_url
+        return kwargs
+
+    def create_client(self) -> OpenAI:
+        """Instantiate an ``OpenAI`` client using the stored configuration.
+
+        Returns
+        -------
+        OpenAI
+        Client initialized with ``client_kwargs``.
+        """
+        return OpenAI(**self.client_kwargs())
+
+
+__all__ = ["OpenAISettings"]

openai_sdk_helpers/enums/__init__.py

@@ -0,0 +1,7 @@
+"""Enum helpers for openai-sdk-helpers."""
+
+from __future__ import annotations
+
+from .base import CrosswalkJSONEnum
+
+__all__ = ["CrosswalkJSONEnum"]

openai_sdk_helpers/enums/base.py

@@ -0,0 +1,29 @@
+"""Base enum helpers."""
+
+from __future__ import annotations
+
+from enum import Enum
+
+
+class CrosswalkJSONEnum(str, Enum):
+    """Enum base class with crosswalk metadata hooks.
+
+    Methods
+    -------
+    CROSSWALK()
+        Return metadata describing enum values keyed by name.
+    """
+
+    @classmethod
+    def CROSSWALK(cls) -> dict[str, dict[str, object]]:
+        """Return metadata describing enum values keyed by name.
+
+        Returns
+        -------
+        dict[str, dict[str, object]]
+            Mapping of enum member names to structured metadata details.
+        """
+        raise NotImplementedError("CROSSWALK must be implemented by subclasses.")
+
+
+__all__ = ["CrosswalkJSONEnum"]

openai_sdk_helpers/environment.py

@@ -0,0 +1,27 @@
+"""Environment helpers for openai-sdk-helpers."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+DATETIME_FMT = "%Y%m%d_%H%M%S"
+
+
+def get_data_path(module_name: str) -> Path:
+    """Return a writable data directory for the given module name.
+
+    Parameters
+    ----------
+    module_name : str
+        Name of the module requesting a data directory.
+
+    Returns
+    -------
+    Path
+        Directory path under ``~/.openai-sdk-helpers`` specific to ``module_name``. The
+        directory is created if it does not already exist.
+    """
+    base = Path.home() / ".openai-sdk-helpers"
+    path = base / module_name
+    path.mkdir(parents=True, exist_ok=True)
+    return path

openai_sdk_helpers/prompt/__init__.py

@@ -0,0 +1,77 @@
+"""Core prompt rendering utilities."""
+
+from __future__ import annotations
+
+import warnings
+from pathlib import Path
+from typing import Any, Mapping, Optional
+
+from dotenv import load_dotenv
+from jinja2 import Environment, FileSystemLoader, Template
+
+load_dotenv()
+warnings.filterwarnings("ignore")
+
+
+class PromptRenderer:
+    """Render prompts using Jinja2 templates.
+
+    The renderer loads templates from a base directory (defaulting to the
+    ``prompt`` package directory) and exposes a rendering helper for
+    injecting context values.
+
+    Methods
+    -------
+    render(template_path, context)
+        Render the template at ``template_path`` with the supplied context.
+    """
+
+    def __init__(self, base_dir: Optional[Path] = None) -> None:
+        """Initialize the renderer with a Jinja2 environment.
+
+        Parameters
+        ----------
+        base_dir : Path or None, default=None
+            Base directory containing Jinja2 templates. Defaults to the
+            ``prompt`` directory adjacent to this file when ``None``.
+
+        Returns
+        -------
+        None
+        """
+        if base_dir is None:
+            # Defaults to the directory containing this file, which also
+            # contains the builtin prompt templates.
+            self.base_dir = Path(__file__).resolve().parent
+        else:
+            self.base_dir = base_dir
+
+        self._env = Environment(
+            loader=FileSystemLoader(str(self.base_dir)),
+            autoescape=False,  # Prompts are plain text
+        )
+
+    def render(
+        self, template_path: str, context: Optional[Mapping[str, Any]] = None
+    ) -> str:
+        """Render a Jinja2 template with the given context.
+
+        Parameters
+        ----------
+        template_path : str
+            Path to the template file, relative to ``base_dir``.
+        context : Mapping[str, Any] or None, default=None
+            Context variables passed to the template.
+
+        Returns
+        -------
+        str
+            Rendered prompt as a string.
+        """
+        template_path_ = Path(self.base_dir, template_path)
+        template_path_text = template_path_.read_text()
+        template = Template(template_path_text)
+        return template.render(context or {})
+
+
+__all__ = ["PromptRenderer"]