scripts/cloud_router.py

"""
Cloud Services Integration for COSMOS

Enables optional Azure, IBM, and other cloud APIs while maintaining local-first design.
Users can input their own credentials and model references.
"""

import os
import json
from pathlib import Path
from typing import Optional, Dict, Any

class CloudConfig:
    """Manage cloud service credentials and configuration."""
    
    def __init__(self, config_path: str = "cloud_config.json"):
        self.config_path = Path(config_path)
        self.config = self._load_config()
    
    def _load_config(self) -> Dict[str, Any]:
        """Load cloud config from file or create template."""
        if self.config_path.exists():
            with open(self.config_path, 'r') as f:
                return json.load(f)
        
        # Default template
        return {
            "enabled": False,
            "default_provider": "ollama",  # ollama, azure, ibm, custom
            "providers": {
                "azure": {
                    "enabled": False,
                    "api_key": "",  # Set via env: COSMOS_AZURE_KEY
                    "api_endpoint": "",  # https://{resource}.openai.azure.com/
                    "deployment_name": "",  # Model deployment name
                    "model_blob": "gpt-4",  # Vision: gpt-4-vision, etc.
                    "temperature": 0.7,
                    "timeout": 30
                },
                "ibm": {
                    "enabled": False,
                    "api_key": "",  # Set via env: COSMOS_IBM_KEY
                    "api_endpoint": "",  # https://api.us-south.watson-platform.net/instances/...
                    "model_name": "granite-13b-chat-v2",  # or custom model ID
                    "model_blob": "ibm/granite",
                    "temperature": 0.7,
                    "timeout": 30
                },
                "ollama": {
                    "enabled": True,  # Local by default
                    "api_endpoint": "http://localhost:11434",
                    "model_name": "cosmos-q4:latest",
                    "timeout": 60
                }
            }
        }
    
    def save(self):
        """Save configuration to file."""
        # Don't save API keys to disk — they must come from env vars
        safe_config = json.loads(json.dumps(self.config))
        safe_config["providers"]["azure"]["api_key"] = "[SET_VIA_ENV]"
        safe_config["providers"]["ibm"]["api_key"] = "[SET_VIA_ENV]"
        
        with open(self.config_path, 'w') as f:
            json.dump(safe_config, f, indent=2)
    
    def load_credentials_from_env(self):
        """Load API credentials from environment variables (secure method)."""
        self.config["providers"]["azure"]["api_key"] = os.getenv("COSMOS_AZURE_KEY", "")
        self.config["providers"]["ibm"]["api_key"] = os.getenv("COSMOS_IBM_KEY", "")
    
    def enable_provider(self, provider: str, enabled: bool = True):
        """Enable/disable a cloud provider."""
        if provider in self.config["providers"]:
            self.config["providers"][provider]["enabled"] = enabled
    
    def set_provider_endpoint(self, provider: str, endpoint: str):
        """Set API endpoint for a provider."""
        if provider in self.config["providers"]:
            self.config["providers"][provider]["api_endpoint"] = endpoint
    
    def set_provider_model(self, provider: str, model_name: str, model_blob: str = None):
        """Set model name and optional blob reference."""
        if provider in self.config["providers"]:
            self.config["providers"][provider]["model_name"] = model_name
            if model_blob:
                self.config["providers"][provider]["model_blob"] = model_blob
    
    def get_active_provider(self) -> str:
        """Get the currently active provider."""
        return self.config.get("default_provider", "ollama")
    
    def set_default_provider(self, provider: str):
        """Set default provider for requests."""
        if provider in self.config["providers"]:
            self.config["default_provider"] = provider
    
    def get_provider_config(self, provider: str) -> Dict[str, Any]:
        """Get full config for a specific provider."""
        return self.config["providers"].get(provider, {})


class CloudRouter:
    """Route requests to appropriate cloud service."""
    
    def __init__(self, config: CloudConfig):
        self.config = config
        self.config.load_credentials_from_env()
    
    def generate(self, prompt: str, provider: Optional[str] = None, **kwargs) -> str:
        """Generate response from configured provider."""
        provider = provider or self.config.get_active_provider()
        
        if provider == "azure":
            return self._generate_azure(prompt, **kwargs)
        elif provider == "ibm":
            return self._generate_ibm(prompt, **kwargs)
        elif provider == "ollama":
            return self._generate_ollama(prompt, **kwargs)
        else:
            raise ValueError(f"Unknown provider: {provider}")
    
    def _generate_azure(self, prompt: str, **kwargs) -> str:
        """Call Azure OpenAI API."""
        try:
            import openai
        except ImportError:
            raise ImportError("Install openai: pip install openai")
        
        cfg = self.config.get_provider_config("azure")
        if not cfg["enabled"] or not cfg["api_key"]:
            raise ValueError("Azure not enabled or API key not set (use COSMOS_AZURE_KEY env var)")
        
        client = openai.AzureOpenAI(
            api_key=cfg["api_key"],
            api_version="2024-02-15-preview",
            azure_endpoint=cfg["api_endpoint"]
        )
        
        response = client.chat.completions.create(
            model=cfg["deployment_name"],
            messages=[{"role": "user", "content": prompt}],
            temperature=cfg.get("temperature", 0.7),
            timeout=cfg.get("timeout", 30)
        )
        
        return response.choices[0].message.content
    
    def _generate_ibm(self, prompt: str, **kwargs) -> str:
        """Call IBM Watsonx API."""
        try:
            from ibm_cloud_sdk_core import Authenticator, IAMAuthenticator
            from ibm_platform_services import WatsonxAiAnalyticsV1
        except ImportError:
            raise ImportError("Install IBM SDK: pip install ibm-cloud-sdk-core ibm-cloud-sdk-watsonx")
        
        cfg = self.config.get_provider_config("ibm")
        if not cfg["enabled"] or not cfg["api_key"]:
            raise ValueError("IBM not enabled or API key not set (use COSMOS_IBM_KEY env var)")
        
        authenticator = IAMAuthenticator(apikey=cfg["api_key"])
        service = WatsonxAiAnalyticsV1(
            version="2024-01-01",
            authenticator=authenticator,
            service_url=cfg["api_endpoint"]
        )
        
        response = service.generate(
            input=prompt,
            model_id=cfg["model_name"],
            parameters={
                "temperature": cfg.get("temperature", 0.7),
                "max_tokens": 512
            }
        ).get_result()
        
        return response["results"][0]["generated_text"]
    
    def _generate_ollama(self, prompt: str, **kwargs) -> str:
        """Call local Ollama API."""
        try:
            import requests
        except ImportError:
            raise ImportError("Install requests: pip install requests")
        
        cfg = self.config.get_provider_config("ollama")
        
        response = requests.post(
            f"{cfg['api_endpoint']}/api/generate",
            json={
                "model": cfg["model_name"],
                "prompt": prompt,
                "stream": False
            },
            timeout=cfg.get("timeout", 60)
        )
        
        if response.status_code == 200:
            return response.json()["response"]
        else:
            raise RuntimeError(f"Ollama error: {response.text}")
    
    def vision(self, image_path: str, prompt: str, provider: Optional[str] = None) -> str:
        """Process image with vision model (Azure/IBM only)."""
        provider = provider or self.config.get_active_provider()
        
        if provider == "azure":
            return self._vision_azure(image_path, prompt)
        elif provider == "ibm":
            return self._vision_ibm(image_path, prompt)
        else:
            raise ValueError(f"Vision not supported on {provider} provider")
    
    def _vision_azure(self, image_path: str, prompt: str) -> str:
        """Azure vision analysis."""
        import base64
        from pathlib import Path
        import openai
        
        cfg = self.config.get_provider_config("azure")
        client = openai.AzureOpenAI(
            api_key=cfg["api_key"],
            api_version="2024-02-15-preview",
            azure_endpoint=cfg["api_endpoint"]
        )
        
        # Read and encode image
        with open(image_path, "rb") as img_file:
            image_data = base64.standard_b64encode(img_file.read()).decode("utf-8")
        
        ext = Path(image_path).suffix.lower()
        media_type = {"jpg": "image/jpeg", "jpeg": "image/jpeg", "png": "image/png", "gif": "image/gif", "webp": "image/webp"}.get(ext[1:], "image/jpeg")
        
        response = client.chat.completions.create(
            model=cfg["deployment_name"],
            messages=[
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": prompt},
                        {
                            "type": "image_url",
                            "image_url": {"url": f"data:{media_type};base64,{image_data}"}
                        }
                    ]
                }
            ]
        )
        
        return response.choices[0].message.content
    
    def _vision_ibm(self, image_path: str, prompt: str) -> str:
        """IBM vision analysis."""
        raise NotImplementedError("IBM vision support coming soon")


# Export
__all__ = ["CloudConfig", "CloudRouter"]