cost-katana 1.0.0__py3-none-any.whl

cost_katana/config.py

@@ -0,0 +1,181 @@
+"""
+Configuration management for Cost Katana
+"""
+
+import json
+import os
+from typing import Dict, Any, Optional
+from dataclasses import dataclass, asdict
+from pathlib import Path
+
+@dataclass
+class Config:
+    """Configuration class for Cost Katana client"""
+    
+    api_key: Optional[str] = None
+    base_url: str = "https://cost-katana-backend.store"
+    timeout: int = 30
+    max_retries: int = 3
+    retry_delay: float = 1.0
+    default_model: str = "nova-lite"
+    default_temperature: float = 0.7
+    default_max_tokens: int = 2000
+    default_chat_mode: str = "balanced"
+    enable_analytics: bool = True
+    enable_optimization: bool = True
+    enable_failover: bool = True
+    cost_limit_per_request: Optional[float] = None
+    cost_limit_per_day: Optional[float] = None
+    
+    def __post_init__(self):
+        """Load from environment variables if not set"""
+        if not self.api_key:
+            self.api_key = os.getenv('COST_KATANA_API_KEY')
+        
+        # Override with environment variables if they exist
+        if os.getenv('COST_KATANA_BASE_URL'):
+            self.base_url = os.getenv('COST_KATANA_BASE_URL')
+        if os.getenv('COST_KATANA_DEFAULT_MODEL'):
+            self.default_model = os.getenv('COST_KATANA_DEFAULT_MODEL')
+        if os.getenv('COST_KATANA_TIMEOUT'):
+            self.timeout = int(os.getenv('COST_KATANA_TIMEOUT'))
+    
+    @classmethod
+    def from_file(cls, config_path: str) -> 'Config':
+        """
+        Load configuration from JSON file.
+        
+        Args:
+            config_path: Path to JSON configuration file
+            
+        Returns:
+            Config instance
+            
+        Example config.json:
+        {
+            "api_key": "dak_your_key_here",
+            "base_url": "https://api.costkatana.com",
+            "default_model": "claude-3-sonnet",
+            "default_temperature": 0.3,
+            "cost_limit_per_day": 100.0,
+            "providers": {
+                "anthropic": {
+                    "priority": 1,
+                    "models": ["claude-3-sonnet", "claude-3-haiku"]
+                },
+                "openai": {
+                    "priority": 2,
+                    "models": ["gpt-4", "gpt-3.5-turbo"]
+                }
+            }
+        }
+        """
+        config_path = Path(config_path).expanduser()
+        
+        if not config_path.exists():
+            raise FileNotFoundError(f"Configuration file not found: {config_path}")
+            
+        try:
+            with open(config_path, 'r', encoding='utf-8') as f:
+                data = json.load(f)
+        except json.JSONDecodeError as e:
+            raise ValueError(f"Invalid JSON in config file: {e}")
+            
+        # Extract known fields
+        config_fields = {
+            field.name for field in cls.__dataclass_fields__.values()
+        }
+        
+        config_data = {k: v for k, v in data.items() if k in config_fields}
+        config = cls(**config_data)
+        
+        # Store additional data
+        config._extra_data = {k: v for k, v in data.items() if k not in config_fields}
+        
+        return config
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert config to dictionary"""
+        result = asdict(self)
+        
+        # Add extra data if it exists
+        if hasattr(self, '_extra_data'):
+            result.update(self._extra_data)
+            
+        return result
+    
+    def save(self, config_path: str):
+        """Save configuration to JSON file"""
+        config_path = Path(config_path).expanduser()
+        config_path.parent.mkdir(parents=True, exist_ok=True)
+        
+        with open(config_path, 'w', encoding='utf-8') as f:
+            json.dump(self.to_dict(), f, indent=2, ensure_ascii=False)
+    
+    def get_provider_config(self, provider: str) -> Dict[str, Any]:
+        """Get configuration for a specific provider"""
+        if hasattr(self, '_extra_data') and 'providers' in self._extra_data:
+            return self._extra_data['providers'].get(provider, {})
+        return {}
+    
+    def get_model_mapping(self, model_name: str) -> str:
+        """
+        Map user-friendly model names to internal model IDs.
+        This allows users to use names like 'gemini-2.0-flash' while
+        the backend uses the actual model IDs.
+        """
+        # Default mapping - can be overridden in config file
+        # Based on actual models available from Cost Katana Backend
+        default_mappings = {
+            # Amazon Nova models (primary recommendation)
+            'nova-micro': 'amazon.nova-micro-v1:0',
+            'nova-lite': 'amazon.nova-lite-v1:0', 
+            'nova-pro': 'amazon.nova-pro-v1:0',
+            'fast': 'amazon.nova-micro-v1:0',
+            'balanced': 'amazon.nova-lite-v1:0',
+            'powerful': 'amazon.nova-pro-v1:0',
+            
+            # Anthropic Claude models
+            'claude-3-haiku': 'anthropic.claude-3-haiku-20240307-v1:0',
+            'claude-3-sonnet': 'anthropic.claude-3-sonnet-20240229-v1:0',
+            'claude-3-opus': 'anthropic.claude-3-opus-20240229-v1:0',
+            'claude-3.5-haiku': 'anthropic.claude-3-5-haiku-20241022-v1:0',
+            'claude-3.5-sonnet': 'anthropic.claude-3-5-sonnet-20241022-v2:0',
+            'claude': 'anthropic.claude-3-5-sonnet-20241022-v2:0',
+            
+            # Meta Llama models
+            'llama-3.1-8b': 'meta.llama3-1-8b-instruct-v1:0',
+            'llama-3.1-70b': 'meta.llama3-1-70b-instruct-v1:0',
+            'llama-3.1-405b': 'meta.llama3-1-405b-instruct-v1:0',
+            'llama-3.2-1b': 'meta.llama3-2-1b-instruct-v1:0',
+            'llama-3.2-3b': 'meta.llama3-2-3b-instruct-v1:0',
+            
+            # Mistral models
+            'mistral-7b': 'mistral.mistral-7b-instruct-v0:2',
+            'mixtral-8x7b': 'mistral.mixtral-8x7b-instruct-v0:1',
+            'mistral-large': 'mistral.mistral-large-2402-v1:0',
+            
+            # Cohere models
+            'command': 'cohere.command-text-v14',
+            'command-light': 'cohere.command-light-text-v14',
+            'command-r': 'cohere.command-r-v1:0',
+            'command-r-plus': 'cohere.command-r-plus-v1:0',
+            
+            # AI21 models
+            'jamba': 'ai21.jamba-instruct-v1:0',
+            'j2-ultra': 'ai21.j2-ultra-v1',
+            'j2-mid': 'ai21.j2-mid-v1',
+            
+            # Backwards compatibility aliases
+            'gemini-2.0-flash': 'amazon.nova-lite-v1:0',  # Map to similar performance
+            'gemini-pro': 'amazon.nova-pro-v1:0',
+            'gpt-4': 'anthropic.claude-3-5-sonnet-20241022-v2:0',
+            'gpt-3.5-turbo': 'anthropic.claude-3-haiku-20240307-v1:0',
+        }
+        
+        # Check for custom mappings in config
+        if hasattr(self, '_extra_data') and 'model_mappings' in self._extra_data:
+            custom_mappings = self._extra_data['model_mappings']
+            default_mappings.update(custom_mappings)
+            
+        return default_mappings.get(model_name, model_name)

cost_katana/exceptions.py

@@ -0,0 +1,39 @@
+"""
+Custom exceptions for Cost Katana
+"""
+
+class CostKatanaError(Exception):
+    """Base exception for Cost Katana errors"""
+    pass
+
+class AuthenticationError(CostKatanaError):
+    """Raised when authentication fails"""
+    pass
+
+class ModelNotAvailableError(CostKatanaError):
+    """Raised when requested model is not available"""
+    pass
+
+class RateLimitError(CostKatanaError):
+    """Raised when rate limit is exceeded"""
+    pass
+
+class CostLimitExceededError(CostKatanaError):
+    """Raised when cost limits are exceeded"""
+    pass
+
+class ConversationNotFoundError(CostKatanaError):
+    """Raised when conversation is not found"""
+    pass
+
+class InvalidConfigurationError(CostKatanaError):
+    """Raised when configuration is invalid"""
+    pass
+
+class NetworkError(CostKatanaError):
+    """Raised when network requests fail"""
+    pass
+
+class ModelTimeoutError(CostKatanaError):
+    """Raised when model request times out"""
+    pass

cost_katana/models.py

@@ -0,0 +1,343 @@
+"""
+Generative AI Models - Simple interface similar to google-generative-ai
+"""
+
+import time
+from typing import Dict, Any, Optional, List, Iterator, Union
+from dataclasses import dataclass
+from .client import CostKatanaClient
+from .exceptions import CostKatanaError, ModelNotAvailableError
+
+@dataclass
+class GenerationConfig:
+    """Configuration for text generation"""
+    temperature: float = 0.7
+    max_output_tokens: int = 2000
+    top_p: Optional[float] = None
+    top_k: Optional[int] = None
+    candidate_count: int = 1
+    stop_sequences: Optional[List[str]] = None
+
+@dataclass
+class UsageMetadata:
+    """Usage metadata returned with responses"""
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+    cost: float
+    latency: float
+    model: str
+    optimizations_applied: Optional[List[str]] = None
+    cache_hit: bool = False
+    agent_path: Optional[List[str]] = None
+    risk_level: Optional[str] = None
+
+class GenerateContentResponse:
+    """Response from generate_content method"""
+    
+    def __init__(self, response_data: Dict[str, Any]):
+        self._data = response_data
+        self._text = response_data.get('data', {}).get('response', '')
+        
+        # Extract usage metadata
+        data = response_data.get('data', {})
+        self.usage_metadata = UsageMetadata(
+            prompt_tokens=data.get('tokenCount', 0),  # This might need adjustment based on actual response
+            completion_tokens=data.get('tokenCount', 0), 
+            total_tokens=data.get('tokenCount', 0),
+            cost=data.get('cost', 0.0),
+            latency=data.get('latency', 0.0),
+            model=data.get('model', ''),
+            optimizations_applied=data.get('optimizationsApplied'),
+            cache_hit=data.get('cacheHit', False),
+            agent_path=data.get('agentPath'),
+            risk_level=data.get('riskLevel')
+        )
+        
+        # Store thinking/reasoning if available
+        self.thinking = data.get('thinking')
+    
+    @property 
+    def text(self) -> str:
+        """Get the response text"""
+        return self._text
+    
+    @property
+    def parts(self) -> List[Dict[str, Any]]:
+        """Get response parts (for compatibility)"""
+        return [{'text': self._text}] if self._text else []
+    
+    def __str__(self) -> str:
+        return self._text
+    
+    def __repr__(self) -> str:
+        return f"GenerateContentResponse(text='{self._text[:50]}...', cost=${self.usage_metadata.cost:.4f})"
+
+class ChatSession:
+    """A chat session for maintaining conversation context"""
+    
+    def __init__(
+        self, 
+        client: CostKatanaClient,
+        model_id: str,
+        generation_config: Optional[GenerationConfig] = None,
+        conversation_id: str = None
+    ):
+        self.client = client
+        self.model_id = model_id
+        self.generation_config = generation_config or GenerationConfig()
+        self.conversation_id = conversation_id
+        self.history: List[Dict[str, Any]] = []
+        
+        # Create conversation if not provided
+        if not self.conversation_id:
+            try:
+                conv_response = self.client.create_conversation(
+                    title=f"Chat with {model_id}",
+                    model_id=model_id
+                )
+                self.conversation_id = conv_response['data']['id']
+            except Exception as e:
+                raise CostKatanaError(f"Failed to create conversation: {str(e)}")
+    
+    def send_message(
+        self, 
+        message: str,
+        **kwargs
+    ) -> GenerateContentResponse:
+        """
+        Send a message in the chat session.
+        
+        Args:
+            message: The message to send
+            **kwargs: Additional parameters to override defaults
+            
+        Returns:
+            GenerateContentResponse with the model's reply
+            
+        Example:
+            response = chat.send_message("What's the weather like?")
+            print(response.text)
+        """
+        # Merge generation config with kwargs
+        params = {
+            'temperature': kwargs.get('temperature', self.generation_config.temperature),
+            'max_tokens': kwargs.get('max_tokens', self.generation_config.max_output_tokens),
+            'chat_mode': kwargs.get('chat_mode', 'balanced'),
+            'use_multi_agent': kwargs.get('use_multi_agent', False),
+        }
+        
+        # Add any additional parameters
+        for key, value in kwargs.items():
+            if key not in params:
+                params[key] = value
+        
+        try:
+            response_data = self.client.send_message(
+                message=message,
+                model_id=self.model_id,
+                conversation_id=self.conversation_id,
+                **params
+            )
+            
+            # Add to history
+            self.history.append({
+                'role': 'user',
+                'content': message,
+                'timestamp': time.time()
+            })
+            
+            response_text = response_data.get('data', {}).get('response', '')
+            self.history.append({
+                'role': 'assistant', 
+                'content': response_text,
+                'timestamp': time.time(),
+                'metadata': response_data.get('data', {})
+            })
+            
+            return GenerateContentResponse(response_data)
+            
+        except Exception as e:
+            if isinstance(e, CostKatanaError):
+                raise
+            raise CostKatanaError(f"Failed to send message: {str(e)}")
+    
+    def get_history(self) -> List[Dict[str, Any]]:
+        """Get the conversation history"""
+        try:
+            history_response = self.client.get_conversation_history(self.conversation_id)
+            return history_response.get('data', [])
+        except Exception as e:
+            # Fall back to local history if API call fails
+            return self.history
+    
+    def clear_history(self):
+        """Clear the local conversation history"""
+        self.history = []
+    
+    def delete_conversation(self):
+        """Delete the conversation from the server"""
+        try:
+            self.client.delete_conversation(self.conversation_id)
+            self.conversation_id = None
+            self.history = []
+        except Exception as e:
+            raise CostKatanaError(f"Failed to delete conversation: {str(e)}")
+
+class GenerativeModel:
+    """
+    A generative AI model with a simple interface similar to google-generative-ai.
+    All requests are routed through Cost Katana for optimization and cost management.
+    """
+    
+    def __init__(
+        self,
+        client: CostKatanaClient,
+        model_name: str,
+        generation_config: Optional[GenerationConfig] = None,
+        **kwargs
+    ):
+        """
+        Initialize a generative model.
+        
+        Args:
+            client: Cost Katana client instance
+            model_name: Name of the model (e.g., 'gemini-2.0-flash', 'claude-3-sonnet')
+            generation_config: Generation configuration
+            **kwargs: Additional model parameters
+        """
+        self.client = client
+        self.model_name = model_name
+        self.model_id = client.config.get_model_mapping(model_name)
+        self.generation_config = generation_config or GenerationConfig()
+        self.model_params = kwargs
+        
+        # Validate model is available
+        self._validate_model()
+    
+    def _validate_model(self):
+        """Validate that the model is available"""
+        try:
+            available_models = self.client.get_available_models()
+            model_ids = [model.get('id', model.get('modelId', '')) for model in available_models]
+            
+            if self.model_id not in model_ids and self.model_name not in model_ids:
+                raise ModelNotAvailableError(
+                    f"Model '{self.model_name}' (ID: {self.model_id}) is not available. "
+                    f"Available models: {', '.join(model_ids[:5])}..."
+                )
+        except ModelNotAvailableError:
+            raise
+        except Exception as e:
+            # If we can't validate, log but don't fail - the model might still work
+            print(f"Warning: Could not validate model availability: {e}")
+    
+    def generate_content(
+        self,
+        prompt: Union[str, List[str]],
+        generation_config: Optional[GenerationConfig] = None,
+        **kwargs
+    ) -> GenerateContentResponse:
+        """
+        Generate content from a prompt.
+        
+        Args:
+            prompt: Text prompt or list of prompts
+            generation_config: Generation configuration (overrides instance config)
+            **kwargs: Additional parameters
+            
+        Returns:
+            GenerateContentResponse with the generated content
+            
+        Example:
+            model = cost_katana.GenerativeModel('gemini-2.0-flash')
+            response = model.generate_content("Tell me about AI")
+            print(response.text)
+            print(f"Cost: ${response.usage_metadata.cost:.4f}")
+        """
+        # Handle multiple prompts
+        if isinstance(prompt, list):
+            prompt = "\n\n".join(str(p) for p in prompt)
+        
+        # Use provided config or instance config
+        config = generation_config or self.generation_config
+        
+        # Prepare parameters
+        params = {
+            'temperature': kwargs.get('temperature', config.temperature),
+            'max_tokens': kwargs.get('max_tokens', config.max_output_tokens),
+            'chat_mode': kwargs.get('chat_mode', 'balanced'),
+            'use_multi_agent': kwargs.get('use_multi_agent', False),
+        }
+        
+        # Add any additional parameters from model_params or kwargs
+        params.update(self.model_params)
+        for key, value in kwargs.items():
+            if key not in params:
+                params[key] = value
+        
+        try:
+            response_data = self.client.send_message(
+                message=prompt,
+                model_id=self.model_id,
+                **params
+            )
+            
+            return GenerateContentResponse(response_data)
+            
+        except Exception as e:
+            if isinstance(e, CostKatanaError):
+                raise
+            raise CostKatanaError(f"Failed to generate content: {str(e)}")
+    
+    def start_chat(
+        self,
+        history: Optional[List[Dict[str, Any]]] = None,
+        **kwargs
+    ) -> ChatSession:
+        """
+        Start a chat session.
+        
+        Args:
+            history: Optional conversation history
+            **kwargs: Additional chat configuration
+            
+        Returns:
+            ChatSession instance
+            
+        Example:
+            model = cost_katana.GenerativeModel('gemini-2.0-flash')
+            chat = model.start_chat()
+            response = chat.send_message("Hello!")
+            print(response.text)
+        """
+        chat_session = ChatSession(
+            client=self.client,
+            model_id=self.model_id,
+            generation_config=self.generation_config,
+            **kwargs
+        )
+        
+        # Add history if provided
+        if history:
+            chat_session.history = history
+            
+        return chat_session
+    
+    def count_tokens(self, prompt: str) -> Dict[str, int]:
+        """
+        Count tokens in a prompt (estimated).
+        Note: This is a client-side estimate. Actual tokenization happens on the server.
+        """
+        # Simple word-based estimation - not accurate but gives an idea
+        words = len(prompt.split())
+        estimated_tokens = int(words * 1.3)  # Rough approximation
+        
+        return {
+            'total_tokens': estimated_tokens,
+            'prompt_tokens': estimated_tokens,
+            'completion_tokens': 0
+        }
+    
+    def __repr__(self) -> str:
+        return f"GenerativeModel(model_name='{self.model_name}', model_id='{self.model_id}')"