code-puppy 0.0.141__py3-none-any.whl → 0.0.143__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,102 @@
+from dataclasses import dataclass, field
+from typing import Any, Callable, AsyncIterator, List
+from contextlib import asynccontextmanager, suppress
+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
+
+@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.141.dist-info → code_puppy-0.0.143.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: code-puppy
-Version: 0.0.141
+Version: 0.0.143
 Summary: Code generation agent
 Project-URL: repository, https://github.com/mpfaffenberger/code_puppy
 Project-URL: HomePage, https://github.com/mpfaffenberger/code_puppy

{code_puppy-0.0.141.dist-info → code_puppy-0.0.143.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=3-WmVDbxrPATHFbmYp_sskn5669jYDtDq9P6lrWrl0o,4300
 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.141.data/data/code_puppy/models.json,sha256=dAfpMMI2EEeOMv0ynHSmMuJAYDLcZrs5gCLX3voC4-A,3252
-code_puppy-0.0.141.dist-info/METADATA,sha256=09qtNO_oD5YGszPMVCR3mZ2YsvcqrXn-PU31tLRE8sE,19873
-code_puppy-0.0.141.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-code_puppy-0.0.141.dist-info/entry_points.txt,sha256=d8YkBvIUxF-dHNJAj-x4fPEqizbY5d_TwvYpc01U5kw,58
-code_puppy-0.0.141.dist-info/licenses/LICENSE,sha256=31u8x0SPgdOq3izJX41kgFazWsM43zPEF9eskzqbJMY,1075
-code_puppy-0.0.141.dist-info/RECORD,,
+code_puppy-0.0.143.data/data/code_puppy/models.json,sha256=dAfpMMI2EEeOMv0ynHSmMuJAYDLcZrs5gCLX3voC4-A,3252
+code_puppy-0.0.143.dist-info/METADATA,sha256=XbQIbiQFjwgUAnpB0_-ObMgoro79rBMivTNi_5Zzebk,19873
+code_puppy-0.0.143.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+code_puppy-0.0.143.dist-info/entry_points.txt,sha256=d8YkBvIUxF-dHNJAj-x4fPEqizbY5d_TwvYpc01U5kw,58
+code_puppy-0.0.143.dist-info/licenses/LICENSE,sha256=31u8x0SPgdOq3izJX41kgFazWsM43zPEF9eskzqbJMY,1075
+code_puppy-0.0.143.dist-info/RECORD,,