agentpool 2.1.9__py3-none-any.whl

agentpool/agents/events/tts_handlers.py

@@ -0,0 +1,186 @@
+"""TTS event handlers for streaming text-to-speech synthesis."""
+
+from __future__ import annotations
+
+from abc import abstractmethod
+from typing import TYPE_CHECKING, Any, Literal
+
+from pydantic_ai import PartDeltaEvent, PartStartEvent, TextPart, TextPartDelta
+
+
+if TYPE_CHECKING:
+    from anyvoice import TTSStream
+    from pydantic_ai import RunContext
+
+    from agentpool.agents.events import RichAgentStreamEvent
+
+TTSMode = Literal["sync_sentence", "sync_run", "async_queue", "async_cancel"]
+TTSModel = Literal["tts-1", "tts-1-hd"]
+TTSVoice = Literal["alloy", "echo", "fable", "onyx", "nova", "shimmer"]
+
+
+class BaseTTSEventHandler:
+    """Base TTS event handler with shared event handling logic.
+
+    This base class handles the common pattern of:
+    - Managing TTS stream lifecycle (create, close)
+    - Translating pydantic-ai stream events to TTSStream.feed() calls
+    - Handling different synchronization modes
+
+    Subclasses only need to implement _create_stream() with their
+    specific provider and session configuration.
+    """
+
+    def __init__(
+        self,
+        *,
+        sample_rate: int = 24000,
+        min_text_length: int = 20,
+        mode: TTSMode = "sync_run",
+    ) -> None:
+        self._sample_rate = sample_rate
+        self._min_text_length = min_text_length
+        self._mode: TTSMode = mode
+        self._tts_stream: TTSStream | None = None
+
+    @abstractmethod
+    async def _create_stream(self) -> TTSStream:
+        """Create and return a configured TTS stream.
+
+        Subclasses implement this to create their provider-specific stream.
+
+        Returns:
+            Configured TTSStream instance (not yet entered).
+        """
+        raise NotImplementedError
+
+    async def _ensure_stream(self) -> TTSStream:
+        """Get or create the TTS stream."""
+        if self._tts_stream is None:
+            self._tts_stream = await self._create_stream()
+            await self._tts_stream.__aenter__()
+        return self._tts_stream
+
+    async def _close_stream(self) -> None:
+        """Close the TTS stream if open."""
+        if self._tts_stream is not None:
+            await self._tts_stream.__aexit__(None, None, None)
+            self._tts_stream = None
+
+    async def __call__(self, ctx: RunContext[Any], event: RichAgentStreamEvent[Any]) -> None:
+        """Handle stream events and trigger TTS synthesis."""
+        from agentpool.agents.events import RunStartedEvent, StreamCompleteEvent
+
+        match event:
+            case RunStartedEvent():
+                # For async_cancel mode, cancel any pending audio from previous run
+                if self._mode == "async_cancel" and self._tts_stream is not None:
+                    await self._tts_stream.cancel()
+
+            case (
+                PartStartEvent(part=TextPart(content=delta))
+                | PartDeltaEvent(delta=TextPartDelta(content_delta=delta))
+            ):
+                stream = await self._ensure_stream()
+                await stream.feed(delta)
+
+            case StreamCompleteEvent():
+                await self._close_stream()
+
+
+class OpenAITTSEventHandler(BaseTTSEventHandler):
+    """TTS event handler using OpenAI's Text-to-Speech API.
+
+    Translates pydantic-ai stream events to anyvoice TTSStream.feed() calls
+    using the OpenAI TTS provider.
+    """
+
+    def __init__(
+        self,
+        *,
+        api_key: str | None = None,
+        model: TTSModel = "tts-1",
+        voice: TTSVoice = "alloy",
+        speed: float = 1.0,
+        chunk_size: int = 1024,
+        sample_rate: int = 24000,
+        min_text_length: int = 20,
+        mode: TTSMode = "sync_run",
+    ) -> None:
+        super().__init__(
+            sample_rate=sample_rate,
+            min_text_length=min_text_length,
+            mode=mode,
+        )
+        self._api_key = api_key
+        self._model: TTSModel = model
+        self._voice: TTSVoice = voice
+        self._speed = speed
+        self._chunk_size = chunk_size
+
+    async def _create_stream(self) -> TTSStream:
+        """Create OpenAI TTS stream."""
+        from anyvoice import OpenAITTSProvider, SoundDeviceSink, TTSStream
+
+        provider = OpenAITTSProvider(api_key=self._api_key)
+        session = provider.session(
+            model=self._model,
+            voice=self._voice,
+            speed=self._speed,
+            chunk_size=self._chunk_size,
+        )
+        sink = SoundDeviceSink(sample_rate=self._sample_rate)
+        return TTSStream(
+            session,
+            sink=sink,
+            mode=self._mode,
+            min_text_length=self._min_text_length,
+        )
+
+
+class EdgeTTSEventHandler(BaseTTSEventHandler):
+    """TTS event handler using Microsoft Edge's free TTS service.
+
+    Translates pydantic-ai stream events to anyvoice TTSStream.feed() calls
+    using the Edge TTS provider (no API key required).
+    """
+
+    def __init__(
+        self,
+        *,
+        voice: str = "en-US-AriaNeural",
+        rate: str = "+0%",
+        volume: str = "+0%",
+        pitch: str = "+0Hz",
+        sample_rate: int = 24000,
+        min_text_length: int = 20,
+        mode: TTSMode = "sync_run",
+    ) -> None:
+        super().__init__(
+            sample_rate=sample_rate,
+            min_text_length=min_text_length,
+            mode=mode,
+        )
+        self._voice = voice
+        self._rate = rate
+        self._volume = volume
+        self._pitch = pitch
+
+    async def _create_stream(self) -> TTSStream:
+        """Create Edge TTS stream."""
+        from anyvoice import EdgeTTSProvider, SoundDeviceSink, TTSStream
+
+        provider = EdgeTTSProvider(sample_rate=self._sample_rate)
+        session = provider.session(
+            voice=self._voice,
+            rate=self._rate,
+            volume=self._volume,
+            pitch=self._pitch,
+        )
+        sink = SoundDeviceSink(sample_rate=self._sample_rate)
+        return TTSStream(
+            session,
+            sink=sink,
+            mode=self._mode,
+            min_text_length=self._min_text_length,
+        )

agentpool/agents/interactions.py

@@ -0,0 +1,419 @@
+"""Agent interaction patterns."""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Mapping
+from typing import TYPE_CHECKING, Any, Literal, cast, overload
+
+from schemez import Schema
+
+from agentpool.log import get_logger
+from agentpool.messaging import ChatMessage
+
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from toprompt import AnyPromptType
+
+    from agentpool import AgentPool, MessageNode
+    from agentpool.agents import Agent
+    from agentpool.common_types import SupportsStructuredOutput
+    from agentpool.delegation.base_team import BaseTeam
+
+
+ExtractionMode = Literal["structured", "tool_calls"]
+
+type EndCondition = Callable[[list[ChatMessage[Any]], ChatMessage[Any]], bool]
+
+logger = get_logger(__name__)
+
+
+class LLMPick(Schema):
+    """Decision format for LLM response."""
+
+    selection: str  # The label/name of the selected option
+    reason: str
+
+
+class Pick[T](Schema):
+    """Type-safe decision with original object."""
+
+    selection: T
+    reason: str
+
+
+class LLMMultiPick(Schema):
+    """Multiple selection format for LLM response."""
+
+    selections: list[str]  # Labels of selected options
+    reason: str
+
+
+class MultiPick[T](Schema):
+    """Type-safe multiple selection with original objects."""
+
+    selections: list[T]
+    reason: str
+
+
+def get_label(item: Any) -> str:
+    """Get label for an item to use in selection.
+
+    Args:
+        item: Item to get label for
+
+    Returns:
+        Label to use for selection
+
+    Strategy:
+        - strings stay as-is
+        - types use __name__
+        - others use __repr__ for unique identifiable string
+    """
+    from agentpool import MessageNode
+
+    match item:
+        case str():
+            return item
+        case type():
+            return item.__name__
+        case MessageNode():
+            return item.name or "unnamed_agent"
+        case _:
+            return repr(item)
+
+
+class Interactions:
+    """Manages agent communication patterns."""
+
+    def __init__(self, agent: SupportsStructuredOutput) -> None:
+        self.agent = agent
+
+    # async def conversation(
+    #     self,
+    #     other: MessageNode[Any, Any],
+    #     initial_message: AnyPromptType,
+    #     *,
+    #     max_rounds: int | None = None,
+    #     end_condition: Callable[[list[ChatMessage[Any]], ChatMessage[Any]], bool] | None = None,
+    # ) -> AsyncIterator[ChatMessage[Any]]:
+    #     """Maintain conversation between two agents.
+
+    #     Args:
+    #         other: Agent to converse with
+    #         initial_message: Message to start conversation with
+    #         max_rounds: Optional maximum number of exchanges
+    #         end_condition: Optional predicate to check for conversation end
+
+    #     Yields:
+    #         Messages from both agents in conversation order
+    #     """
+    #     rounds = 0
+    #     messages: list[ChatMessage[Any]] = []
+    #     current_message = initial_message
+    #     current_node: MessageNode[Any, Any] = self.agent
+
+    #     while True:
+    #         if max_rounds and rounds >= max_rounds:
+    #             logger.debug("Conversation ended", max_rounds=max_rounds)
+    #             return
+
+    #         response = await current_node.run(current_message)
+    #         messages.append(response)
+    #         yield response
+
+    #         if end_condition and end_condition(messages, response):
+    #             logger.debug("Conversation ended: end condition met")
+    #             return
+
+    #         # Switch agents for next round
+    #         current_node = other if current_node == self.agent else self.agent
+    #         current_message = response.content
+    #         rounds += 1
+
+    @overload
+    async def pick(
+        self,
+        selections: AgentPool,
+        task: str,
+        prompt: AnyPromptType | None = None,
+    ) -> Pick[Agent[Any, Any]]: ...
+
+    @overload
+    async def pick(
+        self,
+        selections: BaseTeam[Any, Any],
+        task: str,
+        prompt: AnyPromptType | None = None,
+    ) -> Pick[MessageNode[Any, Any]]: ...
+
+    @overload
+    async def pick[T: AnyPromptType](
+        self,
+        selections: Sequence[T] | Mapping[str, T],
+        task: str,
+        prompt: AnyPromptType | None = None,
+    ) -> Pick[T]: ...
+
+    async def pick[T](
+        self,
+        selections: Sequence[T] | Mapping[str, T] | AgentPool | BaseTeam[Any, Any],
+        task: str,
+        prompt: AnyPromptType | None = None,
+    ) -> Pick[T]:
+        """Pick from available options with reasoning.
+
+        Args:
+            selections: What to pick from:
+                - Sequence of items (auto-labeled)
+                - Dict mapping labels to items
+                - AgentPool
+                - Team
+            task: Task/decision description
+            prompt: Optional custom selection prompt
+
+        Returns:
+            Decision with selected item and reasoning
+
+        Raises:
+            ValueError: If no choices available or invalid selection
+        """
+        # Get items and create label mapping
+        from toprompt import to_prompt
+
+        from agentpool import AgentPool
+        from agentpool.delegation.base_team import BaseTeam
+
+        match selections:
+            case dict():
+                label_map = selections
+                items: list[Any] = list(selections.values())
+            case BaseTeam():
+                items = list(selections.nodes)
+                label_map = {get_label(item): item for item in items}
+            case AgentPool():
+                items = list(selections.agents.values())
+                label_map = {get_label(item): item for item in items}
+            case _:
+                items = list(selections)
+                label_map = {get_label(item): item for item in items}
+
+        if not items:
+            msg = "No choices available"
+            raise ValueError(msg)
+
+        # Get descriptions for all items
+        descriptions = []
+        for label, item in label_map.items():
+            item_desc = await to_prompt(item)
+            descriptions.append(f"{label}:\n{item_desc}")
+
+        default_prompt = f"""Task/Decision: {task}
+
+Available options:
+{"-" * 40}
+{"\n\n".join(descriptions)}
+{"-" * 40}
+
+Select ONE option by its exact label."""
+
+        # Get LLM's string-based decision
+        result = await self.agent.run(prompt or default_prompt, output_type=LLMPick)
+
+        # Convert to type-safe decision
+        if result.content.selection not in label_map:
+            msg = f"Invalid selection: {result.content.selection}"
+            raise ValueError(msg)
+
+        selected = cast(T, label_map[result.content.selection])
+        return Pick(selection=selected, reason=result.content.reason)
+
+    @overload
+    async def pick_multiple(
+        self,
+        selections: BaseTeam[Any, Any],
+        task: str,
+        *,
+        min_picks: int = 1,
+        max_picks: int | None = None,
+        prompt: AnyPromptType | None = None,
+    ) -> MultiPick[MessageNode[Any, Any]]: ...
+
+    @overload
+    async def pick_multiple(
+        self,
+        selections: AgentPool,
+        task: str,
+        *,
+        min_picks: int = 1,
+        max_picks: int | None = None,
+        prompt: AnyPromptType | None = None,
+    ) -> MultiPick[Agent[Any, Any]]: ...
+
+    @overload
+    async def pick_multiple[T: AnyPromptType](
+        self,
+        selections: Sequence[T] | Mapping[str, T],
+        task: str,
+        *,
+        min_picks: int = 1,
+        max_picks: int | None = None,
+        prompt: AnyPromptType | None = None,
+    ) -> MultiPick[T]: ...
+
+    async def pick_multiple[T](
+        self,
+        selections: Sequence[T] | Mapping[str, T] | AgentPool | BaseTeam[Any, Any],
+        task: str,
+        *,
+        min_picks: int = 1,
+        max_picks: int | None = None,
+        prompt: AnyPromptType | None = None,
+    ) -> MultiPick[T]:
+        """Pick multiple options from available choices.
+
+        Args:
+            selections: What to pick from
+            task: Task/decision description
+            min_picks: Minimum number of selections required
+            max_picks: Maximum number of selections (None for unlimited)
+            prompt: Optional custom selection prompt
+        """
+        from toprompt import to_prompt
+
+        from agentpool import AgentPool
+        from agentpool.delegation.base_team import BaseTeam
+
+        match selections:
+            case AgentPool():
+                items: list[Any] = list(selections.agents.values())
+                label_map: Mapping[str, Any] = {get_label(item): item for item in items}
+            case Mapping():
+                label_map = selections
+                items = list(selections.values())
+            case BaseTeam():
+                items = list(selections.nodes)
+                label_map = {get_label(item): item for item in items}
+            case _:
+                items = list(selections)
+                label_map = {get_label(item): item for item in items}
+
+        if not items:
+            msg = "No choices available"
+            raise ValueError(msg)
+
+        if max_picks is not None and max_picks < min_picks:
+            msg = f"max_picks ({max_picks}) cannot be less than min_picks ({min_picks})"
+            raise ValueError(msg)
+
+        descriptions = []
+        for label, item in label_map.items():
+            item_desc = await to_prompt(item)
+            descriptions.append(f"{label}:\n{item_desc}")
+
+        picks_info = (
+            f"Select between {min_picks} and {max_picks}"
+            if max_picks is not None
+            else f"Select at least {min_picks}"
+        )
+
+        default_prompt = f"""Task/Decision: {task}
+
+Available options:
+{"-" * 40}
+{"\n\n".join(descriptions)}
+{"-" * 40}
+
+{picks_info} options by their exact labels.
+List your selections, one per line, followed by your reasoning."""
+
+        result = await self.agent.run(prompt or default_prompt, output_type=LLMMultiPick)
+
+        # Validate selections
+        invalid = [s for s in result.content.selections if s not in label_map]
+        if invalid:
+            msg = f"Invalid selections: {', '.join(invalid)}"
+            raise ValueError(msg)
+        num_picks = len(result.content.selections)
+        if num_picks < min_picks:
+            msg = f"Too few selections: got {num_picks}, need {min_picks}"
+            raise ValueError(msg)
+
+        if max_picks and num_picks > max_picks:
+            msg = f"Too many selections: got {num_picks}, max {max_picks}"
+            raise ValueError(msg)
+
+        selected = [cast(T, label_map[label]) for label in result.content.selections]
+        return MultiPick(selections=selected, reason=result.content.reason)
+
+    async def extract[T](
+        self,
+        text: str,
+        as_type: type[T],
+        *,
+        prompt: AnyPromptType | None = None,
+    ) -> T:
+        """Extract single instance of type from text.
+
+        Args:
+            text: Text to extract from
+            as_type: Type to extract
+            prompt: Optional custom prompt
+        """
+        item_model = Schema.for_class_ctor(as_type)
+        final_prompt = prompt or f"Extract {as_type.__name__} from: {text}"
+
+        class Extraction(Schema):
+            instance: item_model  # type: ignore
+            # explanation: str | None = None
+
+        result = await self.agent.run(final_prompt, output_type=Extraction)
+        return as_type(**result.content.instance.model_dump())
+
+    async def extract_multiple[T](
+        self,
+        text: str,
+        as_type: type[T],
+        *,
+        min_items: int = 1,
+        max_items: int | None = None,
+        prompt: AnyPromptType | None = None,
+    ) -> list[T]:
+        """Extract multiple instances of type from text.
+
+        Args:
+            text: Text to extract from
+            as_type: Type to extract
+            min_items: Minimum number of instances to extract
+            max_items: Maximum number of instances (None=unlimited)
+            prompt: Optional custom prompt
+        """
+        item_model = Schema.for_class_ctor(as_type)
+        final_prompt = prompt or "\n".join([
+            f"Extract {as_type.__name__} instances from text.",
+            # "Requirements:",
+            # f"- Extract at least {min_items} instances",
+            # f"- Extract at most {max_items} instances" if max_items else "",
+            "\nText to analyze:",
+            text,
+        ])
+        # Create model for individual instance
+
+        class Extraction(Schema):
+            instances: list[item_model]  # type: ignore
+            # explanation: str | None = None
+
+        result = await self.agent.run(final_prompt, output_type=Extraction)
+        num_instances = len(result.content.instances)  # Validate counts
+        if len(result.content.instances) < min_items:
+            msg = f"Found only {num_instances} instances, need {min_items}"
+            raise ValueError(msg)
+
+        if max_items and num_instances > max_items:
+            msg = f"Found {num_instances} instances, max is {max_items}"
+            raise ValueError(msg)
+        return [
+            as_type(**instance.data if hasattr(instance, "data") else instance.model_dump())
+            for instance in result.content.instances
+        ]