code-puppy 0.0.142__py3-none-any.whl → 0.0.144__py3-none-any.whl

code_puppy/model_factory.py

@@ -8,6 +8,7 @@ import httpx
 from anthropic import AsyncAnthropic
 from openai import AsyncAzureOpenAI  # For Azure OpenAI client
 from pydantic_ai.models.anthropic import AnthropicModel
+from pydantic_ai.models.fallback import infer_model
 from pydantic_ai.models.gemini import GeminiModel
 from pydantic_ai.models.openai import OpenAIChatModel
 from pydantic_ai.providers.anthropic import AnthropicProvider
@@ -18,6 +19,7 @@ from pydantic_ai.providers.cerebras import CerebrasProvider
 from . import callbacks
 from .config import EXTRA_MODELS_FILE
 from .http_utils import create_async_client
+from .round_robin_model import RoundRobinModel
 
 # Environment variables used in this module:
 # - GEMINI_API_KEY: API key for Google's Gemini models. Required when using Gemini models.
@@ -246,5 +248,22 @@ class ModelFactory:
             model = OpenAIChatModel(model_name=model_config["name"], provider=provider)
             setattr(model, "provider", provider)
             return model
+        
+        elif model_type == "round_robin":
+            # Get the list of model names to use in the round-robin
+            model_names = model_config.get("models")
+            if not model_names or not isinstance(model_names, list):
+                raise ValueError(f"Round-robin model '{model_name}' requires a 'models' list in its configuration.")
+            
+            # Resolve each model name to an actual model instance
+            models = []
+            for name in model_names:
+                # Recursively get each model using the factory
+                model = ModelFactory.get_model(name, config)
+                models.append(model)
+            
+            # Create and return the round-robin model
+            return RoundRobinModel(*models)
+        
         else:
             raise ValueError(f"Unsupported model type: {model_type}")

code_puppy/round_robin_model.py

@@ -0,0 +1,115 @@
+from contextlib import asynccontextmanager, suppress
+from dataclasses import dataclass, field
+from typing import Any, Callable, AsyncIterator, List
+
+from pydantic_ai.models import Model, ModelMessage, ModelSettings, ModelRequestParameters, ModelResponse, StreamedResponse
+from pydantic_ai.models.fallback import KnownModelName, infer_model, merge_model_settings
+from pydantic_ai.result import RunContext
+
+try:
+    from opentelemetry.context import get_current_span
+except ImportError:
+    # If opentelemetry is not installed, provide a dummy implementation
+    def get_current_span():
+        class DummySpan:
+            def is_recording(self):
+                return False
+            def set_attributes(self, attributes):
+                pass
+        return DummySpan()
+
+@dataclass(init=False)
+class RoundRobinModel(Model):
+    """A model that cycles through multiple models in a round-robin fashion.
+    
+    This model distributes requests across multiple candidate models to help
+    overcome rate limits or distribute load.
+    """
+    
+    models: List[Model]
+    _current_index: int = field(default=0, repr=False)
+    _model_name: str = field(repr=False)
+
+    def __init__(
+        self,
+        *models: Model | KnownModelName | str,
+    ):
+        """Initialize a round-robin model instance.
+        
+        Args:
+            models: The names or instances of models to cycle through.
+        """
+        super().__init__()
+        if not models:
+            raise ValueError("At least one model must be provided")
+        self.models = [infer_model(m) for m in models]
+        self._current_index = 0
+
+    @property
+    def model_name(self) -> str:
+        """The model name showing this is a round-robin model with its candidates."""
+        return f'round_robin:{",".join(model.model_name for model in self.models)}'
+
+    @property
+    def system(self) -> str:
+        """System prompt from the current model."""
+        return self.models[self._current_index].system
+
+    @property
+    def base_url(self) -> str | None:
+        """Base URL from the current model."""
+        return self.models[self._current_index].base_url
+
+    def _get_next_model(self) -> Model:
+        """Get the next model in the round-robin sequence and update the index."""
+        model = self.models[self._current_index]
+        self._current_index = (self._current_index + 1) % len(self.models)
+        return model
+
+    async def request(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> ModelResponse:
+        """Make a request using the next model in the round-robin sequence."""
+        current_model = self._get_next_model()
+        merged_settings = merge_model_settings(current_model.settings, model_settings)
+        customized_model_request_parameters = current_model.customize_request_parameters(model_request_parameters)
+        
+        try:
+            response = await current_model.request(messages, merged_settings, customized_model_request_parameters)
+            self._set_span_attributes(current_model)
+            return response
+        except Exception as exc:
+            # Unlike FallbackModel, we don't try other models here
+            # The round-robin strategy is about distribution, not failover
+            raise exc
+
+    @asynccontextmanager
+    async def request_stream(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+        run_context: RunContext[Any] | None = None,
+    ) -> AsyncIterator[StreamedResponse]:
+        """Make a streaming request using the next model in the round-robin sequence."""
+        current_model = self._get_next_model()
+        merged_settings = merge_model_settings(current_model.settings, model_settings)
+        customized_model_request_parameters = current_model.customize_request_parameters(model_request_parameters)
+        
+        async with current_model.request_stream(
+            messages, merged_settings, customized_model_request_parameters, run_context
+        ) as response:
+            self._set_span_attributes(current_model)
+            yield response
+
+    def _set_span_attributes(self, model: Model):
+        """Set span attributes for observability."""
+        with suppress(Exception):
+            span = get_current_span()
+            if span.is_recording():
+                attributes = getattr(span, 'attributes', {})
+                if attributes.get('gen_ai.request.model') == self.model_name:
+                    span.set_attributes(model.model_attributes(model))

{code_puppy-0.0.142.dist-info → code_puppy-0.0.144.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: code-puppy
-Version: 0.0.142
+Version: 0.0.144
 Summary: Code generation agent
 Project-URL: repository, https://github.com/mpfaffenberger/code_puppy
 Project-URL: HomePage, https://github.com/mpfaffenberger/code_puppy
@@ -189,6 +189,68 @@ If you need to run more exotic setups or connect to remote MCPs, just update you
 
 ---
 
+## Round Robin Model Distribution
+
+Code Puppy supports **Round Robin model distribution** to help you overcome rate limits and distribute load across multiple AI models. This feature automatically cycles through configured models with each request, maximizing your API usage while staying within rate limits.
+
+### Configuration
+Add a round-robin model configuration to your `extra_models.json` file:
+
+```bash
+export CEREBRAS_API_KEY1=csk-...
+export CEREBRAS_API_KEY2=csk-...
+export CEREBRAS_API_KEY3=csk-...
+
+```
+
+```json
+{
+  "qwen1": {
+    "type": "cerebras",
+    "name": "qwen-3-coder-480b",
+    "custom_endpoint": {
+      "url": "https://api.cerebras.ai/v1",
+      "api_key": "$CEREBRAS_API_KEY1"
+    },
+    "context_length": 131072
+  },
+  "qwen2": {
+    "type": "cerebras",
+    "name": "qwen-3-coder-480b",
+    "custom_endpoint": {
+      "url": "https://api.cerebras.ai/v1",
+      "api_key": "$CEREBRAS_API_KEY2"
+    },
+    "context_length": 131072
+  },
+  "qwen3": {
+    "type": "cerebras",
+    "name": "qwen-3-coder-480b",
+    "custom_endpoint": {
+      "url": "https://api.cerebras.ai/v1",
+      "api_key": "$CEREBRAS_API_KEY3"
+    },
+    "context_length": 131072
+  },
+  "cerebras_round_robin": {
+    "type": "round_robin",
+    "models": ["qwen1", "qwen2", "qwen3"]
+  }
+}
+```
+
+Then just use /model and tab to select your round-robin model!
+
+### Benefits
+- **Rate Limit Protection**: Automatically distribute requests across multiple models
+- **Load Balancing**: Share workload between different model providers
+- **Fallback Resilience**: Continue working even if one model has temporary issues
+- **Cost Optimization**: Use different models for different types of tasks
+
+**NOTE:** Unlike fallback models, round-robin models distribute load but don't automatically retry with another model on failure. If a request fails, it will raise the exception directly.
+
+---
+
 ## Create your own Agent!!!
 
 Code Puppy features a flexible agent system that allows you to work with specialized AI assistants tailored for different coding tasks. The system supports both built-in Python agents and custom JSON agents that you can create yourself.

{code_puppy-0.0.142.dist-info → code_puppy-0.0.144.dist-info}/RECORD

@@ -6,9 +6,10 @@ code_puppy/config.py,sha256=9yWKHKjLJ2Ddl4frrBI9VRIwPvoWpIx1fAd1YpAvOSQ,15330
 code_puppy/http_utils.py,sha256=BAvt4hed7fVMXglA7eS9gOb08h2YTuOyai6VmQq09fg,3432
 code_puppy/main.py,sha256=Vv5HSJnkgZhCvvOoXrJ2zqM5P-i47-RcYAU00Z1Pfx0,21733
 code_puppy/message_history_processor.py,sha256=O2rKp7W6YeIg93W8b0XySTUEQgIZm0f_06--_kzHugM,16145
-code_puppy/model_factory.py,sha256=NoG9wDTosaaDrFIGtq3oq8gDe0J_7N6CUKuesXz87qM,10878
+code_puppy/model_factory.py,sha256=kTVaHNm6S1cLw6vHE6kH0WS6JZLRoZ8qFGKCp_fdDM4,11756
 code_puppy/models.json,sha256=dAfpMMI2EEeOMv0ynHSmMuJAYDLcZrs5gCLX3voC4-A,3252
 code_puppy/reopenable_async_client.py,sha256=4UJRaMp5np8cbef9F0zKQ7TPKOfyf5U-Kv-0zYUWDho,8274
+code_puppy/round_robin_model.py,sha256=DmbO1_SIWevdhb9nN1eNVh0dNIF-XzLYX-9gra5xVsY,4670
 code_puppy/state_management.py,sha256=o4mNBCPblRyVrNBH-992-1YqffgH6AKHU7iZRqgP1LI,5925
 code_puppy/status_display.py,sha256=F6eEAkGePDp4StM2BWj-uLLQTDGtJrf0IufzCeP1rRg,8336
 code_puppy/summarization_agent.py,sha256=-e6yUGZ22ahSaF0y7QhgVcQBfx5ktNUkPxBIWQfPaA4,3275
@@ -125,9 +126,9 @@ code_puppy/tui/tests/test_sidebar_history_navigation.py,sha256=JGiyua8A2B8dLfwiE
 code_puppy/tui/tests/test_status_bar.py,sha256=nYT_FZGdmqnnbn6o0ZuOkLtNUtJzLSmtX8P72liQ5Vo,1797
 code_puppy/tui/tests/test_timestamped_history.py,sha256=nVXt9hExZZ_8MFP-AZj4L4bB_1Eo_mc-ZhVICzTuw3I,1799
 code_puppy/tui/tests/test_tools.py,sha256=kgzzAkK4r0DPzQwHHD4cePpVNgrHor6cFr05Pg6DBWg,2687
-code_puppy-0.0.142.data/data/code_puppy/models.json,sha256=dAfpMMI2EEeOMv0ynHSmMuJAYDLcZrs5gCLX3voC4-A,3252
-code_puppy-0.0.142.dist-info/METADATA,sha256=NbPyMrnJMdLPcVGsRVp6EF9rvRyWDRi7G7kDfRQrmJI,19873
-code_puppy-0.0.142.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-code_puppy-0.0.142.dist-info/entry_points.txt,sha256=d8YkBvIUxF-dHNJAj-x4fPEqizbY5d_TwvYpc01U5kw,58
-code_puppy-0.0.142.dist-info/licenses/LICENSE,sha256=31u8x0SPgdOq3izJX41kgFazWsM43zPEF9eskzqbJMY,1075
-code_puppy-0.0.142.dist-info/RECORD,,
+code_puppy-0.0.144.data/data/code_puppy/models.json,sha256=dAfpMMI2EEeOMv0ynHSmMuJAYDLcZrs5gCLX3voC4-A,3252
+code_puppy-0.0.144.dist-info/METADATA,sha256=86kIwQ2Vf9hFT7PL6NBbHaMZGBzJ6L-CtXVn3IXULk0,21743
+code_puppy-0.0.144.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+code_puppy-0.0.144.dist-info/entry_points.txt,sha256=d8YkBvIUxF-dHNJAj-x4fPEqizbY5d_TwvYpc01U5kw,58
+code_puppy-0.0.144.dist-info/licenses/LICENSE,sha256=31u8x0SPgdOq3izJX41kgFazWsM43zPEF9eskzqbJMY,1075
+code_puppy-0.0.144.dist-info/RECORD,,