ebk 0.4.4__py3-none-any.whl

ebk/ai/llm_providers/anthropic.py

@@ -0,0 +1,209 @@
+"""
+Anthropic Claude LLM Provider.
+
+Supports Claude models via the Anthropic API.
+"""
+
+import json
+from typing import Dict, Any, List, Optional, AsyncIterator
+import httpx
+
+from .base import BaseLLMProvider, LLMConfig, LLMResponse, ModelCapability
+
+
+class AnthropicProvider(BaseLLMProvider):
+    """
+    Anthropic Claude LLM provider.
+
+    Supports:
+    - Claude 3 models (haiku, sonnet, opus)
+    - Claude 3.5/4 models
+    - Streaming completions
+    - JSON mode (via prompting)
+    """
+
+    API_VERSION = "2023-06-01"
+    DEFAULT_MODEL = "claude-sonnet-4-20250514"
+
+    def __init__(self, config: LLMConfig):
+        """
+        Initialize Anthropic provider.
+
+        Args:
+            config: LLM configuration with api_key for Anthropic
+        """
+        super().__init__(config)
+        if not config.api_key:
+            raise ValueError("Anthropic API key is required")
+
+    @property
+    def name(self) -> str:
+        return "anthropic"
+
+    @property
+    def supported_capabilities(self) -> List[ModelCapability]:
+        return [
+            ModelCapability.TEXT_GENERATION,
+            ModelCapability.JSON_MODE,
+            ModelCapability.STREAMING,
+            ModelCapability.VISION,
+        ]
+
+    @classmethod
+    def create(
+        cls,
+        api_key: str,
+        model: str = DEFAULT_MODEL,
+        **kwargs
+    ) -> 'AnthropicProvider':
+        """
+        Create an Anthropic provider.
+
+        Args:
+            api_key: Anthropic API key
+            model: Model name (e.g., 'claude-sonnet-4-20250514', 'claude-3-5-sonnet-20241022')
+            **kwargs: Additional config parameters
+
+        Returns:
+            Configured AnthropicProvider
+
+        Example:
+            >>> provider = AnthropicProvider.create(
+            ...     api_key="sk-ant-...",
+            ...     model="claude-sonnet-4-20250514"
+            ... )
+        """
+        config = LLMConfig(
+            base_url="https://api.anthropic.com",
+            api_key=api_key,
+            model=model,
+            **kwargs
+        )
+        return cls(config)
+
+    async def initialize(self) -> None:
+        """Initialize HTTP client with Anthropic headers."""
+        self._client = httpx.AsyncClient(
+            base_url=self.config.base_url,
+            timeout=self.config.timeout,
+            headers={
+                "x-api-key": self.config.api_key,
+                "anthropic-version": self.API_VERSION,
+                "content-type": "application/json",
+            }
+        )
+
+    async def cleanup(self) -> None:
+        """Close HTTP client."""
+        if self._client:
+            await self._client.aclose()
+            self._client = None
+
+    async def complete(
+        self,
+        prompt: str,
+        system_prompt: Optional[str] = None,
+        **kwargs
+    ) -> LLMResponse:
+        """
+        Generate completion using Anthropic Claude.
+
+        Args:
+            prompt: User prompt
+            system_prompt: Optional system prompt
+            **kwargs: Additional parameters (temperature, max_tokens, etc.)
+
+        Returns:
+            LLMResponse with generated text
+        """
+        if not self._client:
+            await self.initialize()
+
+        # Build request payload - Anthropic uses messages format
+        data = {
+            "model": self.config.model,
+            "messages": [
+                {"role": "user", "content": prompt}
+            ],
+            "max_tokens": kwargs.get("max_tokens", self.config.max_tokens or 4096),
+            "temperature": kwargs.get("temperature", self.config.temperature),
+            "top_p": kwargs.get("top_p", self.config.top_p),
+        }
+
+        if system_prompt:
+            data["system"] = system_prompt
+
+        # Make request
+        response = await self._client.post("/v1/messages", json=data)
+        response.raise_for_status()
+
+        result = response.json()
+
+        # Extract content from response
+        content = ""
+        if result.get("content"):
+            for block in result["content"]:
+                if block.get("type") == "text":
+                    content += block.get("text", "")
+
+        return LLMResponse(
+            content=content,
+            model=result.get("model", self.config.model),
+            finish_reason=result.get("stop_reason"),
+            usage={
+                "prompt_tokens": result.get("usage", {}).get("input_tokens", 0),
+                "completion_tokens": result.get("usage", {}).get("output_tokens", 0),
+            },
+            raw_response=result,
+        )
+
+    async def complete_streaming(
+        self,
+        prompt: str,
+        system_prompt: Optional[str] = None,
+        **kwargs
+    ) -> AsyncIterator[str]:
+        """
+        Generate streaming completion.
+
+        Args:
+            prompt: User prompt
+            system_prompt: Optional system prompt
+            **kwargs: Additional parameters
+
+        Yields:
+            Text chunks as they are generated
+        """
+        if not self._client:
+            await self.initialize()
+
+        data = {
+            "model": self.config.model,
+            "messages": [
+                {"role": "user", "content": prompt}
+            ],
+            "max_tokens": kwargs.get("max_tokens", self.config.max_tokens or 4096),
+            "temperature": kwargs.get("temperature", self.config.temperature),
+            "top_p": kwargs.get("top_p", self.config.top_p),
+            "stream": True,
+        }
+
+        if system_prompt:
+            data["system"] = system_prompt
+
+        async with self._client.stream("POST", "/v1/messages", json=data) as response:
+            response.raise_for_status()
+
+            async for line in response.aiter_lines():
+                if line.startswith("data: "):
+                    try:
+                        chunk = json.loads(line[6:])
+                        # Handle different event types
+                        if chunk.get("type") == "content_block_delta":
+                            delta = chunk.get("delta", {})
+                            if delta.get("type") == "text_delta":
+                                yield delta.get("text", "")
+                        elif chunk.get("type") == "message_stop":
+                            break
+                    except json.JSONDecodeError:
+                        continue

ebk/ai/llm_providers/base.py

@@ -0,0 +1,295 @@
+"""
+Base abstract LLM provider interface.
+
+This module defines the abstract base class that all LLM providers must implement.
+"""
+
+import json
+import re
+from abc import ABC, abstractmethod
+from typing import Dict, Any, List, Optional, AsyncIterator
+from dataclasses import dataclass
+from enum import Enum
+
+
+class ModelCapability(Enum):
+    """Capabilities that an LLM model might support."""
+    TEXT_GENERATION = "text_generation"
+    JSON_MODE = "json_mode"
+    FUNCTION_CALLING = "function_calling"
+    STREAMING = "streaming"
+    VISION = "vision"
+    EMBEDDINGS = "embeddings"
+
+
+@dataclass
+class LLMConfig:
+    """Configuration for LLM provider."""
+
+    # Connection settings
+    base_url: str
+    api_key: Optional[str] = None
+
+    # Model settings
+    model: str = "default"
+    temperature: float = 0.7
+    max_tokens: Optional[int] = None
+    top_p: float = 0.9
+
+    # Behavior settings
+    timeout: float = 60.0
+    max_retries: int = 3
+
+    # Additional provider-specific settings
+    extra_params: Dict[str, Any] = None
+
+    def __post_init__(self):
+        if self.extra_params is None:
+            self.extra_params = {}
+
+
+@dataclass
+class LLMResponse:
+    """Response from LLM completion."""
+
+    content: str
+    model: str
+    finish_reason: Optional[str] = None
+    usage: Optional[Dict[str, int]] = None  # tokens used
+    raw_response: Optional[Dict[str, Any]] = None
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary."""
+        return {
+            "content": self.content,
+            "model": self.model,
+            "finish_reason": self.finish_reason,
+            "usage": self.usage,
+        }
+
+
+class BaseLLMProvider(ABC):
+    """
+    Abstract base class for LLM providers.
+
+    All LLM providers must implement this interface to ensure consistency
+    across different backends (Ollama, OpenAI, Anthropic, etc.).
+    """
+
+    def __init__(self, config: LLMConfig):
+        """
+        Initialize the provider with configuration.
+
+        Args:
+            config: LLM configuration
+        """
+        self.config = config
+        self._client = None
+
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        """Provider name (e.g., 'ollama', 'openai')."""
+        pass
+
+    @property
+    @abstractmethod
+    def supported_capabilities(self) -> List[ModelCapability]:
+        """List of capabilities supported by this provider."""
+        pass
+
+    @abstractmethod
+    async def initialize(self) -> None:
+        """
+        Initialize the provider (establish connections, etc.).
+
+        This is called once before first use.
+        """
+        pass
+
+    @abstractmethod
+    async def cleanup(self) -> None:
+        """
+        Cleanup resources (close connections, etc.).
+
+        This is called when the provider is no longer needed.
+        """
+        pass
+
+    @abstractmethod
+    async def complete(
+        self,
+        prompt: str,
+        system_prompt: Optional[str] = None,
+        **kwargs
+    ) -> LLMResponse:
+        """
+        Generate a text completion.
+
+        Args:
+            prompt: The user prompt
+            system_prompt: Optional system prompt to set context
+            **kwargs: Additional provider-specific parameters
+
+        Returns:
+            LLMResponse with generated text
+
+        Raises:
+            Exception: If completion fails
+        """
+        pass
+
+    async def complete_json(
+        self,
+        prompt: str,
+        system_prompt: Optional[str] = None,
+        schema: Optional[Dict[str, Any]] = None,
+        **kwargs
+    ) -> Dict[str, Any]:
+        """
+        Generate a JSON completion.
+
+        Builds a JSON-focused system/user prompt, delegates to complete(),
+        and parses the result. Subclasses can override to add provider-specific
+        behavior (e.g., Ollama's format="json" parameter).
+
+        Args:
+            prompt: The user prompt
+            system_prompt: Optional system prompt
+            schema: Optional JSON schema to validate against
+            **kwargs: Additional parameters
+
+        Returns:
+            Parsed JSON object
+
+        Raises:
+            Exception: If completion or parsing fails
+        """
+        # Enhance prompt for JSON output
+        json_system = "You are a helpful assistant that responds only in valid JSON format."
+        if system_prompt:
+            json_system = f"{system_prompt}\n\n{json_system}"
+
+        json_prompt = f"{prompt}\n\nRespond with valid JSON only. Do not include any explanation or markdown formatting."
+
+        if schema:
+            json_prompt += f"\n\nFollow this schema:\n```json\n{json.dumps(schema, indent=2)}\n```"
+
+        response = await self.complete(json_prompt, system_prompt=json_system, **kwargs)
+
+        # Parse JSON from response
+        return self._parse_json_response(response.content)
+
+    def _parse_json_response(self, content: str) -> Dict[str, Any]:
+        """
+        Parse JSON from response content.
+
+        Handles common issues like markdown code blocks.
+        """
+        try:
+            # Try direct parse first
+            return json.loads(content)
+        except json.JSONDecodeError:
+            pass
+
+        # Try to extract JSON from markdown code block
+        cleaned = content.strip()
+        if cleaned.startswith("```json"):
+            cleaned = cleaned[7:]
+        elif cleaned.startswith("```"):
+            cleaned = cleaned[3:]
+
+        if cleaned.endswith("```"):
+            cleaned = cleaned[:-3]
+
+        cleaned = cleaned.strip()
+
+        try:
+            return json.loads(cleaned)
+        except json.JSONDecodeError:
+            pass
+
+        # Try to find JSON object in text
+        json_match = re.search(r'\{.*\}', content, re.DOTALL)
+        if json_match:
+            try:
+                return json.loads(json_match.group())
+            except json.JSONDecodeError:
+                pass
+
+        # Last resort: try to find JSON array
+        json_match = re.search(r'\[.*\]', content, re.DOTALL)
+        if json_match:
+            try:
+                return json.loads(json_match.group())
+            except json.JSONDecodeError:
+                pass
+
+        raise ValueError(f"Failed to parse JSON from response: {content[:200]}...")
+
+    async def complete_streaming(
+        self,
+        prompt: str,
+        system_prompt: Optional[str] = None,
+        **kwargs
+    ) -> AsyncIterator[str]:
+        """
+        Generate a streaming text completion.
+
+        Args:
+            prompt: The user prompt
+            system_prompt: Optional system prompt
+            **kwargs: Additional parameters
+
+        Yields:
+            Text chunks as they are generated
+
+        Raises:
+            NotImplementedError: If streaming is not supported
+        """
+        raise NotImplementedError(
+            f"{self.name} does not support streaming completions"
+        )
+
+    async def get_embeddings(
+        self,
+        texts: List[str],
+        **kwargs
+    ) -> List[List[float]]:
+        """
+        Get embeddings for text inputs.
+
+        Args:
+            texts: List of texts to embed
+            **kwargs: Additional parameters
+
+        Returns:
+            List of embedding vectors
+
+        Raises:
+            NotImplementedError: If embeddings are not supported
+        """
+        raise NotImplementedError(
+            f"{self.name} does not support embeddings"
+        )
+
+    def supports_capability(self, capability: ModelCapability) -> bool:
+        """
+        Check if provider supports a specific capability.
+
+        Args:
+            capability: The capability to check
+
+        Returns:
+            True if supported
+        """
+        return capability in self.supported_capabilities
+
+    async def __aenter__(self):
+        """Async context manager entry."""
+        await self.initialize()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Async context manager exit."""
+        await self.cleanup()