cost-katana 1.0.0__py3-none-any.whl

cost_katana/__init__.py

@@ -0,0 +1,69 @@
+"""
+Cost Katana - Unified AI Interface with Cost Optimization
+
+Simple interface for AI models that routes through Cost Katana for:
+- Cost optimization and tracking
+- Automatic failover between providers
+- Usage analytics and insights
+- No API key management needed in your code
+
+Example:
+    import cost_katana as ck
+    
+    # Configure once
+    ck.configure(config_file='config.json')
+    
+    # Use like any AI library
+    model = ck.GenerativeModel('gemini-2.0-flash')
+    chat = model.start_chat()
+    response = chat.send_message("Hello!")
+    print(response.text)
+"""
+
+from .client import CostKatanaClient, get_global_client
+from .models import GenerativeModel, ChatSession
+from .exceptions import (
+    CostKatanaError,
+    AuthenticationError,
+    ModelNotAvailableError,
+    RateLimitError,
+    CostLimitExceededError
+)
+from .config import Config
+
+__version__ = "1.0.0"
+__all__ = [
+    "configure",
+    "GenerativeModel", 
+    "ChatSession",
+    "CostKatanaClient",
+    "CostKatanaError",
+    "AuthenticationError", 
+    "ModelNotAvailableError",
+    "RateLimitError",
+    "CostLimitExceededError",
+    "Config"
+]
+
+# Import configure function from client
+from .client import configure
+
+def GenerativeModel(model_name: str, **kwargs):
+    """
+    Create a generative model instance.
+    
+    Args:
+        model_name: Name of the model (e.g., 'gemini-2.0-flash', 'claude-3-sonnet', 'gpt-4')
+        **kwargs: Additional model configuration
+        
+    Returns:
+        GenerativeModel instance
+        
+    Example:
+        model = cost_katana.GenerativeModel('gemini-2.0-flash')
+        response = model.generate_content("Hello, world!")
+    """
+    client = get_global_client()
+    
+    from .models import GenerativeModel as GM
+    return GM(client, model_name, **kwargs)

cost_katana/cli.py

@@ -0,0 +1,303 @@
+"""
+Command-line interface for Cost Katana
+"""
+
+import argparse
+import json
+import sys
+from pathlib import Path
+from typing import Optional
+from rich.console import Console
+from rich.table import Table
+from rich.panel import Panel
+from rich.prompt import Prompt, Confirm
+from rich.syntax import Syntax
+
+try:
+    from . import configure, GenerativeModel, CostKatanaClient
+    from .config import Config
+    from .exceptions import CostKatanaError
+except ImportError:
+    # Handle case when running as script
+    import cost_katana as ck
+    from cost_katana.config import Config
+    from cost_katana.exceptions import CostKatanaError
+
+console = Console()
+
+def create_sample_config():
+    """Create a sample configuration file"""
+    sample_config = {
+        "api_key": "dak_your_api_key_here",
+        "base_url": "https://api.costkatana.com",
+        "default_model": "gemini-2.0-flash",
+        "default_temperature": 0.7,
+        "default_max_tokens": 2000,
+        "cost_limit_per_day": 50.0,
+        "enable_analytics": True,
+        "enable_optimization": True,
+        "enable_failover": True,
+        "model_mappings": {
+            "gemini": "gemini-2.0-flash-exp",
+            "claude": "anthropic.claude-3-sonnet-20240229-v1:0",
+            "gpt4": "gpt-4-turbo-preview"
+        },
+        "providers": {
+            "google": {
+                "priority": 1,
+                "models": ["gemini-2.0-flash", "gemini-pro"]
+            },
+            "anthropic": {
+                "priority": 2,
+                "models": ["claude-3-sonnet", "claude-3-haiku"]
+            },
+            "openai": {
+                "priority": 3,
+                "models": ["gpt-4", "gpt-3.5-turbo"]
+            }
+        }
+    }
+    return sample_config
+
+def init_config(args):
+    """Initialize configuration"""
+    config_path = Path(args.config or 'cost_katana_config.json')
+    
+    if config_path.exists() and not args.force:
+        console.print(f"[yellow]Configuration file already exists: {config_path}[/yellow]")
+        if not Confirm.ask("Overwrite existing configuration?"):
+            return
+    
+    console.print("[bold blue]Setting up Cost Katana configuration...[/bold blue]")
+    
+    # Get API key
+    api_key = Prompt.ask(
+        "Enter your Cost Katana API key",
+        default=args.api_key if args.api_key else None
+    )
+    
+    # Get base URL
+    base_url = Prompt.ask(
+        "Enter base URL",
+        default="https://api.costkatana.com"
+    )
+    
+    # Get default model
+    default_model = Prompt.ask(
+        "Enter default model",
+        default="gemini-2.0-flash",
+        choices=["gemini-2.0-flash", "claude-3-sonnet", "gpt-4", "nova-pro"]
+    )
+    
+    # Create configuration
+    config_data = create_sample_config()
+    config_data.update({
+        "api_key": api_key,
+        "base_url": base_url,
+        "default_model": default_model
+    })
+    
+    # Save configuration
+    try:
+        with open(config_path, 'w') as f:
+            json.dump(config_data, f, indent=2)
+        
+        console.print(f"[green]Configuration saved to: {config_path}[/green]")
+        console.print("\n[bold]Next steps:[/bold]")
+        console.print("1. Test the configuration: [cyan]cost-katana test[/cyan]")
+        console.print("2. Start a chat session: [cyan]cost-katana chat[/cyan]")
+        console.print("3. See available models: [cyan]cost-katana models[/cyan]")
+        
+    except Exception as e:
+        console.print(f"[red]Failed to save configuration: {e}[/red]")
+        sys.exit(1)
+
+def test_connection(args):
+    """Test connection to Cost Katana API"""
+    try:
+        config_path = args.config or 'cost_katana_config.json'
+        
+        if Path(config_path).exists():
+            configure(config_file=config_path)
+        elif args.api_key:
+            configure(api_key=args.api_key)
+        else:
+            console.print("[red]No configuration found. Run 'cost-katana init' first.[/red]")
+            return
+        
+        console.print("[bold blue]Testing Cost Katana connection...[/bold blue]")
+        
+        # Test with a simple model
+        model = GenerativeModel('gemini-2.0-flash')
+        response = model.generate_content("Hello! Please respond with just 'OK' to test the connection.")
+        
+        console.print(Panel(
+            f"[green]✓ Connection successful![/green]\n"
+            f"Model: {response.usage_metadata.model}\n"
+            f"Response: {response.text}\n"
+            f"Cost: ${response.usage_metadata.cost:.4f}\n"
+            f"Latency: {response.usage_metadata.latency:.2f}s",
+            title="Test Results"
+        ))
+        
+    except Exception as e:
+        console.print(f"[red]✗ Connection failed: {e}[/red]")
+        sys.exit(1)
+
+def list_models(args):
+    """List available models"""
+    try:
+        config_path = args.config or 'cost_katana_config.json'
+        
+        if Path(config_path).exists():
+            configure(config_file=config_path)
+        elif args.api_key:
+            configure(api_key=args.api_key)
+        else:
+            console.print("[red]No configuration found. Run 'cost-katana init' first.[/red]")
+            return
+        
+        client = CostKatanaClient(config_file=config_path if Path(config_path).exists() else None)
+        models = client.get_available_models()
+        
+        table = Table(title="Available Models")
+        table.add_column("Model ID", style="cyan", no_wrap=True)
+        table.add_column("Display Name", style="magenta")
+        table.add_column("Provider", style="green")
+        table.add_column("Type", style="yellow")
+        
+        for model in models:
+            model_id = model.get('id', model.get('modelId', 'Unknown'))
+            name = model.get('name', model.get('displayName', model_id))
+            provider = model.get('provider', 'Unknown')
+            model_type = model.get('type', 'Text')
+            
+            table.add_row(model_id, name, provider, model_type)
+        
+        console.print(table)
+        
+    except Exception as e:
+        console.print(f"[red]Failed to fetch models: {e}[/red]")
+        sys.exit(1)
+
+def start_chat(args):
+    """Start an interactive chat session"""
+    try:
+        config_path = args.config or 'cost_katana_config.json'
+        
+        if Path(config_path).exists():
+            configure(config_file=config_path)
+            config = Config.from_file(config_path)
+        elif args.api_key:
+            configure(api_key=args.api_key)
+            config = Config(api_key=args.api_key)
+        else:
+            console.print("[red]No configuration found. Run 'cost-katana init' first.[/red]")
+            return
+        
+        model_name = args.model or config.default_model
+        
+        console.print(Panel(
+            f"[bold blue]Cost Katana Chat Session[/bold blue]\n"
+            f"Model: {model_name}\n"
+            f"Type 'quit' to exit, 'clear' to clear history",
+            title="Welcome"
+        ))
+        
+        model = GenerativeModel(model_name)
+        chat = model.start_chat()
+        
+        total_cost = 0.0
+        
+        while True:
+            try:
+                message = Prompt.ask("[bold cyan]You[/bold cyan]")
+                
+                if message.lower() in ['quit', 'exit', 'q']:
+                    break
+                elif message.lower() == 'clear':
+                    chat.clear_history()
+                    console.print("[yellow]Chat history cleared.[/yellow]")
+                    continue
+                elif message.lower() == 'cost':
+                    console.print(f"[green]Total session cost: ${total_cost:.4f}[/green]")
+                    continue
+                
+                console.print("[bold green]Assistant[/bold green]: ", end="")
+                
+                with console.status("Thinking..."):
+                    response = chat.send_message(message)
+                
+                console.print(response.text)
+                
+                # Show cost info
+                total_cost += response.usage_metadata.cost
+                console.print(
+                    f"[dim]Cost: ${response.usage_metadata.cost:.4f} | "
+                    f"Total: ${total_cost:.4f} | "
+                    f"Tokens: {response.usage_metadata.total_tokens}[/dim]\n"
+                )
+                
+            except KeyboardInterrupt:
+                console.print("\n[yellow]Chat session interrupted.[/yellow]")
+                break
+            except Exception as e:
+                console.print(f"[red]Error: {e}[/red]")
+                continue
+        
+        console.print(f"\n[bold]Session Summary:[/bold]")
+        console.print(f"Total Cost: ${total_cost:.4f}")
+        console.print("Thanks for using Cost Katana!")
+        
+    except Exception as e:
+        console.print(f"[red]Failed to start chat: {e}[/red]")
+        sys.exit(1)
+
+def main():
+    """Main CLI entry point"""
+    parser = argparse.ArgumentParser(
+        description="Cost Katana - Unified AI interface with cost optimization"
+    )
+    parser.add_argument(
+        '--config', '-c',
+        help='Configuration file path (default: cost_katana_config.json)'
+    )
+    parser.add_argument(
+        '--api-key', '-k',
+        help='Cost Katana API key'
+    )
+    
+    subparsers = parser.add_subparsers(dest='command', help='Available commands')
+    
+    # Init command
+    init_parser = subparsers.add_parser('init', help='Initialize configuration')
+    init_parser.add_argument('--force', action='store_true', help='Overwrite existing config')
+    
+    # Test command
+    subparsers.add_parser('test', help='Test API connection')
+    
+    # Models command
+    subparsers.add_parser('models', help='List available models')
+    
+    # Chat command
+    chat_parser = subparsers.add_parser('chat', help='Start interactive chat')
+    chat_parser.add_argument('--model', '-m', help='Model to use for chat')
+    
+    args = parser.parse_args()
+    
+    if not args.command:
+        parser.print_help()
+        return
+    
+    # Route to appropriate function
+    if args.command == 'init':
+        init_config(args)
+    elif args.command == 'test':
+        test_connection(args)
+    elif args.command == 'models':
+        list_models(args)
+    elif args.command == 'chat':
+        start_chat(args)
+
+if __name__ == '__main__':
+    main()

cost_katana/client.py

@@ -0,0 +1,235 @@
+"""
+Cost Katana HTTP Client
+Handles communication with the Cost Katana backend API
+"""
+
+import json
+import os
+from typing import Dict, Any, Optional, List
+import httpx
+from .config import Config
+from .exceptions import (
+    CostKatanaError,
+    AuthenticationError,
+    ModelNotAvailableError,
+    RateLimitError,
+    CostLimitExceededError
+)
+
+# Global client instance for the configure function
+_global_client = None
+
+def configure(
+    api_key: str = None,
+    base_url: str = None,
+    config_file: str = None,
+    **kwargs
+):
+    """
+    Configure Cost Katana client globally.
+    
+    Args:
+        api_key: Your Cost Katana API key (starts with 'dak_')
+        base_url: Base URL for Cost Katana API (optional)
+        config_file: Path to JSON configuration file (optional)
+        **kwargs: Additional configuration options
+        
+    Example:
+        # Using API key
+        cost_katana.configure(api_key='dak_your_key_here')
+        
+        # Using config file
+        cost_katana.configure(config_file='config.json')
+    """
+    global _global_client
+    _global_client = CostKatanaClient(
+        api_key=api_key,
+        base_url=base_url, 
+        config_file=config_file,
+        **kwargs
+    )
+    return _global_client
+
+def get_global_client():
+    """Get the global client instance"""
+    if _global_client is None:
+        raise CostKatanaError("Cost Katana not configured. Call cost_katana.configure() first.")
+    return _global_client
+
+class CostKatanaClient:
+    """HTTP client for Cost Katana API"""
+    
+    def __init__(
+        self,
+        api_key: str = None,
+        base_url: str = None,
+        config_file: str = None,
+        timeout: int = 30,
+        **kwargs
+    ):
+        """
+        Initialize Cost Katana client.
+        
+        Args:
+            api_key: Your Cost Katana API key
+            base_url: Base URL for the API
+            config_file: Path to JSON configuration file
+            timeout: Request timeout in seconds
+        """
+        self.config = Config.from_file(config_file) if config_file else Config()
+        
+        # Override with provided parameters
+        if api_key:
+            self.config.api_key = api_key
+        if base_url:
+            self.config.base_url = base_url
+            
+        # Apply additional config
+        for key, value in kwargs.items():
+            setattr(self.config, key, value)
+            
+        # Validate configuration
+        if not self.config.api_key:
+            raise AuthenticationError(
+                "API key is required. Get one from https://costkatana.com/dashboard/api-keys"
+            )
+            
+        # Initialize HTTP client
+        self.client = httpx.Client(
+            base_url=self.config.base_url,
+            timeout=timeout,
+            headers={
+                'Authorization': f'Bearer {self.config.api_key}',
+                'Content-Type': 'application/json',
+                'User-Agent': f'cost-katana-python/1.0.0'
+            }
+        )
+    
+    def __enter__(self):
+        return self
+        
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+        
+    def close(self):
+        """Close the HTTP client"""
+        if hasattr(self, 'client'):
+            self.client.close()
+    
+    def _handle_response(self, response: httpx.Response) -> Dict[str, Any]:
+        """Handle HTTP response and raise appropriate exceptions"""
+        try:
+            data = response.json()
+        except json.JSONDecodeError:
+            raise CostKatanaError(f"Invalid JSON response: {response.text}")
+            
+        if response.status_code == 401:
+            raise AuthenticationError(data.get('message', 'Authentication failed'))
+        elif response.status_code == 403:
+            raise AuthenticationError(data.get('message', 'Access forbidden'))
+        elif response.status_code == 404:
+            raise ModelNotAvailableError(data.get('message', 'Model not found'))
+        elif response.status_code == 429:
+            raise RateLimitError(data.get('message', 'Rate limit exceeded'))
+        elif response.status_code == 400 and 'cost' in data.get('message', '').lower():
+            raise CostLimitExceededError(data.get('message', 'Cost limit exceeded'))
+        elif not response.is_success:
+            raise CostKatanaError(
+                data.get('message', f'API error: {response.status_code}')
+            )
+            
+        return data
+    
+    def get_available_models(self) -> List[Dict[str, Any]]:
+        """Get list of available models"""
+        try:
+            response = self.client.get('/api/chat/models')
+            data = self._handle_response(response)
+            return data.get('data', [])
+        except Exception as e:
+            if isinstance(e, CostKatanaError):
+                raise
+            raise CostKatanaError(f"Failed to get models: {str(e)}")
+    
+    def send_message(
+        self,
+        message: str,
+        model_id: str,
+        conversation_id: str = None,
+        temperature: float = 0.7,
+        max_tokens: int = 2000,
+        chat_mode: str = 'balanced',
+        use_multi_agent: bool = False,
+        **kwargs
+    ) -> Dict[str, Any]:
+        """
+        Send a message to the AI model via Cost Katana.
+        
+        Args:
+            message: The message to send
+            model_id: ID of the model to use
+            conversation_id: Optional conversation ID
+            temperature: Sampling temperature (0.0 to 1.0)
+            max_tokens: Maximum tokens to generate
+            chat_mode: Chat optimization mode ('fastest', 'cheapest', 'balanced')
+            use_multi_agent: Whether to use multi-agent processing
+            
+        Returns:
+            Response data from the API
+        """
+        payload = {
+            'message': message,
+            'modelId': model_id,
+            'temperature': temperature,
+            'maxTokens': max_tokens,
+            'chatMode': chat_mode,
+            'useMultiAgent': use_multi_agent,
+            **kwargs
+        }
+        
+        if conversation_id:
+            payload['conversationId'] = conversation_id
+            
+        try:
+            response = self.client.post('/api/chat/message', json=payload)
+            return self._handle_response(response)
+        except Exception as e:
+            if isinstance(e, CostKatanaError):
+                raise
+            raise CostKatanaError(f"Failed to send message: {str(e)}")
+    
+    def create_conversation(self, title: str = None, model_id: str = None) -> Dict[str, Any]:
+        """Create a new conversation"""
+        payload = {}
+        if title:
+            payload['title'] = title
+        if model_id:
+            payload['modelId'] = model_id
+            
+        try:
+            response = self.client.post('/api/chat/conversations', json=payload)
+            return self._handle_response(response)
+        except Exception as e:
+            if isinstance(e, CostKatanaError):
+                raise
+            raise CostKatanaError(f"Failed to create conversation: {str(e)}")
+    
+    def get_conversation_history(self, conversation_id: str) -> Dict[str, Any]:
+        """Get conversation history"""
+        try:
+            response = self.client.get(f'/api/chat/conversations/{conversation_id}/history')
+            return self._handle_response(response)
+        except Exception as e:
+            if isinstance(e, CostKatanaError):
+                raise
+            raise CostKatanaError(f"Failed to get conversation history: {str(e)}")
+    
+    def delete_conversation(self, conversation_id: str) -> Dict[str, Any]:
+        """Delete a conversation"""
+        try:
+            response = self.client.delete(f'/api/chat/conversations/{conversation_id}')
+            return self._handle_response(response)
+        except Exception as e:
+            if isinstance(e, CostKatanaError):
+                raise
+            raise CostKatanaError(f"Failed to delete conversation: {str(e)}")