custom_chat_model.py — langchain Source File

Source Code

from __future__ import annotations

from typing import TYPE_CHECKING, Any

from langchain_core.language_models import BaseChatModel
from langchain_core.messages import AIMessage, AIMessageChunk, BaseMessage
from langchain_core.messages.ai import UsageMetadata
from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
from pydantic import Field
from typing_extensions import override

if TYPE_CHECKING:
    from collections.abc import Iterator

    from langchain_core.callbacks import CallbackManagerForLLMRun


class ChatParrotLink(BaseChatModel):
    """Chat Parrot Link.

    A custom chat model that echoes the first `parrot_buffer_length` characters
    of the input.

    When contributing an implementation to LangChain, carefully document
    the model including the initialization parameters, include
    an example of how to initialize the model and include any relevant
    links to the underlying models documentation or API.

    Example:
    ```python
    model = ChatParrotLink(parrot_buffer_length=2, model="bird-brain-001")
    result = model.invoke([HumanMessage(content="hello")])
    result = model.batch(
        [
            [HumanMessage(content="hello")],
            [HumanMessage(content="world")],
        ]
    )
    ```
    """

    model_name: str = Field(alias="model")
    """The name of the model"""
    parrot_buffer_length: int
    """The number of characters from the last message of the prompt to be echoed."""
    temperature: float | None = None
    max_tokens: int | None = None
    timeout: int | None = None
    stop: list[str] | None = None
    max_retries: int = 2

    @override
    def _generate(
        self,
        messages: list[BaseMessage],
        stop: list[str] | None = None,
        run_manager: CallbackManagerForLLMRun | None = None,
        **kwargs: Any,
    ) -> ChatResult:
        """Override the _generate method to implement the chat model logic.
// ... (129 more lines)