hyperinfer-client 0.2.0

//! HyperInfer Client Library - Data Plane

pub mod cache;
pub mod http_client;
pub mod mirroring;
pub mod router;
pub mod router_engine;
pub mod telemetry;
pub mod telemetry_otlp;
mod util;

pub use cache::ExactMatchCache;
pub use http_client::HttpCaller;
pub use mirroring::{MirrorConfig, MirrorHandle};
pub use router::Router;
pub use router_engine::RouterEngine;
pub use telemetry::Telemetry;
pub use telemetry_otlp::{
    init_langfuse_telemetry, init_telemetry, init_telemetry_with_headers, set_gen_ai_attributes,
    set_gen_ai_response, set_gen_ai_usage, shutdown_telemetry,
};

use futures::{Stream, StreamExt};
use hyperinfer_core::{
    rate_limiting::RateLimiter, ChatChunk, ChatRequest, ChatResponse, Config, HyperInferError,
    Provider,
};
use hyperinfer_providers::ProviderRegistry;
use std::pin::Pin;
use std::sync::{Arc, LazyLock};
use std::task::{Context, Poll};

static HTTP_CLIENT: LazyLock<reqwest::Client> = LazyLock::new(reqwest::Client::new);
use tokio::sync::RwLock;
use tracing::Instrument as _;

/// Wraps a provider `ChatChunk` stream and performs the same accounting as
/// `chat()` once the stream terminates (naturally or via an error):
///
/// - Fires Redis telemetry off the critical path via `tokio::spawn`.
/// - Records output-token usage in the rate-limiter bucket.
/// - Sets OTel span usage / response attributes.
///
/// The accounting is triggered exactly once, either when a `[DONE]`-equivalent
/// chunk with a `finish_reason` is seen **or** when the stream signals
/// `Poll::Ready(None)`.
struct AccountedStream {
    inner: Pin<Box<dyn Stream<Item = Result<ChatChunk, HyperInferError>> + Send>>,
    telemetry: Telemetry,
    rate_limiter: RateLimiter,
    key: String,
    model: String,
    start: std::time::Instant,
    /// Accumulated token counts from the stream's usage chunk (if any).
    input_tokens: u32,
    output_tokens: u32,
    /// Guards against running the accounting block more than once.
    accounted: bool,
    /// OTel span that lives for the full stream lifetime.
    span: tracing::Span,
}

impl AccountedStream {
    /// Run the accounting side-effects exactly once.
    fn account(&mut self) {
        if self.accounted {
            return;
        }
        self.accounted = true;

        let elapsed = self.start.elapsed().as_millis() as u64;
        let input_tokens = self.input_tokens;
        let output_tokens = self.output_tokens;

        let _enter = self.span.clone().entered();
        crate::telemetry_otlp::set_gen_ai_usage(&self.span, input_tokens, output_tokens);

        // Telemetry write is off the critical path.
        let telemetry = self.telemetry.clone();
        let key = self.key.clone();
        let model = self.model.clone();
        tokio::spawn(async move {
            if let Err(e) = telemetry
                .record_with_tokens(&key, &model, input_tokens, output_tokens, elapsed)
                .await
            {
                tracing::warn!(error = %e, "stream telemetry record failed");
            }
        });

        // Rate-limiter token-bucket update is lightweight and synchronous-ish;
        // run it in a spawn to avoid blocking the poll path.
        let rate_limiter = self.rate_limiter.clone();
        let key2 = self.key.clone();
        let total = (input_tokens + output_tokens) as u64;
        tokio::spawn(async move {
            let _ = rate_limiter.record_usage(&key2, total).await;
        });
    }
}

impl Drop for AccountedStream {
    fn drop(&mut self) {
        self.account();
    }
}

impl Stream for AccountedStream {
    type Item = Result<ChatChunk, HyperInferError>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        // Clone the span so the Entered guard holds no borrow into `self`,
        // allowing subsequent mutable borrows of other fields.
        let _enter = self.span.clone().entered();
        match self.inner.as_mut().poll_next(cx) {
            Poll::Ready(Some(Ok(chunk))) => {
                // Capture usage from whatever chunk carries it (typically the last).
                if let Some(ref u) = chunk.usage {
                    self.input_tokens = u.input_tokens;
                    self.output_tokens = u.output_tokens;
                }
                // If this chunk has a finish_reason the stream is done; account now
                // so the span attributes are set while the span is still open.
                if chunk.finish_reason.is_some() {
                    self.account();
                }
                Poll::Ready(Some(Ok(chunk)))
            }
            Poll::Ready(Some(Err(e))) => {
                self.account();
                Poll::Ready(Some(Err(e)))
            }
            Poll::Ready(None) => {
                self.account();
                Poll::Ready(None)
            }
            Poll::Pending => Poll::Pending,
        }
    }
}

pub struct HyperInferClient {
    config: Arc<RwLock<Config>>,
    http_caller: Arc<HttpCaller>,
    router: Arc<Router>,
    router_engine: Arc<RouterEngine>,
    rate_limiter: RateLimiter,
    telemetry: Telemetry,
    cache: ExactMatchCache,
    mirror: MirrorHandle,
    provider_registry: Arc<RwLock<Arc<ProviderRegistry>>>,
}

impl HyperInferClient {
    pub async fn new(redis_url: &str, config: Config) -> Result<Self, HyperInferError> {
        let http_caller = Arc::new(HttpCaller::new().map_err(HyperInferError::Http)?);
        let router = Arc::new(
            Router::new(config.routing_rules.clone())
                .with_aliases(config.model_aliases.clone())
                .with_default_provider(config.default_provider.clone()),
        );
        let rate_limiter = RateLimiter::new(Some(redis_url))
            .await
            .map_err(|e| HyperInferError::Config(std::io::Error::other(e.to_string())))?;
        let telemetry = Telemetry::new(redis_url)
            .await
            .map_err(|e| HyperInferError::Config(std::io::Error::other(e.to_string())))?;
        let cache = ExactMatchCache::new(redis_url, "default").await;
        let mirror: MirrorHandle = Arc::new(RwLock::new(None));
        let config = Arc::new(RwLock::new(config));

        let provider_registry_inner = Arc::new(ProviderRegistry::new());
        hyperinfer_providers::init_default_registry(&provider_registry_inner);
        let provider_registry = Arc::new(RwLock::new(provider_registry_inner));

        Ok(Self {
            config,
            http_caller,
            router,
            router_engine: Arc::new(RouterEngine::new().await),
            rate_limiter,
            telemetry,
            cache,
            mirror,
            provider_registry,
        })
    }

    /// Configure traffic mirroring.  Pass `None` to disable.
    pub async fn set_mirror(&self, cfg: Option<MirrorConfig>) {
        let mut guard = self.mirror.write().await;
        *guard = cfg;
    }

    pub async fn inject_provider_registry(&self, external_registry: Arc<ProviderRegistry>) {
        let mut guard = self.provider_registry.write().await;
        *guard = external_registry;
    }

    /// Load deployments into the router engine for deployment-based routing
    pub async fn load_deployments(&self, deployments: Vec<hyperinfer_core::Deployment>) {
        self.router_engine.load_deployments(deployments).await;
    }

    /// Subscribe to Redis Pub/Sub for live deployment config changes.
    /// When a message arrives on "hyperinfer:config_updates", calls the provided
    /// fetcher function to re-fetch deployments and rebuilds the routing pool.
    pub async fn subscribe_config_updates<F, Fut>(
        &self,
        redis_url: &str,
        fetcher: F,
    ) -> Result<(), HyperInferError>
    where
        F: Fn() -> Fut + Send + Sync + 'static,
        Fut: std::future::Future<Output = Result<Vec<hyperinfer_core::Deployment>, String>> + Send,
    {
        let client = redis::Client::open(redis_url)
            .map_err(|e| HyperInferError::Config(std::io::Error::other(e.to_string())))?;
        let mut pubsub = client
            .get_async_pubsub()
            .await
            .map_err(|e| HyperInferError::Config(std::io::Error::other(e.to_string())))?;

        pubsub
            .subscribe("hyperinfer:config_updates")
            .await
            .map_err(|e| HyperInferError::Config(std::io::Error::other(e.to_string())))?;

        let engine = self.router_engine.clone();
        let _handle = tokio::spawn(async move {
            let mut stream = pubsub.on_message();
            loop {
                match stream.next().await {
                    Some(_msg) => {
                        tracing::info!(
                            "Received config update notification, re-fetching deployments"
                        );
                        match fetcher().await {
                            Ok(deployments) => {
                                engine.rebuild_pool(deployments).await;
                                tracing::info!("Rebuilt deployment pool after config update");
                            }
                            Err(e) => {
                                tracing::warn!(error = %e, "Failed to re-fetch deployments after config update");
                            }
                        }
                    }
                    None => {
                        tracing::info!("Pub/Sub stream ended");
                        break;
                    }
                }
            }
        });

        Ok(())
    }

    /// Get a reference to the router engine
    pub fn router_engine(&self) -> &Arc<RouterEngine> {
        &self.router_engine
    }

    pub async fn chat(
        &self,
        key: &str,
        request: ChatRequest,
    ) -> Result<ChatResponse, HyperInferError> {
        request.validate()?;

        // 0. Exact-match cache lookup (before rate-limiting to avoid wasting quota).
        if let Some(cached) = self.cache.get(&request).await {
            return Ok(cached);
        }

        // Create a root OTel span following the GenAI Semantic Conventions.
        // We use `.instrument(span)` on the inner async block so the span is
        // properly propagated across every `.await` point (using `span.enter()`
        // in an async function is unsafe — the guard can survive suspension).
        let span = tracing::info_span!(
            "gen_ai.chat",
            gen_ai.operation.name = "chat",
            gen_ai.request.model = %request.model,
        );

        async move {
            let start = std::time::Instant::now();

            // 1. Check rate limit
            let allowed = self.rate_limiter.is_allowed(key, 1).await;
            if let Err(e) = allowed {
                return Err(HyperInferError::RateLimit(e.to_string()));
            }
            if !allowed.unwrap() {
                return Err(HyperInferError::RateLimit(
                    "Rate limit exceeded".to_string(),
                ));
            }

            // 2. Try deployment-based routing first (if deployments are loaded)
            let deployment_result = self.router_engine.select_deployment(&request).await;
            if let Ok(routing_result) = deployment_result {
                let deployment = &routing_result.deployment;
                let default_url = match &deployment.provider {
                    Provider::Anthropic => "https://api.anthropic.com/v1",
                    _ => "https://api.openai.com/v1",
                };
                let base_url = deployment.base_url.as_deref().unwrap_or(default_url);
                let api_key = &deployment.api_key_ref;

                let url = format!("{}/chat/completions", base_url.trim_end_matches('/'));

                let mut headers = reqwest::header::HeaderMap::new();
                headers.insert("content-type", "application/json".parse().unwrap());
                if !api_key.is_empty() {
                    match &deployment.provider {
                        Provider::Anthropic => {
                            headers.insert("x-api-key", api_key.parse().unwrap());
                            headers.insert("anthropic-version", "2023-06-01".parse().unwrap());
                        }
                        _ => {
                            headers.insert(
                                "authorization",
                                format!("Bearer {}", api_key).parse().unwrap(),
                            );
                        }
                    }
                }

                match HTTP_CLIENT
                    .post(&url)
                    .headers(headers)
                    .json(&request)
                    .send()
                    .await
                {
                    Ok(response) => {
                        let status = response.status();
                        if status.is_success() {
                            if let Ok(body) = response.json::<ChatResponse>().await {
                                // Record success metrics
                                let latency = start.elapsed().as_secs_f64() * 1000.0;
                                let tokens =
                                    (body.usage.input_tokens + body.usage.output_tokens) as u64;
                                self.router_engine
                                    .record_success(&deployment.id, latency, tokens)
                                    .await;

                                // Record telemetry
                                let elapsed = start.elapsed().as_millis() as u64;
                                let telemetry = self.telemetry.clone();
                                let key_owned = key.to_string();
                                let model_owned = request.model.clone();
                                tokio::spawn(async move {
                                    let _ = telemetry
                                        .record_with_tokens(
                                            &key_owned,
                                            &model_owned,
                                            body.usage.input_tokens,
                                            body.usage.output_tokens,
                                            elapsed,
                                        )
                                        .await;
                                });

                                return Ok(body);
                            }
                        }
                        // If request failed, record failure and fall through to fallback
                        self.router_engine.record_failure(&deployment.id).await;
                    }
                    Err(_) => {
                        // Request failed, record failure and fall through to fallback
                        self.router_engine.record_failure(&deployment.id).await;
                    }
                }
            }

            // 3. Fallback to existing Router-based flow
            let (model, provider, api_key, config_snapshot) = {
                let config = self.config.read().await;
                let resolved = self.router.resolve(&request.model, &config);

                let (model, provider) = resolved.ok_or_else(|| {
                    HyperInferError::Config(std::io::Error::new(
                        std::io::ErrorKind::NotFound,
                        format!(
                            "Unknown model: '{}'. No routing rule or alias found.",
                            request.model
                        ),
                    ))
                })?;

                let api_key = config
                    .api_keys
                    .get(&provider.to_string())
                    .cloned()
                    .ok_or_else(|| {
                        HyperInferError::Config(std::io::Error::new(
                            std::io::ErrorKind::NotFound,
                            format!("API key not found for provider: {:?}", provider),
                        ))
                    })?;

                (model, provider, api_key, Arc::new(config.clone()))
            };

            // Enrich span with the resolved provider and final model name.
            let provider_name = provider.to_string();
            crate::telemetry_otlp::set_gen_ai_attributes(
                &tracing::Span::current(),
                &provider_name,
                &model,
                "chat",
            );

            // 3. Execute HTTP call via provider registry
            let llm_provider = {
                let registry = self.provider_registry.read().await;
                registry.get(&provider_name).ok_or_else(|| {
                    HyperInferError::Config(std::io::Error::new(
                        std::io::ErrorKind::NotFound,
                        format!("Provider '{}' not found in registry", provider_name),
                    ))
                })?
            };

            let mut resolved_request = request.clone();
            resolved_request.model = model.clone();
            let response = llm_provider.chat(&resolved_request, &api_key).await?;

            // 4. Record OTel usage and response attributes on the span.
            let elapsed = start.elapsed().as_millis() as u64;
            let input_tokens = response.usage.input_tokens;
            let output_tokens = response.usage.output_tokens;

            crate::telemetry_otlp::set_gen_ai_usage(
                &tracing::Span::current(),
                input_tokens,
                output_tokens,
            );

            let finish_reason = response
                .choices
                .first()
                .and_then(|c| c.finish_reason.as_deref())
                .unwrap_or("unknown");
            crate::telemetry_otlp::set_gen_ai_response(
                &tracing::Span::current(),
                &response.id,
                finish_reason,
            );

            // Store successful response in exact-match cache.
            self.cache.set(&request, &response).await;

            // Record async Redis telemetry off the critical path.
            let telemetry = self.telemetry.clone();
            let key_owned = key.to_string();
            let model_owned = model.clone();
            tokio::spawn(async move {
                if let Err(e) = telemetry
                    .record_with_tokens(
                        &key_owned,
                        &model_owned,
                        input_tokens,
                        output_tokens,
                        elapsed,
                    )
                    .await
                {
                    tracing::warn!(error = %e, "telemetry record failed");
                }
            });

            // Record usage for rate-limiter token bucket.
            let total_tokens = response.usage.input_tokens + response.usage.output_tokens;
            let _ = self
                .rate_limiter
                .record_usage(key, total_tokens as u64)
                .await;

            // 5. Fire-and-forget traffic mirror (if configured).
            mirroring::maybe_mirror(
                self.mirror.clone(),
                self.http_caller.clone(),
                self.router.clone(),
                config_snapshot,
                key.to_string(),
                request,
            );

            // 6. Return response
            Ok(response)
        }
        .instrument(span)
        .await
    }

    /// Stream token chunks for a chat request.
    ///
    /// Returns a `Stream` of `ChatChunk` items.  The caller is responsible for
    /// collecting `delta` fields and assembling the final text.  The last chunk
    /// in the stream has a non-`None` `finish_reason` and may carry `usage`.
    ///
    /// Rate-limiting and routing follow the same logic as `chat()`.
    pub async fn chat_stream(
        &self,
        key: &str,
        request: ChatRequest,
    ) -> Result<
        Pin<Box<dyn Stream<Item = Result<ChatChunk, HyperInferError>> + Send>>,
        HyperInferError,
    > {
        request.validate()?;

        // 1. Rate limit check (same as non-streaming path).
        let allowed = self.rate_limiter.is_allowed(key, 1).await;
        if let Err(e) = allowed {
            return Err(HyperInferError::RateLimit(e.to_string()));
        }
        if !allowed.unwrap() {
            return Err(HyperInferError::RateLimit(
                "Rate limit exceeded".to_string(),
            ));
        }

        // 2. Resolve model / provider / api key.
        let (model, provider_name, api_key) = {
            let config = self.config.read().await;
            let resolved = self.router.resolve(&request.model, &config);

            let (model, provider) = resolved.ok_or_else(|| {
                HyperInferError::Config(std::io::Error::new(
                    std::io::ErrorKind::NotFound,
                    format!(
                        "Unknown model: '{}'. No routing rule or alias found.",
                        request.model
                    ),
                ))
            })?;

            let provider_name = provider.to_string();
            let api_key = config
                .api_keys
                .get(&provider_name)
                .cloned()
                .ok_or_else(|| {
                    HyperInferError::Config(std::io::Error::new(
                        std::io::ErrorKind::NotFound,
                        format!("API key not found for provider: {:?}", provider),
                    ))
                })?;

            (model, provider_name, api_key)
        };

        // 3. Get streaming provider from registry (already checks supports_streaming)
        let streaming_provider = {
            let registry = self.provider_registry.read().await;
            registry.get_streaming(&provider_name).ok_or_else(|| {
                HyperInferError::Config(std::io::Error::new(
                    std::io::ErrorKind::NotFound,
                    format!(
                        "Provider '{}' not found in registry or does not support streaming",
                        provider_name
                    ),
                ))
            })?
        };

        let mut resolved_request = request.clone();
        resolved_request.model = model.clone();
        let provider_stream: Pin<
            Box<dyn Stream<Item = Result<ChatChunk, HyperInferError>> + Send>,
        > = streaming_provider.into_stream(&resolved_request, &api_key);
        // Note: streaming responses are not cached — the stream is consumed
        // incrementally by the caller so we cannot inspect it here.

        // 4. Create an OTel span for the stream lifetime and enrich it with
        //    resolved provider / model information (mirrors chat()).
        let span = tracing::info_span!(
            "gen_ai.chat_stream",
            gen_ai.operation.name = "chat_stream",
            gen_ai.request.model = %request.model,
        );
        crate::telemetry_otlp::set_gen_ai_attributes(&span, &provider_name, &model, "chat_stream");

        // 5. Wrap the provider stream so usage/telemetry are recorded on
        //    termination — the same accounting chat() performs, but deferred
        //    to when the last chunk (or an error) is polled.
        //    The span is stored inside the wrapper; poll_next enters it on
        //    every poll so it covers the full stream lifetime.
        let stream = AccountedStream {
            inner: provider_stream,
            telemetry: self.telemetry.clone(),
            rate_limiter: self.rate_limiter.clone(),
            key: key.to_string(),
            model,
            start: std::time::Instant::now(),
            input_tokens: 0,
            output_tokens: 0,
            accounted: false,
            span,
        };

        Ok(Box::pin(stream))
    }
}