mllm_tools/openrouter.py

import os
import re
from typing import List, Dict, Any, Optional, Union
import io
import base64
from PIL import Image
import mimetypes
from litellm import completion, completion_cost
from dotenv import load_dotenv

load_dotenv()

class OpenRouterWrapper:
    """
    OpenRouter wrapper using LiteLLM for various language models.
    Compatible with the existing wrapper interface.
    """
    
    def __init__(
        self, 
        model_name: str = "openrouter/deepseek/deepseek-chat-v3-0324:free",
        temperature: float = 0.7,
        print_cost: bool = False,
        verbose: bool = False,
        use_langfuse: bool = True,
        site_url: str = "",
        app_name: str = "Theory2Manim"
    ):
        """
        Initialize OpenRouter wrapper.
        
        Args:
            model_name: OpenRouter model name (with openrouter/ prefix)
            temperature: Temperature for completion
            print_cost: Whether to print the cost of the completion
            verbose: Whether to print verbose output
            use_langfuse: Whether to enable Langfuse logging
            site_url: Optional site URL for tracking
            app_name: Optional app name for tracking
        """
        self.model_name = model_name
        self.temperature = temperature
        self.print_cost = print_cost
        self.verbose = verbose
        self.accumulated_cost = 0
        
        # Setup OpenRouter environment variables
        api_key = os.getenv("OPENROUTER_API_KEY")
        if not api_key:
            raise ValueError("No OPENROUTER_API_KEY found. Please set the environment variable.")
        
        os.environ["OPENROUTER_API_KEY"] = api_key
        os.environ["OPENROUTER_API_BASE"] = "https://openrouter.ai/api/v1"
        
        if site_url or os.getenv("OR_SITE_URL"):
            os.environ["OR_SITE_URL"] = site_url or os.getenv("OR_SITE_URL", "")
        if app_name:
            os.environ["OR_APP_NAME"] = app_name
            
        if self.verbose:
            os.environ['LITELLM_LOG'] = 'DEBUG'
            
        # Set langfuse callback only if enabled
        if use_langfuse:
            import litellm
            litellm.success_callback = ["langfuse"]
            litellm.failure_callback = ["langfuse"]
    
    def _encode_file(self, file_path: Union[str, Image.Image]) -> str:
        """
        Encode local file or PIL Image to base64 string
        
        Args:
            file_path: Path to local file or PIL Image object
            
        Returns:
            Base64 encoded file string
        """
        if isinstance(file_path, Image.Image):
            buffered = io.BytesIO()
            file_path.save(buffered, format="PNG")
            return base64.b64encode(buffered.getvalue()).decode("utf-8")
        else:
            with open(file_path, "rb") as file:
                return base64.b64encode(file.read()).decode("utf-8")
    
    def _get_mime_type(self, file_path: str) -> str:
        """
        Get the MIME type of a file based on its extension
        
        Args:
            file_path: Path to the file
            
        Returns:
            MIME type as a string (e.g., "image/jpeg", "audio/mp3")
        """
        mime_type, _ = mimetypes.guess_type(file_path)
        if mime_type is None:
            raise ValueError(f"Unsupported file type: {file_path}")
        return mime_type
    
    def __call__(self, messages: List[Dict[str, Any]], metadata: Optional[Dict[str, Any]] = None) -> str:
        """
        Process messages and return completion
        
        Args:
            messages: List of message dictionaries with 'type' and 'content' keys
            metadata: Optional metadata to pass to completion
        
        Returns:
            Generated text response
        """
        if metadata is None:
            metadata = {}
        metadata["trace_name"] = f"openrouter-completion-{self.model_name}"
        
        # Convert messages to LiteLLM format
        formatted_messages = []
        for msg in messages:
            if msg["type"] == "text":
                formatted_messages.append({
                    "role": "user",
                    "content": [{"type": "text", "text": msg["content"]}]
                })
            elif msg["type"] in ["image", "audio", "video"]:
                # Check if content is a local file path or PIL Image
                if isinstance(msg["content"], Image.Image) or os.path.isfile(msg["content"]):
                    try:
                        if isinstance(msg["content"], Image.Image):
                            mime_type = "image/png"
                        else:
                            mime_type = self._get_mime_type(msg["content"])
                        base64_data = self._encode_file(msg["content"])
                        data_url = f"data:{mime_type};base64,{base64_data}"
                    except ValueError as e:
                        print(f"Error processing file {msg['content']}: {e}")
                        continue
                else:
                    data_url = msg["content"]
                
                # Format for vision models
                if msg["type"] == "image":
                    formatted_messages.append({
                        "role": "user",
                        "content": [
                            {
                                "type": "image_url",
                                "image_url": {
                                    "url": data_url,
                                    "detail": "high"
                                }
                            }
                        ]
                    })
                else:
                    # For audio/video, treat as text for now
                    formatted_messages.append({
                        "role": "user",
                        "content": [{"type": "text", "text": f"[{msg['type'].upper()}]: {msg['content']}"}]
                    })

        try:
            response = completion(
                model=self.model_name,
                messages=formatted_messages,
                temperature=self.temperature,
                metadata=metadata,
                max_retries=99
            )
            if self.print_cost:
                # Calculate and print cost
                cost = completion_cost(completion_response=response)
                self.accumulated_cost += cost
                print(f"Accumulated Cost: ${self.accumulated_cost:.10f}")
            
            content = response.choices[0].message.content
            if content is None:
                print(f"Got null response from model. Full response: {response}")
                return "Error: Received null response from model"
            
            # Check if the response contains error messages about unmapped models
            if "This model isn't mapped yet" in content or "model isn't mapped" in content.lower():
                error_msg = f"Error: Model {self.model_name} is not supported by LiteLLM. Please use a supported model."
                print(error_msg)
                return error_msg
            
            return content
        
        except Exception as e:
            print(f"Error in OpenRouter completion: {e}")
            return f"Error: {str(e)}"


class OpenRouterClient:
    """
    Legacy OpenRouter client for backward compatibility.
    """
    
    def __init__(self, api_key: str, site_url: str = "", app_name: str = "Theory2Manim"):
        """
        Initialize OpenRouter client.
        
        Args:
            api_key: OpenRouter API key
            site_url: Optional site URL for tracking
            app_name: Optional app name for tracking
        """
        os.environ["OPENROUTER_API_KEY"] = api_key
        os.environ["OPENROUTER_API_BASE"] = "https://openrouter.ai/api/v1"
        
        if site_url:
            os.environ["OR_SITE_URL"] = site_url
        if app_name:
            os.environ["OR_APP_NAME"] = app_name
    
    def complete(
        self, 
        messages: List[Dict[str, str]], 
        model: str = "openrouter/openai/gpt-3.5-turbo",
        transforms: Optional[List[str]] = None,
        route: Optional[str] = None,
        **kwargs
    ) -> Any:
        """
        Generate completion using OpenRouter model.
        
        Args:
            messages: List of message dictionaries with 'role' and 'content'
            model: Model name (with openrouter/ prefix)
            transforms: Optional transforms to apply
            route: Optional route specification
            **kwargs: Additional parameters for completion
            
        Returns:
            Completion response
        """
        params = {
            "model": model,
            "messages": messages,
            **kwargs
        }
        
        if transforms:
            params["transforms"] = transforms
        if route:
            params["route"] = route
            
        return completion(**params)

# Convenience functions for common models
def ds_r1(messages: List[Dict[str, str]], **kwargs) -> Any:
    """Use GPT-3.5 Turbo via OpenRouter"""
    client = OpenRouterClient(os.environ.get("OPENROUTER_API_KEY", ""))
    return client.complete(messages, "deepseek/deepseek-r1:free", **kwargs)

def ds_v3(messages: List[Dict[str, str]], **kwargs) -> Any:
    """Use GPT-4 via OpenRouter"""
    client = OpenRouterClient(os.environ.get("OPENROUTER_API_KEY", ""))
    return client.complete(messages, "deepseek/deepseek-chat-v3-0324:free", **kwargs)

def qwen3(messages: List[Dict[str, str]], **kwargs) -> Any:
    """Use Claude-2 via OpenRouter"""
    client = OpenRouterClient(os.environ.get("OPENROUTER_API_KEY", ""))
    return client.complete(messages, "qwen/qwen3-235b-a22b:free", **kwargs)