ai-sdk-core 0.3.0

//! Generate structured objects from language models.
//!
//! This module provides the `generate_object` function that generates validated,
//! schema-based outputs from language models.

use super::output_strategy::{OutputStrategy, ValidationContext, ValidationResult};
use super::GenerationObjectConfig;
use crate::error::GenerateObjectError;
use crate::retry::RetryPolicy;
use ai_sdk_provider::language_model::{
    CallOptions, CallWarning, Content, FinishReason, LanguageModel, Message, ResponseFormat,
    ResponseMetadata, TextPart, Usage,
};
use futures::future::BoxFuture;
use serde_json::Value;
use std::future::IntoFuture;
use std::sync::Arc;

/// Builder for generating structured objects.
pub struct GenerateObjectBuilder<M, P, S> {
    model: M,
    prompt: P,
    strategy: S,
    config: GenerationObjectConfig,
}

// 1. Initial State
impl GenerateObjectBuilder<(), (), ()> {
    /// Creates a new GenerateObjectBuilder
    pub fn new() -> Self {
        Self {
            model: (),
            prompt: (),
            strategy: (),
            config: GenerationObjectConfig::default(),
        }
    }
}

impl Default for GenerateObjectBuilder<(), (), ()> {
    fn default() -> Self {
        Self::new()
    }
}

// 2. Configuration Setters (Available in ANY state)
impl<M, P, S> GenerateObjectBuilder<M, P, S> {
    /// Set the schema name
    pub fn schema_name(mut self, name: impl Into<String>) -> Self {
        self.config.schema_name = Some(name.into());
        self
    }

    /// Set the schema description
    pub fn schema_description(mut self, description: impl Into<String>) -> Self {
        self.config.schema_description = Some(description.into());
        self
    }

    /// Set temperature (0.0 to 2.0)
    pub fn temperature(mut self, temperature: f32) -> Self {
        self.config.temperature = Some(temperature);
        self
    }

    /// Set maximum tokens to generate
    pub fn max_tokens(mut self, max_tokens: u32) -> Self {
        self.config.max_tokens = Some(max_tokens);
        self
    }

    /// Set retry policy
    pub fn retry_policy(mut self, retry_policy: RetryPolicy) -> Self {
        self.config.retry_policy = retry_policy;
        self
    }
}

// 3. State Transition: Set Model
impl<P, S> GenerateObjectBuilder<(), P, S> {
    /// Set the language model to use
    pub fn model<Mod: LanguageModel + 'static>(
        self,
        model: Mod,
    ) -> GenerateObjectBuilder<Arc<dyn LanguageModel>, P, S> {
        GenerateObjectBuilder {
            model: Arc::new(model),
            prompt: self.prompt,
            strategy: self.strategy,
            config: self.config,
        }
    }
}

// 4. State Transition: Set Prompt
impl<M, S> GenerateObjectBuilder<M, (), S> {
    /// Set the prompt from a string
    pub fn prompt(self, prompt: impl Into<String>) -> GenerateObjectBuilder<M, Vec<Message>, S> {
        let text = prompt.into();
        let messages = vec![Message::User {
            content: vec![ai_sdk_provider::language_model::UserContentPart::Text { text }],
        }];

        GenerateObjectBuilder {
            model: self.model,
            prompt: messages,
            strategy: self.strategy,
            config: self.config,
        }
    }

    /// Set the prompt from messages
    pub fn messages(self, messages: Vec<Message>) -> GenerateObjectBuilder<M, Vec<Message>, S> {
        GenerateObjectBuilder {
            model: self.model,
            prompt: messages,
            strategy: self.strategy,
            config: self.config,
        }
    }
}

// 5. State Transition: Set Output Strategy
impl<M, P> GenerateObjectBuilder<M, P, ()> {
    /// Set the output strategy
    pub fn output_strategy<Strat: OutputStrategy + 'static>(
        self,
        strategy: Strat,
    ) -> GenerateObjectBuilder<M, P, Arc<Strat>> {
        GenerateObjectBuilder {
            model: self.model,
            prompt: self.prompt,
            strategy: Arc::new(strategy),
            config: self.config,
        }
    }
}

// 6. Execution Logic (Requires all states set)
impl<S> GenerateObjectBuilder<Arc<dyn LanguageModel>, Vec<Message>, Arc<S>>
where
    S: OutputStrategy + 'static,
{
    /// Execute the object generation
    pub async fn execute(self) -> Result<GenerateObjectResult<S::Result>, GenerateObjectError> {
        let model = self.model;
        let messages = self.prompt;
        let strategy = self.strategy;
        let config = self.config;

        // Get schema from strategy
        let schema = strategy.json_schema().await;

        // Prepare call options
        let options = CallOptions {
            prompt: messages.clone().into(),
            response_format: Some(ResponseFormat::Json {
                schema: schema.clone(),
                name: config.schema_name,
                description: config.schema_description,
            }),
            temperature: config.temperature,
            max_output_tokens: config.max_tokens,
            tools: None, // No tools in object mode
            tool_choice: None,
            ..Default::default()
        };

        // Call model with retry
        let response = config
            .retry_policy
            .retry(|| {
                let model = model.clone();
                let options = options.clone();
                async move { model.do_generate(options).await }
            })
            .await?;

        // Extract text from response
        let text = extract_text_from_content(&response.content)
            .ok_or(GenerateObjectError::NoTextContent)?;

        // Parse JSON
        let value: Value = serde_json::from_str(&text)?;

        // Validate final result
        let context = ValidationContext {
            text: text.clone(),
            response: response.response.clone().map(|r| ResponseMetadata {
                id: r.id,
                timestamp: r.timestamp,
                model_id: r.model_id,
            }),
            usage: response.usage.clone(),
        };

        let validation = strategy.validate_final_result(Some(value), context).await;

        match validation {
            ValidationResult::Success { value: result, .. } => Ok(GenerateObjectResult {
                object: result,
                usage: response.usage.clone(),
                finish_reason: response.finish_reason,
                warnings: response.warnings.clone(),
                raw_response: response,
            }),
            ValidationResult::Failure { error, .. } => {
                Err(GenerateObjectError::ValidationFailed(error.to_string()))
            }
        }
    }
}

// 7. IntoFuture Implementation
impl<S> IntoFuture for GenerateObjectBuilder<Arc<dyn LanguageModel>, Vec<Message>, Arc<S>>
where
    S: OutputStrategy + 'static,
{
    type Output = Result<GenerateObjectResult<S::Result>, GenerateObjectError>;
    type IntoFuture = BoxFuture<'static, Self::Output>;

    fn into_future(self) -> Self::IntoFuture {
        Box::pin(self.execute())
    }
}

/// Result of generating a structured object.
#[derive(Debug)]
pub struct GenerateObjectResult<R> {
    /// The generated and validated object
    pub object: R,
    /// Token usage statistics
    pub usage: Usage,
    /// Finish reason
    pub finish_reason: FinishReason,
    /// Warnings from the provider
    pub warnings: Vec<CallWarning>,
    /// Raw response from the model
    pub raw_response: ai_sdk_provider::language_model::GenerateResponse,
}

/// Creates a new GenerateObjectBuilder
pub fn generate_object() -> GenerateObjectBuilder<(), (), ()> {
    GenerateObjectBuilder::new()
}

/// Extracts text content from a list of content items.
fn extract_text_from_content(content: &[Content]) -> Option<String> {
    let mut text = String::new();
    for item in content {
        if let Content::Text(TextPart { text: t, .. }) = item {
            text.push_str(t);
        }
    }
    if text.is_empty() {
        None
    } else {
        Some(text)
    }
}

#[cfg(test)]
mod tests {
    // Note: These tests would require a mock LanguageModel implementation
    // For now, they serve as documentation of the expected API
}