❯ /Users/mehmetoneryalcin/dev/dspy/.venv/bin/python /Users/mehmetoneryalcin/dev/dspy/human.py
18:10:29 - LiteLLM Router:INFO: router.py:659 - Routing strategy: simple-shuffle
INFO:LiteLLM Router:Routing strategy: simple-shuffle
..
LiteLLM completion() model= gpt-4.1-nano-2; provider = openai

INFO:httpx:HTTP Request: POST https://api.openai.com/v1/chat/completions "HTTP/1.1 404 Not Found"
18:10:29 - LiteLLM Router:INFO: router.py:1064 - litellm.acompletion(model=openai/gpt-4.1-nano-2) Exception litellm.NotFoundError: OpenAIException - The model `gpt-4.1-nano-2` does not exist or you do not have access to it.
..
INFO:LiteLLM Router:Trying to fallback b/w models
18:10:29 - LiteLLM Router:INFO: fallback_event_handlers.py:128 - Falling back to model_group = claude-3-5-haiku-latest
..
INFO:httpx:HTTP Request: POST https://api.anthropic.com/v1/messages "HTTP/1.1 200 OK"
..
18:10:30 - LiteLLM Router:INFO: fallback_event_handlers.py:142 - Successful fallback b/w models.
INFO:LiteLLM Router:Successful fallback b/w models.
✅ Async router call successful
   Async response: Four.

🚀 Testing LiteLLM Router Integration with DSPy
==================================================
📋 Step 1: Creating LiteLLM Router...
✅ Router created successfully

📋 Step 2: Creating DSPy LM with router...
✅ DSPy LM with router created successfully
   Router configured: True
   Model: gpt-4.1-nano
   Provider: None

📋 Step 3: Testing basic LM forward call...
✅ Router forward call successful
   Response: Hello, how are?
   Usage: Usage(completion_tokens=5, prompt_tokens=15, total_tokens=20, completion_tokens_details=CompletionTokensDetailsWrapper(accepted_prediction_tokens=0, audio_tokens=0, reasoning_tokens=0, rejected_prediction_tokens=0, text_tokens=None), prompt_tokens_details=PromptTokensDetailsWrapper(audio_tokens=0, cached_tokens=0, text_tokens=None, image_tokens=None))

📋 Step 4: Testing with DSPy Predict module...
✅ DSPy Predict with router successful
   Input: Alice
   Output: Hello, Alice! It's great to meet you. Hope you have a wonderful day!

📋 Step 5: Testing async functionality...
✅ Async router call successful
   Async response: Four

📋 Step 6: Testing ChainOfThought with router...
✅ ChainOfThought with router successful
   Problem: If I have 5 apples and eat 2, how many are left?
   Answer: 3

📋 Step 7: Testing state management...
✅ State dump successful
   State keys: ['temperature', 'max_tokens', 'model', 'model_type', 'cache', 'cache_in_memory', 'num_retries', 'finetuning_model', 'launch_kwargs', 'train_kwargs', 'router_is_configured', 'provider_name']
   Router configured in state: True
   Provider name: None

==================================================
🎉 All tests passed! LiteLLM Router integration is working correctly.

📊 Summary:
   ✓ Router creation
   ✓ DSPy LM initialization with router
   ✓ Basic forward calls
   ✓ DSPy Predict module
   ✓ Async functionality
   ✓ ChainOfThought reasoning
   ✓ State management

   Response: Hello! How can I assist you today?

🎯 Conclusion: LiteLLM Router integration in DSPy is working correctly!

class RouterLM(BaseLM):
    """
    A DSPy-compatible LM that uses LiteLLM Router for load balancing and fallbacks.
    
    This allows you to use LiteLLM Router features like:
    - Load balancing across multiple model deployments
    - Automatic fallbacks when a model fails
    - Rate limit handling with cooldowns
    - Retries with exponential backoff
    
    Example usage:
    ```python
    # Option 1: Pass a model list directly
    router_config = [
        {
            "model_name": "gpt-4",  # model alias
            "litellm_params": {
                "model": "openai/gpt-4",
                "api_key": os.environ["OPENAI_API_KEY"],
            },
        },
        {
            "model_name": "gpt-4",  # same alias for fallback
            "litellm_params": {
                "model": "azure/gpt-4-deployment",
                "api_key": os.environ["AZURE_API_KEY"],
                "api_base": os.environ["AZURE_API_BASE"],
            },
        },
    ]
    lm = RouterLM(model="gpt-4", router_config=router_config)
    
    # Option 2: Use get_default_router_config() for full config with fallbacks
    from notebooks.src.utils.router_config import get_default_router_config
    lm = RouterLM(model="gpt-4", router_config=get_default_router_config())
    
    dspy.configure(lm=lm)
    ```
    """

    def __init__(
        self,
        model: str,
        router_config: list[dict[str, Any]] | dict[str, Any],
        model_type: Literal["chat", "text"] = "chat",
        temperature: float = 0.0,
        max_tokens: int = 1000,
        cache: bool = True,
        num_retries: int | None = None,
        routing_strategy: Literal[
            "simple-shuffle",
            "least-busy", 
            "usage-based-routing",
            "latency-based-routing",
        ] | None = None,
        fallbacks: list[dict[str, list[str]]] | None = None,
        set_verbose: bool = False,
        **kwargs,
    ):
        """
        Create a new RouterLM instance.

        Args:
            model: The model alias to use (must match a model_name in router_config).
            router_config: Either:
                - A list of model deployments (each with model_name and litellm_params)
                - A full router config dict (from get_default_router_config()) with model_list,
                  fallbacks, routing_strategy, etc.
            model_type: Either "chat" or "text" completion mode.
            temperature: Sampling temperature for generation.
            max_tokens: Maximum tokens to generate.
            cache: Whether to enable DSPy's response caching.
            num_retries: Number of retries for failed requests (default: 3).
            routing_strategy: How to route requests across deployments:
                - "simple-shuffle": Random selection (default)
                - "least-busy": Route to least busy deployment
                - "usage-based-routing": Route based on usage/rate limits
                - "latency-based-routing": Route to lowest latency deployment
            fallbacks: Optional list of fallback model mappings.
            set_verbose: Enable verbose logging from LiteLLM Router.
            **kwargs: Additional kwargs passed to router.completion().
        """
        super().__init__(
            model=model,
            model_type=model_type,
            temperature=temperature,
            max_tokens=max_tokens,
            cache=cache,
            **kwargs,
        )
        
        # Handle both full config dict and simple model list
        if isinstance(router_config, dict):
            # Full config dict from get_default_router_config()
            model_list = router_config.get("model_list", [])
            router_kwargs = {
                "model_list": model_list,
                "num_retries": num_retries or router_config.get("num_retries", 3),
                "routing_strategy": routing_strategy or router_config.get("routing_strategy", "simple-shuffle"),
                "fallbacks": fallbacks or router_config.get("fallbacks"),
                "context_window_fallbacks": router_config.get("context_window_fallbacks"),
                "content_policy_fallbacks": router_config.get("content_policy_fallbacks"),
                "default_fallbacks": router_config.get("default_fallbacks"),
                "allowed_fails": router_config.get("allowed_fails"),
                "cooldown_time": router_config.get("cooldown_time"),
                "timeout": router_config.get("timeout"),
                "retry_after": router_config.get("retry_after"),
                "enable_pre_call_checks": router_config.get("enable_pre_call_checks", False),
                "set_verbose": set_verbose,
            }
            # Remove None values
            router_kwargs = {k: v for k, v in router_kwargs.items() if v is not None}
            num_deployments = len(model_list)
        else:
            # Simple model list
            model_list = router_config
            router_kwargs = {
                "model_list": model_list,
                "routing_strategy": routing_strategy or "simple-shuffle",
                "num_retries": num_retries or 3,
                "fallbacks": fallbacks,
                "set_verbose": set_verbose,
            }
            # Remove None values
            router_kwargs = {k: v for k, v in router_kwargs.items() if v is not None}
            num_deployments = len(model_list)
        
        self.num_retries = router_kwargs.get("num_retries", 3)
        self.routing_strategy = router_kwargs.get("routing_strategy", "simple-shuffle")
        
        # Initialize the LiteLLM Router
        self.router = Router(**router_kwargs)
        
        logger.info(
            f"RouterLM initialized with model='{model}', "
            f"routing_strategy='{self.routing_strategy}', "
            f"{num_deployments} deployment(s)"
        )

    def forward(
        self,
        prompt: str | None = None,
        messages: list[dict[str, Any]] | None = None,
        **kwargs,
    ):
        """
        Forward pass using LiteLLM Router's completion method.
        
        Returns response in OpenAI format for compatibility with DSPy.
        """
        messages = messages or [{"role": "user", "content": prompt}]
        merged_kwargs = {**self.kwargs, **kwargs}
        
        # Remove DSPy-specific kwargs that shouldn't go to LiteLLM
        merged_kwargs.pop("rollout_id", None)
        
        if self.model_type == "chat":
            response = self.router.completion(
                model=self.model,
                messages=messages,
                **merged_kwargs,
            )
        else:
            # For text completion, build prompt from messages
            prompt_text = "\n\n".join([m["content"] for m in messages] + ["BEGIN RESPONSE:"])
            response = self.router.text_completion(
                model=self.model,
                prompt=prompt_text,
                **merged_kwargs,
            )
        
        return response

    async def aforward(
        self,
        prompt: str | None = None,
        messages: list[dict[str, Any]] | None = None,
        **kwargs,
    ):
        """
        Async forward pass using LiteLLM Router's acompletion method.
        
        Returns response in OpenAI format for compatibility with DSPy.
        """
        messages = messages or [{"role": "user", "content": prompt}]
        merged_kwargs = {**self.kwargs, **kwargs}
        
        # Remove DSPy-specific kwargs that shouldn't go to LiteLLM
        merged_kwargs.pop("rollout_id", None)
        
        if self.model_type == "chat":
            response = await self.router.acompletion(
                model=self.model,
                messages=messages,
                **merged_kwargs,
            )
        else:
            # For text completion, build prompt from messages
            prompt_text = "\n\n".join([m["content"] for m in messages] + ["BEGIN RESPONSE:"])
            response = await self.router.atext_completion(
                model=self.model,
                prompt=prompt_text,
                **merged_kwargs,
            )
        
        return response

    def get_router_stats(self) -> dict[str, Any]:
        """
        Get statistics about the router's deployment usage.
        
        Returns:
            Dict with deployment stats including request counts, latencies, etc.
        """
        return {
            "model_list": self.router.model_list,
            "routing_strategy": self.routing_strategy,
            "num_retries": self.num_retries,
        }

    def copy(self, **kwargs):
        """Create a copy of this RouterLM with updated parameters."""
        import copy as copy_module
        
        new_instance = copy_module.deepcopy(self)
        new_instance.history = []
        
        for key, value in kwargs.items():
            if hasattr(self, key):
                setattr(new_instance, key, value)
            if (key in self.kwargs) or (not hasattr(self, key)):
                if value is None:
                    new_instance.kwargs.pop(key, None)
                else:
                    new_instance.kwargs[key] = value
        
        return new_instance