agnt5 0.1.0__cp39-abi3-macosx_11_0_arm64.whl

agnt5/llm/base.py

@@ -0,0 +1,315 @@
+"""
+Base classes and interfaces for LLM provider integration.
+
+Defines the common interface that all LLM providers must implement,
+along with shared types and utilities.
+"""
+
+import json
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, AsyncIterator, Dict, List, Optional, Union
+from datetime import datetime
+
+
+class Role(Enum):
+    """Message roles in conversations."""
+    SYSTEM = "system"
+    USER = "user"
+    ASSISTANT = "assistant"
+    TOOL = "tool"
+
+
+class LanguageModelType:
+    """
+    Simple model type that parses provider/model format.
+    
+    Supports formats:
+    - "provider/model" (e.g., "anthropic/claude-3-5-sonnet")
+    - "model" (auto-detects provider, e.g., "gpt-4o")
+    """
+    
+    def __init__(self, model_string: str):
+        """
+        Initialize with a model string.
+        
+        Args:
+            model_string: Either "provider/model" or just "model"
+        """
+        self.original = model_string
+        
+        if "/" in model_string:
+            # Format: "provider/model"
+            self._provider, self.value = model_string.split("/", 1)
+        else:
+            # Format: "model" - auto-detect provider
+            self.value = model_string
+            self._provider = self._detect_provider(model_string)
+    
+    @classmethod
+    def from_string(cls, model_string: str) -> "LanguageModelType":
+        """Create LanguageModelType from string."""
+        return cls(model_string)
+    
+    def _detect_provider(self, model_name: str) -> str:
+        """Auto-detect provider from model name."""
+        model_lower = model_name.lower()
+        
+        # Claude models
+        if any(keyword in model_lower for keyword in ["claude", "sonnet", "haiku", "opus"]):
+            return "anthropic"
+        
+        # OpenAI models
+        elif model_lower.startswith("gpt"):
+            return "openai"
+        
+        # Google models
+        elif model_lower.startswith("gemini"):
+            return "google"
+        
+        # Mistral models
+        elif any(model_lower.startswith(prefix) for prefix in ["mistral", "codestral"]):
+            return "mistral"
+        
+        # Together AI models (usually have namespace/model format, but handle edge cases)
+        elif any(keyword in model_lower for keyword in ["llama", "mixtral", "qwen"]) or model_name.count("/") > 0:
+            return "together"
+        
+        # Default to OpenAI for unknown models
+        else:
+            return "openai"
+    
+    def get_provider(self) -> str:
+        """Get the provider name for this model."""
+        return self._provider
+    
+    @property
+    def provider(self) -> str:
+        """Provider property for easy access."""
+        return self._provider
+    
+    @property
+    def model(self) -> str:
+        """Model name property for easy access."""
+        return self.value
+    
+    def __str__(self):
+        return self.value
+    
+    def __repr__(self):
+        return f"LanguageModelType('{self.original}' -> provider='{self._provider}', model='{self.value}')"
+    
+    def __eq__(self, other):
+        if isinstance(other, LanguageModelType):
+            return self.value == other.value and self._provider == other._provider
+        elif isinstance(other, str):
+            return self.value == other or self.original == other
+        return False
+    
+    def __hash__(self):
+        return hash((self._provider, self.value))
+
+
+@dataclass
+class TokenUsage:
+    """Token usage information."""
+    prompt_tokens: int = 0
+    completion_tokens: int = 0
+    total_tokens: int = 0
+
+
+@dataclass
+class ToolCall:
+    """Represents a tool call made by the LLM."""
+    id: str
+    name: str
+    arguments: Dict[str, Any]
+
+
+@dataclass
+class ToolResult:
+    """Represents the result of a tool call."""
+    tool_call_id: str
+    output: Optional[Any] = None
+    error: Optional[str] = None
+
+
+@dataclass
+class Message:
+    """Represents a message in a conversation."""
+    role: Role
+    content: Union[str, List[Dict[str, Any]]]
+    name: Optional[str] = None
+    tool_calls: Optional[List[ToolCall]] = None
+    tool_call_id: Optional[str] = None
+    metadata: Dict[str, Any] = field(default_factory=dict)
+    timestamp: datetime = field(default_factory=datetime.utcnow)
+
+
+@dataclass
+class LanguageModelResponse:
+    """Response from a language model."""
+    message: str
+    usage: TokenUsage
+    tool_calls: Optional[List[ToolCall]] = None
+    model: Optional[str] = None
+    finish_reason: Optional[str] = None
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+
+class LLMError(Exception):
+    """Base exception for LLM-related errors."""
+    
+    def __init__(self, message: str, provider: str = None, model: str = None, **kwargs):
+        super().__init__(message)
+        self.message = message
+        self.provider = provider
+        self.model = model
+        self.metadata = kwargs
+
+
+class LanguageModel(ABC):
+    """
+    Abstract base class for all language model providers.
+    
+    Provides a unified interface for different LLM providers with support for:
+    - Text generation with tool calling
+    - Streaming responses
+    - Message format conversion
+    - Error handling
+    """
+    
+    def __init__(
+        self,
+        llm_model: LanguageModelType,
+        system_prompt: Optional[str] = None,
+        **kwargs
+    ):
+        self.llm_model = llm_model
+        self.system_prompt = system_prompt
+        self.config = kwargs
+        
+    @property
+    def provider_name(self) -> str:
+        """Get the provider name."""
+        return self.llm_model.get_provider()
+    
+    @property
+    def model_name(self) -> str:
+        """Get the model name."""
+        return self.llm_model.value
+    
+    @abstractmethod
+    async def generate(
+        self,
+        messages: List[Message],
+        tools: Optional[List[Dict[str, Any]]] = None,
+        max_tokens: int = 1024,
+        temperature: float = 0.7,
+        top_p: float = 1.0,
+        stream: bool = False,
+        **kwargs
+    ) -> Union[LanguageModelResponse, AsyncIterator[LanguageModelResponse]]:
+        """
+        Generate a response from the language model.
+        
+        Args:
+            messages: List of conversation messages
+            tools: Optional list of available tools
+            max_tokens: Maximum tokens to generate
+            temperature: Sampling temperature (0.0 to 2.0)
+            top_p: Nucleus sampling parameter
+            stream: Whether to stream the response
+            **kwargs: Provider-specific parameters
+            
+        Returns:
+            LanguageModelResponse or async iterator for streaming
+        """
+        pass
+    
+    @abstractmethod
+    def convert_messages_to_provider_format(self, messages: List[Message]) -> List[Dict[str, Any]]:
+        """Convert internal message format to provider-specific format."""
+        pass
+    
+    @abstractmethod
+    def convert_tools_to_provider_format(self, tools: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """Convert internal tool format to provider-specific format."""
+        pass
+    
+    def prepare_system_message(self, messages: List[Message]) -> List[Message]:
+        """Prepare system message for the conversation."""
+        if not self.system_prompt:
+            return messages
+            
+        # Check if system message already exists
+        if messages and messages[0].role == Role.SYSTEM:
+            # Merge with existing system message
+            existing_content = messages[0].content
+            if isinstance(existing_content, str):
+                combined_content = f"{self.system_prompt}\n\n{existing_content}"
+            else:
+                combined_content = self.system_prompt
+            
+            messages[0].content = combined_content
+            return messages
+        else:
+            # Add new system message
+            system_message = Message(role=Role.SYSTEM, content=self.system_prompt)
+            return [system_message] + messages
+    
+    def validate_messages(self, messages: List[Message]) -> None:
+        """Validate message format and content."""
+        if not messages:
+            raise LLMError("No messages provided", provider=self.provider_name)
+        
+        for i, message in enumerate(messages):
+            if not isinstance(message.role, Role):
+                raise LLMError(f"Invalid role at message {i}: {message.role}", provider=self.provider_name)
+            
+            if not message.content and not message.tool_calls:
+                raise LLMError(f"Empty content at message {i}", provider=self.provider_name)
+    
+    def extract_tool_calls_from_response(self, response: Any) -> List[ToolCall]:
+        """Extract tool calls from provider response."""
+        # Default implementation - override in subclasses
+        return []
+    
+    async def generate_with_retry(
+        self,
+        messages: List[Message],
+        max_retries: int = 3,
+        **kwargs
+    ) -> LanguageModelResponse:
+        """Generate response with automatic retry on failure."""
+        last_error = None
+        
+        for attempt in range(max_retries + 1):
+            try:
+                result = await self.generate(messages, **kwargs)
+                if isinstance(result, AsyncIterator):
+                    # Convert streaming to single response for retry logic
+                    response_text = ""
+                    async for chunk in result:
+                        response_text += chunk.message
+                    return LanguageModelResponse(
+                        message=response_text,
+                        usage=TokenUsage(),
+                        model=self.model_name
+                    )
+                return result
+            except Exception as e:
+                last_error = e
+                if attempt < max_retries:
+                    # Exponential backoff
+                    import asyncio
+                    await asyncio.sleep(2 ** attempt)
+                    continue
+                break
+        
+        raise LLMError(
+            f"Failed after {max_retries + 1} attempts: {last_error}",
+            provider=self.provider_name,
+            model=self.model_name
+        ) from last_error

agnt5/llm/google.py

@@ -0,0 +1,373 @@
+"""
+Google Gemini integration for AGNT5 SDK.
+
+Provides integration with Google's Gemini models including proper message conversion,
+tool calling, and streaming support.
+"""
+
+import json
+import os
+from typing import Any, AsyncIterator, Dict, List, Optional, Union
+
+from .base import (
+    LanguageModel,
+    LanguageModelResponse,
+    LanguageModelType,
+    LLMError,
+    Message,
+    Role,
+    TokenUsage,
+    ToolCall,
+)
+
+try:
+    import google.generativeai as genai
+    from google.generativeai.types import (
+        GenerateContentResponse,
+        Content,
+        Part,
+        FunctionCall,
+        FunctionResponse,
+    )
+    GOOGLE_AVAILABLE = True
+except ImportError:
+    GOOGLE_AVAILABLE = False
+    # Define placeholder types for type hints when Google SDK is not available
+    Content = Any
+    Part = Any
+    FunctionCall = Any
+    FunctionResponse = Any
+    GenerateContentResponse = Any
+
+
+class GoogleError(LLMError):
+    """Google Gemini-specific errors."""
+    pass
+
+
+class GoogleLanguageModel(LanguageModel):
+    """
+    Google Gemini language model implementation.
+    
+    Supports all Gemini models with proper message conversion, tool calling,
+    and streaming capabilities.
+    """
+    
+    def __init__(
+        self,
+        llm_model: LanguageModelType,
+        api_key: Optional[str] = None,
+        system_prompt: Optional[str] = None,
+        **kwargs
+    ):
+        if not GOOGLE_AVAILABLE:
+            raise GoogleError("Google AI library not installed. Install with: pip install google-generativeai")
+        
+        super().__init__(llm_model, system_prompt, **kwargs)
+        
+        # Get API key
+        self.api_key = api_key or os.getenv("GOOGLE_API_KEY")
+        if not self.api_key:
+            raise GoogleError("Google API key required. Set GOOGLE_API_KEY or pass api_key parameter")
+        
+        # Configure the client
+        genai.configure(api_key=self.api_key)
+        
+        # Validate model is supported by Google
+        if not self.model_name.startswith("gemini"):
+            raise GoogleError(f"Model {self.model_name} is not a Google Gemini model")
+        
+        # Initialize the model
+        self.model = genai.GenerativeModel(self.model_name)
+    
+    async def generate(
+        self,
+        messages: List[Message],
+        tools: Optional[List[Dict[str, Any]]] = None,
+        max_tokens: int = 1024,
+        temperature: float = 0.7,
+        top_p: float = 1.0,
+        stream: bool = False,
+        **kwargs
+    ) -> Union[LanguageModelResponse, AsyncIterator[LanguageModelResponse]]:
+        """Generate response using Google Gemini."""
+        try:
+            # Validate and prepare messages
+            self.validate_messages(messages)
+            prepared_messages = self.prepare_system_message(messages)
+            
+            # Convert to Gemini format
+            gemini_contents = self.convert_messages_to_provider_format(prepared_messages)
+            
+            # Prepare generation config
+            generation_config = {
+                "max_output_tokens": max_tokens,
+                "temperature": temperature,
+                "top_p": top_p,
+            }
+            generation_config.update(kwargs)
+            
+            # Prepare tools if provided
+            gemini_tools = None
+            if tools:
+                gemini_tools = self.convert_tools_to_provider_format(tools)
+            
+            if stream:
+                return self._generate_stream(gemini_contents, generation_config, gemini_tools)
+            else:
+                return await self._generate_single(gemini_contents, generation_config, gemini_tools)
+                
+        except Exception as e:
+            # Handle various Google AI exceptions
+            error_msg = str(e)
+            if "API_KEY" in error_msg.upper():
+                raise GoogleError(f"Google API authentication error: {e}", provider="google", model=self.model_name) from e
+            elif "QUOTA" in error_msg.upper() or "RATE_LIMIT" in error_msg.upper():
+                raise GoogleError(f"Google API quota/rate limit error: {e}", provider="google", model=self.model_name) from e
+            else:
+                raise GoogleError(f"Google API error: {e}", provider="google", model=self.model_name) from e
+    
+    async def _generate_single(
+        self, 
+        contents: List[Content], 
+        generation_config: Dict[str, Any],
+        tools: Optional[List[Dict[str, Any]]] = None
+    ) -> LanguageModelResponse:
+        """Generate a single response."""
+        try:
+            # Create generation arguments
+            generate_kwargs = {
+                "contents": contents,
+                "generation_config": generation_config,
+            }
+            
+            if tools:
+                generate_kwargs["tools"] = tools
+            
+            # Generate content
+            response = self.model.generate_content(**generate_kwargs)
+            
+            # Extract text content
+            response_text = ""
+            tool_calls = []
+            
+            if response.candidates:
+                candidate = response.candidates[0]
+                for part in candidate.content.parts:
+                    if hasattr(part, 'text') and part.text:
+                        response_text += part.text
+                    elif hasattr(part, 'function_call') and part.function_call:
+                        # Extract function call
+                        func_call = part.function_call
+                        tool_calls.append(ToolCall(
+                            id=f"call_{hash(func_call.name)}",  # Generate ID since Gemini doesn't provide one
+                            name=func_call.name,
+                            arguments=dict(func_call.args) if func_call.args else {}
+                        ))
+            
+            # Calculate token usage (Gemini provides usage in response)
+            usage = TokenUsage()
+            if hasattr(response, 'usage_metadata') and response.usage_metadata:
+                usage = TokenUsage(
+                    prompt_tokens=response.usage_metadata.prompt_token_count,
+                    completion_tokens=response.usage_metadata.candidates_token_count,
+                    total_tokens=response.usage_metadata.total_token_count
+                )
+            
+            return LanguageModelResponse(
+                message=response_text,
+                usage=usage,
+                tool_calls=tool_calls if tool_calls else None,
+                model=self.model_name,
+                finish_reason=getattr(response.candidates[0], 'finish_reason', None) if response.candidates else None,
+                metadata={
+                    "safety_ratings": getattr(response.candidates[0], 'safety_ratings', []) if response.candidates else []
+                }
+            )
+            
+        except Exception as e:
+            raise GoogleError(f"Error generating single response: {e}", provider="google", model=self.model_name) from e
+    
+    async def _generate_stream(
+        self, 
+        contents: List[Content], 
+        generation_config: Dict[str, Any],
+        tools: Optional[List[Dict[str, Any]]] = None
+    ) -> AsyncIterator[LanguageModelResponse]:
+        """Generate streaming response."""
+        try:
+            # Create generation arguments
+            generate_kwargs = {
+                "contents": contents,
+                "generation_config": generation_config,
+                "stream": True,
+            }
+            
+            if tools:
+                generate_kwargs["tools"] = tools
+            
+            # Generate streaming content
+            response_stream = self.model.generate_content(**generate_kwargs)
+            
+            for chunk in response_stream:
+                if chunk.candidates:
+                    candidate = chunk.candidates[0]
+                    for part in candidate.content.parts:
+                        if hasattr(part, 'text') and part.text:
+                            yield LanguageModelResponse(
+                                message=part.text,
+                                usage=TokenUsage(),
+                                model=self.model_name
+                            )
+                        elif hasattr(part, 'function_call') and part.function_call:
+                            # Handle function calls in streaming
+                            func_call = part.function_call
+                            yield LanguageModelResponse(
+                                message="",
+                                usage=TokenUsage(),
+                                tool_calls=[ToolCall(
+                                    id=f"call_{hash(func_call.name)}",
+                                    name=func_call.name,
+                                    arguments=dict(func_call.args) if func_call.args else {}
+                                )],
+                                model=self.model_name
+                            )
+                            
+        except Exception as e:
+            raise GoogleError(f"Error generating streaming response: {e}", provider="google", model=self.model_name) from e
+    
+    def convert_messages_to_provider_format(self, messages: List[Message]) -> List[Content]:
+        """Convert internal messages to Gemini Content format."""
+        contents = []
+        system_instruction = None
+        
+        for message in messages:
+            # Handle system messages separately
+            if message.role == Role.SYSTEM:
+                system_instruction = message.content
+                continue
+            
+            # Convert role
+            if message.role == Role.USER:
+                role = "user"
+            elif message.role == Role.ASSISTANT:
+                role = "model"  # Gemini uses "model" instead of "assistant"
+            elif message.role == Role.TOOL:
+                # Tool results are handled as function responses
+                role = "function"
+            else:
+                continue  # Skip unsupported roles
+            
+            # Prepare content parts
+            parts = []
+            
+            if isinstance(message.content, str):
+                if message.content:  # Only add non-empty content
+                    parts.append(Part(text=message.content))
+            elif isinstance(message.content, list):
+                parts.extend(self._convert_content_blocks(message.content))
+            
+            # Handle tool calls for assistant messages
+            if message.tool_calls and message.role == Role.ASSISTANT:
+                for tool_call in message.tool_calls:
+                    parts.append(Part(
+                        function_call=FunctionCall(
+                            name=tool_call.name,
+                            args=tool_call.arguments
+                        )
+                    ))
+            
+            # Handle tool results
+            if message.tool_call_id and message.role == Role.TOOL:
+                # Tool results in Gemini are handled as function responses
+                parts.append(Part(
+                    function_response=FunctionResponse(
+                        name=f"function_{message.tool_call_id}",  # Reconstruct function name
+                        response={"result": message.content}
+                    )
+                ))
+            
+            if parts:  # Only add content if there are parts
+                contents.append(Content(role=role, parts=parts))
+        
+        # Add system instruction to the model if present
+        if system_instruction and contents:
+            # Prepend system instruction as user message (Gemini pattern)
+            system_content = Content(
+                role="user",
+                parts=[Part(text=f"System: {system_instruction}")]
+            )
+            contents.insert(0, system_content)
+        
+        return contents
+    
+    def _convert_content_blocks(self, content_blocks: List[Dict[str, Any]]) -> List[Part]:
+        """Convert content blocks to Gemini Part format."""
+        parts = []
+        
+        for block in content_blocks:
+            if isinstance(block, str):
+                parts.append(Part(text=block))
+            elif isinstance(block, dict):
+                block_type = block.get("type", "text")
+                
+                if block_type == "text":
+                    text_content = block.get("text", str(block))
+                    if text_content:
+                        parts.append(Part(text=text_content))
+                elif block_type == "image" or block_type == "image_url":
+                    # Handle image content (Gemini supports images)
+                    # This would require additional image processing
+                    # For now, convert to text description
+                    image_desc = block.get("alt_text", "Image content")
+                    parts.append(Part(text=f"[Image: {image_desc}]"))
+                else:
+                    # Convert unknown blocks to text
+                    parts.append(Part(text=str(block)))
+            else:
+                parts.append(Part(text=str(block)))
+        
+        return parts
+    
+    def convert_tools_to_provider_format(self, tools: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """Convert tools to Gemini function format."""
+        gemini_functions = []
+        
+        for tool in tools:
+            if "function" in tool:
+                # OpenAI-style tool format
+                func = tool["function"]
+                gemini_function = {
+                    "name": func["name"],
+                    "description": func.get("description", ""),
+                    "parameters": func.get("parameters", {})
+                }
+            else:
+                # Direct format or simple format
+                gemini_function = {
+                    "name": tool.get("name", "unknown"),
+                    "description": tool.get("description", ""),
+                    "parameters": tool.get("parameters", tool.get("input_schema", {}))
+                }
+            
+            gemini_functions.append(gemini_function)
+        
+        # Wrap functions in the expected Gemini tools format
+        return [{"function_declarations": gemini_functions}]
+    
+    def extract_tool_calls_from_response(self, response: Any) -> List[ToolCall]:
+        """Extract tool calls from Gemini response."""
+        tool_calls = []
+        
+        if hasattr(response, "candidates") and response.candidates:
+            candidate = response.candidates[0]
+            for part in candidate.content.parts:
+                if hasattr(part, 'function_call') and part.function_call:
+                    func_call = part.function_call
+                    tool_calls.append(ToolCall(
+                        id=f"call_{hash(func_call.name)}",
+                        name=func_call.name,
+                        arguments=dict(func_call.args) if func_call.args else {}
+                    ))
+        
+        return tool_calls