ai_sdk_provider/language_model/

trait_def.rs

use super::*;
use crate::{Error, Result, SharedHeaders, SharedProviderMetadata};
use async_trait::async_trait;
use std::collections::HashMap;
use std::future::{Future, IntoFuture};
use std::pin::Pin;
use std::sync::Arc;
use tokio_stream::Stream;

macro_rules! impl_call_options {
    () => {
        /// Sets the maximum number of tokens to generate in the model output.
        ///
        /// This parameter controls the maximum length of the generated response. The actual
        /// number of tokens generated may be less if the model reaches a natural stopping point
        /// or encounters a stop sequence.
        pub fn max_tokens(mut self, max_tokens: u32) -> Self {
            self.options.max_output_tokens = Some(max_tokens);
            self
        }

        /// Sets the sampling temperature for controlling output randomness.
        ///
        /// Valid range: 0.0 to 2.0
        /// - Lower values (e.g., 0.2): More focused, deterministic outputs
        /// - Higher values (e.g., 1.5): More creative, diverse outputs
        /// - Default: Provider-specific (typically 1.0)
        pub fn temperature(mut self, temperature: f32) -> Self {
            self.options.temperature = Some(temperature);
            self
        }

        /// Sets sequences that will cause generation to stop when encountered.
        ///
        /// When the model generates any of these sequences, generation will terminate
        /// immediately. Useful for structured output or preventing unwanted continuations.
        pub fn stop_sequences(mut self, sequences: Vec<String>) -> Self {
            self.options.stop_sequences = Some(sequences);
            self
        }

        /// Sets the nucleus sampling probability mass.
        ///
        /// Valid range: 0.0 to 1.0
        /// Only tokens whose cumulative probability exceeds this threshold are considered
        /// for sampling. Lower values produce more focused outputs by restricting the
        /// token pool to high-probability candidates.
        pub fn top_p(mut self, top_p: f32) -> Self {
            self.options.top_p = Some(top_p);
            self
        }

        /// Sets the top-k sampling parameter.
        ///
        /// Limits sampling to the k most probable next tokens. Reduces output diversity
        /// by eliminating low-probability tokens from consideration. Provider support varies.
        pub fn top_k(mut self, top_k: u32) -> Self {
            self.options.top_k = Some(top_k);
            self
        }

        /// Sets the presence penalty for reducing topic repetition.
        ///
        /// Valid range: -2.0 to 2.0
        /// Positive values penalize tokens that have already appeared in the generated text,
        /// encouraging the model to explore new topics. Negative values have the opposite effect.
        pub fn presence_penalty(mut self, penalty: f32) -> Self {
            self.options.presence_penalty = Some(penalty);
            self
        }

        /// Sets the frequency penalty for reducing token repetition.
        ///
        /// Valid range: -2.0 to 2.0
        /// Positive values penalize tokens based on their frequency in the generated text,
        /// with stronger penalties for frequently occurring tokens. This reduces repetitive phrasing.
        pub fn frequency_penalty(mut self, penalty: f32) -> Self {
            self.options.frequency_penalty = Some(penalty);
            self
        }

        /// Sets a random seed for deterministic generation.
        ///
        /// Using the same seed with identical parameters should produce consistent outputs
        /// across multiple requests. Useful for reproducible testing and debugging.
        /// Provider support for deterministic generation varies.
        pub fn seed(mut self, seed: i64) -> Self {
            self.options.seed = Some(seed);
            self
        }

        /// Sets the tools available for the model to call during generation.
        ///
        /// Tools enable function calling where the model can request execution of
        /// external functions with structured arguments. The application is responsible
        /// for executing tool calls and providing results back to the model.
        pub fn tools(mut self, tools: Vec<Tool>) -> Self {
            self.options.tools = Some(tools);
            self
        }

        /// Sets the tool selection strategy.
        ///
        /// Controls how the model decides whether and which tools to call:
        /// - `Auto`: Model autonomously decides when to use tools
        /// - `Required`: Model must call at least one tool
        /// - `None`: Tools are disabled for this request
        /// - `Tool { name }`: Model must call the specified tool
        pub fn tool_choice(mut self, choice: ToolChoice) -> Self {
            self.options.tool_choice = Some(choice);
            self
        }

        /// Forces a specific output format for the generated response.
        ///
        /// Constrains the model to produce output in the specified format:
        /// - `ResponseFormat::Text`: Plain text output (default)
        /// - `ResponseFormat::Json`: Valid JSON object or array
        ///
        /// JSON format is useful for structured data extraction and schema-compliant outputs.
        pub fn response_format(mut self, format: ResponseFormat) -> Self {
            self.options.response_format = Some(format);
            self
        }

        /// Sets custom HTTP headers for the provider request.
        ///
        /// Replaces any existing headers. Use this for provider-specific headers such as
        /// organization identifiers, API version selection, or custom authentication schemes.
        pub fn headers(mut self, headers: HashMap<String, String>) -> Self {
            self.options.headers = Some(headers);
            self
        }

        /// Adds a single HTTP header to the provider request.
        ///
        /// Merges with existing headers rather than replacing them. Useful for
        /// incrementally building header sets in a builder pattern.
        pub fn header(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
            let headers = self.options.headers.get_or_insert_with(HashMap::new);
            headers.insert(key.into(), value.into());
            self
        }
    };
}

/// Builder for configuring and executing non-streaming generation requests.
///
/// This builder provides a fluent interface for setting generation parameters before
/// executing the request. It supports both explicit execution via `.send()` and
/// implicit execution by awaiting the builder directly.
#[derive(Clone)]
pub struct GenerateBuilder<'a, M: LanguageModel + ?Sized> {
    model: &'a M,
    options: CallOptions,
    init_error: Option<Arc<Error>>,
}

impl<'a, M: LanguageModel + ?Sized> GenerateBuilder<'a, M> {
    /// Creates a new generation builder with the specified prompt.
    ///
    /// If the prompt conversion fails, the error is captured and will be returned
    /// when the builder is executed.
    pub fn new(model: &'a M, prompt_result: std::result::Result<Prompt, Error>) -> Self {
        match prompt_result {
            Ok(prompt) => Self {
                model,
                options: CallOptions {
                    prompt,
                    ..Default::default()
                },
                init_error: None,
            },
            Err(e) => Self {
                model,
                options: CallOptions::default(),
                init_error: Some(Arc::new(e)),
            },
        }
    }

    impl_call_options!();

    /// Executes the generation request and returns the complete response.
    ///
    /// This method consumes the builder and performs the actual API call to the
    /// language model provider. Use this for explicit execution, or simply await
    /// the builder directly for implicit execution.
    pub async fn send(self) -> Result<GenerateResponse> {
        if let Some(err) = self.init_error {
            // We have to return the error. Since we can't move out of Arc,
            // and Error is a Box<dyn std::error::Error>, we can recreate it
            // from the string representation if we can't clone it.
            // Or better, we just return a new generic error wrapping the old one.
            return Err(format!("Initialization error: {}", err).into());
        }
        self.model.do_generate(self.options).await
    }
}

/// Enables awaiting the builder directly without calling `.send()`.
///
/// This implementation allows the syntax `model.generate(prompt).await` as a
/// shorthand for `model.generate(prompt).send().await`.
impl<'a, M: LanguageModel + ?Sized> IntoFuture for GenerateBuilder<'a, M> {
    type Output = Result<GenerateResponse>;
    type IntoFuture = Pin<Box<dyn Future<Output = Self::Output> + Send + 'a>>;

    fn into_future(self) -> Self::IntoFuture {
        if let Some(err) = self.init_error {
            return Box::pin(async move { Err(format!("Initialization error: {}", err).into()) });
        }
        let model = self.model;
        let options = self.options;
        Box::pin(async move { model.do_generate(options).await })
    }
}

/// Builder for configuring and executing streaming generation requests.
///
/// This builder provides a fluent interface for setting generation parameters before
/// initiating a streaming response. Streaming allows processing partial responses as
/// they arrive, enabling real-time display and reduced time-to-first-token.
#[derive(Clone)]
pub struct StreamBuilder<'a, M: LanguageModel + ?Sized> {
    model: &'a M,
    options: CallOptions,
    init_error: Option<Arc<Error>>,
}

impl<'a, M: LanguageModel + ?Sized> StreamBuilder<'a, M> {
    /// Creates a new streaming builder with the specified prompt.
    ///
    /// If the prompt conversion fails, the error is captured and will be returned
    /// when the builder is executed.
    pub fn new(model: &'a M, prompt_result: std::result::Result<Prompt, Error>) -> Self {
        match prompt_result {
            Ok(prompt) => Self {
                model,
                options: CallOptions {
                    prompt,
                    ..Default::default()
                },
                init_error: None,
            },
            Err(e) => Self {
                model,
                options: CallOptions::default(),
                init_error: Some(Arc::new(e)),
            },
        }
    }

    impl_call_options!();

    /// Initiates the streaming request and returns a stream of response parts.
    ///
    /// This method consumes the builder and establishes the streaming connection.
    /// Use this for explicit execution, or simply await the builder directly.
    pub async fn send(self) -> Result<StreamResponse> {
        if let Some(err) = self.init_error {
            return Err(format!("Initialization error: {}", err).into());
        }
        self.model.do_stream(self.options).await
    }
}

/// Enables awaiting the builder directly without calling `.send()`.
///
/// This implementation allows the syntax `model.stream(prompt).await` as a
/// shorthand for `model.stream(prompt).send().await`.
impl<'a, M: LanguageModel + ?Sized> IntoFuture for StreamBuilder<'a, M> {
    type Output = Result<StreamResponse>;
    type IntoFuture = Pin<Box<dyn Future<Output = Self::Output> + Send + 'a>>;

    fn into_future(self) -> Self::IntoFuture {
        if let Some(err) = self.init_error {
            return Box::pin(async move { Err(format!("Initialization error: {}", err).into()) });
        }
        let model = self.model;
        let options = self.options;
        Box::pin(async move { model.do_stream(options).await })
    }
}

/// Core trait for language model providers.
///
/// This trait defines the standard interface that all language model implementations
/// must satisfy. It provides both high-level builder methods for ergonomic usage and
/// low-level `do_*` methods for direct provider integration.
///
/// # Implementation Requirements
///
/// Implementors must provide:
/// - `provider()`: Provider identifier (e.g., "openai", "anthropic")
/// - `model_id()`: Model identifier (e.g., "gpt-4", "claude-3-opus")
/// - `do_generate()`: Non-streaming generation implementation
/// - `do_stream()`: Streaming generation implementation
///
/// # Example Implementation
///
/// ```rust,ignore
/// use ai_sdk_provider::{LanguageModel, CallOptions, GenerateResponse};
/// use async_trait::async_trait;
///
/// struct MyModel {
///     api_key: String,
/// }
///
/// #[async_trait]
/// impl LanguageModel for MyModel {
///     fn provider(&self) -> &str { "my-provider" }
///     fn model_id(&self) -> &str { "my-model-v1" }
///
///     async fn do_generate(&self, options: CallOptions) -> Result<GenerateResponse> {
///         // Implementation
///     }
///
///     async fn do_stream(&self, options: CallOptions) -> Result<StreamResponse> {
///         // Implementation
///     }
/// }
/// ```
#[async_trait]
pub trait LanguageModel: Send + Sync {
    /// Returns the specification version implemented by this model.
    ///
    /// The default implementation returns "v3". Override this only if implementing
    /// a different specification version.
    fn specification_version(&self) -> &str {
        "v3"
    }

    /// Returns the provider identifier for this model.
    ///
    /// This should be a stable, lowercase identifier for the AI provider
    /// (e.g., "openai", "anthropic", "cohere").
    fn provider(&self) -> &str;

    /// Returns the specific model identifier.
    ///
    /// This should be the exact model name as recognized by the provider
    /// (e.g., "gpt-4-turbo", "claude-3-opus-20240229").
    fn model_id(&self) -> &str;

    /// Returns URLs supported by this model for various operations.
    ///
    /// The default implementation returns an empty map. Providers can override this
    /// to expose endpoint URLs for debugging or advanced configuration.
    async fn supported_urls(&self) -> HashMap<String, Vec<String>> {
        HashMap::new()
    }

    /// Creates a builder for a non-streaming generation request.
    ///
    /// This high-level method provides a fluent interface for configuring and executing
    /// generation requests. The builder supports both explicit execution via `.send()`
    /// and implicit execution by awaiting the builder directly.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// let response = model.generate("Explain photosynthesis")
    ///     .temperature(0.7)
    ///     .max_tokens(500)
    ///     .await?;
    /// println!("{}", response.text);
    /// ```
    fn generate<P>(&self, prompt: P) -> GenerateBuilder<'_, Self>
    where
        Self: Sized,
        P: TryInto<Prompt>,
        Error: From<P::Error>,
    {
        let result = prompt.try_into().map_err(Error::from);
        GenerateBuilder::new(self, result)
    }

    /// Creates a builder for a streaming generation request.
    ///
    /// Streaming enables processing partial responses as they arrive from the provider,
    /// allowing for real-time display and reduced latency to first token. The stream
    /// yields `StreamPart` items containing text deltas, tool calls, and metadata.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// use tokio_stream::StreamExt;
    ///
    /// let mut stream = model.stream("Write a poem")
    ///     .max_tokens(500)
    ///     .await?;
    ///
    /// while let Some(part) = stream.next().await {
    ///     match part? {
    ///         StreamPart::TextDelta(delta) => print!("{}", delta),
    ///         _ => {}
    ///     }
    /// }
    /// ```
    fn stream<P>(&self, prompt: P) -> StreamBuilder<'_, Self>
    where
        Self: Sized,
        P: TryInto<Prompt>,
        Error: From<P::Error>,
    {
        let result = prompt.try_into().map_err(Error::from);
        StreamBuilder::new(self, result)
    }

    /// Executes a non-streaming generation request.
    ///
    /// This method must be implemented by providers to perform the actual API call.
    /// It receives fully configured `CallOptions` and returns a complete `GenerateResponse`
    /// containing all generated content, usage statistics, and metadata.
    ///
    /// Most users should prefer the high-level `generate()` method over calling this directly.
    async fn do_generate(&self, options: CallOptions) -> Result<GenerateResponse>;

    /// Executes a streaming generation request.
    ///
    /// This method must be implemented by providers to establish a streaming connection.
    /// It receives fully configured `CallOptions` and returns a `StreamResponse` containing
    /// an async stream of response parts.
    ///
    /// Most users should prefer the high-level `stream()` method over calling this directly.
    async fn do_stream(&self, options: CallOptions) -> Result<StreamResponse>;
}

#[async_trait]
impl<T: LanguageModel + ?Sized> LanguageModel for Box<T> {
    fn specification_version(&self) -> &str {
        (**self).specification_version()
    }

    fn provider(&self) -> &str {
        (**self).provider()
    }
    fn model_id(&self) -> &str {
        (**self).model_id()
    }

    async fn supported_urls(&self) -> HashMap<String, Vec<String>> {
        (**self).supported_urls().await
    }

    async fn do_generate(&self, options: CallOptions) -> Result<GenerateResponse> {
        (**self).do_generate(options).await
    }

    async fn do_stream(&self, options: CallOptions) -> Result<StreamResponse> {
        (**self).do_stream(options).await
    }
}

#[async_trait]
impl<T: LanguageModel + ?Sized> LanguageModel for std::sync::Arc<T> {
    fn specification_version(&self) -> &str {
        (**self).specification_version()
    }

    fn provider(&self) -> &str {
        (**self).provider()
    }
    fn model_id(&self) -> &str {
        (**self).model_id()
    }

    async fn supported_urls(&self) -> HashMap<String, Vec<String>> {
        (**self).supported_urls().await
    }

    async fn do_generate(&self, options: CallOptions) -> Result<GenerateResponse> {
        (**self).do_generate(options).await
    }

    async fn do_stream(&self, options: CallOptions) -> Result<StreamResponse> {
        (**self).do_stream(options).await
    }
}

#[async_trait]
impl<T: LanguageModel + ?Sized> LanguageModel for &T {
    fn specification_version(&self) -> &str {
        (**self).specification_version()
    }

    fn provider(&self) -> &str {
        (**self).provider()
    }
    fn model_id(&self) -> &str {
        (**self).model_id()
    }

    async fn supported_urls(&self) -> HashMap<String, Vec<String>> {
        (**self).supported_urls().await
    }

    async fn do_generate(&self, options: CallOptions) -> Result<GenerateResponse> {
        (**self).do_generate(options).await
    }

    async fn do_stream(&self, options: CallOptions) -> Result<StreamResponse> {
        (**self).do_stream(options).await
    }
}

/// Complete response from a non-streaming generation request.
///
/// Contains all generated content, usage statistics, and provider metadata
/// from a single generation request. The response includes information about
/// why generation stopped and any warnings from the provider.
#[derive(Debug, Clone)]
pub struct GenerateResponse {
    /// Generated content parts (text, tool calls, etc.).
    pub content: Vec<Content>,

    /// Reason why generation terminated.
    ///
    /// Indicates whether generation stopped naturally, reached token limits,
    /// encountered a stop sequence, or was interrupted for other reasons.
    pub finish_reason: FinishReason,

    /// Token usage statistics for this generation.
    ///
    /// Includes prompt tokens, completion tokens, and total tokens consumed.
    /// Used for cost tracking and quota management.
    pub usage: Usage,

    /// Provider-specific metadata and additional information.
    ///
    /// May include provider-specific fields not covered by the standard response.
    pub provider_metadata: Option<SharedProviderMetadata>,

    /// Information about the request sent to the provider.
    ///
    /// Useful for debugging and auditing. May include the raw request body.
    pub request: Option<RequestInfo>,

    /// Information about the response from the provider.
    ///
    /// Contains response headers, raw body, provider-assigned IDs, and timestamps.
    pub response: Option<ResponseInfo>,

    /// Warnings issued by the provider during generation.
    ///
    /// Non-fatal issues such as unsupported parameters or deprecated features.
    /// The request proceeds despite these warnings.
    pub warnings: Vec<CallWarning>,
}

/// Response from a streaming generation request.
///
/// Contains an async stream that yields response parts as they arrive from the
/// provider. Each item in the stream represents a chunk of generated content,
/// tool calls, or metadata updates.
pub struct StreamResponse {
    /// Async stream of response parts.
    ///
    /// Yields `StreamPart` items containing text deltas, tool calls, usage updates,
    /// and finish reasons. The stream completes when generation finishes.
    pub stream: std::pin::Pin<Box<dyn Stream<Item = Result<StreamPart>> + Send>>,

    /// Information about the request sent to the provider.
    pub request: Option<RequestInfo>,

    /// Information about the initial response from the provider.
    ///
    /// Contains response headers and metadata available before streaming begins.
    pub response: Option<ResponseInfo>,
}

/// Information about the request sent to the provider.
///
/// Captures details about the outbound request for debugging and auditing purposes.
#[derive(Debug, Clone)]
pub struct RequestInfo {
    /// Raw request body sent to the provider's API.
    ///
    /// This is the JSON payload after conversion from `CallOptions` to the
    /// provider's specific request format.
    pub body: Option<serde_json::Value>,
}

/// Information about the response received from the provider.
///
/// Contains metadata and raw response data useful for debugging, logging,
/// and understanding provider behavior.
#[derive(Debug, Clone)]
pub struct ResponseInfo {
    /// HTTP response headers from the provider.
    ///
    /// May include rate limit information, request IDs, and other metadata.
    pub headers: Option<SharedHeaders>,

    /// Raw response body from the provider.
    ///
    /// The JSON payload before conversion to standardized response types.
    pub body: Option<serde_json::Value>,

    /// Unique identifier assigned by the provider to this response.
    ///
    /// Useful for support requests and correlating responses with provider logs.
    pub id: Option<String>,

    /// Timestamp when the provider processed this request.
    ///
    /// Format and timezone depend on the provider implementation.
    pub timestamp: Option<String>,

    /// Actual model used by the provider.
    ///
    /// May differ from the requested model if the provider performed
    /// automatic model selection or substitution.
    pub model_id: Option<String>,
}

#[cfg(test)]
mod tests {
    use super::*;

    struct DummyModel;

    #[async_trait]
    impl LanguageModel for DummyModel {
        fn provider(&self) -> &str {
            "test"
        }
        fn model_id(&self) -> &str {
            "dummy"
        }

        async fn do_generate(&self, _opts: CallOptions) -> Result<GenerateResponse> {
            unimplemented!()
        }

        async fn do_stream(&self, _opts: CallOptions) -> Result<StreamResponse> {
            unimplemented!()
        }
    }

    #[test]
    fn test_trait_implementation() {
        let model = DummyModel;
        assert_eq!(model.provider(), "test");
        assert_eq!(model.specification_version(), "v3");
    }
}