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

aiecs/__init__.py

@@ -5,7 +5,7 @@ A powerful Python middleware framework for building AI-powered applications
 with tool orchestration, task execution, and multi-provider LLM support.
 """
 
-__version__ = "1.7.17"
+__version__ = "1.8.4"
 __author__ = "AIECS Team"
 __email__ = "iretbl@gmail.com"
 

aiecs/application/knowledge_graph/extractors/llm_entity_extractor.py

@@ -162,7 +162,7 @@ class LLMEntityExtractor(EntityExtractor):
         Args:
             text: Input text to extract entities from
             entity_types: Optional filter for specific entity types
-            **kwargs: Additional parameters (e.g., custom prompt, examples)
+            **kwargs: Additional parameters (e.g., custom prompt, examples, context)
 
         Returns:
             List of extracted Entity objects
@@ -174,6 +174,9 @@ class LLMEntityExtractor(EntityExtractor):
         if not text or not text.strip():
             raise ValueError("Input text cannot be empty")
 
+        # Extract context from kwargs if provided
+        context = kwargs.get("context")
+
         # Build extraction prompt
         prompt = self._build_extraction_prompt(text, entity_types)
 
@@ -189,6 +192,7 @@ class LLMEntityExtractor(EntityExtractor):
                     model=self.model,
                     temperature=self.temperature,
                     max_tokens=self.max_tokens,
+                    context=context,
                 )
             # Otherwise use LLM manager with provider
             else:

aiecs/application/knowledge_graph/retrieval/query_intent_classifier.py

@@ -6,7 +6,7 @@ Uses a lightweight LLM to determine the best retrieval approach based on query c
 """
 
 import logging
-from typing import Optional, Dict, TYPE_CHECKING
+from typing import Optional, Dict, Any, TYPE_CHECKING
 from aiecs.application.knowledge_graph.retrieval.strategy_types import RetrievalStrategy
 
 if TYPE_CHECKING:
@@ -58,12 +58,13 @@ class QueryIntentClassifier:
         self.enable_caching = enable_caching
         self._cache: Dict[str, RetrievalStrategy] = {}
 
-    async def classify_intent(self, query: str) -> RetrievalStrategy:
+    async def classify_intent(self, query: str, context: Optional[Dict[str, Any]] = None) -> RetrievalStrategy:
         """
         Classify query intent and return optimal retrieval strategy.
 
         Args:
             query: Query string to classify
+            context: Optional context dictionary for tracking/observability
 
         Returns:
             RetrievalStrategy enum value
@@ -82,7 +83,7 @@ class QueryIntentClassifier:
         # Use LLM classification if client is available
         if self.llm_client is not None:
             try:
-                strategy = await self._classify_with_llm(query)
+                strategy = await self._classify_with_llm(query, context)
             except Exception as e:
                 logger.warning(f"LLM classification failed: {e}, falling back to rule-based")
                 strategy = self._classify_with_rules(query)
@@ -96,7 +97,7 @@ class QueryIntentClassifier:
 
         return strategy
 
-    async def _classify_with_llm(self, query: str) -> RetrievalStrategy:
+    async def _classify_with_llm(self, query: str, context: Optional[Dict[str, Any]] = None) -> RetrievalStrategy:
         """
         Classify query using LLM.
 
@@ -127,11 +128,12 @@ Respond with ONLY the strategy name (e.g., "MULTI_HOP"). No explanation needed."
         if self.llm_client is None:
             # Fallback to rule-based classification if no LLM client
             return self._classify_with_rules(query)
-        
+
         response = await self.llm_client.generate_text(
             messages=messages,
             temperature=0.0,  # Deterministic classification
             max_tokens=20,  # Short response
+            context=context,
         )
 
         # Parse response

aiecs/config/config.py

@@ -47,6 +47,9 @@ class Settings(BaseSettings):
     google_cse_id: str = Field(default="", alias="GOOGLE_CSE_ID")
     xai_api_key: str = Field(default="", alias="XAI_API_KEY")
     grok_api_key: str = Field(default="", alias="GROK_API_KEY")  # Backward compatibility
+    openrouter_api_key: str = Field(default="", alias="OPENROUTER_API_KEY")
+    openrouter_http_referer: str = Field(default="", alias="OPENROUTER_HTTP_REFERER")
+    openrouter_x_title: str = Field(default="", alias="OPENROUTER_X_TITLE")
 
     # LLM Models Configuration
     llm_models_config_path: str = Field(

aiecs/domain/agent/hybrid_agent.py

@@ -242,7 +242,7 @@ class HybridAgent(BaseAIAgent):
         config: AgentConfiguration,
         description: Optional[str] = None,
         version: str = "1.0.0",
-        max_iterations: int = 10,
+        max_iterations: Optional[int] = None,
         config_manager: Optional["ConfigManagerProtocol"] = None,
         checkpointer: Optional["CheckpointerProtocol"] = None,
         context_engine: Optional[Any] = None,
@@ -262,7 +262,7 @@ class HybridAgent(BaseAIAgent):
             config: Agent configuration
             description: Optional description
             version: Agent version
-            max_iterations: Maximum ReAct iterations
+            max_iterations: Maximum ReAct iterations (if None, uses config.max_iterations)
             config_manager: Optional configuration manager for dynamic config
             checkpointer: Optional checkpointer for state persistence
             context_engine: Optional context engine for persistent storage
@@ -316,7 +316,17 @@ class HybridAgent(BaseAIAgent):
 
         # Store LLM client reference (from BaseAIAgent or local)
         self.llm_client = self._llm_client if self._llm_client else llm_client
-        self._max_iterations = max_iterations
+        
+        # Use config.max_iterations if constructor parameter is None
+        # This makes max_iterations consistent with max_tokens (both configurable via config)
+        # If max_iterations is explicitly provided, it takes precedence over config
+        if max_iterations is None:
+            # Use config value (defaults to 10 if not set in config)
+            self._max_iterations = config.max_iterations
+        else:
+            # Constructor parameter explicitly provided, use it
+            self._max_iterations = max_iterations
+        
         self._system_prompt: Optional[str] = None
         self._conversation_history: List[LLMMessage] = []
         self._tool_schemas: List[Dict[str, Any]] = []
@@ -456,6 +466,24 @@ class HybridAgent(BaseAIAgent):
                     agent_id=self.agent_id,
                 )
 
+            # Extract images from task dict and merge into context
+            task_images = task.get("images")
+            if task_images:
+                # Merge images from task into context
+                # If context already has images, combine them
+                if "images" in context:
+                    existing_images = context["images"]
+                    if isinstance(existing_images, list) and isinstance(task_images, list):
+                        context["images"] = existing_images + task_images
+                    elif isinstance(existing_images, list):
+                        context["images"] = existing_images + [task_images]
+                    elif isinstance(task_images, list):
+                        context["images"] = [existing_images] + task_images
+                    else:
+                        context["images"] = [existing_images, task_images]
+                else:
+                    context["images"] = task_images
+
             # Transition to busy state
             self._transition_state(self.state.__class__.BUSY)
             self._current_task_id = task.get("task_id")
@@ -572,6 +600,24 @@ class HybridAgent(BaseAIAgent):
                 }
                 return
 
+            # Extract images from task dict and merge into context
+            task_images = task.get("images")
+            if task_images:
+                # Merge images from task into context
+                # If context already has images, combine them
+                if "images" in context:
+                    existing_images = context["images"]
+                    if isinstance(existing_images, list) and isinstance(task_images, list):
+                        context["images"] = existing_images + task_images
+                    elif isinstance(existing_images, list):
+                        context["images"] = existing_images + [task_images]
+                    elif isinstance(task_images, list):
+                        context["images"] = [existing_images] + task_images
+                    else:
+                        context["images"] = [existing_images, task_images]
+                else:
+                    context["images"] = task_images
+
             # Transition to busy state
             self._transition_state(self.state.__class__.BUSY)
             self._current_task_id = task.get("task_id")
@@ -712,6 +758,7 @@ class HybridAgent(BaseAIAgent):
                     model=self._config.llm_model,
                     temperature=self._config.temperature,
                     max_tokens=self._config.max_tokens,
+                    context=context,
                     tools=tools,
                     tool_choice="auto",
                     return_chunks=True,  # Enable tool_calls accumulation
@@ -723,6 +770,7 @@ class HybridAgent(BaseAIAgent):
                     model=self._config.llm_model,
                     temperature=self._config.temperature,
                     max_tokens=self._config.max_tokens,
+                    context=context,
                 )
 
             # Stream tokens and collect tool calls
@@ -1064,6 +1112,7 @@ class HybridAgent(BaseAIAgent):
                     model=self._config.llm_model,
                     temperature=self._config.temperature,
                     max_tokens=self._config.max_tokens,
+                    context=context,
                     tools=tools,
                     tool_choice="auto",
                 )
@@ -1074,6 +1123,7 @@ class HybridAgent(BaseAIAgent):
                     model=self._config.llm_model,
                     temperature=self._config.temperature,
                     max_tokens=self._config.max_tokens,
+                    context=context,
                 )
 
             thought_raw = response.content or ""
@@ -1330,6 +1380,9 @@ class HybridAgent(BaseAIAgent):
                 )
             )
 
+        # Collect images from context to attach to task message
+        task_images = []
+
         # Add context if provided
         if context:
             # Special handling: if context contains 'history' as a list of messages,
@@ -1340,18 +1393,40 @@ class HybridAgent(BaseAIAgent):
                 for msg in history:
                     if isinstance(msg, dict) and "role" in msg and "content" in msg:
                         # Valid message format - add as separate message
-                        messages.append(
-                            LLMMessage(
-                                role=msg["role"],
-                                content=msg["content"],
+                        # Extract images if present
+                        msg_images = msg.get("images", [])
+                        if msg_images:
+                            messages.append(
+                                LLMMessage(
+                                    role=msg["role"],
+                                    content=msg["content"],
+                                    images=msg_images if isinstance(msg_images, list) else [msg_images],
+                                )
+                            )
+                        else:
+                            messages.append(
+                                LLMMessage(
+                                    role=msg["role"],
+                                    content=msg["content"],
+                                )
                             )
-                        )
                     elif isinstance(msg, LLMMessage):
-                        # Already an LLMMessage instance
+                        # Already an LLMMessage instance (may already have images)
                         messages.append(msg)
             
-            # Format remaining context fields (excluding history) as Additional Context
-            context_without_history = {k: v for k, v in context.items() if k != "history"}
+            # Extract images from context if present
+            context_images = context.get("images")
+            if context_images:
+                if isinstance(context_images, list):
+                    task_images.extend(context_images)
+                else:
+                    task_images.append(context_images)
+            
+            # Format remaining context fields (excluding history and images) as Additional Context
+            context_without_history = {
+                k: v for k, v in context.items() 
+                if k not in ("history", "images")
+            }
             if context_without_history:
                 context_str = self._format_context(context_without_history)
                 if context_str:
@@ -1367,7 +1442,13 @@ class HybridAgent(BaseAIAgent):
             f"Task: {task}\n\n"
             f"[Iteration 1/{self._max_iterations}, remaining: {self._max_iterations - 1}]"
         )
-        messages.append(LLMMessage(role="user", content=task_message))
+        messages.append(
+            LLMMessage(
+                role="user",
+                content=task_message,
+                images=task_images if task_images else [],
+            )
+        )
 
         return messages
 

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/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,7 @@ 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
 
@@ -21,6 +22,7 @@ class AIProvider(str, Enum):
     VERTEX = "Vertex"
     GOOGLEAI = "GoogleAI"
     XAI = "xAI"
+    OPENROUTER = "OpenRouter"
 
 
 class LLMClientFactory:
@@ -133,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}")
 
@@ -262,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(
@@ -372,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(

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
@@ -163,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(
@@ -174,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):
@@ -221,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"),