agent-infra-sdk 0.2.0

//! Model Infra data-plane client.
//!
//! The Model service exposes OpenAI-compatible JSON at `/v1/*`. Keeping this
//! SDK boundary JSON-shaped avoids coupling runtime consumers to one provider
//! protocol crate while still centralising credentials, deadlines, retries,
//! response limits, and Gateway routing in [`HttpTransport`].

use bytes::Bytes;
use futures_util::Stream;
use reqwest::Client;
use serde_json::Value;
use std::pin::Pin;
use std::task::{Context, Poll};

use crate::transport::{
    CallOptions, ClientOptions, HttpTransport, InfraClientError, ServiceEndpoint,
};

const CHAT_COMPLETIONS_PATH: &str = "/v1/model/chat/completions";
const RESPONSES_PATH: &str = "/v1/model/responses";
const MODELS_PATH: &str = "/v1/model/models";
const SERVICE_NAME: &str = "agent-model-infra";
const AUDIENCE: &str = "agent-model-infra";

/// Typed transport owner for Model Infra's JSON data plane.
#[derive(Clone, Debug)]
pub struct ModelClient {
    transport: HttpTransport,
}

/// Raw SSE bytes from a Model Infra streaming response.
///
/// Chunks preserve wire order but are transport chunks, not guaranteed SSE
/// event boundaries. Consumers may proxy them directly or apply their own SSE
/// decoder. Dropping the stream cancels the in-flight response.
pub struct ModelStream {
    inner: crate::transport::HttpByteStream,
}

impl std::fmt::Debug for ModelStream {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter.write_str("ModelStream(..)")
    }
}

impl ModelStream {
    /// Read the next raw SSE transport chunk without importing `StreamExt`.
    pub async fn next_chunk(&mut self) -> Option<Result<Bytes, InfraClientError>> {
        futures_util::StreamExt::next(self).await
    }
}

impl Stream for ModelStream {
    type Item = Result<Bytes, InfraClientError>;

    fn poll_next(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        self.inner.as_mut().poll_next(context)
    }
}

impl ModelClient {
    pub(crate) fn new_with_endpoint(
        http: Client,
        endpoint: ServiceEndpoint,
        options: ClientOptions,
    ) -> Self {
        let endpoint = endpoint.with_default_credential_audience(AUDIENCE);
        Self {
            transport: HttpTransport::new_with_options(http, SERVICE_NAME, endpoint, options),
        }
    }

    /// Invoke the OpenAI-compatible chat-completions endpoint.
    pub async fn chat_completions(
        &self,
        request: &Value,
        options: CallOptions,
    ) -> Result<Value, InfraClientError> {
        self.transport
            .post_json_with_options(CHAT_COMPLETIONS_PATH, request, options)
            .await
    }

    /// Invoke Chat Completions with `stream=true` and return raw data-only SSE.
    pub async fn chat_completions_stream(
        &self,
        request: &Value,
        options: CallOptions,
    ) -> Result<ModelStream, InfraClientError> {
        self.stream(CHAT_COMPLETIONS_PATH, request, options).await
    }

    /// Invoke the OpenAI-compatible Responses endpoint.
    pub async fn responses(
        &self,
        request: &Value,
        options: CallOptions,
    ) -> Result<Value, InfraClientError> {
        self.transport
            .post_json_with_options(RESPONSES_PATH, request, options)
            .await
    }

    /// Invoke Responses with `stream=true` and return raw named-event SSE.
    pub async fn responses_stream(
        &self,
        request: &Value,
        options: CallOptions,
    ) -> Result<ModelStream, InfraClientError> {
        self.stream(RESPONSES_PATH, request, options).await
    }

    pub async fn models(&self) -> Result<Value, InfraClientError> {
        self.transport.get_json(MODELS_PATH).await
    }

    async fn stream(
        &self,
        path: &str,
        request: &Value,
        options: CallOptions,
    ) -> Result<ModelStream, InfraClientError> {
        let mut request = request.clone();
        let object = request
            .as_object_mut()
            .ok_or_else(|| InfraClientError::Protocol {
                service: SERVICE_NAME,
                message: "streaming model request must be a JSON object".into(),
            })?;
        object.insert("stream".into(), Value::Bool(true));
        let inner = self
            .transport
            .post_json_stream_with_options(path, &request, options)
            .await?;
        Ok(ModelStream { inner })
    }
}