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

aiecs/domain/agent/knowledge_aware_agent.py

@@ -95,7 +95,7 @@ class KnowledgeAwareAgent(HybridAgent):
         graph_store: Optional[GraphStore] = None,
         description: Optional[str] = None,
         version: str = "1.0.0",
-        max_iterations: int = 10,
+        max_iterations: Optional[int] = None,
         enable_graph_reasoning: bool = True,
         config_manager: Optional["ConfigManagerProtocol"] = None,
         checkpointer: Optional["CheckpointerProtocol"] = None,
@@ -118,7 +118,7 @@ class KnowledgeAwareAgent(HybridAgent):
             graph_store: Optional knowledge graph store
             description: Optional description
             version: Agent version
-            max_iterations: Maximum ReAct iterations
+            max_iterations: Maximum ReAct iterations (if None, uses config.max_iterations)
             enable_graph_reasoning: Whether to enable graph reasoning capabilities
             config_manager: Optional configuration manager for dynamic config
             checkpointer: Optional checkpointer for state persistence
@@ -745,6 +745,7 @@ Use graph reasoning proactively when questions involve:
                 model=self._config.llm_model,
                 temperature=self._config.temperature,
                 max_tokens=self._config.max_tokens,
+                context=context,
             )
 
             thought = response.content

aiecs/domain/agent/llm_agent.py

@@ -376,6 +376,7 @@ class LLMAgent(BaseAIAgent):
                 model=self._config.llm_model,
                 temperature=self._config.temperature,
                 max_tokens=self._config.max_tokens,
+                context=context,
             )
 
             # Extract result
@@ -513,6 +514,7 @@ class LLMAgent(BaseAIAgent):
                 model=self._config.llm_model,
                 temperature=self._config.temperature,
                 max_tokens=self._config.max_tokens,
+                context=context,
             ):
                 output_tokens.append(token)
                 yield {

aiecs/domain/agent/models.py

@@ -381,6 +381,16 @@ class AgentMetrics(BaseModel):
     p95_operation_time: Optional[float] = Field(None, ge=0, description="95th percentile operation time in seconds")
     p99_operation_time: Optional[float] = Field(None, ge=0, description="99th percentile operation time in seconds")
 
+    # Prompt cache metrics (for LLM provider-level caching observability)
+    total_llm_requests: int = Field(default=0, ge=0, description="Total number of LLM requests made")
+    cache_hits: int = Field(default=0, ge=0, description="Number of LLM requests with cache hits")
+    cache_misses: int = Field(default=0, ge=0, description="Number of LLM requests without cache hits (cache creation)")
+    cache_hit_rate: float = Field(default=0.0, ge=0.0, le=1.0, description="Prompt cache hit rate (0-1)")
+    total_cache_read_tokens: int = Field(default=0, ge=0, description="Total tokens read from prompt cache")
+    total_cache_creation_tokens: int = Field(default=0, ge=0, description="Total tokens used to create cache entries")
+    estimated_cache_savings_tokens: int = Field(default=0, ge=0, description="Estimated tokens saved from cache (cache_read_tokens * 0.9)")
+    estimated_cache_savings_cost: float = Field(default=0.0, ge=0, description="Estimated cost saved from cache in USD")
+
     # Timestamps
     last_reset_at: Optional[datetime] = Field(None, description="When metrics were last reset")
     updated_at: datetime = Field(default_factory=datetime.utcnow, description="Last metrics update")

aiecs/domain/agent/tools/schema_generator.py

@@ -207,6 +207,12 @@ class ToolSchemaGenerator:
         if not hasattr(schema_class, "model_fields"):
             return properties, required
 
+        # Import PydanticUndefined for v2 compatibility
+        try:
+            from pydantic_core import PydanticUndefined
+        except ImportError:
+            PydanticUndefined = type(None)  # Fallback for Pydantic v1
+
         for field_name, field_info in schema_class.model_fields.items():
             # Build property schema
             prop_schema: Dict[str, Any] = {}
@@ -219,11 +225,18 @@ class ToolSchemaGenerator:
             if hasattr(field_info, "description") and field_info.description:
                 prop_schema["description"] = field_info.description
 
-            # Get default value
-            if field_info.default is not None and field_info.default != inspect.Parameter.empty:
-                prop_schema["default"] = field_info.default
+            # Check if required using Pydantic v2 API (preferred)
+            if hasattr(field_info, "is_required") and callable(field_info.is_required):
+                if field_info.is_required():
+                    required.append(field_name)
+                elif field_info.default is not None and field_info.default is not PydanticUndefined:
+                    prop_schema["default"] = field_info.default
             else:
-                required.append(field_name)
+                # Fallback for Pydantic v1
+                if field_info.default is None or field_info.default == inspect.Parameter.empty:
+                    required.append(field_name)
+                else:
+                    prop_schema["default"] = field_info.default
 
             properties[field_name] = prop_schema
 

aiecs/llm/callbacks/custom_callbacks.py

@@ -33,7 +33,9 @@ class RedisTokenCallbackHandler(CustomAsyncCallbackHandler):
         self.start_time = time.time()
         self.messages = messages
 
-        logger.info(f"[Callback] LLM call started for user '{self.user_id}' with {len(messages)} messages")
+        # Defensive check for None messages
+        message_count = len(messages) if messages is not None else 0
+        logger.info(f"[Callback] LLM call started for user '{self.user_id}' with {message_count} messages")
 
     async def on_llm_end(self, response: dict, **kwargs: Any) -> None:
         """Triggered when LLM call ends successfully"""
@@ -93,8 +95,8 @@ class DetailedRedisTokenCallbackHandler(CustomAsyncCallbackHandler):
         self.start_time = time.time()
         self.messages = messages
 
-        # Estimate input token count
-        self.prompt_tokens = self._estimate_prompt_tokens(messages)
+        # Estimate input token count with None check
+        self.prompt_tokens = self._estimate_prompt_tokens(messages) if messages else 0
 
         logger.info(f"[DetailedCallback] LLM call started for user '{self.user_id}' with estimated {self.prompt_tokens} prompt tokens")
 
@@ -144,7 +146,10 @@ class DetailedRedisTokenCallbackHandler(CustomAsyncCallbackHandler):
 
     def _estimate_prompt_tokens(self, messages: List[dict]) -> int:
         """Estimate token count for input messages"""
-        total_chars = sum(len(msg.get("content", "")) for msg in messages)
+        if not messages:
+            return 0
+        # Use `or ""` to handle both missing key AND None value
+        total_chars = sum(len(msg.get("content") or "") for msg in messages)
         # Rough estimation: 4 characters ≈ 1 token
         return total_chars // 4
 

aiecs/llm/client_factory.py

@@ -7,6 +7,8 @@ from .clients.openai_client import OpenAIClient
 from .clients.vertex_client import VertexAIClient
 from .clients.googleai_client import GoogleAIClient
 from .clients.xai_client import XAIClient
+from .clients.openrouter_client import OpenRouterClient
+from .clients.openai_compatible_mixin import StreamChunk
 from .callbacks.custom_callbacks import CustomAsyncCallbackHandler
 
 if TYPE_CHECKING:
@@ -20,6 +22,7 @@ class AIProvider(str, Enum):
     VERTEX = "Vertex"
     GOOGLEAI = "GoogleAI"
     XAI = "xAI"
+    OPENROUTER = "OpenRouter"
 
 
 class LLMClientFactory:
@@ -132,6 +135,8 @@ class LLMClientFactory:
             return GoogleAIClient()
         elif provider == AIProvider.XAI:
             return XAIClient()
+        elif provider == AIProvider.OPENROUTER:
+            return OpenRouterClient()
         else:
             raise ValueError(f"Unsupported provider: {provider}")
 
@@ -261,14 +266,16 @@ class LLMClientManager:
         final_provider = context_provider or provider or AIProvider.OPENAI
         final_model = context_model or model
 
-        # Convert string prompt to messages format
-        if isinstance(messages, str):
+        # Convert string prompt to messages format and handle None
+        if messages is None:
+            messages = []
+        elif isinstance(messages, str):
             messages = [LLMMessage(role="user", content=messages)]
 
         # Execute on_llm_start callbacks
         if callbacks:
             # Convert LLMMessage objects to dictionaries for callbacks
-            messages_dict = [{"role": msg.role, "content": msg.content} for msg in messages]
+            messages_dict = [{"role": msg.role, "content": msg.content} for msg in messages] if messages else []
             for callback in callbacks:
                 try:
                     await callback.on_llm_start(
@@ -371,14 +378,16 @@ class LLMClientManager:
         final_provider = context_provider or provider or AIProvider.OPENAI
         final_model = context_model or model
 
-        # Convert string prompt to messages format
-        if isinstance(messages, str):
+        # Convert string prompt to messages format and handle None
+        if messages is None:
+            messages = []
+        elif isinstance(messages, str):
             messages = [LLMMessage(role="user", content=messages)]
 
         # Execute on_llm_start callbacks
         if callbacks:
             # Convert LLMMessage objects to dictionaries for callbacks
-            messages_dict = [{"role": msg.role, "content": msg.content} for msg in messages]
+            messages_dict = [{"role": msg.role, "content": msg.content} for msg in messages] if messages else []
             for callback in callbacks:
                 try:
                     await callback.on_llm_start(
@@ -407,7 +416,11 @@ class LLMClientManager:
                 max_tokens=max_tokens,
                 **kwargs,
             ):
-                collected_content += chunk
+                # Handle StreamChunk objects (when return_chunks=True or function calling)
+                if hasattr(chunk, 'content') and chunk.content:
+                    collected_content += chunk.content
+                elif isinstance(chunk, str):
+                    collected_content += chunk
                 yield chunk
 
             # Create a response object for callbacks (streaming doesn't return LLMResponse directly)

aiecs/llm/clients/base_client.py

@@ -1,6 +1,6 @@
 from abc import ABC, abstractmethod
-from typing import Dict, Any, Optional, List, AsyncGenerator
-from dataclasses import dataclass
+from typing import Dict, Any, Optional, List, AsyncGenerator, Union
+from dataclasses import dataclass, field
 import logging
 
 logger = logging.getLogger(__name__)
@@ -35,6 +35,7 @@ class LLMMessage:
     Attributes:
         role: Message role - "system", "user", "assistant", or "tool"
         content: Text content of the message (None when using tool calls)
+        images: List of image sources (URLs, base64 data URIs, or file paths) for vision support
         tool_calls: Tool call information for assistant messages
         tool_call_id: Tool call ID for tool response messages
         cache_control: Cache control marker for prompt caching support
@@ -42,6 +43,7 @@ class LLMMessage:
 
     role: str  # "system", "user", "assistant", "tool"
     content: Optional[str] = None  # None when using tool calls
+    images: List[Union[str, Dict[str, Any]]] = field(default_factory=list)  # Image sources for vision support
     tool_calls: Optional[List[Dict[str, Any]]] = None  # For assistant messages with tool calls
     tool_call_id: Optional[str] = None  # For tool messages
     cache_control: Optional[CacheControl] = None  # Cache control for prompt caching
@@ -139,7 +141,11 @@ class SafetyBlockError(LLMClientError):
         if self.block_type:
             msg += f" (Block type: {self.block_type})"
         if self.safety_ratings:
-            categories = [r.get("category", "UNKNOWN") for r in self.safety_ratings if r.get("blocked")]
+            # Safely extract categories, handling potential non-dict elements
+            categories = []
+            for r in self.safety_ratings:
+                if isinstance(r, dict) and r.get("blocked"):
+                    categories.append(r.get("category", "UNKNOWN"))
             if categories:
                 msg += f" (Categories: {', '.join(categories)})"
         return msg
@@ -159,9 +165,28 @@ class BaseLLMClient(ABC):
         model: Optional[str] = None,
         temperature: float = 0.7,
         max_tokens: Optional[int] = None,
+        context: Optional[Dict[str, Any]] = None,
         **kwargs,
     ) -> LLMResponse:
-        """Generate text using the provider's API"""
+        """
+        Generate text using the provider's API.
+
+        Args:
+            messages: List of conversation messages
+            model: Model name (optional, uses default if not provided)
+            temperature: Sampling temperature (0.0 to 1.0)
+            max_tokens: Maximum tokens to generate
+            context: Optional context dictionary containing metadata such as:
+                - user_id: User identifier for tracking/billing
+                - tenant_id: Tenant identifier for multi-tenant setups
+                - request_id: Request identifier for tracing
+                - session_id: Session identifier
+                - Any other custom metadata for observability or middleware
+            **kwargs: Additional provider-specific parameters
+
+        Returns:
+            LLMResponse with generated text and metadata
+        """
 
     @abstractmethod
     async def stream_text(
@@ -170,9 +195,28 @@ class BaseLLMClient(ABC):
         model: Optional[str] = None,
         temperature: float = 0.7,
         max_tokens: Optional[int] = None,
+        context: Optional[Dict[str, Any]] = None,
         **kwargs,
     ) -> AsyncGenerator[str, None]:
-        """Stream text generation using the provider's API"""
+        """
+        Stream text generation using the provider's API.
+
+        Args:
+            messages: List of conversation messages
+            model: Model name (optional, uses default if not provided)
+            temperature: Sampling temperature (0.0 to 1.0)
+            max_tokens: Maximum tokens to generate
+            context: Optional context dictionary containing metadata such as:
+                - user_id: User identifier for tracking/billing
+                - tenant_id: Tenant identifier for multi-tenant setups
+                - request_id: Request identifier for tracing
+                - session_id: Session identifier
+                - Any other custom metadata for observability or middleware
+            **kwargs: Additional provider-specific parameters
+
+        Yields:
+            Text tokens as they are generated
+        """
 
     @abstractmethod
     async def close(self):
@@ -217,6 +261,7 @@ class BaseLLMClient(ABC):
                     LLMMessage(
                         role=msg.role,
                         content=msg.content,
+                        images=msg.images,
                         tool_calls=msg.tool_calls,
                         tool_call_id=msg.tool_call_id,
                         cache_control=CacheControl(type="ephemeral"),

aiecs/llm/clients/google_function_calling_mixin.py

@@ -5,6 +5,7 @@ Provides shared implementation for Google providers (Vertex AI, Google AI)
 that use FunctionDeclaration format for Function Calling.
 """
 
+import json
 import logging
 from typing import Dict, Any, Optional, List, Union, AsyncGenerator
 from dataclasses import dataclass
@@ -12,8 +13,6 @@ from vertexai.generative_models import (
     FunctionDeclaration,
     Tool,
 )
-from google.genai.types import Schema, Type
-
 from .base_client import LLMMessage, LLMResponse
 
 logger = logging.getLogger(__name__)
@@ -32,13 +31,46 @@ except ImportError:
         tool_calls: Optional[List[Dict[str, Any]]] = None
 
 
+def _serialize_function_args(args) -> str:
+    """
+    Safely serialize function call arguments to JSON string.
+
+    Handles MapComposite/protobuf objects from Vertex AI by converting
+    them to regular dicts before JSON serialization.
+
+    Args:
+        args: Function call arguments (may be MapComposite, dict, or other)
+
+    Returns:
+        JSON string representation of the arguments
+    """
+    if args is None:
+        return "{}"
+
+    # Handle MapComposite/protobuf objects (they have items() method)
+    if hasattr(args, 'items'):
+        # Convert to regular dict
+        args_dict = dict(args)
+    elif isinstance(args, dict):
+        args_dict = args
+    else:
+        # Try to convert to dict if possible
+        try:
+            args_dict = dict(args)
+        except (TypeError, ValueError):
+            # Last resort: use str() but this should rarely happen
+            return str(args)
+
+    return json.dumps(args_dict, ensure_ascii=False)
+
+
 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, ...):
@@ -71,15 +103,13 @@ class GoogleFunctionCallingMixin:
                 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
+
+                # Create FunctionDeclaration with raw dict parameters
+                # Let Vertex SDK handle the schema conversion internally
                 function_declaration = FunctionDeclaration(
                     name=func_name,
                     description=func_description,
-                    parameters=google_schema,
+                    parameters=func_parameters,
                 )
                 
                 function_declarations.append(function_declaration)
@@ -91,78 +121,6 @@ class GoogleFunctionCallingMixin:
             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]]]:
@@ -191,10 +149,10 @@ class GoogleFunctionCallingMixin:
                     "type": "function",
                     "function": {
                         "name": func_call.name,
-                        "arguments": str(func_call.args) if hasattr(func_call, "args") else "{}",
+                        "arguments": _serialize_function_args(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:
@@ -205,7 +163,7 @@ class GoogleFunctionCallingMixin:
                             "type": "function",
                             "function": {
                                 "name": func_call.name,
-                                "arguments": str(func_call.args) if hasattr(func_call, "args") else "{}",
+                                "arguments": _serialize_function_args(func_call.args) if hasattr(func_call, "args") else "{}",
                             },
                         })
         
@@ -300,10 +258,10 @@ class GoogleFunctionCallingMixin:
                             "type": "function",
                             "function": {
                                 "name": func_call.name,
-                                "arguments": str(func_call.args) if hasattr(func_call, "args") else "{}",
+                                "arguments": _serialize_function_args(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
@@ -312,7 +270,7 @@ class GoogleFunctionCallingMixin:
                     "type": "function",
                     "function": {
                         "name": func_call.name,
-                        "arguments": str(func_call.args) if hasattr(func_call, "args") else "{}",
+                        "arguments": _serialize_function_args(func_call.args) if hasattr(func_call, "args") else "{}",
                     },
                 })