File size: 1,790 Bytes
6172a47 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 | """Base provider interface - extend this to implement your own provider."""
from abc import ABC, abstractmethod
from collections.abc import AsyncIterator
from typing import Any
from pydantic import BaseModel
class ProviderConfig(BaseModel):
"""Configuration for a provider.
Base fields apply to all providers. Provider-specific parameters
(e.g. NIM temperature, top_p) are passed by the provider constructor.
"""
api_key: str
base_url: str | None = None
rate_limit: int | None = None
rate_window: int = 60
max_concurrency: int = 5
http_read_timeout: float = 300.0
http_write_timeout: float = 10.0
http_connect_timeout: float = 2.0
enable_thinking: bool = True
proxy: str = ""
class BaseProvider(ABC):
"""Base class for all providers. Extend this to add your own."""
def __init__(self, config: ProviderConfig):
self._config = config
def _is_thinking_enabled(self, request: Any) -> bool:
"""Return whether thinking should be enabled for this request."""
thinking = getattr(request, "thinking", None)
request_enabled = (
thinking.enabled
if thinking is not None and hasattr(thinking, "enabled")
else True
)
return self._config.enable_thinking and request_enabled
@abstractmethod
async def cleanup(self) -> None:
"""Release any resources held by this provider."""
@abstractmethod
async def stream_response(
self,
request: Any,
input_tokens: int = 0,
*,
request_id: str | None = None,
) -> AsyncIterator[str]:
"""Stream response in Anthropic SSE format."""
if False:
yield "" # Required for ty/mypy to accept abstract async generator
|