agent_chain_core/language_models/

chat_models.rs

//! Chat models for conversational AI.
//!
//! This module provides the base abstraction for chat models,
//! following the LangChain pattern of having a common interface
//! for different providers.
//!
//! Mirrors `langchain_core.language_models.chat_models`.

use std::collections::HashMap;
use std::pin::Pin;
use std::sync::Arc;

use async_trait::async_trait;
use futures::Stream;
use serde::{Deserialize, Serialize};
use serde_json::Value;

use super::base::{BaseLanguageModel, LangSmithParams, LanguageModelConfig, LanguageModelInput};
use super::model_profile::ModelProfile;
use crate::GenerationType;
use crate::callbacks::{
    AsyncCallbackManagerForLLMRun, BaseCallbackHandler, CallbackManagerForLLMRun, Callbacks,
};
use crate::error::{Error, Result};
use crate::messages::{AIMessage, AIMessageChunk, BaseMessage, ChunkPosition, UsageMetadata};
use crate::outputs::{ChatGeneration, ChatGenerationChunk, ChatResult, Generation, LLMResult};
use crate::rate_limiters::BaseRateLimiter;
use crate::tools::{BaseTool, ToolDefinition};

/// Type alias for streaming output.
pub type ChatStream = Pin<Box<dyn Stream<Item = Result<ChatChunk>> + Send>>;

/// Type alias for a streaming chat generation output.
pub type ChatGenerationStream = Pin<Box<dyn Stream<Item = Result<ChatGenerationChunk>> + Send>>;

/// Type alias for streaming AIMessageChunk output.
pub type AIMessageChunkStream = Pin<Box<dyn Stream<Item = Result<AIMessageChunk>> + Send>>;

/// A chunk of output from streaming.
///
/// This struct carries content deltas during streaming, along with optional
/// metadata that is typically attached to the final chunk.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatChunk {
    /// The content delta.
    pub content: String,
    /// Whether this is the final chunk.
    pub is_final: bool,
    /// Usage metadata (token counts) - typically present on the final chunk.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub usage_metadata: Option<UsageMetadata>,
    /// The reason the model stopped generating (e.g., "stop", "length", "tool_calls").
    /// Typically present on the final chunk.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub finish_reason: Option<String>,
}

impl ChatChunk {
    /// Create a new content chunk (non-final).
    pub fn new(content: impl Into<String>) -> Self {
        Self {
            content: content.into(),
            is_final: false,
            usage_metadata: None,
            finish_reason: None,
        }
    }

    /// Create a final chunk with optional metadata.
    pub fn final_chunk(
        usage_metadata: Option<UsageMetadata>,
        finish_reason: Option<String>,
    ) -> Self {
        Self {
            content: String::new(),
            is_final: true,
            usage_metadata,
            finish_reason,
        }
    }

    /// Set usage metadata on this chunk.
    pub fn with_usage_metadata(mut self, usage: UsageMetadata) -> Self {
        self.usage_metadata = Some(usage);
        self
    }

    /// Set finish reason on this chunk.
    pub fn with_finish_reason(mut self, reason: impl Into<String>) -> Self {
        self.finish_reason = Some(reason.into());
        self
    }
}

/// Configuration for tool choice.
///
/// Mirrors Python's tool_choice parameter patterns.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(untagged)]
pub enum ToolChoice {
    /// String value like "auto", "any", "none", or a specific tool name.
    String(String),
    /// Structured tool choice with type and optional name.
    Structured {
        /// Type of tool choice.
        #[serde(rename = "type")]
        choice_type: String,
        /// Optional tool name.
        #[serde(skip_serializing_if = "Option::is_none")]
        name: Option<String>,
    },
}

impl ToolChoice {
    /// Create an "auto" tool choice - let the model decide.
    pub fn auto() -> Self {
        ToolChoice::String("auto".to_string())
    }

    /// Create an "any" tool choice - model must use at least one tool.
    pub fn any() -> Self {
        ToolChoice::String("any".to_string())
    }

    /// Create a "none" tool choice - model should not use any tools.
    pub fn none() -> Self {
        ToolChoice::String("none".to_string())
    }

    /// Create a tool choice for a specific tool by name.
    pub fn tool(name: impl Into<String>) -> Self {
        ToolChoice::Structured {
            choice_type: "tool".to_string(),
            name: Some(name.into()),
        }
    }
}

/// Disable streaming options.
///
/// Mirrors Python's `disable_streaming: bool | Literal["tool_calling"]` field.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum DisableStreaming {
    /// Boolean value: true = always disable, false = never disable.
    Bool(bool),
    /// Literal "tool_calling": disable only when tools are present.
    ToolCalling,
}

impl Default for DisableStreaming {
    fn default() -> Self {
        DisableStreaming::Bool(false)
    }
}

impl DisableStreaming {
    /// Check if streaming should be bypassed.
    ///
    /// # Arguments
    ///
    /// * `has_tools` - Whether tools are present in the current call.
    pub fn should_disable(&self, has_tools: bool) -> bool {
        match self {
            DisableStreaming::Bool(b) => *b,
            DisableStreaming::ToolCalling => has_tools,
        }
    }
}

impl From<bool> for DisableStreaming {
    fn from(b: bool) -> Self {
        DisableStreaming::Bool(b)
    }
}

/// Configuration specific to chat models.
#[derive(Clone, Default)]
pub struct ChatModelConfig {
    /// Base language model configuration.
    pub base: LanguageModelConfig,

    /// Rate limiter for limiting API requests.
    pub rate_limiter: Option<Arc<dyn BaseRateLimiter>>,

    /// Whether to disable streaming for this model.
    ///
    /// If streaming is bypassed, then `stream`/`astream` will defer to `invoke`/`ainvoke`.
    ///
    /// - If `Bool(true)`, will always bypass streaming case.
    /// - If `ToolCalling`, will bypass streaming case only when tools are present.
    /// - If `Bool(false)` (default), will always use streaming case if available.
    pub disable_streaming: DisableStreaming,

    /// Version of `AIMessage` output format.
    ///
    /// - `"v0"`: provider-specific format in content
    /// - `"v1"`: standardized format in content
    ///
    /// Can also be set via `LC_OUTPUT_VERSION` environment variable.
    pub output_version: Option<String>,

    /// Profile detailing model capabilities.
    pub profile: Option<ModelProfile>,
}

impl std::fmt::Debug for ChatModelConfig {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ChatModelConfig")
            .field("base", &self.base)
            .field(
                "rate_limiter",
                &self.rate_limiter.as_ref().map(|_| "<rate_limiter>"),
            )
            .field("disable_streaming", &self.disable_streaming)
            .field("output_version", &self.output_version)
            .field("profile", &self.profile)
            .finish()
    }
}

impl ChatModelConfig {
    /// Create a new chat model configuration.
    pub fn new() -> Self {
        Self::default()
    }

    /// Set the rate limiter.
    pub fn with_rate_limiter(mut self, rate_limiter: Arc<dyn BaseRateLimiter>) -> Self {
        self.rate_limiter = Some(rate_limiter);
        self
    }

    /// Set whether to disable streaming.
    pub fn with_disable_streaming(mut self, disable: impl Into<DisableStreaming>) -> Self {
        self.disable_streaming = disable.into();
        self
    }

    /// Set the output version.
    pub fn with_output_version(mut self, version: impl Into<String>) -> Self {
        self.output_version = Some(version.into());
        self
    }

    /// Set the model profile.
    pub fn with_profile(mut self, profile: ModelProfile) -> Self {
        self.profile = Some(profile);
        self
    }

    /// Enable caching.
    pub fn with_cache(mut self, cache: bool) -> Self {
        self.base.cache = Some(cache);
        self
    }

    /// Enable verbose mode.
    pub fn with_verbose(mut self, verbose: bool) -> Self {
        self.base.verbose = verbose;
        self
    }

    /// Set tags.
    pub fn with_tags(mut self, tags: Vec<String>) -> Self {
        self.base.tags = Some(tags);
        self
    }

    /// Set metadata.
    pub fn with_metadata(mut self, metadata: HashMap<String, Value>) -> Self {
        self.base.metadata = Some(metadata);
        self
    }
}

/// Base trait for all chat models.
///
/// This trait follows the LangChain pattern where each provider implements
/// the core generation methods. The trait provides both sync-style (via async)
/// and streaming interfaces.
///
/// # Implementation Guide
///
/// Custom chat model implementations should override these methods:
///
/// | Method/Property           | Description                                        | Required |
/// |--------------------------|----------------------------------------------------|---------:|
/// | `_generate`              | Use to generate a chat result from messages        | Required |
/// | `_llm_type` (property)   | Used to uniquely identify the type of the model    | Required |
/// | `_identifying_params`    | Represent model parameterization for tracing       | Optional |
/// | `_stream`                | Use to implement streaming                         | Optional |
/// | `_agenerate`             | Use to implement a native async method             | Optional |
/// | `_astream`               | Use to implement async version of `_stream`        | Optional |
#[async_trait]
pub trait BaseChatModel: BaseLanguageModel {
    /// Get the chat model configuration.
    fn chat_config(&self) -> &ChatModelConfig;

    /// Get the model profile, if available.
    fn profile(&self) -> Option<&ModelProfile> {
        self.chat_config().profile.as_ref()
    }

    /// Core abstract method to generate a chat result.
    ///
    /// Implementations must override this method.
    ///
    /// # Arguments
    ///
    /// * `messages` - The messages to generate from.
    /// * `stop` - Optional list of stop words to use when generating.
    /// * `run_manager` - Optional callback manager to use for this call.
    ///
    /// # Returns
    ///
    /// The output chat result containing generations.
    async fn _generate(
        &self,
        messages: Vec<BaseMessage>,
        stop: Option<Vec<String>>,
        run_manager: Option<&CallbackManagerForLLMRun>,
    ) -> Result<ChatResult>;

    /// Async version of `_generate`.
    ///
    /// Default implementation calls `_generate`.
    async fn _agenerate(
        &self,
        messages: Vec<BaseMessage>,
        stop: Option<Vec<String>>,
        _run_manager: Option<&AsyncCallbackManagerForLLMRun>,
    ) -> Result<ChatResult> {
        self._generate(messages, stop, None).await
    }

    /// Stream the output of the model.
    ///
    /// Default implementation raises NotImplementedError.
    ///
    /// # Arguments
    ///
    /// * `messages` - The messages to generate from.
    /// * `stop` - Optional list of stop words to use when generating.
    /// * `run_manager` - Optional callback manager to use for this call.
    ///
    /// # Yields
    ///
    /// The chat generation chunks.
    fn _stream(
        &self,
        _messages: Vec<BaseMessage>,
        _stop: Option<Vec<String>>,
        _run_manager: Option<&CallbackManagerForLLMRun>,
    ) -> Result<ChatGenerationStream> {
        Err(Error::NotImplemented("Streaming not implemented".into()))
    }

    /// Async stream the output of the model.
    ///
    /// Default implementation calls `_stream`.
    async fn _astream(
        &self,
        messages: Vec<BaseMessage>,
        stop: Option<Vec<String>>,
        _run_manager: Option<&AsyncCallbackManagerForLLMRun>,
    ) -> Result<ChatGenerationStream> {
        self._stream(messages, stop, None)
    }

    /// Get the first AI message from a chat result.
    ///
    /// Helper method to extract the first generation's message as an AIMessage.
    fn get_first_message(&self, result: &ChatResult) -> Result<AIMessage> {
        if result.generations.is_empty() {
            return Err(Error::Other("No generations returned".into()));
        }

        match result.generations[0].message.clone() {
            BaseMessage::AI(message) => Ok(message),
            other => Ok(AIMessage::new(other.content())),
        }
    }

    /// Combine LLM outputs from multiple results.
    ///
    /// This method is called after generating results from multiple prompts
    /// to combine any LLM-specific output information.
    ///
    /// Default implementation returns an empty HashMap.
    /// Subclasses can override to combine provider-specific output data.
    fn _combine_llm_outputs(
        &self,
        _llm_outputs: &[Option<HashMap<String, Value>>],
    ) -> HashMap<String, Value> {
        HashMap::new()
    }

    /// Convert cached Generation objects to ChatGeneration objects.
    ///
    /// Handle case where cache contains Generation objects instead of
    /// ChatGeneration objects. This can happen due to serialization/deserialization
    /// issues or legacy cache data.
    fn _convert_cached_generations(&self, cache_val: Vec<Generation>) -> Vec<ChatGeneration> {
        cache_val
            .into_iter()
            .map(|cached_gen| {
                // Convert Generation to ChatGeneration by creating AIMessage from text
                let message = AIMessage::new(&cached_gen.text);
                match cached_gen.generation_info {
                    Some(info) => ChatGeneration::with_info(message.into(), info),
                    None => ChatGeneration::new(message.into()),
                }
            })
            .collect()
    }

    /// Get invocation parameters for tracing.
    ///
    /// Returns a HashMap containing the model configuration and stop sequences.
    fn _get_invocation_params(
        &self,
        stop: Option<&[String]>,
        kwargs: Option<&HashMap<String, Value>>,
    ) -> HashMap<String, Value> {
        let mut params = self.get_identifying_params();
        if let Some(stop) = stop {
            params.insert(
                "stop".to_string(),
                Value::Array(stop.iter().map(|s| Value::String(s.clone())).collect()),
            );
        }
        if let Some(kw) = kwargs {
            params.extend(kw.clone());
        }
        params
    }

    /// Get the LLM string for cache key generation.
    ///
    /// This string uniquely identifies the model configuration for caching purposes.
    fn _get_llm_string(
        &self,
        stop: Option<&[String]>,
        kwargs: Option<&HashMap<String, Value>>,
    ) -> String {
        let params = self._get_invocation_params(stop, kwargs);

        // Sort params for deterministic key
        let mut sorted_items: Vec<_> = params.iter().collect();
        sorted_items.sort_by_key(|(k, _)| *k);

        format!("{:?}", sorted_items)
    }

    /// Check if `_stream` is implemented (not the default).
    ///
    /// This is used by `_should_stream` to determine if streaming is available.
    /// Implementations that override `_stream` should also override this to return `true`.
    fn has_stream_impl(&self) -> bool {
        false
    }

    /// Check if `_astream` is implemented (not the default).
    ///
    /// This is used by `_should_stream` to determine if async streaming is available.
    /// Implementations that override `_astream` should also override this to return `true`.
    fn has_astream_impl(&self) -> bool {
        false
    }

    /// Check if streaming is enabled via a model field.
    ///
    /// Override this if the model has a `streaming` field that should be checked.
    fn has_streaming_field(&self) -> Option<bool> {
        None
    }

    /// Determine if a given model call should hit the streaming API.
    ///
    /// This method mirrors Python's `_should_stream` behavior:
    /// 1. Check if streaming is implemented (either sync or async)
    /// 2. Check if streaming has been disabled on this instance
    /// 3. Check if streaming is disabled for tool calling and tools are present
    /// 4. Check if streaming field is set on the model
    /// 5. Check if any streaming callback handlers are present
    ///
    /// # Arguments
    ///
    /// * `async_api` - Whether this is an async API call
    /// * `has_tools` - Whether tools are present in the call
    /// * `stream_kwarg` - Optional explicit stream kwarg from caller
    /// * `run_manager` - Optional callback manager for checking streaming handlers
    ///
    /// # Returns
    ///
    /// `true` if streaming should be used, `false` otherwise.
    fn _should_stream(
        &self,
        async_api: bool,
        has_tools: bool,
        stream_kwarg: Option<bool>,
        run_manager: Option<&[Arc<dyn BaseCallbackHandler>]>,
    ) -> bool {
        // Check if streaming is implemented
        let sync_not_implemented = !self.has_stream_impl();
        let async_not_implemented = !self.has_astream_impl();

        // Check if streaming is implemented
        if !async_api && sync_not_implemented {
            return false;
        }
        // Note: since async falls back to sync, we check both here
        if async_api && async_not_implemented && sync_not_implemented {
            return false;
        }

        // Check if streaming has been disabled on this instance
        if self
            .chat_config()
            .disable_streaming
            .should_disable(has_tools)
        {
            return false;
        }

        // Check if a runtime streaming flag has been passed in
        if let Some(stream) = stream_kwarg {
            return stream;
        }

        // Check if streaming field is set on the model
        if let Some(streaming) = self.has_streaming_field() {
            return streaming;
        }

        // Check if any streaming callback handlers are present
        if let Some(handlers) = run_manager {
            // In Python, this checks for `_StreamingCallbackHandler` instances
            // For Rust, we can check if any handler implements StreamingCallbackHandler
            // This is a simplified check - in practice, you'd want a more sophisticated check
            if !handlers.is_empty() {
                // If there are any handlers, we assume streaming might be wanted
                return true;
            }
        }

        // Default: streaming is available and not disabled
        true
    }

    /// Generate from a batch of message lists.
    ///
    /// This method should make use of batched calls for models that expose a batched API.
    ///
    /// Use this method when you want to:
    /// 1. Take advantage of batched calls
    /// 2. Need more output from the model than just the top generated value
    /// 3. Are building chains that are agnostic to the underlying language model type
    ///
    /// # Arguments
    ///
    /// * `messages` - List of message lists.
    /// * `stop` - Stop words to use when generating.
    /// * `callbacks` - Callbacks to pass through.
    ///
    /// # Returns
    ///
    /// An `LLMResult` containing a list of candidate `ChatGeneration` objects.
    async fn generate(
        &self,
        messages: Vec<Vec<BaseMessage>>,
        stop: Option<Vec<String>>,
        _callbacks: Option<Callbacks>,
    ) -> Result<LLMResult> {
        let mut all_generations: Vec<Vec<GenerationType>> = Vec::new();

        for message_list in messages {
            let result = self._generate(message_list, stop.clone(), None).await?;
            all_generations.push(result.generations.into_iter().map(|e| e.into()).collect());
        }

        Ok(LLMResult::new(all_generations))
    }

    /// Async version of `generate`.
    async fn agenerate(
        &self,
        messages: Vec<Vec<BaseMessage>>,
        stop: Option<Vec<String>>,
        _callbacks: Option<Callbacks>,
    ) -> Result<LLMResult> {
        let mut all_generations: Vec<Vec<GenerationType>> = Vec::new();

        for message_list in messages {
            let result = self._agenerate(message_list, stop.clone(), None).await?;
            all_generations.push(result.generations.into_iter().map(|e| e.into()).collect());
        }

        Ok(LLMResult::new(all_generations))
    }

    /// Async call helper.
    ///
    /// This is a convenience method that wraps `agenerate` for single-message calls.
    async fn _call_async(
        &self,
        messages: Vec<BaseMessage>,
        stop: Option<Vec<String>>,
        callbacks: Option<Callbacks>,
    ) -> Result<BaseMessage> {
        let result = self.agenerate(vec![messages], stop, callbacks).await?;

        if result.generations.is_empty() || result.generations[0].is_empty() {
            return Err(Error::Other("No generations returned".into()));
        }

        match &result.generations[0][0] {
            GenerationType::ChatGeneration(chat_gen) => Ok(chat_gen.message.clone()),
            _ => Err(Error::Other("Unexpected generation type".into())),
        }
    }

    /// Generate a response from the model with tools.
    ///
    /// This is the preferred method when tool calling is needed.
    /// Default implementation ignores tools and calls `_generate`.
    ///
    /// # Arguments
    ///
    /// * `messages` - The conversation history.
    /// * `tools` - Tool definitions for the model to use.
    /// * `tool_choice` - Optional configuration for tool selection.
    /// * `stop` - Optional stop sequences.
    ///
    /// # Returns
    ///
    /// An `AIMessage` containing the generated response.
    async fn generate_with_tools(
        &self,
        messages: Vec<BaseMessage>,
        _tools: &[ToolDefinition],
        _tool_choice: Option<&ToolChoice>,
        stop: Option<Vec<String>>,
    ) -> Result<AIMessage> {
        let result = self._generate(messages, stop, None).await?;

        if result.generations.is_empty() {
            return Err(Error::Other("No generations returned".into()));
        }

        match result.generations[0].message.clone() {
            BaseMessage::AI(message) => Ok(message),
            _ => Err(Error::Other("Unexpected message type".into())),
        }
    }

    /// Convert input to messages.
    fn convert_input(&self, input: LanguageModelInput) -> Result<Vec<BaseMessage>> {
        Ok(input.to_messages())
    }

    /// Invoke the model with input.
    async fn invoke(&self, input: LanguageModelInput) -> Result<AIMessage> {
        let messages = self.convert_input(input)?;
        let result = self._generate(messages, None, None).await?;

        if result.generations.is_empty() {
            return Err(Error::Other("No generations returned".into()));
        }

        match result.generations[0].message.clone() {
            BaseMessage::AI(message) => Ok(message),
            _ => Err(Error::Other("Unexpected message type".into())),
        }
    }

    /// Async invoke the model.
    async fn ainvoke(&self, input: LanguageModelInput) -> Result<AIMessage> {
        let messages = self.convert_input(input)?;
        let result = self._agenerate(messages, None, None).await?;

        if result.generations.is_empty() {
            return Err(Error::Other("No generations returned".into()));
        }

        match result.generations[0].message.clone() {
            BaseMessage::AI(message) => Ok(message),
            _ => Err(Error::Other("Unexpected message type".into())),
        }
    }

    /// Bind tools to the model.
    ///
    /// This method returns a Runnable that can call tools. The default
    /// implementation raises NotImplementedError.
    ///
    /// # Arguments
    ///
    /// * `tools` - Sequence of tools to bind to the model.
    /// * `tool_choice` - Optional tool choice configuration.
    ///
    /// # Returns
    ///
    /// A Result with error indicating tools are not supported.
    ///
    /// # Note
    ///
    /// Provider implementations should override this method to return a configured model.
    fn bind_tools(
        &self,
        _tools: &[Arc<dyn BaseTool>],
        _tool_choice: Option<ToolChoice>,
    ) -> Result<()> {
        Err(Error::NotImplemented(
            "bind_tools is not implemented for this model".into(),
        ))
    }

    /// Get tool definitions from tools.
    ///
    /// Helper method to convert tools to their definitions.
    fn get_tool_definitions(&self, tools: &[Arc<dyn BaseTool>]) -> Vec<ToolDefinition> {
        tools.iter().map(|t| t.definition()).collect()
    }

    /// Generate a streaming response from the model.
    ///
    /// This is the main streaming API. It yields `AIMessageChunk`s.
    /// Providers should override `_stream` for native streaming support.
    ///
    /// # Arguments
    ///
    /// * `input` - The input to the model (string, messages, or PromptValue).
    /// * `stop` - Optional stop sequences.
    ///
    /// # Returns
    ///
    /// A stream of `AIMessageChunk`s.
    async fn stream(
        &self,
        input: LanguageModelInput,
        stop: Option<Vec<String>>,
    ) -> Result<AIMessageChunkStream> {
        let messages = self.convert_input(input)?;
        let has_tools = false;

        // Check if streaming should be used
        if !self._should_stream(false, has_tools, Some(true), None) {
            // Model doesn't implement streaming or streaming is disabled,
            // fall back to invoke and yield the result as a single chunk
            let result = self._generate(messages, stop, None).await?;
            let message = self.get_first_message(&result)?;
            let chunk = AIMessageChunk::new(message.content());
            return Ok(Box::pin(futures::stream::once(async move { Ok(chunk) })));
        }

        // Acquire rate limiter if configured (blocking until token available)
        if let Some(ref rate_limiter) = self.chat_config().rate_limiter {
            rate_limiter.acquire(true);
        }

        // Use the _stream method and convert ChatGenerationChunk to AIMessageChunk
        let generation_stream = self._stream(messages, stop, None)?;

        // Transform the stream to yield AIMessageChunk instead of ChatGenerationChunk
        let chunk_stream = async_stream::stream! {
            use futures::StreamExt;

            let mut pinned_stream = generation_stream;
            let mut yielded = false;

            while let Some(result) = pinned_stream.next().await {
                match result {
                    Ok(generation_chunk) => {
                        // Extract AIMessageChunk from the generation chunk
                        let ai_chunk = match generation_chunk.message {
                            BaseMessage::AI(ai_msg) => AIMessageChunk::new(ai_msg.content()),
                            other => AIMessageChunk::new(other.content()),
                        };
                        yielded = true;
                        yield Ok(ai_chunk);
                    }
                    Err(e) => {
                        yield Err(e);
                        return;
                    }
                }
            }

            // Yield a final empty chunk with chunk_position="last" if we yielded anything
            if yielded {
                let mut final_chunk = AIMessageChunk::new("");
                final_chunk.set_chunk_position(Some(ChunkPosition::Last));
                yield Ok(final_chunk);
            }
        };

        Ok(Box::pin(chunk_stream))
    }

    /// Async stream the model output.
    ///
    /// This is the async version of `stream`. It yields `AIMessageChunk`s.
    /// Providers should override `_astream` for native async streaming support.
    ///
    /// # Arguments
    ///
    /// * `input` - The input to the model (string, messages, or PromptValue).
    /// * `stop` - Optional stop sequences.
    ///
    /// # Returns
    ///
    /// A stream of `AIMessageChunk`s.
    async fn astream(
        &self,
        input: LanguageModelInput,
        stop: Option<Vec<String>>,
    ) -> Result<AIMessageChunkStream> {
        let messages = self.convert_input(input)?;
        let has_tools = false;

        // Check if streaming should be used
        if !self._should_stream(true, has_tools, Some(true), None) {
            // No async or sync stream is implemented, fall back to ainvoke
            let result = self._agenerate(messages, stop, None).await?;
            let message = self.get_first_message(&result)?;
            let chunk = AIMessageChunk::new(message.content());
            return Ok(Box::pin(futures::stream::once(async move { Ok(chunk) })));
        }

        // Acquire rate limiter if configured (blocking until token available)
        if let Some(ref rate_limiter) = self.chat_config().rate_limiter {
            rate_limiter.aacquire(true).await;
        }

        // Use the _astream method
        let generation_stream = self._astream(messages, stop, None).await?;

        // Transform the stream to yield AIMessageChunk instead of ChatGenerationChunk
        let chunk_stream = async_stream::stream! {
            use futures::StreamExt;

            let mut pinned_stream = generation_stream;
            let mut yielded = false;

            while let Some(result) = pinned_stream.next().await {
                match result {
                    Ok(generation_chunk) => {
                        // Extract AIMessageChunk from the generation chunk
                        let ai_chunk = match generation_chunk.message {
                            BaseMessage::AI(ai_msg) => AIMessageChunk::new(ai_msg.content()),
                            other => AIMessageChunk::new(other.content()),
                        };
                        yielded = true;
                        yield Ok(ai_chunk);
                    }
                    Err(e) => {
                        yield Err(e);
                        return;
                    }
                }
            }

            // Yield a final empty chunk with chunk_position="last" if we yielded anything
            if yielded {
                let mut final_chunk = AIMessageChunk::new("");
                final_chunk.set_chunk_position(Some(ChunkPosition::Last));
                yield Ok(final_chunk);
            }
        };

        Ok(Box::pin(chunk_stream))
    }

    /// Stream ChatGenerationChunk objects from the model.
    ///
    /// This is a lower-level streaming API that yields `ChatGenerationChunk`s directly.
    /// Most users should use `stream()` or `astream()` instead.
    ///
    /// # Arguments
    ///
    /// * `messages` - The conversation history.
    /// * `stop` - Optional stop sequences.
    /// * `run_manager` - Optional callback manager for the run.
    ///
    /// # Returns
    ///
    /// A stream of `ChatGenerationChunk`s.
    async fn stream_generations(
        &self,
        messages: Vec<BaseMessage>,
        stop: Option<Vec<String>>,
        run_manager: Option<&CallbackManagerForLLMRun>,
    ) -> Result<ChatGenerationStream> {
        let has_tools = false;

        // Check if streaming should be used
        if !self._should_stream(false, has_tools, None, None) {
            // Fall back to non-streaming
            let result = self._generate(messages, stop, run_manager).await?;
            if result.generations.is_empty() {
                return Err(Error::Other("No generations returned".into()));
            }

            let message = result.generations[0].message.clone();
            let chunk = ChatGenerationChunk::new(message);
            return Ok(Box::pin(futures::stream::once(async move { Ok(chunk) })));
        }

        // Try to use streaming
        self._stream(messages, stop, run_manager)
    }

    /// Get standard params for tracing.
    fn get_chat_ls_params(&self, stop: Option<&[String]>) -> LangSmithParams {
        let mut params = self.get_ls_params(stop);
        params.ls_model_type = Some("chat".to_string());
        params
    }

    /// Get a dictionary representation of the model.
    ///
    /// Returns identifying parameters plus the model type.
    fn to_dict(&self) -> HashMap<String, Value> {
        let mut result = self.get_identifying_params();
        result.insert(
            "_type".to_string(),
            Value::String(self.llm_type().to_string()),
        );
        result
    }

    /// Create a wrapper that structures model output using a schema.
    ///
    /// This method returns a Runnable that formats outputs to match the given schema.
    /// The default implementation raises NotImplementedError.
    ///
    /// # Arguments
    ///
    /// * `schema` - The output schema (as a JSON value).
    /// * `include_raw` - If true, include raw model response in output.
    ///
    /// # Returns
    ///
    /// A Result with error indicating structured output is not supported.
    ///
    /// # Note
    ///
    /// Provider implementations should override `bind_tools` first, as the default
    /// implementation uses `bind_tools` internally.
    fn with_structured_output(&self, _schema: Value, _include_raw: bool) -> Result<()> {
        Err(Error::NotImplemented(
            "with_structured_output is not implemented for this model".into(),
        ))
    }

    /// Get the identifying parameters for this model.
    ///
    /// Returns a map of parameters that uniquely identify this model instance.
    fn get_identifying_params(&self) -> HashMap<String, Value> {
        let mut params = HashMap::new();
        params.insert(
            "_type".to_string(),
            Value::String(self.llm_type().to_string()),
        );
        params.insert(
            "model".to_string(),
            Value::String(self.model_name().to_string()),
        );
        params
    }
}

/// Simplified implementation for a chat model to inherit from.
///
/// This implementation is primarily here for backwards compatibility.
/// For new implementations, please use `BaseChatModel` directly.
#[async_trait]
pub trait SimpleChatModel: BaseChatModel {
    /// Simple call method that takes messages and returns a string.
    ///
    /// Implementations should override this method.
    async fn _call(
        &self,
        messages: Vec<BaseMessage>,
        stop: Option<Vec<String>>,
        run_manager: Option<&CallbackManagerForLLMRun>,
    ) -> Result<String>;
}

#[async_trait]
impl<T: SimpleChatModel> BaseChatModel for T {
    fn chat_config(&self) -> &ChatModelConfig {
        <T as BaseChatModel>::chat_config(self)
    }

    async fn _generate(
        &self,
        messages: Vec<BaseMessage>,
        stop: Option<Vec<String>>,
        run_manager: Option<&CallbackManagerForLLMRun>,
    ) -> Result<ChatResult> {
        let output_str = self._call(messages, stop, run_manager).await?;
        let message = AIMessage::new(output_str);
        let generation = ChatGeneration::new(message.into());
        Ok(ChatResult::new(vec![generation]))
    }
}

/// Generate from a stream of chunks.
///
/// Collects all chunks from the stream and generates a final ChatResult.
///
/// This corresponds to `generate_from_stream` in LangChain Python.
///
/// # Arguments
///
/// * `stream` - An iterator of `ChatGenerationChunk` objects.
///
/// # Returns
///
/// A `ChatResult` containing the merged generation.
///
/// # Errors
///
/// Returns an error if no generations are found in the stream.
pub fn generate_from_stream<I>(mut stream: I) -> Result<ChatResult>
where
    I: Iterator<Item = ChatGenerationChunk>,
{
    let first = stream.next();
    if first.is_none() {
        return Err(Error::Other("No generations found in stream.".into()));
    }

    let mut generation = first.unwrap();

    // Merge remaining chunks
    for chunk in stream {
        generation = generation + chunk;
    }

    // Convert ChatGenerationChunk to ChatGeneration
    let chat_generation: ChatGeneration = generation.into();
    Ok(ChatResult::new(vec![chat_generation]))
}

/// Async generate from a stream of chunks.
///
/// Collects all chunks from an async stream and generates a final ChatResult.
///
/// This corresponds to `agenerate_from_stream` in LangChain Python.
///
/// # Arguments
///
/// * `stream` - An async stream of `ChatGenerationChunk` objects.
///
/// # Returns
///
/// A `ChatResult` containing the merged generation.
///
/// # Errors
///
/// Returns an error if no generations are found in the stream.
pub async fn agenerate_from_stream(
    stream: impl futures::Stream<Item = Result<ChatGenerationChunk>> + Unpin,
) -> Result<ChatResult> {
    use futures::StreamExt;

    let chunks: Vec<ChatGenerationChunk> = stream
        .filter_map(|result| async { result.ok() })
        .collect()
        .await;

    if chunks.is_empty() {
        return Err(Error::Other("No generations found in stream.".into()));
    }

    generate_from_stream(chunks.into_iter())
}

/// Collect a stream of ChatGenerationChunks and merge them.
///
/// This is a convenience function that collects all chunks from a stream
/// and returns the merged result.
///
/// # Arguments
///
/// * `stream` - An async stream of `ChatGenerationChunk` results.
///
/// # Returns
///
/// The merged `ChatGenerationChunk`, or `None` if the stream was empty.
pub async fn collect_and_merge_stream(
    mut stream: impl futures::StreamExt<Item = Result<ChatGenerationChunk>> + Unpin,
) -> Result<Option<ChatGenerationChunk>> {
    let mut chunks = Vec::new();
    while let Some(chunk_result) = stream.next().await {
        chunks.push(chunk_result?);
    }

    if chunks.is_empty() {
        return Ok(None);
    }

    Ok(crate::outputs::merge_chat_generation_chunks(chunks))
}

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

    #[test]
    fn test_chat_model_config_builder() {
        let config = ChatModelConfig::new()
            .with_cache(true)
            .with_verbose(true)
            .with_disable_streaming(true)
            .with_output_version("v1");

        assert_eq!(config.base.cache, Some(true));
        assert!(config.base.verbose);
        assert_eq!(config.disable_streaming, DisableStreaming::Bool(true));
        assert_eq!(config.output_version, Some("v1".to_string()));
    }

    #[test]
    fn test_tool_choice_auto() {
        let choice = ToolChoice::auto();
        assert_eq!(choice, ToolChoice::String("auto".to_string()));
    }

    #[test]
    fn test_tool_choice_any() {
        let choice = ToolChoice::any();
        assert_eq!(choice, ToolChoice::String("any".to_string()));
    }

    #[test]
    fn test_tool_choice_none() {
        let choice = ToolChoice::none();
        assert_eq!(choice, ToolChoice::String("none".to_string()));
    }

    #[test]
    fn test_tool_choice_tool() {
        let choice = ToolChoice::tool("my_tool");
        assert_eq!(
            choice,
            ToolChoice::Structured {
                choice_type: "tool".to_string(),
                name: Some("my_tool".to_string()),
            }
        );
    }

    #[test]
    fn test_tool_choice_serialization() {
        let auto = ToolChoice::auto();
        let json = serde_json::to_string(&auto).unwrap();
        assert_eq!(json, "\"auto\"");

        let tool = ToolChoice::tool("my_tool");
        let json = serde_json::to_string(&tool).unwrap();
        assert!(json.contains("my_tool"));
        assert!(json.contains("tool"));
    }

    #[test]
    fn test_disable_streaming() {
        let bool_false = DisableStreaming::Bool(false);
        assert!(!bool_false.should_disable(true));
        assert!(!bool_false.should_disable(false));

        let bool_true = DisableStreaming::Bool(true);
        assert!(bool_true.should_disable(true));
        assert!(bool_true.should_disable(false));

        let tool_calling = DisableStreaming::ToolCalling;
        assert!(tool_calling.should_disable(true));
        assert!(!tool_calling.should_disable(false));
    }

    #[test]
    fn test_generate_from_stream() {
        let chunks = vec![
            ChatGenerationChunk::new(AIMessage::new("Hello, ").into()),
            ChatGenerationChunk::new(AIMessage::new("world!").into()),
        ];

        let result = generate_from_stream(chunks.into_iter()).unwrap();
        assert_eq!(result.generations.len(), 1);
        assert_eq!(result.generations[0].message.content(), "Hello, world!");
    }

    #[test]
    fn test_generate_from_stream_empty() {
        let chunks: Vec<ChatGenerationChunk> = vec![];
        let result = generate_from_stream(chunks.into_iter());
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_agenerate_from_stream() {
        let chunks = vec![
            Ok(ChatGenerationChunk::new(AIMessage::new("Hello, ").into())),
            Ok(ChatGenerationChunk::new(AIMessage::new("world!").into())),
        ];

        let stream = futures::stream::iter(chunks);
        let result = agenerate_from_stream(stream).await.unwrap();
        assert_eq!(result.generations.len(), 1);
        assert_eq!(result.generations[0].message.content(), "Hello, world!");
    }

    #[tokio::test]
    async fn test_collect_and_merge_stream() {
        let chunks = vec![
            Ok(ChatGenerationChunk::new(AIMessage::new("a").into())),
            Ok(ChatGenerationChunk::new(AIMessage::new("b").into())),
            Ok(ChatGenerationChunk::new(AIMessage::new("c").into())),
        ];

        let stream = futures::stream::iter(chunks);
        let merged = collect_and_merge_stream(stream).await.unwrap();

        assert!(merged.is_some());
        assert_eq!(merged.unwrap().text, "abc");
    }

    #[tokio::test]
    async fn test_collect_and_merge_stream_empty() {
        let chunks: Vec<Result<ChatGenerationChunk>> = vec![];
        let stream = futures::stream::iter(chunks);
        let merged = collect_and_merge_stream(stream).await.unwrap();
        assert!(merged.is_none());
    }
}