BaseChatModel: a model of your own

A chat model takes messages and returns one AIMessage. Subclass BaseChatModel to write your own, and LangChain treats it like any hosted model.

Every chat model in LangChain, from OpenAI's to Groq's, is a subclass of BaseChatModel. A subclass has to supply two things: a name for its type, and _generate, which turns messages into a reply. Everything else, including invoke, comes from the base class.

A name for the model

import re

from langchain.chat_models import BaseChatModel
from langchain.messages import AIMessage
from langchain_core.outputs import ChatGeneration, ChatResult


class ShopModel(BaseChatModel):
    @property
    def _llm_type(self):
        return "shop"

_llm_type is a label LangChain uses in logs and traces. The imports are the base class, the message class the model returns, and two small wrappers that _generate has to put its reply in.

Deciding on a reply

    def _generate(self, messages, stop=None, run_manager=None, **kwargs):
        text = messages[-1].text
        orders = re.findall(r"\b[A-Z]\d+\b", text)
        if orders:
            reply = f"I have no way to look up {orders[0]} yet."
        else:
            reply = "Hello. Which order is this about?"
        message = AIMessage(reply)
        return ChatResult(generations=[ChatGeneration(message=message)])

The model reads the last message and looks for an order id, a capital letter followed by digits. With one, it admits it cannot look the order up yet. Without one, it asks which order the customer means. The reply goes into an AIMessage, wrapped in the result type every chat model returns.

Save both pieces as shop_model.py. Lesson 4 imports it, and lesson 6 teaches it to use tools.

Calling it like any other model

model = ShopModel()

reply = model.invoke("Where is my order A17?")
print(type(reply).__name__)
print(reply.text)

invoke came from the base class. It turned the string into a HumanMessage, called your _generate, and returned the AIMessage inside the result.

model = ShopModel()

reply = model.invoke([
    {"role": "system", "content": "You help customers of a small online shop."},
    {"role": "user", "content": "Hello"},
])
print(reply.text)

A list of dictionaries works too, converted to message objects before _generate sees them. The model reads only the last one, the customer's "Hello", so it asks which order this is about.