daita-agents 0.1.0__py3-none-any.whl

daita/llm/base.py

@@ -0,0 +1,373 @@
+"""
+Updated BaseLLMProvider with Unified Tracing Integration
+
+This replaces the old BaseLLMProvider to use the unified tracing system.
+All LLM calls are automatically traced without user configuration.
+
+Key Changes:
+- Removed old token tracking system completely
+- Integrated automatic LLM call tracing
+- Simple cost estimation  
+- Automatic provider/model/token capture
+- Zero configuration required
+"""
+
+from abc import ABC, abstractmethod
+from typing import Dict, Any, Optional, List
+from contextlib import asynccontextmanager
+import logging
+import time
+import asyncio
+
+from ..core.tracing import get_trace_manager, TraceType, TraceStatus
+from ..core.interfaces import LLMProvider
+
+logger = logging.getLogger(__name__)
+
+class BaseLLMProvider(LLMProvider, ABC):
+    """
+    Base class for LLM providers with automatic call tracing.
+    
+    Every LLM call is automatically traced with:
+    - Provider and model details
+    - Token usage and costs
+    - Latency and performance
+    - Input/output content (preview)
+    - Error tracking
+    
+    Users get full LLM observability without any configuration.
+    """
+    
+    def __init__(
+        self,
+        model: str,
+        api_key: Optional[str] = None,
+        **kwargs
+    ):
+        """
+        Initialize the LLM provider with automatic tracing.
+        
+        Args:
+            model: Model identifier
+            api_key: API key for authentication
+            **kwargs: Additional provider-specific options
+        """
+        self.model = model
+        self.api_key = api_key
+        self.config = kwargs
+        
+        # Default parameters
+        self.default_params = {
+            'temperature': kwargs.get('temperature', 0.7),
+            'max_tokens': kwargs.get('max_tokens', 1000),
+            'top_p': kwargs.get('top_p', 1.0),
+        }
+        
+        # Agent ID for tracing (set by agent)
+        self.agent_id = kwargs.get('agent_id')
+        
+        # Get trace manager for automatic tracing
+        self.trace_manager = get_trace_manager()
+        
+        # Provider name for tracing
+        self.provider_name = self.__class__.__name__.replace('Provider', '').lower()
+        
+        # Last usage for cost estimation
+        self._last_usage = None
+        
+        logger.debug(f"Initialized {self.__class__.__name__} with model {model} (automatic tracing enabled)")
+    
+    async def generate(self, prompt: str, **kwargs) -> str:
+        """
+        Generate text from prompt with automatic LLM call tracing.
+        
+        Every call is automatically traced with full metadata.
+        
+        Args:
+            prompt: Input prompt
+            **kwargs: Optional parameters to override defaults
+            
+        Returns:
+            Generated text response
+        """
+        # Merge parameters
+        params = self._merge_params(kwargs)
+        
+        # Automatically trace the LLM call
+        async with self.trace_manager.span(
+            operation_name=f"llm_{self.provider_name}_{self.model}",
+            trace_type=TraceType.LLM_CALL,
+            agent_id=self.agent_id,
+            input_data=prompt,
+            llm_provider=self.provider_name,
+            llm_model=self.model,
+            temperature=str(params.get('temperature', 0.7)),
+            max_tokens=str(params.get('max_tokens', 1000))
+        ) as span_id:
+            
+            try:
+                start_time = time.time()
+                
+                # Execute the actual LLM call
+                response = await self._generate_impl(prompt, **params)
+                
+                end_time = time.time()
+                duration_ms = (end_time - start_time) * 1000
+                
+                # Get token usage from the call
+                token_usage = self._get_last_token_usage()
+                
+                # Record LLM call details in the trace
+                self.trace_manager.record_llm_call(
+                    span_id=span_id,
+                    provider=self.provider_name,
+                    model=self.model,
+                    tokens=token_usage
+                )
+                
+                # Add cost estimation if available
+                cost = self._estimate_cost(token_usage)
+                if cost and span_id in self.trace_manager._active_spans:
+                    span = self.trace_manager._active_spans[span_id]
+                    span.metadata['estimated_cost_usd'] = cost
+                    span.metadata['duration_ms'] = duration_ms
+                
+                logger.debug(f"LLM call completed: {token_usage.get('total_tokens', 0)} tokens in {duration_ms:.1f}ms")
+                return response
+                
+            except Exception as e:
+                # LLM call failed - error automatically recorded by span context
+                logger.warning(f"LLM call failed: {str(e)}")
+                raise
+    
+    @abstractmethod
+    async def _generate_impl(self, prompt: str, **kwargs) -> str:
+        """
+        Provider-specific implementation of text generation.
+        
+        This method must be implemented by each provider (OpenAI, Anthropic, etc.)
+        and contains the actual LLM API call logic.
+        
+        Args:
+            prompt: Input prompt
+            **kwargs: Optional parameters
+            
+        Returns:
+            Generated text response
+        """
+        pass
+    
+    def _get_last_token_usage(self) -> Dict[str, int]:
+        """
+        Get token usage from the last API call.
+        
+        This should be overridden by each provider to return actual
+        token usage from their last API response.
+        """
+        if self._last_usage:
+            # Try to extract from stored usage object
+            if hasattr(self._last_usage, 'total_tokens'):
+                # OpenAI format
+                return {
+                    'total_tokens': getattr(self._last_usage, 'total_tokens', 0),
+                    'prompt_tokens': getattr(self._last_usage, 'prompt_tokens', 0),
+                    'completion_tokens': getattr(self._last_usage, 'completion_tokens', 0)
+                }
+            elif hasattr(self._last_usage, 'input_tokens'):
+                # Anthropic format
+                input_tokens = getattr(self._last_usage, 'input_tokens', 0)
+                output_tokens = getattr(self._last_usage, 'output_tokens', 0)
+                return {
+                    'total_tokens': input_tokens + output_tokens,
+                    'prompt_tokens': input_tokens,
+                    'completion_tokens': output_tokens
+                }
+        
+        # Default fallback
+        return {
+            'total_tokens': 0,
+            'prompt_tokens': 0,
+            'completion_tokens': 0
+        }
+    
+    def _estimate_cost(self, token_usage: Dict[str, int]) -> Optional[float]:
+        """
+        Simple cost estimation for MVP.
+        
+        Providers can override with specific pricing.
+        """
+        total_tokens = token_usage.get('total_tokens', 0)
+        if total_tokens == 0:
+            return None
+        
+        # Generic estimation for MVP - providers should override
+        cost_per_1k_tokens = 0.002  # $0.002 per 1K tokens
+        return (total_tokens / 1000) * cost_per_1k_tokens
+    
+    def _merge_params(self, override_params: Dict[str, Any]) -> Dict[str, Any]:
+        """Merge default parameters with overrides."""
+        params = self.default_params.copy()
+        params.update(override_params)
+        return params
+    
+    def _validate_api_key(self) -> None:
+        """Validate that API key is available."""
+        if not self.api_key:
+            raise ValueError(f"API key required for {self.__class__.__name__}")
+    
+    def set_agent_id(self, agent_id: str):
+        """
+        Set the agent ID for tracing context.
+        
+        This is called automatically by BaseAgent during initialization.
+        """
+        self.agent_id = agent_id
+        logger.debug(f"Set agent ID {agent_id} for {self.provider_name} provider")
+    
+    def get_recent_calls(self, limit: int = 10) -> List[Dict[str, Any]]:
+        """Get recent LLM calls for this provider's agent from unified tracing."""
+        if not self.agent_id:
+            return []
+        
+        operations = self.trace_manager.get_recent_operations(agent_id=self.agent_id, limit=limit * 2)
+        
+        # Filter for LLM calls from this provider
+        llm_calls = [
+            op for op in operations 
+            if (op.get('type') == 'llm_call' and 
+                op.get('metadata', {}).get('llm_provider') == self.provider_name)
+        ]
+        
+        return llm_calls[:limit]
+    
+    def get_token_stats(self) -> Dict[str, Any]:
+        """Get token usage statistics from unified tracing."""
+        if not self.agent_id:
+            return {
+                'total_calls': 0,
+                'total_tokens': 0,
+                'estimated_cost': 0.0
+            }
+        
+        metrics = self.trace_manager.get_agent_metrics(self.agent_id)
+        
+        # Get actual token usage from the most recent call
+        last_usage = self._get_last_token_usage()
+        
+        return {
+            'total_calls': metrics.get('total_operations', 0),  # All operations
+            'total_tokens': last_usage.get('total_tokens', 0),  # From last API call
+            'prompt_tokens': last_usage.get('prompt_tokens', 0),
+            'completion_tokens': last_usage.get('completion_tokens', 0),
+            'estimated_cost': self._estimate_cost(last_usage) or 0.0,
+            'success_rate': metrics.get('success_rate', 0),
+            'avg_latency_ms': metrics.get('avg_latency_ms', 0)
+        }
+    
+    @property
+    def info(self) -> Dict[str, Any]:
+        """Get information about this LLM provider."""
+        return {
+            'provider': self.provider_name,
+            'model': self.model,
+            'agent_id': self.agent_id,
+            'config': {k: v for k, v in self.config.items() if 'key' not in k.lower()},
+            'default_params': self.default_params,
+            'tracing_enabled': True
+        }
+
+
+
+# Context manager for batch LLM operations
+@asynccontextmanager
+async def traced_llm_batch(llm_provider: BaseLLMProvider, batch_name: str = "llm_batch"):
+    """
+    Context manager for tracing batch LLM operations.
+    
+    Usage:
+        async with traced_llm_batch(llm, "document_analysis"):
+            summary = await llm.generate("Summarize: " + doc1)
+            analysis = await llm.generate("Analyze: " + doc2)
+    """
+    trace_manager = get_trace_manager()
+    
+    async with trace_manager.span(
+        operation_name=f"llm_batch_{batch_name}",
+        trace_type=TraceType.LLM_CALL,
+        agent_id=llm_provider.agent_id,
+        llm_provider=llm_provider.provider_name,
+        batch_operation=batch_name
+    ):
+        try:
+            yield llm_provider
+        except Exception as e:
+            logger.error(f"LLM batch {batch_name} failed: {e}")
+            raise
+
+
+# Utility functions
+def get_llm_traces(agent_id: Optional[str] = None, provider: Optional[str] = None, limit: int = 20) -> List[Dict[str, Any]]:
+    """Get recent LLM call traces from unified tracing."""
+    trace_manager = get_trace_manager()
+    operations = trace_manager.get_recent_operations(agent_id=agent_id, limit=limit * 2)
+    
+    # Filter for LLM operations
+    llm_ops = [
+        op for op in operations 
+        if op.get('type') == 'llm_call'
+    ]
+    
+    # Filter by provider if specified
+    if provider:
+        llm_ops = [
+            op for op in llm_ops
+            if op.get('metadata', {}).get('llm_provider') == provider
+        ]
+    
+    return llm_ops[:limit]
+
+
+def get_llm_stats(agent_id: Optional[str] = None, provider: Optional[str] = None) -> Dict[str, Any]:
+    """Get LLM usage statistics from unified tracing."""
+    traces = get_llm_traces(agent_id, provider, limit=50)
+    
+    if not traces:
+        return {"total_calls": 0, "success_rate": 0, "avg_latency_ms": 0}
+    
+    total_calls = len(traces)
+    successful_calls = len([t for t in traces if t.get('status') == 'success'])
+    
+    # Calculate averages
+    latencies = [t.get('duration_ms', 0) for t in traces if t.get('duration_ms')]
+    avg_latency = sum(latencies) / len(latencies) if latencies else 0
+    
+    # Token aggregation (simplified for MVP)
+    total_tokens = 0
+    for trace in traces:
+        metadata = trace.get('metadata', {})
+        total_tokens += metadata.get('tokens_total', 0)
+    
+    return {
+        "total_calls": total_calls,
+        "successful_calls": successful_calls,
+        "failed_calls": total_calls - successful_calls,
+        "success_rate": successful_calls / total_calls if total_calls > 0 else 0,
+        "avg_latency_ms": avg_latency,
+        "total_tokens": total_tokens,
+        "agent_id": agent_id,
+        "provider": provider
+    }
+
+
+# Export everything
+__all__ = [
+    # Base class
+    "BaseLLMProvider",
+    
+    # Context managers
+    "traced_llm_batch",
+    
+    # Utility functions
+    "get_llm_traces",
+    "get_llm_stats"
+]

daita/llm/factory.py

@@ -0,0 +1,101 @@
+"""
+Factory for creating LLM provider instances.
+"""
+import logging
+from typing import Optional
+
+from ..core.exceptions import LLMError
+from .base import BaseLLMProvider
+from .openai import OpenAIProvider
+from .anthropic import AnthropicProvider
+from .grok import GrokProvider
+from .gemini import GeminiProvider
+from .mock import MockLLMProvider
+
+logger = logging.getLogger(__name__)
+
+# Registry of available providers
+PROVIDER_REGISTRY = {
+    'openai': OpenAIProvider,
+    'anthropic': AnthropicProvider,
+    'grok': GrokProvider,
+    'gemini': GeminiProvider,
+    'mock': MockLLMProvider,
+}
+
+def create_llm_provider(
+    provider: str,
+    model: str,
+    api_key: Optional[str] = None,
+    agent_id: Optional[str] = None,
+    **kwargs
+) -> BaseLLMProvider:
+    """
+    Factory function to create LLM provider instances.
+    
+    Args:
+        provider: Provider name ('openai', 'anthropic', 'grok', 'gemini', 'mock')
+        model: Model identifier
+        api_key: API key for authentication
+        agent_id: Agent ID for token tracking
+        **kwargs: Additional provider-specific parameters
+        
+    Returns:
+        LLM provider instance
+        
+    Raises:
+        LLMError: If provider is not supported
+        
+    Examples:
+        >>> # Create OpenAI provider with token tracking
+        >>> llm = create_llm_provider('openai', 'gpt-4', api_key='sk-...', agent_id='my_agent')
+        
+        >>> # Create Anthropic provider with token tracking
+        >>> llm = create_llm_provider('anthropic', 'claude-3-sonnet-20240229', agent_id='my_agent')
+        
+        >>> # Create Grok provider
+        >>> llm = create_llm_provider('grok', 'grok-beta', api_key='xai-...', agent_id='my_agent')
+        
+        >>> # Create Gemini provider
+        >>> llm = create_llm_provider('gemini', 'gemini-1.5-pro', api_key='AIza...', agent_id='my_agent')
+        
+        >>> # Create mock provider for testing
+        >>> llm = create_llm_provider('mock', 'test-model', agent_id='test_agent')
+    """
+    provider_name = provider.lower()
+    
+    if provider_name not in PROVIDER_REGISTRY:
+        available_providers = list(PROVIDER_REGISTRY.keys())
+        raise LLMError(
+            f"Unsupported LLM provider: {provider}. "
+            f"Available providers: {available_providers}"
+        )
+    
+    provider_class = PROVIDER_REGISTRY[provider_name]
+    
+    try:
+        # Pass agent_id to provider for token tracking
+        return provider_class(
+            model=model, 
+            api_key=api_key, 
+            agent_id=agent_id,
+            **kwargs
+        )
+    except Exception as e:
+        logger.error(f"Failed to create {provider} provider: {str(e)}")
+        raise LLMError(f"Failed to create {provider} provider: {str(e)}")
+
+def register_llm_provider(name: str, provider_class) -> None:
+    """
+    Register a custom LLM provider.
+    
+    Args:
+        name: Provider name
+        provider_class: Provider class that implements LLMProvider interface
+    """
+    PROVIDER_REGISTRY[name.lower()] = provider_class
+    logger.info(f"Registered custom LLM provider: {name}")
+
+def list_available_providers() -> list:
+    """Get list of available LLM providers."""
+    return list(PROVIDER_REGISTRY.keys())

daita/llm/gemini.py

@@ -0,0 +1,152 @@
+"""
+Google Gemini LLM provider implementation with integrated tracing.
+"""
+import os
+import logging
+import asyncio
+from typing import Dict, Any, Optional
+
+from ..core.exceptions import LLMError
+from .base import BaseLLMProvider
+
+logger = logging.getLogger(__name__)
+
+class GeminiProvider(BaseLLMProvider):
+    """Google Gemini LLM provider implementation with automatic call tracing."""
+    
+    def __init__(
+        self,
+        model: str = "gemini-1.5-flash",
+        api_key: Optional[str] = None,
+        **kwargs
+    ):
+        """
+        Initialize Gemini provider.
+        
+        Args:
+            model: Gemini model name (e.g., "gemini-1.5-flash", "gemini-1.5-pro", "gemini-1.0-pro")
+            api_key: Google AI API key
+            **kwargs: Additional Gemini-specific parameters
+        """
+        # Get API key from parameter or environment
+        api_key = api_key or os.getenv("GOOGLE_API_KEY") or os.getenv("GEMINI_API_KEY")
+        
+        super().__init__(model=model, api_key=api_key, **kwargs)
+        
+        # Gemini-specific default parameters
+        self.default_params.update({
+            'timeout': kwargs.get('timeout', 60),
+            'safety_settings': kwargs.get('safety_settings', None),
+            'generation_config': kwargs.get('generation_config', None)
+        })
+        
+        # Lazy-load Gemini client
+        self._client = None
+    
+    @property
+    def client(self):
+        """Lazy-load Google Generative AI client."""
+        if self._client is None:
+            try:
+                import google.generativeai as genai
+                self._validate_api_key()
+                
+                # Configure the API key
+                genai.configure(api_key=self.api_key)
+                
+                # Create the generative model
+                self._client = genai.GenerativeModel(self.model)
+                logger.debug("Gemini client initialized")
+            except ImportError:
+                raise LLMError(
+                    "Google Generative AI package not installed. Install with: pip install google-generativeai"
+                )
+        return self._client
+    
+    async def _generate_impl(self, prompt: str, **kwargs) -> str:
+        """
+        Provider-specific implementation of text generation for Gemini.
+        
+        This method contains the actual Gemini 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)
+            
+            # Prepare generation config
+            generation_config = params.get('generation_config', {})
+            if not generation_config:
+                generation_config = {
+                    'max_output_tokens': params.get('max_tokens'),
+                    'temperature': params.get('temperature'),
+                    'top_p': params.get('top_p')
+                }
+            
+            # Make API call (Gemini's generate_content can be sync or async)
+            # For consistency with other providers, we'll run in executor if needed
+            if asyncio.iscoroutinefunction(self.client.generate_content):
+                response = await self.client.generate_content(
+                    prompt,
+                    generation_config=generation_config,
+                    safety_settings=params.get('safety_settings')
+                )
+            else:
+                # Run synchronous method in executor
+                loop = asyncio.get_event_loop()
+                response = await loop.run_in_executor(
+                    None,
+                    lambda: self.client.generate_content(
+                        prompt,
+                        generation_config=generation_config,
+                        safety_settings=params.get('safety_settings')
+                    )
+                )
+            
+            # Store usage info if available (Gemini's usage tracking varies)
+            if hasattr(response, 'usage_metadata'):
+                self._last_usage = response.usage_metadata
+            
+            return response.text
+            
+        except Exception as e:
+            logger.error(f"Gemini generation failed: {str(e)}")
+            raise LLMError(f"Gemini generation failed: {str(e)}")
+    
+    def _get_last_token_usage(self) -> Dict[str, int]:
+        """
+        Override base class method to handle Gemini's token format.
+        
+        Gemini uses different token field names in usage_metadata.
+        """
+        if self._last_usage:
+            # Gemini format varies, try to extract what we can
+            prompt_tokens = getattr(self._last_usage, 'prompt_token_count', 0)
+            completion_tokens = getattr(self._last_usage, 'candidates_token_count', 0)
+            total_tokens = getattr(self._last_usage, 'total_token_count', prompt_tokens + completion_tokens)
+            
+            return {
+                'total_tokens': total_tokens,
+                'prompt_tokens': prompt_tokens,
+                'completion_tokens': completion_tokens
+            }
+        
+        # Fallback to base class estimation
+        return super()._get_last_token_usage()
+    
+    @property
+    def info(self) -> Dict[str, Any]:
+        """Get information about the Gemini provider."""
+        base_info = super().info
+        base_info.update({
+            'provider_name': 'Google Gemini',
+            'api_compatible': 'Google AI'
+        })
+        return base_info