feat: create dynamic model registration for Anthropic remote inferenc… #2879
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -4,11 +4,17 @@ | |
| # This source code is licensed under the terms described in the LICENSE file in | ||
| # the root directory of this source tree. | ||
|
|
||
| import logging | ||
|
|
||
| from anthropic import AsyncAnthropic, NotFoundError | ||
|
|
||
| from llama_stack.providers.utils.inference.litellm_openai_mixin import LiteLLMOpenAIMixin | ||
|
|
||
| from .config import AnthropicConfig | ||
| from .models import MODEL_ENTRIES | ||
|
|
||
| logger = logging.getLogger(__name__) | ||
|
|
||
|
|
||
| class AnthropicInferenceAdapter(LiteLLMOpenAIMixin): | ||
| def __init__(self, config: AnthropicConfig) -> None: | ||
|
|
@@ -19,9 +25,35 @@ def __init__(self, config: AnthropicConfig) -> None: | |
| provider_data_api_key_field="anthropic_api_key", | ||
| ) | ||
| self.config = config | ||
| self._client: AsyncAnthropic | None = None | ||
|
|
||
| async def initialize(self) -> None: | ||
| await super().initialize() | ||
|
|
||
| async def shutdown(self) -> None: | ||
| # Clean up the client connection pool | ||
| if self._client: | ||
| await self._client.aclose() | ||
| self._client = None | ||
| await super().shutdown() | ||
|
|
||
| @property | ||
| def client(self) -> AsyncAnthropic: | ||
| if self._client is None: | ||
| api_key = self.config.api_key if self.config.api_key else "no-key" | ||
| self._client = AsyncAnthropic(api_key=api_key) | ||
| return self._client | ||
|
Comment on lines
+41
to
+45
Collaborator
There was a problem hiding this comment. This will only work if the API key for anthropic is provided in the config. However, Llama Stack users can also provide their own API key in each request for this (any many other) providers. Our general pattern for providers that extend LiteLLM and any that support per-request credential passthrough via the I wonder more generally if the scope of this PR should be adjusted since #2835 landed? It provides a way to fetch clients and check model availability that should work for any of our LiteLLM based providers, I believe?
Contributor
Author
There was a problem hiding this comment. Thanks @bbrowning for clarification. However, Anthropic is not fully OpenAI compatible in terms of retrieve specific model or list models api endpoints. OpenAI requires Bearer token while Anthropic has slightly different structure for API Key to use during calling these endpoints. That's why I used
Collaborator
Contributor
Author
There was a problem hiding this comment. @mattf your PR is nice! Great job 👏 Should I again cherry pickup from your PR and use your infrastructure? I didn’t realise litellm is so powerful 😁 |
||
|
|
||
| async def check_model_availability(self, model: str) -> bool: | ||
| try: | ||
| retrieved_model = await self.client.models.retrieve(model) | ||
| logger.info(f"Model {retrieved_model.id} is available on Anthropic") | ||
| return True | ||
|
|
||
| except NotFoundError: | ||
| logger.info(f"Model {model} was not found on Anthropic") | ||
|
|
||
| except Exception as e: | ||
| logger.error(f"Failed to check model availability for {model} on Anthropic: {e}") | ||
|
|
||
| return False | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
wrong logger