llm-connect/llm_connect/adapter.py

"""
LLM adapter interface for pluggable model providers.

Implements abstraction layer for LLM integration, supporting
multiple providers (OpenAI, Anthropic, local models, etc.).
"""

import asyncio
from abc import ABC, abstractmethod
from typing import Dict, Any

from llm_connect.models import RunConfig, LLMResponse, BudgetTracker
from llm_connect.exceptions import LLMBudgetExceededError


class LLMAdapter(ABC):
    """
    Abstract base class for LLM providers.

    Enables pluggable LLM backends without prescribing implementation.
    Implementations can wrap OpenAI, Anthropic, or other APIs.
    """

    @abstractmethod
    def execute_prompt(
        self,
        prompt: str,
        config: RunConfig,
    ) -> LLMResponse:
        """
        Execute a prompt with the LLM.

        Args:
            prompt: Compiled prompt text
            config: Execution configuration

        Returns:
            LLMResponse with generated content

        Raises:
            Exception: On LLM API errors
        """
        pass

    async def async_execute_prompt(
        self,
        prompt: str,
        config: RunConfig,
    ) -> LLMResponse:
        """Execute a prompt asynchronously.

        Default implementation runs :meth:`execute_prompt` in a thread
        executor so that the event loop is not blocked. Subclasses may
        override with a native ``asyncio``-based implementation.

        Args:
            prompt: Compiled prompt text
            config: Execution configuration

        Returns:
            LLMResponse with generated content
        """
        return await asyncio.to_thread(self.execute_prompt, prompt, config)

    @abstractmethod
    def validate_config(self, config: RunConfig) -> bool:
        """
        Validate that configuration is supported.

        Args:
            config: Configuration to validate

        Returns:
            True if valid, False otherwise
        """
        pass

    # ── Budget helpers (call in execute_prompt implementations) ─────

    def _preflight_budget(self, config: RunConfig) -> None:
        """Raise ``LLMBudgetExceededError`` if the budget is already exhausted."""
        if config.budget_tracker is not None and config.budget_tracker.remaining() == 0:
            tracker = config.budget_tracker
            raise LLMBudgetExceededError(
                "Token budget exhausted before making request",
                total=tracker.total,
                spent=tracker.spent,
                requested=0,
            )

    def _consume_budget(self, config: RunConfig, response: LLMResponse) -> None:
        """Consume tokens from the budget tracker after a successful call."""
        if config.budget_tracker is not None:
            tokens = response.usage.get("total_tokens", 0)
            config.budget_tracker.consume(tokens)


class MockLLMAdapter(LLMAdapter):
    """
    Mock LLM adapter for testing.

    Returns deterministic responses without calling external APIs.
    """

    def __init__(self, mock_response: str = "Mock LLM response"):
        """
        Initialize mock adapter.

        Args:
            mock_response: Response to return
        """
        self.mock_response = mock_response
        self.call_count = 0
        self.last_prompt = None
        self.last_config = None

    def execute_prompt(
        self,
        prompt: str,
        config: RunConfig,
    ) -> LLMResponse:
        """
        Return mock response.

        Args:
            prompt: Prompt (stored for inspection)
            config: Config (stored for inspection)

        Returns:
            Mock LLMResponse
        """
        self._preflight_budget(config)
        self.call_count += 1
        self.last_prompt = prompt
        self.last_config = config

        prompt_tokens = len(prompt.split())
        completion_tokens = len(self.mock_response.split())
        response = LLMResponse(
            content=self.mock_response,
            model=config.model_name,
            usage={
                "prompt_tokens": prompt_tokens,
                "completion_tokens": completion_tokens,
                "total_tokens": prompt_tokens + completion_tokens,
            },
            finish_reason="stop",
            metadata={"mock": True},
        )
        self._consume_budget(config, response)
        return response

    def validate_config(self, config: RunConfig) -> bool:
        """
        Mock validation always succeeds.

        Args:
            config: Configuration

        Returns:
            Always True
        """
        return True

    def reset(self) -> None:
        """Reset mock state."""
        self.call_count = 0
        self.last_prompt = None
        self.last_config = None


class ErrorLLMAdapter(LLMAdapter):
    """
    Mock adapter that always raises an error.

    Useful for testing error handling.
    """

    def __init__(self, error_message: str = "Mock LLM error"):
        """
        Initialize error adapter.

        Args:
            error_message: Error message to raise
        """
        self.error_message = error_message

    def execute_prompt(
        self,
        prompt: str,
        config: RunConfig,
    ) -> LLMResponse:
        """
        Raise error.

        Args:
            prompt: Prompt
            config: Config

        Raises:
            RuntimeError: Always
        """
        raise RuntimeError(self.error_message)

    def validate_config(self, config: RunConfig) -> bool:
        """
        Validation succeeds.

        Args:
            config: Configuration

        Returns:
            True
        """
        return True