import litellm
import os
from typing import List, Optional, Dict, Any

class LLMProvider:
    """Wrapper around LLM clients utilizing litellm for extensive provider coverage."""
    
    def __init__(self, model_name: str = "gpt-4o-mini", **model_kwargs):
        """
        Task: Initializes the LLMProvider class with a specified model and optional keyword arguments.
        
        Input Parameters:
            model_name (str): The name of the LLM model to use (default is "gpt-4o-mini").
            **model_kwargs: Additional keyword arguments to pass to the underlying litellm generator.
            
        Output Parameters / Returns: None
        
        Raises: None
        """
        self.model_name = model_name
        self.model_kwargs = model_kwargs
        
    async def agenerate(
        self, 
        messages: List[Dict[str, Any]], 
        tools: Optional[List[Any]] = None, 
        response_model: Optional[Any] = None,
        **kwargs
    ):
        """
        Task: Asynchronously calls the appropriate LLM via litellm given conversation history and available tools.
        
        Input Parameters:
            messages (List[Dict[str, Any]]): The chronological list of message dictionaries (role, content) for the model context.
            tools (Optional[List[Any]]): A list of Tool objects specifying the schemas available for the model to use.
            response_model (Optional[Any]): A Pydantic BaseClass to enforce structured JSON output from the model.
            **kwargs: Additional runtime keyword arguments passed to litellm.acompletion.
            
        Output Parameters / Returns:
            Any: The raw response object from litellm (usually a ModelResponse containing choices and message details).
            
        Raises:
            Exception: If litellm throws any errors during connection or inference.
        """
        params = {
            "model": self.model_name,
            "messages": messages,
            **self.model_kwargs,
            **kwargs
        }
        
        if response_model:
            # litellm natively supports passing pydantic classes to response_format
            params["response_format"] = response_model
        
        if tools:
            # Inject available tools schema
            params["tools"] = [t.schema for t in tools]
            
        response = await litellm.acompletion(**params)
        return response
