from __future__ import annotations
import os
import httpx
from openai import AsyncOpenAI
from pydantic_ai import ModelProfile
from pydantic_ai.exceptions import UserError
from pydantic_ai.models import cached_async_http_client
from pydantic_ai.profiles.deepseek import deepseek_model_profile
from pydantic_ai.profiles.meta import meta_model_profile
from pydantic_ai.profiles.mistral import mistral_model_profile
from pydantic_ai.profiles.openai import OpenAIJsonSchemaTransformer, OpenAIModelProfile
from pydantic_ai.profiles.qwen import qwen_model_profile
from pydantic_ai.providers import Provider
try:
from openai import AsyncOpenAI
except ImportError as _import_error: # pragma: no cover
raise ImportError(
'Please install the `openai` package to use the SambaNova provider, '
'you can use the `openai` optional group — `pip install "pydantic-ai-slim[openai]"`'
) from _import_error
__all__ = ['SambaNovaProvider']
class SambaNovaProvider(Provider[AsyncOpenAI]):
"""Provider for SambaNova AI models.
SambaNova uses an OpenAI-compatible API.
"""
@property
def name(self) -> str:
"""Return the provider name."""
return 'sambanova'
@property
def base_url(self) -> str:
"""Return the base URL."""
return self._base_url
@property
def client(self) -> AsyncOpenAI:
"""Return the AsyncOpenAI client."""
return self._client
def model_profile(self, model_name: str) -> ModelProfile | None:
"""Get model profile for SambaNova models.
SambaNova serves models from multiple families including Meta Llama,
DeepSeek, Qwen, and Mistral. Model profiles are matched based on
model name prefixes.
"""
prefix_to_profile = {
'deepseek-': deepseek_model_profile,
'meta-llama-': meta_model_profile,
'llama-': meta_model_profile,
'qwen': qwen_model_profile,
'mistral': mistral_model_profile,
}
profile = None
model_name_lower = model_name.lower()
for prefix, profile_func in prefix_to_profile.items():
if model_name_lower.startswith(prefix):
profile = profile_func(model_name)
break
# Wrap into OpenAIModelProfile since SambaNova is OpenAI-compatible
return OpenAIModelProfile(json_schema_transformer=OpenAIJsonSchemaTransformer).update(profile)
def __init__(
self,
*,
api_key: str | None = None,
base_url: str | None = None,
openai_client: AsyncOpenAI | None = None,
http_client: httpx.AsyncClient | None = None,
) -> None:
"""Initialize SambaNova provider.
Args:
api_key: SambaNova API key. If not provided, reads from SAMBANOVA_API_KEY env var.
base_url: Custom API base URL. Defaults to https://api.sambanova.ai/v1
openai_client: Optional pre-configured OpenAI client
http_client: Optional custom httpx.AsyncClient for making HTTP requests
Raises:
UserError: If API key is not provided and SAMBANOVA_API_KEY env var is not set
"""
if openai_client is not None:
self._client = openai_client
self._base_url = str(openai_client.base_url)
else:
# Get API key from parameter or environment
api_key = api_key or os.getenv('SAMBANOVA_API_KEY')
if not api_key:
raise UserError(
'Set the `SAMBANOVA_API_KEY` environment variable or pass it via '
'`SambaNovaProvider(api_key=...)` to use the SambaNova provider.'
)
# Set base URL (default to SambaNova API endpoint)
self._base_url = base_url or os.getenv('SAMBANOVA_BASE_URL', 'https://api.sambanova.ai/v1')
# Create http client and AsyncOpenAI client
http_client = http_client or cached_async_http_client(provider='sambanova')
self._client = AsyncOpenAI(base_url=self._base_url, api_key=api_key, http_client=http_client)
