daita-agents 0.2.0__py3-none-any.whl

daita/llm/mock.py

@@ -0,0 +1,172 @@
+"""
+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]}..."
+    
+    async def _generate_with_tools_single(
+        self,
+        messages: list,
+        tools: list,
+        **kwargs
+    ) -> dict:
+        """
+        Mock implementation of single LLM call with tools.
+
+        Returns a mock response without tool calls (final answer).
+        Override this in tests if you need to test tool calling behavior.
+        """
+        # Extract the last user message
+        user_message = ""
+        for msg in reversed(messages):
+            if msg.get("role") == "user":
+                user_message = msg.get("content", "")
+                break
+
+        # Record the call
+        self.call_history.append({
+            'messages': messages,
+            'tools': tools,
+            'params': kwargs,
+            'timestamp': asyncio.get_event_loop().time()
+        })
+
+        # Simulate API delay
+        if self.delay > 0:
+            await asyncio.sleep(self.delay)
+
+        # Return final answer (no tool calls for basic mock)
+        return {
+            "content": f"Mock response for: {user_message[:50]}...",
+            "tool_calls": None
+        }
+
+    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,220 @@
+"""
+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)}")
+
+    def _convert_messages_to_openai(
+        self,
+        messages: list[Dict[str, Any]]
+    ) -> list[Dict[str, Any]]:
+        """
+        Convert universal flat format to OpenAI's nested format.
+
+        OpenAI expects tool_calls in nested format:
+        {"id": "x", "type": "function", "function": {"name": "...", "arguments": "..."}}
+
+        Our internal format is flat:
+        {"id": "x", "name": "...", "arguments": {...}}
+        """
+        import json
+
+        openai_messages = []
+        for msg in messages:
+            if msg.get("role") == "assistant" and msg.get("tool_calls"):
+                # Convert flat format to OpenAI's nested format
+                converted_tool_calls = []
+                for tc in msg["tool_calls"]:
+                    converted_tool_calls.append({
+                        "id": tc.get("id", ""),
+                        "type": "function",
+                        "function": {
+                            "name": tc["name"],
+                            "arguments": json.dumps(tc["arguments"]) if isinstance(tc["arguments"], dict) else tc["arguments"]
+                        }
+                    })
+
+                openai_messages.append({
+                    "role": "assistant",
+                    "tool_calls": converted_tool_calls
+                })
+            else:
+                # Pass through other messages unchanged
+                openai_messages.append(msg)
+
+        return openai_messages
+
+    async def _generate_with_tools_single(
+        self,
+        messages: list[Dict[str, Any]],
+        tools: list[Dict[str, Any]],
+        **kwargs
+    ) -> Dict[str, Any]:
+        """
+        OpenAI-specific tool calling implementation.
+
+        Args:
+            messages: Conversation history in universal flat format
+            tools: Tool specifications in OpenAI format
+            **kwargs: Optional parameters
+
+        Returns:
+            {
+                "tool_calls": [...],  # If LLM wants to call tools
+                "content": "...",      # If LLM has final answer
+            }
+        """
+        import json
+
+        try:
+            # Merge parameters
+            params = self._merge_params(kwargs)
+
+            # Convert flat format to OpenAI's nested format
+            openai_messages = self._convert_messages_to_openai(messages)
+
+            # Make API call with tools
+            response = await self.client.chat.completions.create(
+                model=self.model,
+                messages=openai_messages,
+                tools=tools,
+                tool_choice="auto",
+                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')
+            )
+
+            message = response.choices[0].message
+
+            # Store usage for token tracking
+            if hasattr(response, 'usage'):
+                self._last_usage = response.usage
+
+            if message.tool_calls:
+                # LLM wants to call tools
+                return {
+                    "tool_calls": [
+                        {
+                            "id": tc.id,
+                            "name": tc.function.name,
+                            "arguments": json.loads(tc.function.arguments)
+                        }
+                        for tc in message.tool_calls
+                    ]
+                }
+            else:
+                # LLM has final answer
+                return {
+                    "content": message.content
+                }
+
+        except Exception as e:
+            logger.error(f"OpenAI tool calling failed: {str(e)}")
+            raise LLMError(f"OpenAI tool calling 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