orchestrata_ultimate.providers.provider_manager

Reference › orchestrata_ultimate › providers › provider_manager

provider_manager is the single entry point of the providers package (specification 011): ProviderManager presents one execute(request) surface over many LLM backends. The LLM client (client) is its main caller; the package __init__ re-exports it. It composes the registry, the router, the cost tracker and the health monitor, and owns the per-provider circuit breakers and rate limiters. The curated narrative is providers.

register_provider adds an adapter with its ProviderConfig and creates its breaker and limiter; deregister_provider removes it and closes its HTTP client. The constructor takes optional Redis connection settings, which are passed to the breakers, limiters and cost tracker so their state is shared and durable. As an IComponent it also has initialize, start, stop, get_health and get_metadata.

execute asks the router for a provider and model, using current health data; if none is eligible it raises ProviderError with the rejection reasons. For each attempt it prices the model, enforces the budget (_enforce_budget raises BudgetExceededError when the estimated cost does not fit the volk's budget), calls the adapter (_invoke_provider) and on success records cost, health, breaker and limiter state (_record_success). On an exception it records the failure (_record_failure), then takes the next provider from the router's fallback chain, excluding those already tried; when the chain is exhausted it raises ProviderError listing every attempt. Other methods: estimate_cost, get_usage, get_remaining_budget, set_budget, get_rate_limit_backoff, get_provider_health and cancel.

At a glance

Class diagram

Dependency diagram

API

Classes

ProviderError

Bases: Exception · Decorators: traces(spec='011', story='US-011-001', reqs=['REQ-011-009']) · Line 50

Raised when all providers fail to serve a request.

BudgetExceededError

Bases: Exception · Decorators: traces(spec='011', story='US-011-001', reqs=['REQ-011-046']) · Line 60

Raised when a request would exceed the budget.

ProviderManager (protocol)

Bases: ProviderManagerProtocol · Decorators: traces(spec='011', story='US-001', reqs=['REQ-011-003']) · Line 75

Unified LLM provider manager.

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_ultimate/providers/provider_manager.py

# mypy: ignore-errors
"""Provider Manager — unified LLM provider abstraction.

SPEC-011: Orchestrates provider registry, request routing, circuit breakers,
rate limiters, cost tracking, and health monitoring into a single interface.
"""

import asyncio
import logging
import time
from collections.abc import Callable, Coroutine
from datetime import datetime
from typing import Any

import redis.asyncio as redis

from orchestrata_ultimate.core.base.traces import traces
from orchestrata_ultimate.core.protocols import (
    BudgetStatus,
    ComponentMetadata,
    ComponentType,
    CostEstimate,
    CostRecord,
    HealthStatus,
    ModelPricing,
    ProviderConfig,
    ProviderManagerProtocol,
    ProviderProtocol,
    ProviderRequest,
    ProviderResponse,
    TokenUsage,
    UsageReport,
)
from orchestrata_ultimate.core.redis_factory import (
    resolve_redis_host,
    resolve_redis_port,
)
from orchestrata_ultimate.memory.token_counter import TiktokenCounter
from orchestrata_ultimate.providers.circuit_breaker import CircuitBreaker
from orchestrata_ultimate.providers.cost_tracker import CostTracker
from orchestrata_ultimate.providers.health_monitor import HealthMonitor
from orchestrata_ultimate.providers.rate_limiter import RateLimiter
from orchestrata_ultimate.providers.registry import ProviderRegistry
from orchestrata_ultimate.providers.router import RequestRouter

logger = logging.getLogger(__name__)


@traces(spec="011", story="US-011-001", reqs=["REQ-011-009"])
class ProviderError(Exception):
    """Raised when all providers fail to serve a request."""

    @traces(spec="011", story="US-011-001", reqs=["REQ-011-009"])
    def __init__(self, message: str, attempts: list[dict] | None = None):
        super().__init__(message)
        self.attempts = attempts or []


@traces(spec="011", story="US-011-001", reqs=["REQ-011-046"])
class BudgetExceededError(Exception):
    """Raised when a request would exceed the budget."""

    @traces(spec="011", story="US-011-001", reqs=["REQ-011-046"])
    def __init__(self, scope_key: str, estimated_cost: float, remaining: float):
        super().__init__(
            f"Budget exceeded for {scope_key}: "
            f"estimated ${estimated_cost:.4f}, remaining ${remaining:.4f}"
        )
        self.scope_key = scope_key
        self.estimated_cost = estimated_cost
        self.remaining = remaining


@traces(spec="011", story="US-001", reqs=["REQ-011-003"])
class ProviderManager(ProviderManagerProtocol):
    """Unified LLM provider manager.

    Orchestrates all provider subsystems: registry, routing,
    circuit breakers, rate limiting, cost tracking, and health monitoring.
    """

    def __init__(
        self,
        redis_host: str | None = None,
        redis_port: int | None = None,
        redis_db: int = 0,
        redis_password: str | None = None,
    ):
        self._redis_host = (
            redis_host if redis_host is not None else resolve_redis_host()
        )
        self._redis_port = (
            redis_port if redis_port is not None else resolve_redis_port()
        )
        self._redis_db = redis_db
        self._redis_password = redis_password
        self._redis: redis.Redis | None = None
        self._registry = ProviderRegistry()
        self._breakers: dict[str, CircuitBreaker] = {}
        self._limiters: dict[str, RateLimiter] = {}
        self._cost_tracker = CostTracker()
        self._token_counters: dict[str, TiktokenCounter] = {}
        self._health = HealthMonitor()
        self._router = RequestRouter(self._registry, self._breakers, self._limiters)
        self._running = False
        self._in_flight: dict[str, asyncio.Task] = {}
        self._metadata = ComponentMetadata(
            id="orchestrata-provider",
            name="ProviderManager",
            version="0.3.0",
            component_type=ComponentType.PROVIDER_MANAGER,
            description="Unified LLM provider abstraction",
            capabilities=[
                "execute",
                "route",
                "fallback",
                "cost_track",
                "rate_limit",
            ],
        )

    @traces(spec="011", story="US-011-001", reqs=["REQ-011-003"])
    async def initialize(self, config: dict[str, Any]) -> None:
        """Initialize the provider manager."""
        from orchestrata_ultimate.core.redis_factory import create_async_redis

        self._redis = create_async_redis(
            host=self._redis_host,
            port=self._redis_port,
            db=self._redis_db,
            password=self._redis_password,
        )
        await self._cost_tracker.set_redis(self._redis)
        # Load circuit breaker states from Redis
        for name, breaker in self._breakers.items():
            await breaker.set_redis(self._redis)
        for name, limiter in self._limiters.items():
            await limiter.set_redis(self._redis)
        self._running = True
        logger.info("ProviderManager initialized")

    @traces(spec="011", story="US-011-001", reqs=["REQ-011-003"])
    async def start(self) -> None:
        """Start the provider manager."""
        self._running = True
        logger.info("ProviderManager started")

    @traces(spec="011", story="US-011-001", reqs=["REQ-011-003"])
    @traces(spec="015", story="US-011-001", reqs=["REQ-015-057"])
    async def stop(self) -> None:
        """Stop the provider manager and release resources."""
        self._running = False
        # Cancel in-flight requests
        for request_id, task in list(self._in_flight.items()):
            if not task.done():
                task.cancel()
        self._in_flight.clear()
        # Close every registered provider's HTTP client so we don't leak
        # httpx AsyncClient connections/sockets on shutdown.
        for name in self._registry.provider_names:
            provider = self._registry.get_provider(name)
            close = getattr(provider, "close", None)
            if close is None:
                continue
            try:
                await close()
            except Exception as exc:  # pragma: no cover - defensive
                logger.warning("Error closing provider %s: %s", name, exc)
        if self._redis:
            await self._redis.aclose()
            self._redis = None
        logger.info("ProviderManager stopped")

    @traces(spec="011", story="US-011-001", reqs=["REQ-011-003"])
    async def get_health(self) -> HealthStatus:
        """Return overall component health."""
        if not self._running:
            return HealthStatus.STOPPED
        # Healthy if at least one provider is available
        for breaker in self._breakers.values():
            if breaker.is_available:
                return HealthStatus.HEALTHY
        if not self._breakers:
            return HealthStatus.HEALTHY
        return HealthStatus.DEGRADED

    def get_metadata(self) -> ComponentMetadata:
        """Return component metadata."""
        return self._metadata

    @traces(spec="011", story="US-001", reqs=["REQ-011-024", "REQ-011-025"])
    def register_provider(
        self,
        name: str,
        provider: ProviderProtocol,
        config: ProviderConfig,
    ) -> None:
        """Register a provider adapter with configuration."""
        self._registry.register(name, provider, config)
        self._breakers[name] = CircuitBreaker(
            provider_name=name,
            failure_threshold=config.max_retries + 2,
        )
        self._limiters[name] = RateLimiter(
            provider_name=name,
            rpm_limit=config.rate_limit_rpm,
            tpm_limit=config.rate_limit_tpm,
        )
        # Update router with new breakers/limiters
        self._router = RequestRouter(self._registry, self._breakers, self._limiters)

    @traces(spec="011", story="US-001", reqs=["REQ-011-024"])
    def deregister_provider(self, name: str) -> bool:
        """Remove a provider and close its HTTP client."""
        provider = self._registry.get_provider(name)
        result = self._registry.deregister(name)
        if result:
            self._breakers.pop(name, None)
            self._limiters.pop(name, None)
            self._router = RequestRouter(self._registry, self._breakers, self._limiters)
            # Close the adapter's HTTP client to avoid leaking connections.
            # deregister_provider is sync (per ProviderManagerProtocol), so dispatch
            # the async close onto the running loop if there is one, else run it.
            close = getattr(provider, "close", None)
            if close is not None:
                self._run_provider_close(name, close)
        return result

    @traces(spec="011", story="US-001", reqs=["REQ-011-024"])
    def _run_provider_close(
        self, name: str, close: Callable[[], Coroutine[Any, Any, object]]
    ) -> None:
        """Run a provider's async ``close()`` from a sync context."""
        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            loop = None
        try:
            if loop is not None:
                # Inside a running loop: schedule the close as a background task.
                loop.create_task(close())
            else:
                asyncio.run(close())
        except Exception as exc:  # pragma: no cover - defensive
            logger.warning("Error closing provider %s: %s", name, exc)

    @traces(
        spec="011",
        story="US-001",
        reqs=[
            "REQ-011-006",
            "REQ-011-007",
            "REQ-011-008",
            "REQ-011-009",
            "REQ-011-010",
        ],
    )
    async def execute(self, request: ProviderRequest) -> ProviderResponse:
        """Execute an LLM request with routing, fallback, and cost tracking."""
        health_data = self._health.get_all_health()
        route = self._router.route(request, health_data)
        if route is None:
            reasons = self._router.get_rejection_reasons(request)
            reasons_text = "; ".join(reasons)
            raise ProviderError(
                f"No eligible provider found for request: {reasons_text}",
                attempts=[{"reasons": reasons}],
            )
        attempts: list[dict] = []
        tried_providers: list[str] = []
        provider_name, model_name = route
        while True:
            tried_providers.append(provider_name)
            pricing = self._registry.get_pricing(provider_name, model_name)
            await self._enforce_budget(request, pricing)

            provider = self._registry.get_provider(provider_name)
            if provider is None:
                attempts.append(
                    {"provider": provider_name, "error": "Provider not found"}
                )
                break

            start_time = time.monotonic()
            call_params = self._build_call_params(request)
            try:
                raw = await self._invoke_provider(
                    provider, request, model_name, call_params
                )
                elapsed_ms = (time.monotonic() - start_time) * 1000.0
                return await self._record_success(
                    request, provider_name, model_name, pricing, raw, elapsed_ms
                )
            except Exception as exc:
                elapsed_ms = (time.monotonic() - start_time) * 1000.0
                await self._record_failure(
                    attempts, provider_name, model_name, exc, elapsed_ms
                )
                fallback_chain = self._router.get_fallback_chain(
                    request, exclude=tried_providers
                )
                if not fallback_chain:
                    break
                provider_name, model_name = fallback_chain[0]
                continue
        raise ProviderError(
            f"All providers failed after {len(attempts)} attempts",
            attempts=attempts,
        )

    @traces(spec="011", story="US-001", reqs=["REQ-011-006", "REQ-011-009"])
    async def _enforce_budget(
        self, request: ProviderRequest, pricing: ModelPricing | None
    ) -> None:
        """Raise BudgetExceededError if the estimated cost exceeds the budget."""
        if not (pricing and request.volk_id):
            return
        est = self._estimate_cost_for_pricing(request, pricing)
        allowed = await self._cost_tracker.check_budget(
            request.volk_id, est.estimated_cost.total_cost
        )
        if allowed:
            return
        budget = await self._cost_tracker.get_remaining_budget(request.volk_id)
        raise BudgetExceededError(
            request.volk_id,
            est.estimated_cost.total_cost,
            budget.remaining_usd,
        )

    def _build_call_params(self, request: ProviderRequest) -> dict[str, Any]:
        """Merge generation/tool params onto a copy of request.parameters.

        Forwards max_tokens/tools/tool_choice set on the request but living
        outside request.parameters, without mutating the caller's dict.
        Explicit keys in request.parameters win.
        """
        call_params = dict(request.parameters or {})
        if request.max_tokens is not None:
            call_params.setdefault("max_tokens", request.max_tokens)
        if request.tools is not None:
            call_params.setdefault("tools", request.tools)
        if request.tool_choice is not None:
            call_params.setdefault("tool_choice", request.tool_choice)
        return call_params

    @traces(spec="015", story="US-011-001", reqs=["REQ-015-057"])
    async def _invoke_provider(
        self,
        provider: ProviderProtocol,
        request: ProviderRequest,
        model_name: str,
        call_params: dict[str, Any],
    ) -> Any:
        """Run provider.complete as a tracked task so cancel()/stop() can interrupt."""
        complete_task = asyncio.ensure_future(
            provider.complete(
                request.messages,
                model_name,
                call_params,
            )
        )
        self._in_flight[request.request_id] = complete_task
        try:
            return await complete_task
        finally:
            self._in_flight.pop(request.request_id, None)

    @traces(spec="011", story="US-001", reqs=["REQ-011-008", "REQ-011-010"])
    async def _record_success(
        self,
        request: ProviderRequest,
        provider_name: str,
        model_name: str,
        pricing: ModelPricing | None,
        raw: Any,
        elapsed_ms: float,
    ) -> ProviderResponse:
        """Record breaker/limiter/cost/health for a successful call and build a response."""
        breaker = self._breakers.get(provider_name)
        limiter = self._limiters.get(provider_name)
        if breaker:
            await breaker.record_success()
        usage = TokenUsage(
            input_tokens=raw.input_tokens,
            output_tokens=raw.output_tokens,
            total_tokens=raw.input_tokens + raw.output_tokens,
        )
        if limiter:
            await limiter.record_request(usage.total_tokens)
        cost = CostRecord()
        if pricing:
            cost = self._cost_tracker.calculate_cost(usage, pricing)
            await self._cost_tracker.record_cost(
                cost,
                volk_id=request.volk_id,
                session_id=request.session_id,
                project_id=request.project_id,
                provider=provider_name,
                model=model_name,
            )
        self._health.record(provider_name, elapsed_ms, success=True)
        return ProviderResponse(
            content=raw.content,
            model_used=raw.model or model_name,
            provider_used=provider_name,
            usage=usage,
            cost=cost,
            latency_ms=elapsed_ms,
            metadata=raw.metadata,
            tool_calls=raw.tool_calls,
            reasoning_content=raw.metadata.get("reasoning_content"),
        )

    @traces(spec="011", story="US-001", reqs=["REQ-011-007", "REQ-011-010"])
    async def _record_failure(
        self,
        attempts: list[dict],
        provider_name: str,
        model_name: str,
        exc: Exception,
        elapsed_ms: float,
    ) -> None:
        """Record breaker/health for a failed call and append an attempt entry."""
        error_type = type(exc).__name__
        self._health.record(
            provider_name, elapsed_ms, success=False, error_type=error_type
        )
        breaker = self._breakers.get(provider_name)
        if breaker:
            await breaker.record_failure()
        attempts.append(
            {
                "provider": provider_name,
                "model": model_name,
                "error": str(exc),
                "error_type": error_type,
                "latency_ms": elapsed_ms,
            }
        )
        logger.warning(
            "Provider %s failed: %s, trying fallback",
            provider_name,
            exc,
        )

    @traces(spec="011", story="US-002", reqs=["REQ-011-015", "REQ-011-016"])
    def get_rate_limit_backoff(self, request: ProviderRequest) -> float:
        """Return the router's recommended rate-limit backoff for this request."""
        return self._router.get_rate_limit_backoff(request)

    @traces(spec="011", story="US-002", reqs=["REQ-011-018"])
    async def estimate_cost(self, request: ProviderRequest) -> CostEstimate:
        """Estimate cost for a request without executing it."""
        health_data = self._health.get_all_health()
        route = self._router.route(request, health_data)
        if route is None:
            return CostEstimate()
        provider_name, model_name = route
        pricing = self._registry.get_pricing(provider_name, model_name)
        if not pricing:
            return CostEstimate(provider=provider_name, model=model_name)
        return self._estimate_cost_for_pricing(
            request, pricing, provider_name, model_name
        )

    @traces(
        spec="011",
        story="US-002",
        reqs=["REQ-011-020", "REQ-011-021", "REQ-011-022", "REQ-011-023"],
    )
    async def get_usage(
        self,
        *,
        volk_id: str | None = None,
        session_id: str | None = None,
        project_id: str | None = None,
        date_from: datetime | None = None,
        date_to: datetime | None = None,
    ) -> UsageReport:
        """Query cumulative usage and cost."""
        return await self._cost_tracker.get_usage(
            volk_id=volk_id,
            session_id=session_id,
            project_id=project_id,
        )

    @traces(spec="011", story="US-002", reqs=["REQ-011-045"])
    async def get_remaining_budget(
        self,
        *,
        volk_id: str | None = None,
        project_id: str | None = None,
        portfolio_id: str | None = None,
    ) -> BudgetStatus:
        """Query remaining budget for a scope."""
        scope_key = volk_id or project_id or portfolio_id or ""
        if not scope_key:
            return BudgetStatus()
        return await self._cost_tracker.get_remaining_budget(scope_key)

    async def cancel(self, request_id: str) -> bool:
        """Cancel an in-flight request."""
        task = self._in_flight.get(request_id)
        if task and not task.done():
            task.cancel()
            del self._in_flight[request_id]
            return True
        return False

    @traces(spec="015", story="US-011-001", reqs=["REQ-015-055"])
    def get_provider_health(self) -> dict[str, Any]:
        """Return health status of all registered providers."""
        result = {}
        for name in self._registry.provider_names:
            health = self._health.get_health(name)
            breaker = self._breakers.get(name)
            limiter = self._limiters.get(name)
            result[name] = {
                **health,
                "circuit_state": breaker.state.value if breaker else "unknown",
                "rate_limit": limiter.get_status() if limiter else {},
            }
        return result

    @traces(spec="011", story="US-011-001", reqs=["REQ-011-045"])
    def set_budget(self, scope_key: str, limit_usd: float) -> None:
        """Set a budget limit for a scope (volk_id, project_id, etc.)."""
        self._cost_tracker.set_budget(scope_key, limit_usd)

    @traces(spec="011", story="US-011-001", reqs=["REQ-011-015"])
    def _token_counter_for(self, model_name: str) -> TiktokenCounter:
        """Return a cached token counter for ``model_name`` (SPEC-011 FR-015).

        OpenAI-family names get their exact tiktoken encoding; other providers
        are approximated with ``cl100k_base`` (via the ``gpt-4`` encoding).
        ``TiktokenCounter`` itself degrades to word counting if tiktoken is not
        installed. Encodings are expensive to build, so counters are memoised
        per model.
        """
        cache = self._token_counters
        if model_name not in cache:
            name = (model_name or "").lower()
            is_openai = any(
                marker in name
                for marker in ("gpt", "o1-", "o3-", "text-embedding", "openai")
            )
            cache[model_name] = TiktokenCounter(
                model_name if (is_openai and model_name) else "gpt-4"
            )
        return cache[model_name]

    @traces(spec="011", story="US-011-001", reqs=["REQ-011-018"])
    def _estimate_cost_for_pricing(
        self,
        request: ProviderRequest,
        pricing: ModelPricing,
        provider_name: str = "",
        model_name: str = "",
    ) -> CostEstimate:
        """Estimate cost using a real tokenizer (SPEC-011 FR-015).

        Replaces the previous flat ``chars/4`` heuristic with tiktoken: an
        OpenAI-family model uses its exact encoding, every other provider is
        approximated with ``cl100k_base`` (modern BPE density tracks
        deepseek/claude/kimi far closer than ``chars/4``). This shares the
        SPEC-010 ``TiktokenCounter`` rather than inventing a third heuristic, so
        budget enforcement sees a materially more accurate input-token count.
        """
        counter = self._token_counter_for(model_name)
        est_input = sum(counter.count(m.content) for m in request.messages)
        est_output = request.max_tokens or 1024
        est_cost = CostRecord(
            input_cost=(est_input / 1_000_000) * pricing.input_per_1m,
            output_cost=(est_output / 1_000_000) * pricing.output_per_1m,
        )
        est_cost.total_cost = est_cost.input_cost + est_cost.output_cost
        return CostEstimate(
            estimated_input_tokens=est_input,
            estimated_output_tokens=est_output,
            estimated_cost=est_cost,
            provider=provider_name,
            model=model_name,
        )