diff --git a/src/celeste/client.py b/src/celeste/client.py index 9596fce2..bee3d9ec 100644 --- a/src/celeste/client.py +++ b/src/celeste/client.py @@ -1,8 +1,10 @@ """Base client and client registry for AI capabilities.""" from abc import ABC, abstractmethod +from json import JSONDecodeError from typing import Any, Unpack +import httpx from pydantic import BaseModel, ConfigDict, Field, SecretStr from celeste.core import Capability, Provider @@ -58,6 +60,23 @@ def _parse_content( """Parse content from provider response.""" pass + def _handle_error_response(self, response: httpx.Response) -> None: + """Handle error responses from provider APIs""" + if not response.is_success: + # Try to extract error message from JSON response + try: + error_data = response.json() + error_msg = error_data.get("error", {}).get("message", response.text) + except JSONDecodeError: + error_msg = response.text or f"HTTP {response.status_code}" + + # Raise HTTPStatusError with provider context + raise httpx.HTTPStatusError( + f"{self.provider.value} API error: {error_msg}", + request=response.request, + response=response, + ) + def _transform_output( self, content: object, **parameters: Unpack[Parameters] ) -> object: diff --git a/src/celeste/core.py b/src/celeste/core.py index 91df5caf..ccbd4a24 100644 --- a/src/celeste/core.py +++ b/src/celeste/core.py @@ -18,6 +18,7 @@ class Provider(StrEnum): LUMA = "luma" TOPAZLABS = "topazlabs" PERPLEXITY = "perplexity" + BYTEDANCE = "bytedance" class Capability(StrEnum): @@ -29,6 +30,7 @@ class Capability(StrEnum): # Image IMAGE_GENERATION = "image_generation" + IMAGE_INTELLIGENCE = "image_intelligence" # Video VIDEO_INTELLIGENCE = "video_intelligence" diff --git a/src/celeste/credentials.py b/src/celeste/credentials.py index 664d240d..a1735988 100644 --- a/src/celeste/credentials.py +++ b/src/celeste/credentials.py @@ -20,6 +20,7 @@ Provider.LUMA: "luma_api_key", Provider.TOPAZLABS: "topazlabs_api_key", Provider.PERPLEXITY: "perplexity_api_key", + Provider.BYTEDANCE: "bytedance_api_key", } @@ -38,6 +39,7 @@ class Credentials(BaseSettings): luma_api_key: SecretStr | None = Field(None, alias="LUMA_API_KEY") topazlabs_api_key: SecretStr | None = Field(None, alias="TOPAZLABS_API_KEY") perplexity_api_key: SecretStr | None = Field(None, alias="PERPLEXITY_API_KEY") + bytedance_api_key: SecretStr | None = Field(None, alias="BYTEDANCE_API_KEY") model_config = { "env_file": find_dotenv(), diff --git a/src/celeste/models.py b/src/celeste/models.py index 49f48f03..8ed08794 100644 --- a/src/celeste/models.py +++ b/src/celeste/models.py @@ -11,8 +11,8 @@ class Model(BaseModel): id: str provider: Provider - capabilities: set[Capability] = Field(default_factory=set) display_name: str + capabilities: set[Capability] = Field(default_factory=set) parameter_constraints: dict[str, Constraint] = Field(default_factory=dict) @property @@ -25,25 +25,45 @@ def supported_parameters(self) -> set[str]: _models: dict[tuple[str, Provider], Model] = {} -def register_models(models: Model | list[Model]) -> None: +def register_models(models: Model | list[Model], capability: Capability) -> None: """Register one or more models in the global registry. Args: models: Single Model instance or list of Models to register. Each model is indexed by (model_id, provider) tuple. + capability: The capability these models are being registered for. Raises: - ValueError: If a model with the same (id, provider) is already registered. + ValueError: If display_name differs for duplicate (id, provider) pairs. """ if isinstance(models, Model): models = [models] for model in models: key = (model.id, model.provider) - if key in _models: - msg = f"Model '{model.id}' for provider {model.provider.value} is already registered" - raise ValueError(msg) - _models[key] = model + + # Get existing or create new model with empty capabilities/constraints + registered = _models.setdefault( + key, + Model( + id=model.id, + provider=model.provider, + display_name=model.display_name, + capabilities=set(), + parameter_constraints={}, + ), + ) + + # Validate display name consistency + if registered.display_name != model.display_name: + raise ValueError( + f"Inconsistent display_name for {model.id}: " + f"'{registered.display_name}' vs '{model.display_name}'" + ) + + # Update capabilities and constraints (single code path) + registered.capabilities.add(capability) + registered.parameter_constraints.update(model.parameter_constraints) def get_model(model_id: str, provider: Provider) -> Model | None: @@ -72,15 +92,15 @@ def list_models( Returns: List of Model instances matching the filters. """ - filtered = list(_models.values()) + models = list(_models.values()) if provider is not None: - filtered = [m for m in filtered if m.provider == provider] + models = [m for m in models if m.provider == provider] if capability is not None: - filtered = [m for m in filtered if capability in m.capabilities] + models = [m for m in models if capability in m.capabilities] - return filtered + return models def clear() -> None: diff --git a/tests/unit_tests/test_models.py b/tests/unit_tests/test_models.py index a6ede149..34d96e47 100644 --- a/tests/unit_tests/test_models.py +++ b/tests/unit_tests/test_models.py @@ -14,19 +14,16 @@ Model( id="gpt-4", provider=Provider.OPENAI, - capabilities={Capability.TEXT_GENERATION}, display_name="GPT-4", ), Model( id="dall-e-3", provider=Provider.OPENAI, - capabilities={Capability.IMAGE_GENERATION}, display_name="DALL-E 3", ), Model( id="claude-3", provider=Provider.ANTHROPIC, - capabilities={Capability.TEXT_GENERATION}, display_name="Claude 3", ), ] @@ -47,33 +44,70 @@ class TestRegisterModels: def test_register_models_accepts_single_or_list(self) -> None: """Registering models works with both single model and list.""" single_model = SAMPLE_MODELS[0] - register_models(single_model) - assert get_model(single_model.id, single_model.provider) == single_model + register_models(single_model, Capability.TEXT_GENERATION) + retrieved = get_model(single_model.id, single_model.provider) + assert retrieved is not None + assert retrieved.id == single_model.id + assert retrieved.provider == single_model.provider + assert retrieved.display_name == single_model.display_name + assert Capability.TEXT_GENERATION in retrieved.capabilities clear() - register_models(SAMPLE_MODELS) + register_models(SAMPLE_MODELS, Capability.TEXT_GENERATION) assert len(list_models()) == 3 for model in SAMPLE_MODELS: - assert get_model(model.id, model.provider) == model + retrieved = get_model(model.id, model.provider) + assert retrieved is not None + assert model.id == retrieved.id + assert model.provider == retrieved.provider + assert Capability.TEXT_GENERATION in retrieved.capabilities def test_reregistering_same_key_raises_error(self) -> None: - """Re-registering with same (id, provider) raises ValueError.""" + """Re-registering with same (id, provider) but different display_name raises ValueError.""" original = SAMPLE_MODELS[0] - register_models(original) + register_models(original, Capability.TEXT_GENERATION) duplicate = Model( id=original.id, provider=original.provider, - capabilities={Capability.IMAGE_GENERATION}, display_name="Duplicate GPT-4", ) - with pytest.raises(ValueError, match="already registered"): - register_models(duplicate) + with pytest.raises(ValueError, match="Inconsistent display_name"): + register_models(duplicate, Capability.IMAGE_GENERATION) result = get_model(original.id, original.provider) - assert result == original + assert result is not None + assert result.display_name == original.display_name + assert len(list_models()) == 1 + + def test_registering_same_model_for_multiple_capabilities_merges(self) -> None: + """Registering the same model for multiple capabilities merges capabilities.""" + model = Model( + id="multi-cap-model", + provider=Provider.OPENAI, + display_name="Multi-Cap Model", + ) + + register_models(model, Capability.TEXT_GENERATION) + retrieved = get_model("multi-cap-model", Provider.OPENAI) + assert retrieved is not None + assert Capability.TEXT_GENERATION in retrieved.capabilities + assert Capability.EMBEDDINGS not in retrieved.capabilities + + # Register same model for different capability + embeddings_model = Model( + id="multi-cap-model", + provider=Provider.OPENAI, + display_name="Multi-Cap Model", + ) + register_models(embeddings_model, Capability.EMBEDDINGS) + + retrieved = get_model("multi-cap-model", Provider.OPENAI) + assert retrieved is not None + assert Capability.TEXT_GENERATION in retrieved.capabilities + assert Capability.EMBEDDINGS in retrieved.capabilities assert len(list_models()) == 1 @@ -83,7 +117,9 @@ class TestListModels: @pytest.fixture(autouse=True) def setup_models(self) -> None: """Set up test models for filtering tests.""" - register_models(SAMPLE_MODELS) + register_models(SAMPLE_MODELS[0], Capability.TEXT_GENERATION) + register_models(SAMPLE_MODELS[1], Capability.IMAGE_GENERATION) + register_models(SAMPLE_MODELS[2], Capability.TEXT_GENERATION) def test_list_all_models(self) -> None: """Listing all models without filters.""" @@ -141,10 +177,12 @@ class TestGetModel: def test_get_existing_model(self) -> None: """Retrieving an existing model by id and provider.""" model = SAMPLE_MODELS[0] - register_models(model) + register_models(model, Capability.TEXT_GENERATION) result = get_model(model.id, model.provider) - assert result == model + assert result is not None + assert result.id == model.id + assert result.provider == model.provider def test_get_nonexistent_model_from_empty_registry_returns_none(self) -> None: """Getting a model from empty registry returns None.""" @@ -159,7 +197,9 @@ def test_get_model_from_populated_registry_with_wrong_key( self, model_id: str, provider: Provider ) -> None: """get_model returns None for non-existent model in populated registry.""" - register_models(SAMPLE_MODELS) + register_models(SAMPLE_MODELS[0], Capability.TEXT_GENERATION) + register_models(SAMPLE_MODELS[1], Capability.IMAGE_GENERATION) + register_models(SAMPLE_MODELS[2], Capability.TEXT_GENERATION) assert get_model(model_id, provider) is None def test_same_id_different_providers_are_distinct(self) -> None: @@ -167,20 +207,26 @@ def test_same_id_different_providers_are_distinct(self) -> None: model1 = Model( id="shared-id", provider=Provider.OPENAI, - capabilities={Capability.TEXT_GENERATION}, display_name="OpenAI Model", ) model2 = Model( id="shared-id", provider=Provider.ANTHROPIC, - capabilities={Capability.TEXT_GENERATION}, display_name="Anthropic Model", ) - register_models([model1, model2]) + register_models([model1, model2], Capability.TEXT_GENERATION) - assert get_model("shared-id", Provider.OPENAI) == model1 - assert get_model("shared-id", Provider.ANTHROPIC) == model2 + retrieved1 = get_model("shared-id", Provider.OPENAI) + retrieved2 = get_model("shared-id", Provider.ANTHROPIC) + assert retrieved1 is not None + assert retrieved1.id == model1.id + assert retrieved1.provider == model1.provider + assert retrieved1.display_name == model1.display_name + assert retrieved2 is not None + assert retrieved2.id == model2.id + assert retrieved2.provider == model2.provider + assert retrieved2.display_name == model2.display_name class TestEntryPoints: @@ -196,10 +242,11 @@ def test_entry_point_loading_success( test_model = Model( id="ep-test-model", provider=Provider.OPENAI, - capabilities={Capability.TEXT_GENERATION}, display_name="Entry Point Test Model", ) - mock_ep.load.return_value = lambda: register_models(test_model) + mock_ep.load.return_value = lambda: register_models( + test_model, Capability.TEXT_GENERATION + ) mock_entry_points.return_value = [mock_ep] @@ -242,12 +289,11 @@ def test_register_model_with_parameters(self) -> None: model = Model( id="param-model", provider=Provider.OPENAI, - capabilities={Capability.TEXT_GENERATION}, display_name="Model with Params", parameter_constraints={"temperature": Str(), "max_tokens": Str()}, ) - register_models(model) + register_models(model, Capability.TEXT_GENERATION) retrieved = get_model("param-model", Provider.OPENAI) assert retrieved is not None @@ -261,19 +307,18 @@ def test_list_models_includes_parameters(self) -> None: Model( id="with-params", provider=Provider.ANTHROPIC, - capabilities={Capability.TEXT_GENERATION}, display_name="With Params", parameter_constraints={"feature_a": Str(), "feature_b": Str()}, ), Model( id="without-params", provider=Provider.GOOGLE, - capabilities={Capability.IMAGE_GENERATION}, display_name="Without Params", ), ] - register_models(models) + register_models(models[0], Capability.TEXT_GENERATION) + register_models(models[1], Capability.IMAGE_GENERATION) all_models = list_models() with_params = next((m for m in all_models if m.id == "with-params"), None) @@ -292,7 +337,9 @@ class TestClear: def test_clear_removes_all_models(self) -> None: """clear removes all registered models.""" - register_models(SAMPLE_MODELS) + register_models(SAMPLE_MODELS[0], Capability.TEXT_GENERATION) + register_models(SAMPLE_MODELS[1], Capability.IMAGE_GENERATION) + register_models(SAMPLE_MODELS[2], Capability.TEXT_GENERATION) assert len(list_models()) == 3 clear() @@ -308,12 +355,11 @@ def test_supported_parameters_computed_from_constraints(self) -> None: model = Model( id="test-model", provider=Provider.OPENAI, - capabilities={Capability.TEXT_GENERATION}, display_name="Test Model", parameter_constraints={"param_a": Str(), "param_b": Str()}, ) - register_models(model) + register_models(model, Capability.TEXT_GENERATION) retrieved = get_model("test-model", Provider.OPENAI) assert retrieved is not None @@ -327,11 +373,10 @@ def test_supported_parameters_empty_by_default(self) -> None: model = Model( id="basic-model", provider=Provider.OPENAI, - capabilities={Capability.TEXT_GENERATION}, display_name="Basic Model", ) - register_models(model) + register_models(model, Capability.TEXT_GENERATION) retrieved = get_model("basic-model", Provider.OPENAI) assert retrieved is not None