isa-model 0.2.0__py3-none-any.whl → 0.3.1__py3-none-any.whl

isa_model/inference/providers/openai_provider.py

@@ -7,51 +7,67 @@ import os
 logger = logging.getLogger(__name__)
 
 class OpenAIProvider(BaseProvider):
-    """Provider for OpenAI API"""
+    """Provider for OpenAI API with proper API key management"""
     
     def __init__(self, config=None):
-        """
-        Initialize the OpenAI Provider
+        """Initialize the OpenAI Provider with centralized config management"""
+        super().__init__(config)
+        self.name = "openai"
+        
+        logger.info(f"Initialized OpenAIProvider with URL: {self.config.get('base_url', 'https://api.openai.com/v1')}")
         
-        Args:
-            config (dict, optional): Configuration for the provider
-                - api_key: OpenAI API key (can be passed here or via environment variable)
-                - api_base: Base URL for OpenAI API (default: https://api.openai.com/v1)
-                - timeout: Timeout for API calls in seconds
-        """
-        default_config = {
-            "api_key": "",  # Will be set from config or environment
-            "api_base": os.environ.get("OPENAI_API_BASE", "https://api.openai.com/v1"),
+        if not self.has_valid_credentials():
+            logger.warning("OpenAI API key not found. Set OPENAI_API_KEY environment variable or pass api_key in config.")
+    
+    def _load_provider_env_vars(self):
+        """Load OpenAI-specific environment variables"""
+        # Set defaults first
+        defaults = {
+            "base_url": "https://api.openai.com/v1",
             "timeout": 60,
-            "stream": True,
             "temperature": 0.7,
             "top_p": 0.9,
             "max_tokens": 1024
         }
         
-        # Merge default config with provided config
-        merged_config = {**default_config, **(config or {})}
+        # Apply defaults only if not already set
+        for key, value in defaults.items():
+            if key not in self.config:
+                self.config[key] = value
         
-        # Set API key from config first, then fallback to environment variable
-        if not merged_config["api_key"]:
-            merged_config["api_key"] = os.environ.get("OPENAI_API_KEY", "")
-        
-        super().__init__(config=merged_config)
-        self.name = "openai"
-        
-        logger.info(f"Initialized OpenAIProvider with URL: {self.config['api_base']}")
+        # Load from environment variables (override config if present)
+        env_mappings = {
+            "api_key": "OPENAI_API_KEY",
+            "base_url": "OPENAI_API_BASE",
+            "organization": "OPENAI_ORGANIZATION"
+        }
         
-        # Only warn if no API key is provided at all
-        if not self.config["api_key"]:
-            logger.info("OpenAI API key not provided. You can set it via OPENAI_API_KEY environment variable or pass it in the config when creating services.")
+        for config_key, env_var in env_mappings.items():
+            env_value = os.getenv(env_var)
+            if env_value:
+                self.config[config_key] = env_value
+    
+    def _validate_config(self):
+        """Validate OpenAI configuration"""
+        if not self.config.get("api_key"):
+            logger.debug("OpenAI API key not set - some functionality may not work")
+    
+    def get_model_pricing(self, model_name: str) -> Dict[str, float]:
+        """Get pricing information for a model - delegated to ModelManager"""
+        # Import here to avoid circular imports
+        from isa_model.core.model_manager import ModelManager
+        model_manager = ModelManager()
+        return model_manager.get_model_pricing("openai", model_name)
+    
+    def calculate_cost(self, model_name: str, input_tokens: int, output_tokens: int) -> float:
+        """Calculate cost for a request - delegated to ModelManager"""
+        # Import here to avoid circular imports
+        from isa_model.core.model_manager import ModelManager
+        model_manager = ModelManager()
+        return model_manager.calculate_cost("openai", model_name, input_tokens, output_tokens)
     
     def set_api_key(self, api_key: str):
-        """
-        Set the API key after initialization
-        
-        Args:
-            api_key: OpenAI API key
-        """
+        """Set the API key after initialization"""
         self.config["api_key"] = api_key
         logger.info("OpenAI API key updated")
     
@@ -77,16 +93,29 @@ class OpenAIProvider(BaseProvider):
     def get_models(self, model_type: ModelType) -> List[str]:
         """Get available models for given type"""
         if model_type == ModelType.LLM:
-            return ["gpt-4o", "gpt-4o-mini", "gpt-4-turbo", "gpt-4", "gpt-3.5-turbo"]
+            return ["gpt-4.1-nano", "gpt-4.1-mini", "gpt-4o-mini", "gpt-4o", "gpt-4-turbo", "gpt-4", "gpt-3.5-turbo"]
         elif model_type == ModelType.EMBEDDING:
             return ["text-embedding-3-large", "text-embedding-3-small", "text-embedding-ada-002"]
         elif model_type == ModelType.VISION:
-            return ["gpt-4o", "gpt-4-vision-preview"]
+            return ["gpt-4.1-nano", "gpt-4.1-mini", "gpt-4o-mini", "gpt-4o", "gpt-4-vision-preview"]
         elif model_type == ModelType.AUDIO:
-            return ["whisper-1"]
+            return ["whisper-1", "gpt-4o-transcribe", "tts-1", "tts-1-hd"]
         else:
             return []
     
+    def get_default_model(self, model_type: ModelType) -> str:
+        """Get default model for a given type"""
+        if model_type == ModelType.LLM:
+            return "gpt-4.1-nano"  # Cheapest and most cost-effective
+        elif model_type == ModelType.EMBEDDING:
+            return "text-embedding-3-small"
+        elif model_type == ModelType.VISION:
+            return "gpt-4.1-nano"
+        elif model_type == ModelType.AUDIO:
+            return "whisper-1"
+        else:
+            return ""
+    
     def get_config(self) -> Dict[str, Any]:
         """Get provider configuration"""
         # Return a copy without sensitive information
@@ -97,5 +126,5 @@ class OpenAIProvider(BaseProvider):
     
     def is_reasoning_model(self, model_name: str) -> bool:
         """Check if the model is optimized for reasoning tasks"""
-        reasoning_models = ["gpt-4", "gpt-4o", "gpt-4-turbo"]
+        reasoning_models = ["gpt-4", "gpt-4o", "gpt-4-turbo", "gpt-4.1"]
         return any(rm in model_name.lower() for rm in reasoning_models) 

isa_model/inference/providers/replicate_provider.py

@@ -7,47 +7,56 @@ import os
 logger = logging.getLogger(__name__)
 
 class ReplicateProvider(BaseProvider):
-    """Provider for Replicate API"""
+    """Provider for Replicate API with proper API key management"""
     
     def __init__(self, config=None):
-        """
-        Initialize the Replicate Provider
+        """Initialize the Replicate Provider with centralized config management"""
+        super().__init__(config)
+        self.name = "replicate"
+        
+        logger.info("Initialized ReplicateProvider")
         
-        Args:
-            config (dict, optional): Configuration for the provider
-                - api_token: Replicate API token (can be passed here or via environment variable)
-                - timeout: Timeout for API calls in seconds
-        """
-        default_config = {
-            "api_token": "",  # Will be set from config or environment
+        if not self.has_valid_credentials():
+            logger.warning("Replicate API token not found. Set REPLICATE_API_TOKEN environment variable or pass api_token in config.")
+    
+    def _load_provider_env_vars(self):
+        """Load Replicate-specific environment variables"""
+        # Set defaults first
+        defaults = {
             "timeout": 60,
-            "stream": True,
             "max_tokens": 1024
         }
         
-        # Merge default config with provided config
-        merged_config = {**default_config, **(config or {})}
-        
-        # Set API token from config first, then fallback to environment variable
-        if not merged_config["api_token"]:
-            merged_config["api_token"] = os.environ.get("REPLICATE_API_TOKEN", "")
-        
-        super().__init__(config=merged_config)
-        self.name = "replicate"
+        # Apply defaults only if not already set
+        for key, value in defaults.items():
+            if key not in self.config:
+                self.config[key] = value
         
-        logger.info(f"Initialized ReplicateProvider")
+        # Load from environment variables (override config if present)
+        env_mappings = {
+            "api_token": "REPLICATE_API_TOKEN",
+        }
         
-        # Only warn if no API token is provided at all
-        if not self.config["api_token"]:
-            logger.info("Replicate API token not provided. You can set it via REPLICATE_API_TOKEN environment variable or pass it in the config when creating services.")
+        for config_key, env_var in env_mappings.items():
+            env_value = os.getenv(env_var)
+            if env_value:
+                self.config[config_key] = env_value
+    
+    def _validate_config(self):
+        """Validate Replicate configuration"""
+        if not self.config.get("api_token"):
+            logger.debug("Replicate API token not set - some functionality may not work")
+    
+    def get_api_key(self) -> str:
+        """Get the API token for this provider (override for Replicate naming)"""
+        return self.config.get("api_token", "")
+    
+    def has_valid_credentials(self) -> bool:
+        """Check if provider has valid credentials (override for Replicate naming)"""
+        return bool(self.config.get("api_token"))
     
     def set_api_token(self, api_token: str):
-        """
-        Set the API token after initialization
-        
-        Args:
-            api_token: Replicate API token
-        """
+        """Set the API token after initialization"""
         self.config["api_token"] = api_token
         logger.info("Replicate API token updated")
     
@@ -79,6 +88,8 @@ class ReplicateProvider(BaseProvider):
             ]
         elif model_type == ModelType.VISION:
             return [
+                "black-forest-labs/flux-schnell",
+                "black-forest-labs/flux-kontext-pro",
                 "stability-ai/sdxl",
                 "stability-ai/stable-diffusion-3-medium",
                 "meta/llama-3-70b-vision",
@@ -87,6 +98,7 @@ class ReplicateProvider(BaseProvider):
             ]
         elif model_type == ModelType.AUDIO:
             return [
+                "jaaari/kokoro-82m",
                 "openai/whisper",
                 "suno-ai/bark"
             ]

isa_model/inference/services/audio/base_stt_service.py

@@ -6,7 +6,7 @@ class BaseSTTService(BaseService):
     """Base class for Speech-to-Text services"""
     
     @abstractmethod
-    async def transcribe_audio(
+    async def transcribe(
         self, 
         audio_file: Union[str, BinaryIO], 
         language: Optional[str] = None,
@@ -30,7 +30,26 @@ class BaseSTTService(BaseService):
         pass
     
     @abstractmethod
-    async def transcribe_audio_batch(
+    async def translate(
+        self, 
+        audio_file: Union[str, BinaryIO]
+    ) -> Dict[str, Any]:
+        """
+        Translate audio file to English text
+        
+        Args:
+            audio_file: Path to audio file or file-like object
+            
+        Returns:
+            Dict containing translation results with keys:
+            - text: The translated text (in English)
+            - detected_language: Original language detected
+            - confidence: Confidence score (if available)
+        """
+        pass
+    
+    @abstractmethod
+    async def transcribe_batch(
         self, 
         audio_files: List[Union[str, BinaryIO]], 
         language: Optional[str] = None,

isa_model/inference/services/audio/openai_realtime_service.py

@@ -0,0 +1,353 @@
+import logging
+import json
+import asyncio
+from typing import Dict, Any, List, Optional, Callable, AsyncGenerator
+import aiohttp
+from tenacity import retry, stop_after_attempt, wait_exponential
+
+from isa_model.inference.services.base_service import BaseService
+from isa_model.inference.providers.base_provider import BaseProvider
+from isa_model.inference.billing_tracker import ServiceType
+
+logger = logging.getLogger(__name__)
+
+class OpenAIRealtimeService(BaseService):
+    """
+    OpenAI Realtime API service for real-time audio conversations.
+    Uses gpt-4o-mini-realtime-preview model for interactive audio chat.
+    """
+    
+    def __init__(self, provider: 'BaseProvider', model_name: str = "gpt-4o-mini-realtime-preview"):
+        super().__init__(provider, model_name)
+        
+        self.api_key = self.config.get('api_key')
+        self.base_url = self.config.get('api_base', 'https://api.openai.com/v1')
+        
+        # Default session configuration
+        self.default_config = {
+            "model": self.model_name,
+            "modalities": ["audio", "text"],
+            "voice": "alloy",
+            "input_audio_format": "pcm16",
+            "output_audio_format": "pcm16",
+            "input_audio_transcription": {
+                "model": "whisper-1"
+            },
+            "turn_detection": None,
+            "tools": [],
+            "tool_choice": "none",
+            "temperature": 0.7,
+            "max_response_output_tokens": 200,
+            "speed": 1.1,
+            "tracing": "auto"
+        }
+        
+        logger.info(f"Initialized OpenAIRealtimeService with model '{self.model_name}'")
+    
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, min=4, max=10),
+        reraise=True
+    )
+    async def create_session(
+        self, 
+        instructions: str = "You are a friendly assistant.",
+        modalities: Optional[List[str]] = None,
+        voice: str = "alloy",
+        **kwargs
+    ) -> Dict[str, Any]:
+        """Create a new realtime session"""
+        try:
+            # Prepare session configuration
+            session_config = self.default_config.copy()
+            session_config.update({
+                "instructions": instructions,
+                "modalities": modalities if modalities is not None else ["audio", "text"],
+                "voice": voice,
+                **kwargs
+            })
+            
+            # Create session via REST API
+            url = f"{self.base_url}/realtime/sessions"
+            headers = {
+                "Authorization": f"Bearer {self.api_key}",
+                "Content-Type": "application/json"
+            }
+            
+            async with aiohttp.ClientSession() as session:
+                async with session.post(url, headers=headers, json=session_config) as response:
+                    if response.status == 200:
+                        result = await response.json()
+                        
+                        # Track usage for billing
+                        self._track_usage(
+                            service_type=ServiceType.AUDIO_STT,  # Realtime combines STT/TTS
+                            operation="create_session",
+                            metadata={
+                                "session_id": result.get("id"),
+                                "model": self.model_name,
+                                "modalities": session_config["modalities"]
+                            }
+                        )
+                        
+                        return result
+                    else:
+                        error_text = await response.text()
+                        raise Exception(f"Failed to create session: {response.status} - {error_text}")
+                        
+        except Exception as e:
+            logger.error(f"Error creating realtime session: {e}")
+            raise
+    
+    async def connect_websocket(self, session_id: str) -> aiohttp.ClientWebSocketResponse:
+        """Connect to the realtime WebSocket for a session"""
+        try:
+            ws_url = f"wss://api.openai.com/v1/realtime/sessions/{session_id}/ws"
+            headers = {
+                "Authorization": f"Bearer {self.api_key}",
+                "OpenAI-Beta": "realtime=v1"
+            }
+            
+            session = aiohttp.ClientSession()
+            ws = await session.ws_connect(ws_url, headers=headers)
+            
+            logger.info(f"Connected to realtime WebSocket for session {session_id}")
+            return ws
+            
+        except Exception as e:
+            logger.error(f"Error connecting to WebSocket: {e}")
+            raise
+    
+    async def send_audio_message(
+        self, 
+        ws: aiohttp.ClientWebSocketResponse, 
+        audio_data: bytes,
+        format: str = "pcm16"
+    ):
+        """Send audio data to the realtime session"""
+        try:
+            message = {
+                "type": "input_audio_buffer.append",
+                "audio": audio_data.hex() if format == "pcm16" else audio_data
+            }
+            
+            await ws.send_str(json.dumps(message))
+            
+            # Commit the audio buffer
+            commit_message = {"type": "input_audio_buffer.commit"}
+            await ws.send_str(json.dumps(commit_message))
+            
+        except Exception as e:
+            logger.error(f"Error sending audio message: {e}")
+            raise
+    
+    async def send_text_message(
+        self, 
+        ws: aiohttp.ClientWebSocketResponse, 
+        text: str
+    ):
+        """Send text message to the realtime session"""
+        try:
+            message = {
+                "type": "conversation.item.create",
+                "item": {
+                    "type": "message",
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "input_text",
+                            "text": text
+                        }
+                    ]
+                }
+            }
+            
+            await ws.send_str(json.dumps(message))
+            
+            # Trigger response
+            response_message = {"type": "response.create"}
+            await ws.send_str(json.dumps(response_message))
+            
+        except Exception as e:
+            logger.error(f"Error sending text message: {e}")
+            raise
+    
+    async def listen_for_responses(
+        self, 
+        ws: aiohttp.ClientWebSocketResponse,
+        message_handler: Optional[Callable] = None
+    ) -> AsyncGenerator[Dict[str, Any], None]:
+        """Listen for responses from the realtime session"""
+        try:
+            async for msg in ws:
+                if msg.type == aiohttp.WSMsgType.TEXT:
+                    try:
+                        data = json.loads(msg.data)
+                        
+                        # Handle different message types
+                        if data.get("type") == "response.audio.delta":
+                            # Audio response chunk
+                            yield {
+                                "type": "audio",
+                                "data": data.get("delta", ""),
+                                "format": "pcm16"
+                            }
+                        elif data.get("type") == "response.text.delta":
+                            # Text response chunk
+                            yield {
+                                "type": "text",
+                                "data": data.get("delta", "")
+                            }
+                        elif data.get("type") == "response.done":
+                            # Response completed
+                            usage = data.get("response", {}).get("usage", {})
+                            
+                            # Track usage for billing
+                            self._track_usage(
+                                service_type=ServiceType.AUDIO_STT,
+                                operation="realtime_response",
+                                input_tokens=usage.get("input_tokens", 0),
+                                output_tokens=usage.get("output_tokens", 0),
+                                metadata={
+                                    "response_id": data.get("response", {}).get("id"),
+                                    "model": self.model_name
+                                }
+                            )
+                            
+                            yield {
+                                "type": "done",
+                                "usage": usage
+                            }
+                        
+                        # Call custom message handler if provided
+                        if message_handler:
+                            await message_handler(data)
+                            
+                    except json.JSONDecodeError as e:
+                        logger.error(f"Error parsing WebSocket message: {e}")
+                        continue
+                        
+                elif msg.type == aiohttp.WSMsgType.ERROR:
+                    logger.error(f"WebSocket error: {ws.exception()}")
+                    break
+                    
+        except Exception as e:
+            logger.error(f"Error listening for responses: {e}")
+            raise
+    
+    async def simple_audio_chat(
+        self, 
+        audio_data: bytes, 
+        instructions: str = "You are a helpful assistant. Respond in audio.",
+        voice: str = "alloy"
+    ) -> Dict[str, Any]:
+        """Simple audio chat - send audio, get audio response"""
+        try:
+            # Create session
+            session = await self.create_session(
+                instructions=instructions,
+                modalities=["audio"],
+                voice=voice
+            )
+            session_id = session["id"]
+            
+            # Connect to WebSocket
+            ws = await self.connect_websocket(session_id)
+            
+            try:
+                # Send audio
+                await self.send_audio_message(ws, audio_data)
+                
+                # Collect response
+                audio_chunks = []
+                usage_info = {}
+                
+                async for response in self.listen_for_responses(ws):
+                    if response["type"] == "audio":
+                        audio_chunks.append(response["data"])
+                    elif response["type"] == "done":
+                        usage_info = response["usage"]
+                        break
+                
+                # Combine audio chunks
+                full_audio = "".join(audio_chunks)
+                
+                return {
+                    "audio_response": full_audio,
+                    "session_id": session_id,
+                    "usage": usage_info
+                }
+                
+            finally:
+                await ws.close()
+                
+        except Exception as e:
+            logger.error(f"Error in simple audio chat: {e}")
+            raise
+    
+    async def simple_text_chat(
+        self, 
+        text: str, 
+        instructions: str = "You are a helpful assistant.",
+        voice: str = "alloy"
+    ) -> Dict[str, Any]:
+        """Simple text chat - send text, get audio response"""
+        try:
+            # Create session
+            session = await self.create_session(
+                instructions=instructions,
+                modalities=["text", "audio"],
+                voice=voice
+            )
+            session_id = session["id"]
+            
+            # Connect to WebSocket
+            ws = await self.connect_websocket(session_id)
+            
+            try:
+                # Send text
+                await self.send_text_message(ws, text)
+                
+                # Collect response
+                text_response = ""
+                audio_chunks = []
+                usage_info = {}
+                
+                async for response in self.listen_for_responses(ws):
+                    if response["type"] == "text":
+                        text_response += response["data"]
+                    elif response["type"] == "audio":
+                        audio_chunks.append(response["data"])
+                    elif response["type"] == "done":
+                        usage_info = response["usage"]
+                        break
+                
+                # Combine audio chunks
+                full_audio = "".join(audio_chunks)
+                
+                return {
+                    "text_response": text_response,
+                    "audio_response": full_audio,
+                    "session_id": session_id,
+                    "usage": usage_info
+                }
+                
+            finally:
+                await ws.close()
+                
+        except Exception as e:
+            logger.error(f"Error in simple text chat: {e}")
+            raise
+    
+    def get_supported_voices(self) -> List[str]:
+        """Get list of supported voice options"""
+        return ["alloy", "echo", "fable", "onyx", "nova", "shimmer"]
+    
+    def get_supported_formats(self) -> List[str]:
+        """Get list of supported audio formats"""
+        return ["pcm16", "g711_ulaw", "g711_alaw"]
+    
+    async def close(self):
+        """Cleanup resources"""
+        # No persistent connections to close for REST API
+        pass