autobyteus 1.1.0__py3-none-any.whl → 1.1.2__py3-none-any.whl

autobyteus/llm/ollama_provider.py

@@ -2,6 +2,7 @@ from autobyteus.llm.models import LLMModel
 from autobyteus.llm.api.ollama_llm import OllamaLLM
 from autobyteus.llm.providers import LLMProvider
 from autobyteus.llm.utils.llm_config import LLMConfig, TokenPricingConfig
+from autobyteus.llm.ollama_provider_resolver import OllamaProviderResolver
 from typing import TYPE_CHECKING
 import os
 import logging
@@ -39,7 +40,7 @@ class OllamaModelProvider:
         try:
             from autobyteus.llm.llm_factory import LLMFactory  # Local import to avoid circular dependency
             
-            ollama_host = os.getenv('OLLAMA_HOST', OllamaLLM.DEFAULT_OLLAMA_HOST)
+            ollama_host = os.getenv('DEFAULT_OLLAMA_HOST', OllamaLLM.DEFAULT_OLLAMA_HOST)
             
             if not OllamaModelProvider.is_valid_url(ollama_host):
                 logger.error(f"Invalid Ollama host URL: {ollama_host}")
@@ -73,11 +74,14 @@ class OllamaModelProvider:
                     model_name = model_info.get('model')
                     if not model_name:
                         continue
+
+                    # Determine the provider based on the model name
+                    provider = OllamaProviderResolver.resolve(model_name)
                         
                     llm_model = LLMModel(
                         name=model_name,
                         value=model_name,
-                        provider=LLMProvider.OLLAMA,
+                        provider=provider,
                         llm_class=OllamaLLM,
                         canonical_name=model_name,  # Use model_name as the canonical_name
                         default_config=LLMConfig(

autobyteus/llm/ollama_provider_resolver.py

@@ -0,0 +1,44 @@
+from autobyteus.llm.providers import LLMProvider
+import logging
+
+logger = logging.getLogger(__name__)
+
+class OllamaProviderResolver:
+    """
+    A utility class to resolve the correct LLMProvider for Ollama models
+    based on keywords in their names. This helps attribute models to their
+    original creators (e.g., Google for 'gemma').
+    """
+    
+    # A mapping from keywords to providers. The list is ordered to handle
+    # potential overlaps, though current keywords are distinct.
+    KEYWORD_PROVIDER_MAP = [
+        (['gemma', 'gemini'], LLMProvider.GEMINI),
+        (['llama'], LLMProvider.GROQ),
+        (['mistral'], LLMProvider.MISTRAL),
+        (['deepseek'], LLMProvider.DEEPSEEK),
+    ]
+
+    @staticmethod
+    def resolve(model_name: str) -> LLMProvider:
+        """
+        Resolves the LLMProvider for a given model name from Ollama.
+        It checks for keywords in the model name and returns the corresponding
+        provider. If no specific provider is found, it defaults to OLLAMA.
+
+        Args:
+            model_name (str): The name of the model discovered from Ollama (e.g., 'gemma:7b').
+
+        Returns:
+            LLMProvider: The resolved provider for the model.
+        """
+        lower_model_name = model_name.lower()
+        
+        for keywords, provider in OllamaProviderResolver.KEYWORD_PROVIDER_MAP:
+            for keyword in keywords:
+                if keyword in lower_model_name:
+                    logger.debug(f"Resolved provider for model '{model_name}' to '{provider.name}' based on keyword '{keyword}'.")
+                    return provider
+                    
+        logger.debug(f"Model '{model_name}' did not match any specific provider keywords. Defaulting to OLLAMA provider.")
+        return LLMProvider.OLLAMA

autobyteus/llm/providers.py

@@ -12,3 +12,4 @@ class LLMProvider(Enum):
     DEEPSEEK = "deepseek"
     GROK = "grok"
     AUTOBYTEUS = "autobyteus"
+    KIMI = "kimi"

autobyteus/llm/token_counter/kimi_token_counter.py

@@ -0,0 +1,24 @@
+from typing import TYPE_CHECKING
+from autobyteus.llm.token_counter.openai_token_counter import OpenAITokenCounter
+from autobyteus.llm.models import LLMModel
+
+if TYPE_CHECKING:
+    from autobyteus.llm.base_llm import BaseLLM
+
+class KimiTokenCounter(OpenAITokenCounter):
+    """
+    Token counter for Kimi (Moonshot AI) models. Uses the same token counting implementation as OpenAI.
+    
+    This implementation inherits from OpenAITokenCounter as Kimi uses the same tokenization
+    approach as OpenAI's models.
+    """
+    
+    def __init__(self, model: LLMModel, llm: 'BaseLLM' = None):
+        """
+        Initialize the Kimi token counter.
+        
+        Args:
+            model (LLMModel): The Kimi model to count tokens for.
+            llm (BaseLLM, optional): The LLM instance. Defaults to None.
+        """
+        super().__init__(model, llm)

autobyteus/llm/token_counter/token_counter_factory.py

@@ -3,6 +3,7 @@ from autobyteus.llm.token_counter.openai_token_counter import OpenAITokenCounter
 from autobyteus.llm.token_counter.claude_token_counter import ClaudeTokenCounter
 from autobyteus.llm.token_counter.mistral_token_counter import MistralTokenCounter
 from autobyteus.llm.token_counter.deepseek_token_counter import DeepSeekTokenCounter
+from autobyteus.llm.token_counter.kimi_token_counter import KimiTokenCounter
 from autobyteus.llm.token_counter.base_token_counter import BaseTokenCounter
 from autobyteus.llm.models import LLMModel
 from autobyteus.llm.providers import LLMProvider
@@ -31,6 +32,8 @@ def get_token_counter(model: LLMModel, llm: 'BaseLLM') -> BaseTokenCounter:
         return DeepSeekTokenCounter(model, llm)
     elif model.provider == LLMProvider.GROK:
         return DeepSeekTokenCounter(model, llm)
+    elif model.provider == LLMProvider.KIMI:
+        return KimiTokenCounter(model, llm)
     elif model.provider == LLMProvider.OLLAMA:
         return OpenAITokenCounter(model, llm)
     elif model.provider == LLMProvider.GEMINI:

autobyteus/llm/utils/messages.py

@@ -20,9 +20,9 @@ class Message:
         self.content = content
         self.reasoning_content = reasoning_content  # Optional field for reasoning content
 
-    def to_dict(self) -> Dict[str, Union[str, None]]:
-        result = {"role": self.role.value, "content": self.content}
-        if self.reasoning_content is not None:
+    def to_dict(self) -> Dict[str, Union[str, List[Dict]]]:
+        result: Dict[str, Union[str, List[Dict]]] = {"role": self.role.value, "content": self.content}
+        if self.reasoning_content:
             result["reasoning_content"] = self.reasoning_content
         return result
 

autobyteus/tools/__init__.py

@@ -10,6 +10,7 @@ from .base_tool import BaseTool
 from .functional_tool import tool # The @tool decorator
 from .parameter_schema import ParameterSchema, ParameterDefinition, ParameterType
 from .tool_config import ToolConfig # Configuration data object, primarily for class-based tools
+from .tool_category import ToolCategory
 
 # --- Re-export specific tools for easier access ---
 
@@ -47,6 +48,7 @@ __all__ = [
     "ParameterDefinition",
     "ParameterType",
     "ToolConfig",
+    "ToolCategory",
 
     # Re-exported functional tool instances
     "ask_user_input",

autobyteus/tools/base_tool.py

@@ -9,10 +9,13 @@ from autobyteus.events.event_emitter import EventEmitter
 from autobyteus.events.event_types import EventType
 
 from .tool_meta import ToolMeta
+from .tool_state import ToolState
+
 if TYPE_CHECKING:
     from autobyteus.agent.context import AgentContext
     from autobyteus.tools.parameter_schema import ParameterSchema
     from autobyteus.tools.tool_config import ToolConfig
+    from .tool_state import ToolState
 
 logger = logging.getLogger('autobyteus')
 
@@ -25,7 +28,10 @@ class BaseTool(ABC, EventEmitter, metaclass=ToolMeta):
         self.agent_id: Optional[str] = None
         # The config is stored primarily for potential use by subclasses or future base features.
         self._config = config
-        logger.debug(f"BaseTool instance initializing for potential class {self.__class__.__name__}")
+        # Add a dedicated state dictionary for the tool instance
+        # CHANGED: Use ToolState class for explicit state management.
+        self.tool_state: 'ToolState' = ToolState()
+        logger.debug(f"BaseTool instance initializing for potential class {self.__class__.__name__}. tool_state initialized.")
 
     @classmethod
     def get_name(cls) -> str:

autobyteus/tools/functional_tool.py

@@ -8,6 +8,7 @@ from autobyteus.tools.base_tool import BaseTool
 from autobyteus.tools.parameter_schema import ParameterSchema, ParameterDefinition, ParameterType
 from autobyteus.tools.tool_config import ToolConfig
 from autobyteus.tools.registry import default_tool_registry, ToolDefinition
+from autobyteus.tools.tool_category import ToolCategory
 
 if TYPE_CHECKING:
     from autobyteus.agent.context import AgentContext 
@@ -28,15 +29,20 @@ class FunctionalTool(BaseTool):
                  config_schema: Optional[ParameterSchema],
                  is_async: bool,
                  expects_context: bool,
+                 expects_tool_state: bool,
                  func_param_names: TypingList[str],
                  instantiation_config: Optional[Dict[str, Any]] = None):
         super().__init__(config=ToolConfig(params=instantiation_config) if instantiation_config else None)
         self._original_func = original_func
         self._is_async = is_async
         self._expects_context = expects_context
+        self._expects_tool_state = expects_tool_state
         self._func_param_names = func_param_names
         self._instantiation_config = instantiation_config or {}
         
+        # This instance has its own state dictionary, inherited from BaseTool's __init__
+        # self.tool_state: Dict[str, Any] = {} # This is now handled by super().__init__()
+        
         # Override instance methods to provide specific schema info
         self.get_name = lambda: name
         self.get_description = lambda: description
@@ -65,6 +71,9 @@ class FunctionalTool(BaseTool):
         
         if self._expects_context:
             call_args['context'] = context
+        
+        if self._expects_tool_state:
+            call_args['tool_state'] = self.tool_state
             
         if self._is_async:
             return await self._original_func(**call_args)
@@ -143,15 +152,19 @@ def _get_parameter_type_from_hint(py_type: Any, param_name: str) -> Tuple[Parame
     logger.warning(f"Unmapped type hint {py_type} (actual_type: {actual_type}) for param '{param_name}'. Defaulting to ParameterType.STRING.")
     return ParameterType.STRING, None
 
-def _parse_signature(sig: inspect.Signature, tool_name: str) -> Tuple[TypingList[str], bool, ParameterSchema]:
+def _parse_signature(sig: inspect.Signature, tool_name: str) -> Tuple[TypingList[str], bool, bool, ParameterSchema]:
     func_param_names = []
     expects_context = False
+    expects_tool_state = False
     generated_arg_schema = ParameterSchema()
 
     for param_name, param_obj in sig.parameters.items():
         if param_name == "context":
             expects_context = True
-            continue 
+            continue
+        if param_name == "tool_state":
+            expects_tool_state = True
+            continue
         
         func_param_names.append(param_name)
         
@@ -177,7 +190,7 @@ def _parse_signature(sig: inspect.Signature, tool_name: str) -> Tuple[TypingList
         )
         generated_arg_schema.add_parameter(schema_param)
         
-    return func_param_names, expects_context, generated_arg_schema
+    return func_param_names, expects_context, expects_tool_state, generated_arg_schema
 
 # --- The refactored @tool decorator ---
 
@@ -196,7 +209,7 @@ def tool(
         
         sig = inspect.signature(func)
         is_async = inspect.iscoroutinefunction(func)
-        func_param_names, expects_context, gen_arg_schema = _parse_signature(sig, tool_name)
+        func_param_names, expects_context, expects_tool_state, gen_arg_schema = _parse_signature(sig, tool_name)
         
         final_arg_schema = argument_schema if argument_schema is not None else gen_arg_schema
 
@@ -209,6 +222,7 @@ def tool(
                 config_schema=config_schema,
                 is_async=is_async,
                 expects_context=expects_context,
+                expects_tool_state=expects_tool_state,
                 func_param_names=func_param_names,
                 instantiation_config=inst_config.params if inst_config else None
             )
@@ -221,7 +235,8 @@ def tool(
             argument_schema=final_arg_schema,
             config_schema=config_schema,
             custom_factory=factory,
-            tool_class=None
+            tool_class=None,
+            category=ToolCategory.LOCAL
         )
         default_tool_registry.register_tool(tool_def)
         

autobyteus/tools/mcp/call_handlers/stdio_handler.py

@@ -1,5 +1,6 @@
 # file: autobyteus/autobyteus/tools/mcp/call_handlers/stdio_handler.py
 import logging
+import asyncio
 from typing import Dict, Any, cast, TYPE_CHECKING
 
 from .base_handler import McpCallHandler
@@ -11,6 +12,9 @@ if TYPE_CHECKING:
 
 logger = logging.getLogger(__name__)
 
+# A default timeout for STDIO subprocesses to prevent indefinite hangs.
+DEFAULT_STDIO_TIMEOUT = 30  # seconds
+
 class StdioMcpCallHandler(McpCallHandler):
     """Handles MCP tool calls over a stateless STDIO transport."""
 
@@ -23,6 +27,7 @@ class StdioMcpCallHandler(McpCallHandler):
         """
         Creates a new subprocess, establishes a session, and executes the
         requested tool call. It handles 'list_tools' as a special case.
+        Includes a timeout to prevent hanging on unresponsive subprocesses.
         """
         logger.debug(f"Handling STDIO call to tool '{remote_tool_name}' on server '{config.server_id}'.")
         
@@ -39,7 +44,8 @@ class StdioMcpCallHandler(McpCallHandler):
             cwd=stdio_config.cwd
         )
 
-        try:
+        async def _perform_call():
+            """Inner function to be wrapped by the timeout."""
             # The stdio_client context manager provides the read/write streams.
             async with stdio_client(mcp_lib_stdio_params) as (read_stream, write_stream):
                 # The ClientSession is its own context manager that handles initialization.
@@ -54,6 +60,14 @@ class StdioMcpCallHandler(McpCallHandler):
                     
                     logger.debug(f"STDIO call to tool '{remote_tool_name}' on server '{config.server_id}' completed.")
                     return result
+
+        try:
+            return await asyncio.wait_for(_perform_call(), timeout=DEFAULT_STDIO_TIMEOUT)
+        except asyncio.TimeoutError:
+            error_message = (f"MCP call to '{remote_tool_name}' on server '{config.server_id}' timed out "
+                             f"after {DEFAULT_STDIO_TIMEOUT} seconds. The subprocess may have hung.")
+            logger.error(error_message)
+            raise RuntimeError(error_message) from None
         except Exception as e:
             logger.error(
                 f"An error occurred during STDIO tool call to '{remote_tool_name}' on server '{config.server_id}': {e}",