daita-agents 0.1.0__py3-none-any.whl

daita/llm/grok.py

@@ -0,0 +1,114 @@
+"""
+Grok (xAI) LLM provider implementation with integrated tracing.
+"""
+import os
+import logging
+from typing import Dict, Any, Optional
+
+from ..core.exceptions import LLMError
+from .base import BaseLLMProvider
+
+logger = logging.getLogger(__name__)
+
+class GrokProvider(BaseLLMProvider):
+    """Grok (xAI) LLM provider implementation with automatic call tracing."""
+    
+    def __init__(
+        self,
+        model: str = "grok-beta",
+        api_key: Optional[str] = None,
+        **kwargs
+    ):
+        """
+        Initialize Grok provider.
+        
+        Args:
+            model: Grok model name (e.g., "grok-beta", "grok-vision-beta")
+            api_key: xAI API key
+            **kwargs: Additional Grok-specific parameters
+        """
+        # Get API key from parameter or environment
+        api_key = api_key or os.getenv("XAI_API_KEY") or os.getenv("GROK_API_KEY")
+        
+        super().__init__(model=model, api_key=api_key, **kwargs)
+        
+        # Grok-specific default parameters
+        self.default_params.update({
+            'stream': kwargs.get('stream', False),
+            'timeout': kwargs.get('timeout', 60)
+        })
+        
+        # Base URL for xAI API
+        self.base_url = kwargs.get('base_url', 'https://api.x.ai/v1')
+        
+        # Lazy-load OpenAI client (Grok uses OpenAI-compatible API)
+        self._client = None
+    
+    @property
+    def client(self):
+        """Lazy-load OpenAI client configured for xAI."""
+        if self._client is None:
+            try:
+                import openai
+                self._validate_api_key()
+                self._client = openai.AsyncOpenAI(
+                    api_key=self.api_key,
+                    base_url=self.base_url
+                )
+                logger.debug("Grok client initialized")
+            except ImportError:
+                raise LLMError(
+                    "OpenAI package not installed. Install with: pip install openai"
+                )
+        return self._client
+    
+    async def _generate_impl(self, prompt: str, **kwargs) -> str:
+        """
+        Provider-specific implementation of text generation for Grok.
+        
+        This method contains the actual Grok API call logic and is automatically
+        wrapped with tracing by the base class generate() method.
+        
+        Args:
+            prompt: Input prompt
+            **kwargs: Optional parameters
+            
+        Returns:
+            Generated text response
+        """
+        try:
+            # Merge parameters
+            params = self._merge_params(kwargs)
+            
+            # Make API call using OpenAI-compatible interface
+            response = await self.client.chat.completions.create(
+                model=self.model,
+                messages=[
+                    {"role": "user", "content": prompt}
+                ],
+                max_tokens=params.get('max_tokens'),
+                temperature=params.get('temperature'),
+                top_p=params.get('top_p'),
+                stream=params.get('stream'),
+                timeout=params.get('timeout')
+            )
+            
+            # Store usage for base class token extraction
+            self._last_usage = response.usage
+            
+            return response.choices[0].message.content
+            
+        except Exception as e:
+            logger.error(f"Grok generation failed: {str(e)}")
+            raise LLMError(f"Grok generation failed: {str(e)}")
+    
+    @property
+    def info(self) -> Dict[str, Any]:
+        """Get information about the Grok provider."""
+        base_info = super().info
+        base_info.update({
+            'base_url': self.base_url,
+            'provider_name': 'Grok (xAI)',
+            'api_compatible': 'OpenAI'
+        })
+        return base_info

daita/llm/mock.py

@@ -0,0 +1,135 @@
+"""
+Mock LLM provider for testing with integrated tracing.
+"""
+import asyncio
+import logging
+from typing import Dict, Any, Optional
+
+from .base import BaseLLMProvider
+
+logger = logging.getLogger(__name__)
+
+class MockLLMProvider(BaseLLMProvider):
+    """Mock LLM provider for testing purposes with automatic call tracing."""
+    
+    def __init__(
+        self,
+        model: str = "mock-model",
+        responses: Optional[Dict[str, str]] = None,
+        delay: float = 0.1,
+        **kwargs
+    ):
+        """
+        Initialize mock provider.
+        
+        Args:
+            model: Mock model name
+            responses: Dictionary mapping prompts to responses
+            delay: Artificial delay to simulate API calls
+            **kwargs: Additional parameters
+        """
+        # Remove api_key from kwargs to avoid conflict, then pass it explicitly
+        kwargs.pop('api_key', None)  # Remove if exists
+        super().__init__(model=model, api_key="mock-key", **kwargs)
+        
+        # Predefined responses
+        self.responses = responses or {}
+        self.delay = delay
+        
+        # Default responses
+        self.default_responses = {
+            "default": "This is a mock response from the LLM.",
+            "analyze": "Based on the data provided, here are the key insights: [mock analysis]",
+            "summarize": "Summary: [mock summary of the content]",
+            "error": "This is an error response for testing."
+        }
+        
+        # Track calls for testing
+        self.call_history = []
+    
+    async def _generate_impl(self, prompt: str, **kwargs) -> str:
+        """
+        Provider-specific implementation of mock text generation.
+        
+        This method contains the mock generation logic and is automatically
+        wrapped with tracing by the base class generate() method.
+        
+        Args:
+            prompt: Input prompt
+            **kwargs: Optional parameters
+            
+        Returns:
+            Mock response
+        """
+        # Record the call
+        self.call_history.append({
+            'prompt': prompt,
+            'params': kwargs,
+            'timestamp': asyncio.get_event_loop().time()
+        })
+        
+        # Simulate API delay
+        if self.delay > 0:
+            await asyncio.sleep(self.delay)
+        
+        # Check for specific response
+        if prompt in self.responses:
+            return self.responses[prompt]
+        
+        # Check for keyword-based responses
+        prompt_lower = prompt.lower()
+        for keyword, response in self.default_responses.items():
+            if keyword in prompt_lower:
+                return response
+        
+        # Default response
+        return f"Mock response for: {prompt[:50]}..."
+    
+    def _get_last_token_usage(self) -> Dict[str, int]:
+        """
+        Override base class method to return mock token usage.
+        
+        Provides realistic but fake token counts for testing.
+        """
+        if self.call_history:
+            # Get the last call to estimate tokens
+            last_call = self.call_history[-1]
+            prompt = last_call.get('prompt', '')
+            
+            # Mock realistic token counts
+            estimated_prompt_tokens = max(5, len(prompt) // 4)  # Rough estimate
+            # Assume a moderate response length for mocking
+            estimated_completion_tokens = max(10, estimated_prompt_tokens // 2)
+            
+            return {
+                'total_tokens': estimated_prompt_tokens + estimated_completion_tokens,
+                'prompt_tokens': estimated_prompt_tokens,
+                'completion_tokens': estimated_completion_tokens
+            }
+        
+        # Fallback to default
+        return super()._get_last_token_usage()
+    
+    def set_response(self, prompt: str, response: str) -> None:
+        """Set a specific response for a prompt."""
+        self.responses[prompt] = response
+    
+    def clear_history(self) -> None:
+        """Clear call history."""
+        self.call_history.clear()
+    
+    def get_last_call(self) -> Optional[Dict[str, Any]]:
+        """Get the last call made to the provider."""
+        return self.call_history[-1] if self.call_history else None
+    
+    @property
+    def info(self) -> Dict[str, Any]:
+        """Get information about the mock provider."""
+        base_info = super().info
+        base_info.update({
+            'provider_name': 'Mock LLM (Testing)',
+            'call_count': len(self.call_history),
+            'configured_responses': len(self.responses),
+            'delay_seconds': self.delay
+        })
+        return base_info

daita/llm/openai.py

@@ -0,0 +1,109 @@
+"""
+OpenAI LLM provider implementation with integrated tracing.
+"""
+import os
+import logging
+from typing import Dict, Any, Optional
+
+from ..core.exceptions import LLMError
+from .base import BaseLLMProvider
+
+logger = logging.getLogger(__name__)
+
+class OpenAIProvider(BaseLLMProvider):
+    """OpenAI LLM provider implementation with automatic call tracing."""
+    
+    def __init__(
+        self,
+        model: str = "gpt-4",
+        api_key: Optional[str] = None,
+        **kwargs
+    ):
+        """
+        Initialize OpenAI provider.
+        
+        Args:
+            model: OpenAI model name (e.g., "gpt-4", "gpt-3.5-turbo")
+            api_key: OpenAI API key
+            **kwargs: Additional OpenAI-specific parameters
+        """
+        # Get API key from parameter or environment
+        api_key = api_key or os.getenv("OPENAI_API_KEY")
+        
+        super().__init__(model=model, api_key=api_key, **kwargs)
+        
+        # OpenAI-specific default parameters
+        self.default_params.update({
+            'frequency_penalty': kwargs.get('frequency_penalty', 0.0),
+            'presence_penalty': kwargs.get('presence_penalty', 0.0),
+            'timeout': kwargs.get('timeout', 60)
+        })
+        
+        # Lazy-load OpenAI client
+        self._client = None
+    
+    @property
+    def client(self):
+        """Lazy-load OpenAI client."""
+        if self._client is None:
+            try:
+                import openai
+                self._validate_api_key()
+                self._client = openai.AsyncOpenAI(api_key=self.api_key)
+                logger.debug("OpenAI client initialized")
+            except ImportError:
+                raise LLMError(
+                    "OpenAI package not installed. Install with: pip install openai"
+                )
+        return self._client
+    
+    async def _generate_impl(self, prompt: str, **kwargs) -> str:
+        """
+        Provider-specific implementation of text generation for OpenAI.
+        
+        This method contains the actual OpenAI API call logic and is automatically
+        wrapped with tracing by the base class generate() method.
+        
+        Args:
+            prompt: Input prompt
+            **kwargs: Optional parameters
+            
+        Returns:
+            Generated text response
+        """
+        try:
+            # Merge parameters
+            params = self._merge_params(kwargs)
+            
+            # Make API call
+            response = await self.client.chat.completions.create(
+                model=self.model,
+                messages=[
+                    {"role": "user", "content": prompt}
+                ],
+                max_tokens=params.get('max_tokens'),
+                temperature=params.get('temperature'),
+                top_p=params.get('top_p'),
+                frequency_penalty=params.get('frequency_penalty'),
+                presence_penalty=params.get('presence_penalty'),
+                timeout=params.get('timeout')
+            )
+            
+            # Store usage for base class token extraction
+            self._last_usage = response.usage
+            
+            return response.choices[0].message.content
+            
+        except Exception as e:
+            logger.error(f"OpenAI generation failed: {str(e)}")
+            raise LLMError(f"OpenAI generation failed: {str(e)}")
+    
+    @property
+    def info(self) -> Dict[str, Any]:
+        """Get information about the OpenAI provider."""
+        base_info = super().info
+        base_info.update({
+            'provider_name': 'OpenAI',
+            'api_compatible': 'OpenAI'
+        })
+        return base_info

daita/plugins/__init__.py

@@ -0,0 +1,141 @@
+"""
+Plugin system for Daita Agents.
+
+This module provides database, API, cloud storage, search, collaboration, and MCP integrations:
+- PostgreSQL plugin for async database operations
+- MySQL plugin for async database operations
+- MongoDB plugin for async document database operations
+- REST API plugin for HTTP client functionality
+- AWS S3 plugin for cloud object storage operations
+- Slack plugin for team collaboration and notifications
+- Elasticsearch plugin for search and analytics
+- MCP plugin for Model Context Protocol server integration
+
+All plugins follow async patterns and provide simple, clean interfaces
+without over-engineering.
+
+Usage:
+    ```python
+    from daita.plugins import postgresql, mysql, mongodb, rest, s3, slack, elasticsearch, mcp
+    from daita import SubstrateAgent
+
+    # Database plugins
+    async with postgresql(host="localhost", database="mydb") as db:
+        results = await db.query("SELECT * FROM users")
+
+    # REST API plugin
+    async with rest(base_url="https://api.example.com") as api:
+        data = await api.get("/users")
+
+    # S3 plugin
+    async with s3(bucket="my-bucket", region="us-west-2") as storage:
+        data = await storage.get_object("data/file.csv", format="pandas")
+
+    # Slack plugin
+    async with slack(token="xoxb-token") as slack_client:
+        await slack_client.send_agent_summary("#alerts", agent_results)
+
+    # Elasticsearch plugin
+    async with elasticsearch(hosts=["localhost:9200"]) as es:
+        results = await es.search("logs", {"match": {"level": "ERROR"}}, focus=["timestamp", "message"])
+
+    # MCP plugin with agent integration
+    agent = SubstrateAgent(
+        name="file_analyzer",
+        mcp=mcp.server(command="uvx", args=["mcp-server-filesystem", "/data"])
+    )
+    result = await agent.process("Read report.csv and calculate totals")
+    ```
+"""
+
+# Database plugins
+from .postgresql import PostgreSQLPlugin, postgresql
+from .mysql import MySQLPlugin, mysql
+from .mongodb import MongoDBPlugin, mongodb
+
+# API plugins  
+from .rest import RESTPlugin, rest
+
+# Cloud storage plugins
+from .s3 import S3Plugin, s3
+
+# Collaboration plugins
+from .slack import SlackPlugin, slack
+
+# Search and analytics plugins
+from .elasticsearch import ElasticsearchPlugin, elasticsearch
+
+# Messaging plugins
+from .redis_messaging import RedisMessagingPlugin, redis_messaging
+
+# MCP plugin
+from . import mcp
+
+# Simple plugin access class for SDK
+class PluginAccess:
+    """
+    Simple plugin access for the SDK.
+    
+    Provides clean interface: sdk.plugins.postgresql(...)
+    """
+    
+    def postgresql(self, **kwargs) -> PostgreSQLPlugin:
+        """Create PostgreSQL plugin."""
+        return postgresql(**kwargs)
+    
+    def mysql(self, **kwargs) -> MySQLPlugin:
+        """Create MySQL plugin."""
+        return mysql(**kwargs)
+    
+    def mongodb(self, **kwargs) -> MongoDBPlugin:
+        """Create MongoDB plugin."""
+        return mongodb(**kwargs)
+    
+    def rest(self, **kwargs) -> RESTPlugin:
+        """Create REST API plugin."""
+        return rest(**kwargs)
+    
+    def s3(self, **kwargs) -> S3Plugin:
+        """Create S3 plugin."""
+        return s3(**kwargs)
+    
+    def slack(self, **kwargs) -> SlackPlugin:
+        """Create Slack plugin."""
+        return slack(**kwargs)
+    
+    def elasticsearch(self, **kwargs) -> ElasticsearchPlugin:
+        """Create Elasticsearch plugin."""
+        return elasticsearch(**kwargs)
+    
+    def redis_messaging(self, **kwargs) -> RedisMessagingPlugin:
+        """Create Redis messaging plugin."""
+        return redis_messaging(**kwargs)
+
+# Export everything needed
+__all__ = [
+    # Plugin classes
+    'PostgreSQLPlugin',
+    'MySQLPlugin',
+    'MongoDBPlugin',
+    'RESTPlugin',
+    'S3Plugin',
+    'SlackPlugin',
+    'ElasticsearchPlugin',
+    'RedisMessagingPlugin',
+
+    # Factory functions
+    'postgresql',
+    'mysql',
+    'mongodb',
+    'rest',
+    's3',
+    'slack',
+    'elasticsearch',
+    'redis_messaging',
+
+    # MCP module
+    'mcp',
+
+    # SDK access class
+    'PluginAccess',
+]

daita/plugins/base.py

@@ -0,0 +1,37 @@
+"""
+Base classes for Daita plugins.
+
+Plugins are infrastructure utilities (databases, APIs, storage) that can
+optionally expose their capabilities as agent tools.
+"""
+
+from abc import ABC
+from typing import List, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ..core.tools import AgentTool
+
+
+class BasePlugin(ABC):
+    """
+    Base class for all Daita plugins.
+
+    Plugins provide infrastructure utilities (S3, Slack, REST APIs, etc).
+    They can optionally expose their capabilities as agent tools via get_tools().
+    """
+
+    def get_tools(self) -> List['AgentTool']:
+        """
+        Get agent-usable tools from this plugin.
+
+        Override in subclasses to expose plugin capabilities as LLM tools.
+
+        Returns:
+            List of AgentTool instances
+        """
+        return []
+
+    @property
+    def has_tools(self) -> bool:
+        """Check if plugin exposes any tools"""
+        return len(self.get_tools()) > 0

daita/plugins/base_db.py

@@ -0,0 +1,167 @@
+"""
+Base class for database plugins.
+
+Provides common connection management, error handling, and context manager
+support for all database plugins in the Daita framework.
+"""
+import logging
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Optional, TYPE_CHECKING
+from ..core.exceptions import PluginError, ConnectionError as DaitaConnectionError, ValidationError
+
+if TYPE_CHECKING:
+    from ..core.tools import AgentTool
+
+logger = logging.getLogger(__name__)
+
+class BaseDatabasePlugin(ABC):
+    """
+    Base class for all database plugins with common connection management.
+    
+    This class provides:
+    - Standardized connection/disconnection lifecycle
+    - Context manager support for automatic cleanup
+    - Common error handling patterns
+    - Consistent configuration patterns
+    
+    Database-specific plugins should inherit from this class and implement
+    the abstract methods for their specific database requirements.
+    """
+    
+    def __init__(self, **kwargs):
+        """
+        Initialize base database plugin.
+        
+        Args:
+            **kwargs: Database-specific configuration parameters
+        """
+        # Common connection state
+        self._connection = None
+        self._pool = None
+        self._client = None
+        self._db = None
+        
+        # Connection configuration
+        self.config = kwargs
+        self.timeout = kwargs.get('timeout', 30)
+        self.max_retries = kwargs.get('max_retries', 3)
+        
+        logger.debug(f"{self.__class__.__name__} initialized with config keys: {list(kwargs.keys())}")
+    
+    @abstractmethod
+    async def connect(self) -> None:
+        """
+        Connect to the database.
+        
+        This method must be implemented by each database plugin to handle
+        the specific connection logic for that database type.
+        """
+        pass
+    
+    @abstractmethod
+    async def disconnect(self) -> None:
+        """
+        Disconnect from the database and clean up resources.
+        
+        This method must be implemented by each database plugin to handle
+        the specific disconnection and cleanup logic for that database type.
+        """
+        pass
+    
+    @property
+    def is_connected(self) -> bool:
+        """
+        Check if the database connection is active.
+        
+        Returns:
+            True if connected, False otherwise
+        """
+        return (
+            self._connection is not None or 
+            self._pool is not None or 
+            self._client is not None
+        )
+    
+    async def __aenter__(self):
+        """Async context manager entry - automatically connect."""
+        await self.connect()
+        return self
+    
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Async context manager exit - automatically disconnect."""
+        await self.disconnect()
+    
+    def _validate_connection(self) -> None:
+        """
+        Validate that the database connection is available.
+        
+        Raises:
+            ValidationError: If not connected to database
+        """
+        if not self.is_connected:
+            raise ValidationError(
+                f"{self.__class__.__name__} is not connected to database",
+                field="connection_state"
+            )
+    
+    def _handle_connection_error(self, error: Exception, operation: str) -> None:
+        """
+        Handle database connection errors with consistent logging.
+        
+        Args:
+            error: The exception that occurred
+            operation: Description of the operation that failed
+            
+        Raises:
+            PluginError: Wrapped database error with context
+        """
+        error_msg = f"{self.__class__.__name__} {operation} failed: {str(error)}"
+        logger.error(error_msg)
+        
+        # Choose appropriate exception type based on the original error
+        if isinstance(error, ImportError):
+            # Missing dependency - permanent error
+            raise PluginError(
+                error_msg,
+                plugin_name=self.__class__.__name__,
+                retry_hint="permanent",
+                context={"operation": operation, "original_error": str(error)}
+            ) from error
+        else:
+            # Connection issues - typically transient
+            raise DaitaConnectionError(
+                error_msg,
+                context={"plugin": self.__class__.__name__, "operation": operation}
+            ) from error
+    
+    def get_tools(self) -> List['AgentTool']:
+        """
+        Get agent-usable tools from this database plugin.
+
+        Override in subclasses to expose database operations as LLM tools.
+
+        Returns:
+            List of AgentTool instances
+        """
+        return []
+
+    @property
+    def has_tools(self) -> bool:
+        """Check if plugin exposes any tools"""
+        return len(self.get_tools()) > 0
+
+    @property
+    def info(self) -> Dict[str, Any]:
+        """
+        Get information about the database plugin.
+
+        Returns:
+            Dictionary with plugin information
+        """
+        return {
+            'plugin_type': self.__class__.__name__,
+            'connected': self.is_connected,
+            'timeout': self.timeout,
+            'max_retries': self.max_retries,
+            'config_keys': list(self.config.keys())
+        }