aiecs 1.0.1__py3-none-any.whl → 1.7.6__py3-none-any.whl

aiecs/llm/clients/google_function_calling_mixin.py

@@ -0,0 +1,457 @@
+"""
+Google Function Calling Mixin
+
+Provides shared implementation for Google providers (Vertex AI, Google AI)
+that use FunctionDeclaration format for Function Calling.
+"""
+
+import logging
+from typing import Dict, Any, Optional, List, Union, AsyncGenerator
+from dataclasses import dataclass
+from vertexai.generative_models import (
+    FunctionDeclaration,
+    Tool,
+)
+from google.genai.types import Schema, Type
+
+from .base_client import LLMMessage, LLMResponse
+
+logger = logging.getLogger(__name__)
+
+# Import StreamChunk from OpenAI mixin for compatibility
+try:
+    from .openai_compatible_mixin import StreamChunk
+except ImportError:
+    # Fallback if not available
+    @dataclass
+    class StreamChunk:
+        """Fallback StreamChunk definition"""
+        type: str
+        content: Optional[str] = None
+        tool_call: Optional[Dict[str, Any]] = None
+        tool_calls: Optional[List[Dict[str, Any]]] = None
+
+
+class GoogleFunctionCallingMixin:
+    """
+    Mixin class providing Google Function Calling implementation.
+    
+    This mixin can be used by Google providers (Vertex AI, Google AI)
+    that use FunctionDeclaration format for Function Calling.
+    
+    Usage:
+        class VertexAIClient(BaseLLMClient, GoogleFunctionCallingMixin):
+            async def generate_text(self, messages, tools=None, ...):
+                if tools:
+                    vertex_tools = self._convert_openai_to_google_format(tools)
+                    # Use in API call
+    """
+
+    def _convert_openai_to_google_format(
+        self, tools: List[Dict[str, Any]]
+    ) -> List[Tool]:
+        """
+        Convert OpenAI tools format to Google FunctionDeclaration format.
+        
+        Args:
+            tools: List of OpenAI-format tool dictionaries
+            
+        Returns:
+            List of Google Tool objects containing FunctionDeclaration
+        """
+        function_declarations = []
+        
+        for tool in tools:
+            if tool.get("type") == "function":
+                func = tool.get("function", {})
+                func_name = func.get("name", "")
+                func_description = func.get("description", "")
+                func_parameters = func.get("parameters", {})
+                
+                if not func_name:
+                    logger.warning(f"Skipping tool without name: {tool}")
+                    continue
+                
+                # Convert parameters schema
+                google_schema = self._convert_json_schema_to_google_schema(func_parameters)
+                
+                # Create FunctionDeclaration
+                function_declaration = FunctionDeclaration(
+                    name=func_name,
+                    description=func_description,
+                    parameters=google_schema,
+                )
+                
+                function_declarations.append(function_declaration)
+            else:
+                logger.warning(f"Unsupported tool type: {tool.get('type')}")
+        
+        # Wrap in Tool objects (Google format requires tools to be wrapped)
+        if function_declarations:
+            return [Tool(function_declarations=function_declarations)]
+        return []
+
+    def _convert_json_schema_to_google_schema(
+        self, json_schema: Dict[str, Any]
+    ) -> Schema:
+        """
+        Convert JSON Schema to Google Schema format.
+        
+        Args:
+            json_schema: JSON Schema dictionary
+            
+        Returns:
+            Google Schema object
+        """
+        schema_type = json_schema.get("type", "object")
+        properties = json_schema.get("properties", {})
+        required = json_schema.get("required", [])
+        
+        # Convert type
+        google_type = self._convert_json_type_to_google_type(schema_type)
+        
+        # Convert properties (only for object types)
+        google_properties = None
+        if schema_type == "object" and properties:
+            google_properties = {}
+            for prop_name, prop_schema in properties.items():
+                google_properties[prop_name] = self._convert_json_schema_to_google_schema(
+                    prop_schema
+                )
+        
+        # Handle array items
+        items = None
+        if schema_type == "array" and "items" in json_schema:
+            items = self._convert_json_schema_to_google_schema(json_schema["items"])
+        
+        # Create Schema
+        schema_kwargs = {
+            "type": google_type,
+        }
+        
+        if google_properties is not None:
+            schema_kwargs["properties"] = google_properties
+        
+        if required:
+            schema_kwargs["required"] = required
+        
+        if items is not None:
+            schema_kwargs["items"] = items
+        
+        schema = Schema(**schema_kwargs)
+        
+        return schema
+
+    def _convert_json_type_to_google_type(self, json_type: str) -> Type:
+        """
+        Convert JSON Schema type to Google Type enum.
+        
+        Args:
+            json_type: JSON Schema type string
+            
+        Returns:
+            Google Type enum value
+        """
+        type_mapping = {
+            "string": Type.STRING,
+            "number": Type.NUMBER,
+            "integer": Type.NUMBER,  # Google uses NUMBER for both
+            "boolean": Type.BOOLEAN,
+            "array": Type.ARRAY,
+            "object": Type.OBJECT,
+        }
+        
+        return type_mapping.get(json_type.lower(), Type.OBJECT)
+
+    def _extract_function_calls_from_google_response(
+        self, response: Any
+    ) -> Optional[List[Dict[str, Any]]]:
+        """
+        Extract function calls from Google Vertex AI response.
+        
+        Args:
+            response: Response object from Google Vertex AI API
+            
+        Returns:
+            List of function call dictionaries in OpenAI-compatible format,
+            or None if no function calls found
+        """
+        function_calls = []
+        
+        # Check for function calls in response
+        # Google Vertex AI returns function calls in different places depending on API version
+        if hasattr(response, "candidates") and response.candidates:
+            candidate = response.candidates[0]
+            
+            # Check for function_call attribute (older API)
+            if hasattr(candidate, "function_call") and candidate.function_call:
+                func_call = candidate.function_call
+                function_calls.append({
+                    "id": f"call_{len(function_calls)}",
+                    "type": "function",
+                    "function": {
+                        "name": func_call.name,
+                        "arguments": str(func_call.args) if hasattr(func_call, "args") else "{}",
+                    },
+                })
+            
+            # Check for content.parts with function_call (newer API)
+            elif hasattr(candidate, "content") and hasattr(candidate.content, "parts"):
+                for part in candidate.content.parts:
+                    if hasattr(part, "function_call") and part.function_call:
+                        func_call = part.function_call
+                        function_calls.append({
+                            "id": f"call_{len(function_calls)}",
+                            "type": "function",
+                            "function": {
+                                "name": func_call.name,
+                                "arguments": str(func_call.args) if hasattr(func_call, "args") else "{}",
+                            },
+                        })
+        
+        return function_calls if function_calls else None
+
+    def _attach_function_calls_to_response(
+        self,
+        response: LLMResponse,
+        function_calls: Optional[List[Dict[str, Any]]] = None,
+    ) -> LLMResponse:
+        """
+        Attach function call information to LLMResponse.
+        
+        Args:
+            response: LLMResponse object
+            function_calls: List of function call dictionaries
+            
+        Returns:
+            LLMResponse with function call info attached
+        """
+        if function_calls:
+            setattr(response, "tool_calls", function_calls)
+        return response
+
+    def _convert_messages_to_google_format(
+        self, messages: List[LLMMessage]
+    ) -> List[Dict[str, Any]]:
+        """
+        Convert LLMMessage list to Google message format.
+        
+        Args:
+            messages: List of LLMMessage objects
+            
+        Returns:
+            List of Google-format message dictionaries
+        """
+        google_messages = []
+        
+        for msg in messages:
+            # Google format uses "role" and "parts" structure
+            parts = []
+            
+            if msg.content:
+                parts.append({"text": msg.content})
+            
+            # Handle tool responses (role="tool")
+            if msg.role == "tool" and msg.tool_call_id:
+                # Google format uses function_response
+                # Note: This may need adjustment based on actual API format
+                if msg.content:
+                    parts.append({
+                        "function_response": {
+                            "name": msg.tool_call_id,  # May need mapping
+                            "response": {"result": msg.content},
+                        }
+                    })
+            
+            if parts:
+                google_messages.append({
+                    "role": msg.role if msg.role != "tool" else "model",  # Adjust role mapping
+                    "parts": parts,
+                })
+        
+        return google_messages
+
+    def _extract_function_calls_from_google_chunk(
+        self, chunk: Any
+    ) -> Optional[List[Dict[str, Any]]]:
+        """
+        Extract function calls from Google Vertex AI streaming chunk.
+        
+        Args:
+            chunk: Streaming chunk object from Google Vertex AI API
+            
+        Returns:
+            List of function call dictionaries in OpenAI-compatible format,
+            or None if no function calls found
+        """
+        function_calls = []
+        
+        # Check for function calls in chunk
+        if hasattr(chunk, "candidates") and chunk.candidates:
+            candidate = chunk.candidates[0]
+            
+            # Check for content.parts with function_call
+            if hasattr(candidate, "content") and hasattr(candidate.content, "parts"):
+                for part in candidate.content.parts:
+                    if hasattr(part, "function_call") and part.function_call:
+                        func_call = part.function_call
+                        function_calls.append({
+                            "id": f"call_{len(function_calls)}",
+                            "type": "function",
+                            "function": {
+                                "name": func_call.name,
+                                "arguments": str(func_call.args) if hasattr(func_call, "args") else "{}",
+                            },
+                        })
+            
+            # Check for function_call attribute directly on candidate
+            elif hasattr(candidate, "function_call") and candidate.function_call:
+                func_call = candidate.function_call
+                function_calls.append({
+                    "id": f"call_{len(function_calls)}",
+                    "type": "function",
+                    "function": {
+                        "name": func_call.name,
+                        "arguments": str(func_call.args) if hasattr(func_call, "args") else "{}",
+                    },
+                })
+        
+        return function_calls if function_calls else None
+
+    async def _stream_text_with_function_calling(
+        self,
+        model_instance: Any,
+        contents: Any,
+        generation_config: Any,
+        safety_settings: List[Any],
+        tools: Optional[List[Tool]] = None,
+        return_chunks: bool = False,
+        **kwargs,
+    ) -> AsyncGenerator[Union[str, StreamChunk], None]:
+        """
+        Stream text with Function Calling support (Google Vertex AI format).
+
+        Args:
+            model_instance: GenerativeModel instance (should include system_instruction)
+            contents: Input contents (string or list of Content objects)
+            generation_config: GenerationConfig object
+            safety_settings: List of SafetySetting objects
+            tools: List of Tool objects (Google format)
+            return_chunks: If True, returns StreamChunk objects; if False, returns str tokens only
+            **kwargs: Additional arguments
+
+        Yields:
+            str or StreamChunk: Text tokens or StreamChunk objects
+        """
+        # Build API call parameters
+        api_params = {
+            "contents": contents,
+            "generation_config": generation_config,
+            "safety_settings": safety_settings,
+            "stream": True,
+        }
+        
+        # Add tools if available
+        if tools:
+            api_params["tools"] = tools
+        
+        # Add any additional kwargs
+        api_params.update(kwargs)
+        
+        # Get streaming response
+        stream_response = model_instance.generate_content(**api_params)
+        
+        # Accumulator for tool calls
+        tool_calls_accumulator: Dict[str, Dict[str, Any]] = {}
+        
+        # Stream chunks
+        import asyncio
+        first_chunk_checked = False
+        
+        for chunk in stream_response:
+            # Yield control to event loop
+            await asyncio.sleep(0)
+            
+            # Check for prompt-level safety blocks
+            if not first_chunk_checked and hasattr(chunk, "prompt_feedback"):
+                pf = chunk.prompt_feedback
+                if hasattr(pf, "block_reason") and pf.block_reason:
+                    block_reason = str(pf.block_reason)
+                    if block_reason not in ["BLOCKED_REASON_UNSPECIFIED", "OTHER"]:
+                        from .base_client import SafetyBlockError
+                        raise SafetyBlockError(
+                            "Prompt blocked by safety filters",
+                            block_reason=block_reason,
+                            block_type="prompt",
+                        )
+                first_chunk_checked = True
+            
+            # Extract text content and function calls
+            if hasattr(chunk, "candidates") and chunk.candidates:
+                candidate = chunk.candidates[0]
+                
+                # Check for safety blocks in response
+                if hasattr(candidate, "finish_reason"):
+                    finish_reason = candidate.finish_reason
+                    if finish_reason in ["SAFETY", "RECITATION"]:
+                        from .base_client import SafetyBlockError
+                        raise SafetyBlockError(
+                            "Response blocked by safety filters",
+                            block_type="response",
+                        )
+                
+                # Extract text from chunk parts
+                if hasattr(candidate, "content") and hasattr(candidate.content, "parts"):
+                    for part in candidate.content.parts:
+                        if hasattr(part, "text") and part.text:
+                            text_content = part.text
+                            if return_chunks:
+                                yield StreamChunk(type="token", content=text_content)
+                            else:
+                                yield text_content
+                
+                # Also check if text is directly available
+                elif hasattr(candidate, "text") and candidate.text:
+                    text_content = candidate.text
+                    if return_chunks:
+                        yield StreamChunk(type="token", content=text_content)
+                    else:
+                        yield text_content
+                
+                # Extract and accumulate function calls
+                function_calls = self._extract_function_calls_from_google_chunk(chunk)
+                if function_calls:
+                    for func_call in function_calls:
+                        call_id = func_call["id"]
+                        
+                        # Initialize accumulator if needed
+                        if call_id not in tool_calls_accumulator:
+                            tool_calls_accumulator[call_id] = func_call.copy()
+                        else:
+                            # Update accumulator (merge arguments if needed)
+                            existing_call = tool_calls_accumulator[call_id]
+                            if func_call["function"]["name"]:
+                                existing_call["function"]["name"] = func_call["function"]["name"]
+                            if func_call["function"]["arguments"]:
+                                # Merge arguments (Google may send partial arguments)
+                                existing_args = existing_call["function"].get("arguments", "{}")
+                                new_args = func_call["function"]["arguments"]
+                                # Simple merge: append new args (may need JSON parsing for proper merge)
+                                if new_args and new_args != "{}":
+                                    existing_call["function"]["arguments"] = new_args
+                        
+                        # Yield tool call update if return_chunks=True
+                        if return_chunks:
+                            yield StreamChunk(
+                                type="tool_call",
+                                tool_call=tool_calls_accumulator[call_id].copy(),
+                            )
+        
+        # At the end of stream, yield complete tool_calls if any
+        if tool_calls_accumulator and return_chunks:
+            complete_tool_calls = list(tool_calls_accumulator.values())
+            yield StreamChunk(
+                type="tool_calls",
+                tool_calls=complete_tool_calls,
+            )
+

aiecs/llm/clients/googleai_client.py

@@ -0,0 +1,241 @@
+import logging
+import os
+from typing import Optional, List, AsyncGenerator
+
+from google import genai
+from google.genai import types
+
+from aiecs.llm.clients.base_client import (
+    BaseLLMClient,
+    LLMMessage,
+    LLMResponse,
+    ProviderNotAvailableError,
+    RateLimitError,
+)
+from aiecs.config.config import get_settings
+
+logger = logging.getLogger(__name__)
+
+
+class GoogleAIClient(BaseLLMClient):
+    """Google AI (Gemini) provider client using the new google.genai SDK"""
+
+    def __init__(self):
+        super().__init__("GoogleAI")
+        self.settings = get_settings()
+        self._initialized = False
+        self._client: Optional[genai.Client] = None
+
+    def _init_google_ai(self) -> genai.Client:
+        """Lazy initialization of Google AI SDK and return the client"""
+        if not self._initialized or self._client is None:
+            api_key = self.settings.googleai_api_key or os.environ.get("GOOGLEAI_API_KEY")
+            if not api_key:
+                raise ProviderNotAvailableError("Google AI API key not configured. Set GOOGLEAI_API_KEY.")
+
+            try:
+                self._client = genai.Client(api_key=api_key)
+                self._initialized = True
+                self.logger.info("Google AI SDK (google.genai) initialized successfully.")
+            except Exception as e:
+                raise ProviderNotAvailableError(f"Failed to initialize Google AI SDK: {str(e)}")
+
+        return self._client
+
+    def _convert_messages_to_contents(
+        self, messages: List[LLMMessage]
+    ) -> List[types.Content]:
+        """Convert LLMMessage list to Google GenAI Content objects."""
+        contents = []
+        for msg in messages:
+            # Map 'assistant' role to 'model' for Google AI
+            role = "model" if msg.role == "assistant" else msg.role
+            contents.append(
+                types.Content(role=role, parts=[types.Part(text=msg.content)])
+            )
+        return contents
+
+    async def generate_text(
+        self,
+        messages: List[LLMMessage],
+        model: Optional[str] = None,
+        temperature: float = 0.7,
+        max_tokens: Optional[int] = None,
+        system_instruction: Optional[str] = None,
+        **kwargs,
+    ) -> LLMResponse:
+        """Generate text using Google AI (google.genai SDK)"""
+        client = self._init_google_ai()
+
+        # Get model name from config if not provided
+        model_name = model or self._get_default_model() or "gemini-2.5-pro"
+
+        # Get model config for default parameters
+        model_config = self._get_model_config(model_name)
+        if model_config and max_tokens is None:
+            max_tokens = model_config.default_params.max_tokens
+
+        try:
+            # Extract system message from messages if not provided
+            system_msg = None
+            user_messages = []
+            for msg in messages:
+                if msg.role == "system":
+                    system_msg = msg.content
+                else:
+                    user_messages.append(msg)
+
+            # Use provided system_instruction or extracted system message
+            final_system_instruction = system_instruction or system_msg
+
+            # Convert messages to Content objects
+            contents = self._convert_messages_to_contents(user_messages)
+
+            # Create GenerateContentConfig with all settings
+            config = types.GenerateContentConfig(
+                system_instruction=final_system_instruction,
+                temperature=temperature,
+                max_output_tokens=max_tokens or 8192,
+                top_p=kwargs.get("top_p", 0.95),
+                top_k=kwargs.get("top_k", 40),
+                safety_settings=[
+                    types.SafetySetting(
+                        category="HARM_CATEGORY_HARASSMENT",
+                        threshold="OFF",
+                    ),
+                    types.SafetySetting(
+                        category="HARM_CATEGORY_HATE_SPEECH",
+                        threshold="OFF",
+                    ),
+                    types.SafetySetting(
+                        category="HARM_CATEGORY_SEXUALLY_EXPLICIT",
+                        threshold="OFF",
+                    ),
+                    types.SafetySetting(
+                        category="HARM_CATEGORY_DANGEROUS_CONTENT",
+                        threshold="OFF",
+                    ),
+                ],
+            )
+
+            # Use async client for async operations
+            response = await client.aio.models.generate_content(
+                model=model_name,
+                contents=contents,
+                config=config,
+            )
+
+            content = response.text
+            prompt_tokens = response.usage_metadata.prompt_token_count
+            completion_tokens = response.usage_metadata.candidates_token_count
+            total_tokens = response.usage_metadata.total_token_count
+
+            # Extract cache metadata from Google AI response
+            cache_read_tokens = None
+            cache_hit = None
+            if hasattr(response.usage_metadata, "cached_content_token_count"):
+                cache_read_tokens = response.usage_metadata.cached_content_token_count
+                cache_hit = cache_read_tokens is not None and cache_read_tokens > 0
+
+            # Use config-based cost estimation
+            cost = self._estimate_cost_from_config(model_name, prompt_tokens, completion_tokens)
+
+            return LLMResponse(
+                content=content,
+                provider=self.provider_name,
+                model=model_name,
+                tokens_used=total_tokens,
+                prompt_tokens=prompt_tokens,
+                completion_tokens=completion_tokens,
+                cost_estimate=cost,
+                cache_read_tokens=cache_read_tokens,
+                cache_hit=cache_hit,
+            )
+
+        except Exception as e:
+            if "quota" in str(e).lower():
+                raise RateLimitError(f"Google AI quota exceeded: {str(e)}")
+            self.logger.error(f"Error generating text with Google AI: {e}")
+            raise
+
+    async def stream_text(  # type: ignore[override]
+        self,
+        messages: List[LLMMessage],
+        model: Optional[str] = None,
+        temperature: float = 0.7,
+        max_tokens: Optional[int] = None,
+        system_instruction: Optional[str] = None,
+        **kwargs,
+    ) -> AsyncGenerator[str, None]:
+        """Stream text generation using Google AI (google.genai SDK)"""
+        client = self._init_google_ai()
+
+        # Get model name from config if not provided
+        model_name = model or self._get_default_model() or "gemini-2.5-pro"
+
+        # Get model config for default parameters
+        model_config = self._get_model_config(model_name)
+        if model_config and max_tokens is None:
+            max_tokens = model_config.default_params.max_tokens
+
+        try:
+            # Extract system message from messages if not provided
+            system_msg = None
+            user_messages = []
+            for msg in messages:
+                if msg.role == "system":
+                    system_msg = msg.content
+                else:
+                    user_messages.append(msg)
+
+            # Use provided system_instruction or extracted system message
+            final_system_instruction = system_instruction or system_msg
+
+            # Convert messages to Content objects
+            contents = self._convert_messages_to_contents(user_messages)
+
+            # Create GenerateContentConfig with all settings
+            config = types.GenerateContentConfig(
+                system_instruction=final_system_instruction,
+                temperature=temperature,
+                max_output_tokens=max_tokens or 8192,
+                top_p=kwargs.get("top_p", 0.95),
+                top_k=kwargs.get("top_k", 40),
+                safety_settings=[
+                    types.SafetySetting(
+                        category="HARM_CATEGORY_HARASSMENT",
+                        threshold="OFF",
+                    ),
+                    types.SafetySetting(
+                        category="HARM_CATEGORY_HATE_SPEECH",
+                        threshold="OFF",
+                    ),
+                    types.SafetySetting(
+                        category="HARM_CATEGORY_SEXUALLY_EXPLICIT",
+                        threshold="OFF",
+                    ),
+                    types.SafetySetting(
+                        category="HARM_CATEGORY_DANGEROUS_CONTENT",
+                        threshold="OFF",
+                    ),
+                ],
+            )
+
+            # Use async streaming with the new SDK
+            async for chunk in client.aio.models.generate_content_stream(
+                model=model_name,
+                contents=contents,
+                config=config,
+            ):
+                if chunk.text:
+                    yield chunk.text
+
+        except Exception as e:
+            self.logger.error(f"Error streaming text with Google AI: {e}")
+            raise
+
+    async def close(self):
+        """Clean up resources"""
+        # Google GenAI SDK does not require explicit closing of a client
+        self._initialized = False
+        self._client = None