flashlite 0.1.0__py3-none-any.whl

flashlite/conversation/multi_agent.py

@@ -0,0 +1,378 @@
+"""Multi-agent conversation support for agent-to-agent interactions."""
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+
+from ..types import CompletionResponse
+
+if TYPE_CHECKING:
+    from ..client import Flashlite
+
+
+@dataclass
+class Agent:
+    """
+    An agent with a name, persona, and optional model override.
+
+    Attributes:
+        name: Display name for the agent (used in transcript and message attribution)
+        system_prompt: The agent's personality, instructions, and behavior guidelines
+        model: Optional model override (uses MultiAgentChat default if None)
+
+    Example:
+        agent = Agent(
+            name="Scientist",
+            system_prompt="You are a curious scientist who loves experiments.",
+            model="gpt-4o",  # Optional: use specific model for this agent
+        )
+    """
+
+    name: str
+    system_prompt: str
+    model: str | None = None
+
+
+@dataclass
+class ChatMessage:
+    """A message in the multi-agent conversation."""
+
+    agent_name: str
+    content: str
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+
+class MultiAgentChat:
+    """
+    Manages conversations between multiple AI agents.
+
+    This class enables agent-to-agent conversations where multiple AI agents
+    can discuss, debate, or collaborate. Each agent maintains its own persona
+    and sees the conversation from its perspective.
+
+    Key features:
+    - Multiple agents with different personas and optionally different models
+    - Automatic context building from each agent's perspective
+    - Round-robin or directed turn-taking
+    - Full conversation transcript with metadata
+    - Support for injecting external messages (moderator, user input)
+
+    How it works:
+    - Each agent has a system prompt defining their persona
+    - When an agent speaks, they see:
+      - Their own previous messages as "assistant" role
+      - Other agents' messages as "user" role with name attribution
+    - This creates natural back-and-forth conversation
+
+    Example:
+        client = Flashlite(default_model="gpt-4o-mini")
+        chat = MultiAgentChat(client)
+
+        # Add agents with different personas
+        chat.add_agent(Agent(
+            name="Optimist",
+            system_prompt="You see the positive side of everything. Be concise."
+        ))
+        chat.add_agent(Agent(
+            name="Skeptic",
+            system_prompt="You question assumptions. Be concise."
+        ))
+
+        # Start with a topic
+        chat.add_message("Moderator", "Discuss: Will AI help or hurt jobs?")
+
+        # Have agents take turns
+        await chat.speak("Optimist")  # Optimist responds
+        await chat.speak("Skeptic")   # Skeptic responds to Optimist
+        await chat.speak("Optimist")  # Continue the debate
+
+        # Or use round-robin for structured turns
+        await chat.round_robin(rounds=2)
+
+        # Get formatted transcript
+        print(chat.format_transcript())
+    """
+
+    def __init__(
+        self,
+        client: "Flashlite",
+        default_model: str | None = None,
+    ):
+        """
+        Initialize a multi-agent chat.
+
+        Args:
+            client: Flashlite client for making completions
+            default_model: Default model for agents (uses client default if None)
+        """
+        self._client = client
+        self._default_model = default_model
+        self._agents: dict[str, Agent] = {}
+        self._transcript: list[ChatMessage] = []
+
+    def add_agent(self, agent: Agent) -> "MultiAgentChat":
+        """
+        Add an agent to the chat.
+
+        Args:
+            agent: Agent to add
+
+        Returns:
+            Self for method chaining
+
+        Example:
+            chat.add_agent(Agent("Alice", "You are helpful."))
+                .add_agent(Agent("Bob", "You are curious."))
+        """
+        self._agents[agent.name] = agent
+        return self
+
+    def remove_agent(self, agent_name: str) -> bool:
+        """
+        Remove an agent from the chat.
+
+        Args:
+            agent_name: Name of agent to remove
+
+        Returns:
+            True if agent was removed, False if not found
+        """
+        if agent_name in self._agents:
+            del self._agents[agent_name]
+            return True
+        return False
+
+    def add_message(
+        self,
+        agent_name: str,
+        content: str,
+        metadata: dict[str, Any] | None = None,
+    ) -> "MultiAgentChat":
+        """
+        Manually add a message to the transcript.
+
+        Useful for:
+        - Injecting moderator or facilitator prompts
+        - Adding user input to the conversation
+        - Simulating agent messages for testing
+
+        Args:
+            agent_name: Name to attribute the message to
+            content: Message content
+            metadata: Optional metadata to attach
+
+        Returns:
+            Self for method chaining
+        """
+        self._transcript.append(
+            ChatMessage(
+                agent_name=agent_name,
+                content=content,
+                metadata=metadata or {},
+            )
+        )
+        return self
+
+    async def speak(
+        self,
+        agent_name: str,
+        additional_context: str | None = None,
+        **kwargs: Any,
+    ) -> str:
+        """
+        Have an agent respond to the conversation.
+
+        The agent sees the full conversation history from their perspective:
+        - Their own previous messages appear as "assistant" messages
+        - Other agents' messages appear as "user" messages with name attribution
+
+        Args:
+            agent_name: Name of the agent to speak
+            additional_context: Optional extra context/instruction for this turn
+            **kwargs: Additional kwargs passed to client.complete()
+
+        Returns:
+            The agent's response content
+
+        Raises:
+            ValueError: If agent_name is not found
+        """
+        if agent_name not in self._agents:
+            raise ValueError(
+                f"Unknown agent: {agent_name}. Available agents: {list(self._agents.keys())}"
+            )
+
+        agent = self._agents[agent_name]
+
+        # Build messages from this agent's perspective
+        messages = self._build_messages_for(agent)
+
+        # Add any additional context as a user message
+        if additional_context:
+            messages.append({"role": "user", "content": additional_context})
+
+        # Make completion
+        response: CompletionResponse = await self._client.complete(
+            model=agent.model or self._default_model,
+            messages=messages,
+            **kwargs,
+        )
+
+        # Record in transcript with metadata
+        self._transcript.append(
+            ChatMessage(
+                agent_name=agent_name,
+                content=response.content,
+                metadata={
+                    "model": response.model,
+                    "tokens": response.usage.total_tokens if response.usage else None,
+                },
+            )
+        )
+
+        return response.content
+
+    def _build_messages_for(self, agent: Agent) -> list[dict[str, str]]:
+        """
+        Build the message history from a specific agent's perspective.
+
+        The agent's own messages become "assistant" role (what they said).
+        Other agents' messages become "user" role with speaker attribution.
+        """
+        messages: list[dict[str, str]] = []
+
+        # System prompt for this agent
+        messages.append({"role": "system", "content": agent.system_prompt})
+
+        # Add conversation history
+        for msg in self._transcript:
+            if msg.agent_name == agent.name:
+                # Agent's own previous messages
+                messages.append({"role": "assistant", "content": msg.content})
+            else:
+                # Other agents'/sources' messages - prefix with speaker name
+                messages.append({"role": "user", "content": f"[{msg.agent_name}]: {msg.content}"})
+
+        return messages
+
+    async def round_robin(
+        self,
+        rounds: int = 1,
+        **kwargs: Any,
+    ) -> list[str]:
+        """
+        Have all agents speak in turn for the specified number of rounds.
+
+        Each agent speaks once per round, in the order they were added.
+
+        Args:
+            rounds: Number of complete rounds (each agent speaks once per round)
+            **kwargs: Additional kwargs passed to speak()
+
+        Returns:
+            List of all responses in order
+        """
+        responses = []
+        agent_names = list(self._agents.keys())
+
+        for _ in range(rounds):
+            for name in agent_names:
+                response = await self.speak(name, **kwargs)
+                responses.append(response)
+
+        return responses
+
+    async def speak_sequence(
+        self,
+        agent_sequence: list[str],
+        **kwargs: Any,
+    ) -> list[str]:
+        """
+        Have agents speak in a specific sequence.
+
+        Args:
+            agent_sequence: List of agent names in desired speaking order
+            **kwargs: Additional kwargs passed to speak()
+
+        Returns:
+            List of responses in order
+        """
+        responses = []
+        for name in agent_sequence:
+            response = await self.speak(name, **kwargs)
+            responses.append(response)
+        return responses
+
+    @property
+    def transcript(self) -> list[ChatMessage]:
+        """Get a copy of the conversation transcript."""
+        return list(self._transcript)
+
+    @property
+    def agents(self) -> dict[str, Agent]:
+        """Get the registered agents."""
+        return dict(self._agents)
+
+    @property
+    def agent_names(self) -> list[str]:
+        """Get list of agent names."""
+        return list(self._agents.keys())
+
+    def format_transcript(self, include_metadata: bool = False) -> str:
+        """
+        Format the transcript as a readable string.
+
+        Args:
+            include_metadata: Whether to include metadata like tokens used
+
+        Returns:
+            Formatted transcript string
+        """
+        lines = []
+        for msg in self._transcript:
+            lines.append(f"[{msg.agent_name}]:")
+            # Indent content for readability
+            for line in msg.content.split("\n"):
+                lines.append(f"  {line}")
+            if include_metadata and msg.metadata:
+                meta_str = ", ".join(f"{k}={v}" for k, v in msg.metadata.items() if v)
+                if meta_str:
+                    lines.append(f"  ({meta_str})")
+            lines.append("")
+        return "\n".join(lines)
+
+    def get_messages_for(self, agent_name: str) -> list[dict[str, str]]:
+        """
+        Get the messages list as a specific agent would see it.
+
+        Useful for debugging or custom processing.
+
+        Args:
+            agent_name: Name of agent to get perspective for
+
+        Returns:
+            Messages list from that agent's perspective
+        """
+        if agent_name not in self._agents:
+            raise ValueError(f"Unknown agent: {agent_name}")
+        return self._build_messages_for(self._agents[agent_name])
+
+    def clear(self) -> "MultiAgentChat":
+        """
+        Clear the conversation transcript.
+
+        Does not remove agents.
+
+        Returns:
+            Self for method chaining
+        """
+        self._transcript = []
+        return self
+
+    def __len__(self) -> int:
+        """Number of messages in transcript."""
+        return len(self._transcript)
+
+    def __repr__(self) -> str:
+        return (
+            f"MultiAgentChat(agents={list(self._agents.keys())}, messages={len(self._transcript)})"
+        )

flashlite/core/__init__.py

@@ -0,0 +1,13 @@
+"""Core completion functionality."""
+
+from .completion import complete, complete_sync
+from .messages import assistant_message, format_messages, system_message, user_message
+
+__all__ = [
+    "complete",
+    "complete_sync",
+    "format_messages",
+    "user_message",
+    "system_message",
+    "assistant_message",
+]

flashlite/core/completion.py

@@ -0,0 +1,145 @@
+"""Core completion logic using litellm."""
+
+import asyncio
+from typing import Any
+
+import litellm
+
+from ..types import (
+    CompletionError,
+    CompletionRequest,
+    CompletionResponse,
+    Messages,
+    UsageInfo,
+)
+
+
+async def complete(request: CompletionRequest) -> CompletionResponse:
+    """
+    Execute a completion request using litellm.
+
+    This is the lowest-level async completion function.
+    It handles the litellm API call and response parsing.
+
+    Args:
+        request: The completion request
+
+    Returns:
+        CompletionResponse with the model's response
+
+    Raises:
+        CompletionError: If the API call fails
+    """
+    kwargs = request.to_litellm_kwargs()
+
+    try:
+        # litellm.acompletion is the async version
+        response = await litellm.acompletion(**kwargs)
+
+        # Extract content from response
+        content = ""
+        finish_reason = None
+
+        if response.choices:
+            choice = response.choices[0]
+            if choice.message and choice.message.content:
+                content = choice.message.content
+            finish_reason = getattr(choice, "finish_reason", None)
+
+        # Parse usage info
+        usage = UsageInfo.from_litellm(
+            response.usage.model_dump() if response.usage else None
+        )
+
+        return CompletionResponse(
+            content=content,
+            model=response.model or request.model,
+            finish_reason=finish_reason,
+            usage=usage,
+            raw_response=response,
+        )
+
+    except litellm.exceptions.APIConnectionError as e:
+        raise CompletionError(f"API connection error: {e}", response=e) from e
+    except litellm.exceptions.RateLimitError as e:
+        raise CompletionError(
+            f"Rate limit exceeded: {e}",
+            status_code=429,
+            response=e,
+        ) from e
+    except litellm.exceptions.APIError as e:
+        raise CompletionError(
+            f"API error: {e}",
+            status_code=getattr(e, "status_code", None),
+            response=e,
+        ) from e
+    except Exception as e:
+        raise CompletionError(f"Completion failed: {e}", response=e) from e
+
+
+def complete_sync(request: CompletionRequest) -> CompletionResponse:
+    """
+    Synchronous version of complete().
+
+    Runs the async completion in an event loop.
+    """
+    try:
+        loop = asyncio.get_running_loop()
+    except RuntimeError:
+        loop = None
+
+    if loop and loop.is_running():
+        # We're already in an async context - need to use a new thread
+        import concurrent.futures
+
+        with concurrent.futures.ThreadPoolExecutor() as executor:
+            future = executor.submit(asyncio.run, complete(request))
+            return future.result()
+    else:
+        # No running loop, we can use asyncio.run
+        return asyncio.run(complete(request))
+
+
+async def complete_simple(
+    model: str,
+    messages: Messages,
+    **kwargs: Any,
+) -> CompletionResponse:
+    """
+    Simplified completion function for quick usage.
+
+    Args:
+        model: Model identifier
+        messages: List of messages
+        **kwargs: Additional parameters passed to CompletionRequest
+
+    Returns:
+        CompletionResponse
+    """
+    # Separate known kwargs from extra kwargs
+    known_params = {
+        "temperature",
+        "max_tokens",
+        "max_completion_tokens",
+        "top_p",
+        "stop",
+        "reasoning_effort",
+    }
+
+    request_kwargs: dict[str, Any] = {}
+    extra_kwargs: dict[str, Any] = {}
+
+    for key, value in kwargs.items():
+        if key in known_params:
+            request_kwargs[key] = value
+        else:
+            extra_kwargs[key] = value
+
+    request = CompletionRequest(
+        model=model,
+        messages=messages,
+        extra_kwargs=extra_kwargs,
+        **request_kwargs,
+    )
+
+    return await complete(request)

flashlite/core/messages.py

@@ -0,0 +1,130 @@
+"""Message formatting utilities."""
+
+
+from ..types import Message, Messages, Role
+
+
+def user_message(content: str, name: str | None = None) -> Message:
+    """Create a user message."""
+    msg: Message = {"role": "user", "content": content}
+    if name:
+        msg["name"] = name
+    return msg
+
+
+def system_message(content: str) -> Message:
+    """Create a system message."""
+    return {"role": "system", "content": content}
+
+
+def assistant_message(content: str, name: str | None = None) -> Message:
+    """Create an assistant message."""
+    msg: Message = {"role": "assistant", "content": content}
+    if name:
+        msg["name"] = name
+    return msg
+
+
+def tool_message(content: str, tool_call_id: str) -> Message:
+    """Create a tool result message."""
+    return {
+        "role": "tool",
+        "content": content,
+        "tool_call_id": tool_call_id,
+    }
+
+
+def format_messages(
+    messages: Messages | str | None = None,
+    system: str | None = None,
+    user: str | None = None,
+) -> list[Message]:
+    """
+    Flexibly format messages into a standard list format.
+
+    This allows multiple ways to specify messages:
+    - As a list of message dicts (pass through)
+    - As a single string (converted to user message)
+    - Using system/user kwargs for simple single-turn
+
+    Args:
+        messages: Existing messages list or single string
+        system: System prompt to prepend
+        user: User message to append
+
+    Returns:
+        Normalized list of message dicts
+    """
+    result: list[Message] = []
+
+    # Add system message if provided
+    if system:
+        result.append(system_message(system))
+
+    # Handle messages parameter
+    if messages is not None:
+        if isinstance(messages, str):
+            # Single string becomes user message
+            result.append(user_message(messages))
+        else:
+            # List of messages - copy to avoid mutation
+            result.extend(list(messages))
+
+    # Add user message if provided separately
+    if user:
+        result.append(user_message(user))
+
+    return result
+
+
+def extract_content(message: Message) -> str:
+    """Extract text content from a message."""
+    content = message.get("content", "")
+    if isinstance(content, str):
+        return content
+    elif isinstance(content, list):
+        # Handle multimodal content (text parts)
+        text_parts = []
+        for part in content:
+            if isinstance(part, dict) and part.get("type") == "text":
+                text_parts.append(part.get("text", ""))
+            elif isinstance(part, str):
+                text_parts.append(part)
+        return "".join(text_parts)
+    return str(content) if content else ""
+
+
+def validate_messages(messages: Messages) -> list[str]:
+    """
+    Validate a list of messages.
+
+    Returns list of validation errors (empty if valid).
+    """
+    errors: list[str] = []
+    valid_roles: set[Role] = {"system", "user", "assistant", "tool"}
+
+    for i, msg in enumerate(messages):
+        if not isinstance(msg, dict):
+            errors.append(f"Message {i}: must be a dict, got {type(msg).__name__}")
+            continue
+
+        role = msg.get("role")
+        if role not in valid_roles:
+            errors.append(f"Message {i}: invalid role '{role}', must be one of {valid_roles}")
+
+        if "content" not in msg and "tool_calls" not in msg:
+            errors.append(f"Message {i}: must have 'content' or 'tool_calls'")
+
+        if role == "tool" and "tool_call_id" not in msg:
+            errors.append(f"Message {i}: tool message must have 'tool_call_id'")
+
+    return errors
+
+
+def count_messages_by_role(messages: Messages) -> dict[str, int]:
+    """Count messages by role."""
+    counts: dict[str, int] = {}
+    for msg in messages:
+        role = msg.get("role", "unknown")
+        counts[role] = counts.get(role, 0) + 1
+    return counts

flashlite/middleware/__init__.py

@@ -0,0 +1,18 @@
+"""Middleware for request/response processing."""
+
+from .base import Middleware, MiddlewareChain
+from .cache import CacheConfig, CacheMiddleware
+from .logging import CostTrackingMiddleware, LoggingMiddleware
+from .rate_limit import RateLimitMiddleware
+from .retry import RetryMiddleware
+
+__all__ = [
+    "Middleware",
+    "MiddlewareChain",
+    "RetryMiddleware",
+    "RateLimitMiddleware",
+    "CacheMiddleware",
+    "CacheConfig",
+    "LoggingMiddleware",
+    "CostTrackingMiddleware",
+]