[PYTHON/LANGCHAIN] BaseChatModel 클래스 : 커스텀 채팅 모델 만들기

■ BaseChatModel 클래스를 사용해 커스텀 채팅 모델을 만드는 방법을 보여준다.

▶ main.py


import asyncio

from langchain_core.language_models import BaseChatModel
from typing                         import List
from langchain_core.messages        import BaseMessage
from typing                         import Optional
from langchain_core.callbacks       import CallbackManagerForLLMRun
from typing                         import Any
from langchain_core.outputs         import ChatResult
from langchain_core.messages        import AIMessage
from langchain_core.outputs         import ChatGeneration
from typing                         import Iterator
from langchain_core.outputs         import ChatGenerationChunk
from langchain_core.messages        import AIMessageChunk
from typing                         import Dict
from langchain_core.messages        import HumanMessage

class CustomChatModel(BaseChatModel):
    """A custom chat model that echoes the first `n` 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:

        .. code-block:: python

            customChatModel = CustomChatModel(n = 2)

            responseAIMessage     = customChatModel.invoke([HumanMessage(content = "hello")])
            responseAIMessageList = customChatModel.batch([[HumanMessage(content = "hello")], [HumanMessage(content = "world")]])
    """

    model_name : str
    """The name of the model"""
    n : int
    """The number of characters from the last message of the prompt to be echoed."""

    def _generate(
        self,
        messages    : List[BaseMessage],
        stop        : Optional[List[str]] = None,
        run_manager : Optional[CallbackManagerForLLMRun] = None,
        **kwargs    : Any,
    ) -> ChatResult:
        """Override the _generate method to implement the chat model logic.

        This can be a call to an API, a call to a local model, or any other
        implementation that generates a response to the input prompt.

        Args:
            messages : the prompt composed of a list of messages.
            stop : a list of strings on which the model should stop generating.
                   If generation stops due to a stop token, the stop token itself
                   SHOULD BE INCLUDED as part of the output. This is not enforced
                   across models right now, but it's a good practice to follow since
                   it makes it much easier to parse the output of the model
                   downstream and understand why generation stopped.
            run_manager : A run manager with callbacks for the LLM.
        """
        # 이것을 실제 논리로 바꿔서 메시지 목록에서 응답을 생성한다.
        lastMessage = messages[-1]
        tokenCount  = lastMessage.content[: self.n]
        aiMessage   = AIMessage(
            content           = tokenCount,
            additional_kwargs = {}, # 추가 페이로드(예 : 함수 호출 요청)를 추가하는 데 사용한다.
            response_metadata = {   # 응답 메타 데이터에 사용한다.
                "time_in_seconds" : 3
            }
        )

        chatGeneration = ChatGeneration(message = aiMessage)

        return ChatResult(generations = [chatGeneration])

    def _stream(
        self,
        messages    : List[BaseMessage],
        stop        : Optional[List[str]] = None,
        run_manager : Optional[CallbackManagerForLLMRun] = None,
        **kwargs    : Any,
    ) -> Iterator[ChatGenerationChunk]:
        """Stream the output of the model.

        This method should be implemented if the model can generate output in a streaming fashion.
        If the model does not support streaming, do not implement it.
        In that case streaming requests will be automatically handled by the _generate method.

        Args :
            messages : the prompt composed of a list of messages.
            stop : a list of strings on which the model should stop generating.
                   If generation stops due to a stop token, the stop token itself
                   SHOULD BE INCLUDED as part of the output. This is not enforced
                   across models right now, but it's a good practice to follow since
                   it makes it much easier to parse the output of the model
                   downstream and understand why generation stopped.
            run_manager : A run manager with callbacks for the LLM.
        """
        lastMessage = messages[-1]
        tokenCount  = lastMessage.content[: self.n]

        for token in tokenCount:
            chatGenerationChunk = ChatGenerationChunk(message = AIMessageChunk(content = token))

            if run_manager:
                # 이것은 LangChain의 최신 버전에서는 선택 사항이다.
                # on_llm_new_token이 자동으로 호출된다.
                run_manager.on_llm_new_token(token, chunk = chatGenerationChunk)

            yield chatGenerationChunk

        # 다른 정보(예 : 응답 메타 데이터)를 추가한다.
        chatGenerationChunk = ChatGenerationChunk(message = AIMessageChunk(content = "", response_metadata = {"time_in_sec" : 3}))

        if run_manager:
            # 이것은 LangChain의 새로운 버전에서는 선택 사항이다.
            # on_llm_new_token이 자동으로 호출된다.
            run_manager.on_llm_new_token(token, chunk = chatGenerationChunk)
        yield chatGenerationChunk

    @property
    def _llm_type(self) -> str:
        """Get the type of language model used by this chat model."""
        return "echoing-chat-model-advanced"

    @property
    def _identifying_params(self) -> Dict[str, Any]:
        """Return a dictionary of identifying parameters.

        This information is used by the LangChain callback system, which
        is used for tracing purposes make it possible to monitor LLMs.
        """
        return {
            # 모델 이름을 사용하면 사용자가 사용자 정의 토큰 계산을 지정할 수 있다.
            # LLM 모니터링 애플리케이션의 규칙(예 : LangSmith에서 사용자는 모델에 대한 토큰당 가격을 제공하고 지정된 LLM에 대한 비용을 모니터링할 수 있다)
            "model_name": self.model_name,
        }

customChatModel = CustomChatModel(n = 3, model_name = "custom_chat_model")

print("-" * 50)

responseAIMessage1 = customChatModel.invoke(
    [
        HumanMessage(content = "hello!"),
        AIMessage(content = "Hi there human!"),
        HumanMessage(content = "Meow!")
    ]
)

print(responseAIMessage1)
print("-" * 50)

responseAIMessage2 = customChatModel.invoke("hello")

print(responseAIMessage2)
print("-" * 50)

responseAIMessageList = customChatModel.batch(["hello", "goodbye"])

print(responseAIMessageList)
print("-" * 50)

for responseAIMessageChunk in customChatModel.stream("cat"):
    print(responseAIMessageChunk.content, end = "|")
print()
print("-" * 50)

async def main():
    async for eventDictionary in customChatModel.astream_events("cat", version = "v2"):
        print(eventDictionary)
    print("-" * 50)

asyncio.run(main())

"""
---------------------------------------------------
content='Meo' additional_kwargs={} response_metadata={'time_in_seconds': 3} id='run-dc5a9152-b42f-4767-b788-be2d845e5b6c-0'
---------------------------------------------------
content='hel' additional_kwargs={} response_metadata={'time_in_seconds': 3} id='run-27021883-ebfc-4419-b132-69a83f64bb17-0'
---------------------------------------------------
[AIMessage(content='hel', additional_kwargs={}, response_metadata={'time_in_seconds': 3}, id='run-e74822d3-b1b5-44f5-9bbf-190984467b03-0'), AIMessage(content='goo', additional_kwargs={}, response_metadata={'time_in_seconds': 3}, id='run-41dfe740-9519-49af-bb74-e3fab27539f0-0')]
---------------------------------------------------
c|a|t||
---------------------------------------------------
{'event': 'on_chat_model_start', 'data': {'input': 'cat'}, 'name': 'CustomChatModel', 'tags': [], 'run_id': '9c1a640c-8abc-4010-8aac-b950cc4cf3a8', 'metadata': {'ls_provider': 'customchatmodel', 'ls_model_type': 'chat', 'ls_model_name': 'custom_chat_model'}, 'parent_ids': []}
{'event': 'on_chat_model_stream', 'run_id': '9c1a640c-8abc-4010-8aac-b950cc4cf3a8', 'name': 'CustomChatModel', 'tags': [], 'metadata': {'ls_provider': 'customchatmodel', 'ls_model_type': 'chat', 'ls_model_name': 'custom_chat_model'}, 'data': {'chunk': AIMessageChunk(content='c', additional_kwargs={}, response_metadata={}, id='run-9c1a640c-8abc-4010-8aac-b950cc4cf3a8')}, 'parent_ids': []}
{'event': 'on_chat_model_stream', 'run_id': '9c1a640c-8abc-4010-8aac-b950cc4cf3a8', 'name': 'CustomChatModel', 'tags': [], 'metadata': {'ls_provider': 'customchatmodel', 'ls_model_type': 'chat', 'ls_model_name': 'custom_chat_model'}, 'data': {'chunk': AIMessageChunk(content='a', additional_kwargs={}, response_metadata={}, id='run-9c1a640c-8abc-4010-8aac-b950cc4cf3a8')}, 'parent_ids': []}
{'event': 'on_chat_model_stream', 'run_id': '9c1a640c-8abc-4010-8aac-b950cc4cf3a8', 'name': 'CustomChatModel', 'tags': [], 'metadata': {'ls_provider': 'customchatmodel', 'ls_model_type': 'chat', 'ls_model_name': 'custom_chat_model'}, 'data': {'chunk': AIMessageChunk(content='t', additional_kwargs={}, response_metadata={}, id='run-9c1a640c-8abc-4010-8aac-b950cc4cf3a8')}, 'parent_ids': []}
{'event': 'on_chat_model_stream', 'run_id': '9c1a640c-8abc-4010-8aac-b950cc4cf3a8', 'name': 'CustomChatModel', 'tags': [], 'metadata': {'ls_provider': 'customchatmodel', 'ls_model_type': 'chat', 'ls_model_name': 'custom_chat_model'}, 'data': {'chunk': AIMessageChunk(content='', additional_kwargs={}, response_metadata={'time_in_sec': 3}, id='run-9c1a640c-8abc-4010-8aac-b950cc4cf3a8')}, 'parent_ids': []}
{'event': 'on_chat_model_end', 'data': {'output': AIMessageChunk(content='cat', additional_kwargs={}, response_metadata={'time_in_sec': 3}, id='run-9c1a640c-8abc-4010-8aac-b950cc4cf3a8')}, 'run_id': '9c1a640c-8abc-4010-8aac-b950cc4cf3a8', 'name': 'CustomChatModel', 'tags': [], 'metadata': {'ls_provider': 'customchatmodel', 'ls_model_type': 'chat', 'ls_model_name': 'custom_chat_model'}, 'parent_ids': []}
---------------------------------------------------
"""

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

import asyncio

from langchain_core.language_models import BaseChatModel

from typing import List

from langchain_core.messages import BaseMessage

from typing import Optional

from langchain_core.callbacks import CallbackManagerForLLMRun

from typing import Any

from langchain_core.outputs import ChatResult

from langchain_core.messages import AIMessage

from langchain_core.outputs import ChatGeneration

from typing import Iterator

from langchain_core.outputs import ChatGenerationChunk

from langchain_core.messages import AIMessageChunk

from typing import Dict

from langchain_core.messages import HumanMessage

class CustomChatModel(BaseChatModel):

"""A custom chat model that echoes the first `n` 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:

.. code-block:: python

customChatModel = CustomChatModel(n = 2)

responseAIMessage = customChatModel.invoke([HumanMessage(content = "hello")])

responseAIMessageList = customChatModel.batch([[HumanMessage(content = "hello")], [HumanMessage(content = "world")]])

"""

model_name : str

"""The name of the model"""

n : int

"""The number of characters from the last message of the prompt to be echoed."""

def _generate(

self,

messages : List[BaseMessage],

stop : Optional[List[str]] = None,

run_manager : Optional[CallbackManagerForLLMRun] = None,

**kwargs : Any,

) -> ChatResult:

"""Override the _generate method to implement the chat model logic.

This can be a call to an API, a call to a local model, or any other

implementation that generates a response to the input prompt.

Args:

messages : the prompt composed of a list of messages.

stop : a list of strings on which the model should stop generating.

If generation stops due to a stop token, the stop token itself

SHOULD BE INCLUDED as part of the output. This is not enforced

across models right now, but it's a good practice to follow since

it makes it much easier to parse the output of the model

downstream and understand why generation stopped.

run_manager : A run manager with callbacks for the LLM.

"""

# 이것을 실제 논리로 바꿔서 메시지 목록에서 응답을 생성한다.

lastMessage = messages[-1]

tokenCount = lastMessage.content[: self.n]

aiMessage = AIMessage(

content = tokenCount,

additional_kwargs = {}, # 추가 페이로드(예 : 함수 호출 요청)를 추가하는 데 사용한다.

response_metadata = { # 응답 메타 데이터에 사용한다.

"time_in_seconds" : 3

}

)

chatGeneration = ChatGeneration(message = aiMessage)

return ChatResult(generations = [chatGeneration])

def _stream(

self,

messages : List[BaseMessage],

stop : Optional[List[str]] = None,

run_manager : Optional[CallbackManagerForLLMRun] = None,

**kwargs : Any,

) -> Iterator[ChatGenerationChunk]:

"""Stream the output of the model.

This method should be implemented if the model can generate output in a streaming fashion.

If the model does not support streaming, do not implement it.

In that case streaming requests will be automatically handled by the _generate method.

Args :

messages : the prompt composed of a list of messages.

stop : a list of strings on which the model should stop generating.

If generation stops due to a stop token, the stop token itself

SHOULD BE INCLUDED as part of the output. This is not enforced

across models right now, but it's a good practice to follow since

it makes it much easier to parse the output of the model

downstream and understand why generation stopped.

run_manager : A run manager with callbacks for the LLM.

"""

lastMessage = messages[-1]

tokenCount = lastMessage.content[: self.n]

for token in tokenCount:

chatGenerationChunk = ChatGenerationChunk(message = AIMessageChunk(content = token))

if run_manager:

# 이것은 LangChain의 최신 버전에서는 선택 사항이다.

# on_llm_new_token이 자동으로 호출된다.

run_manager.on_llm_new_token(token, chunk = chatGenerationChunk)

yield chatGenerationChunk

# 다른 정보(예 : 응답 메타 데이터)를 추가한다.

chatGenerationChunk = ChatGenerationChunk(message = AIMessageChunk(content = "", response_metadata = {"time_in_sec" : 3}))

if run_manager:

# 이것은 LangChain의 새로운 버전에서는 선택 사항이다.

# on_llm_new_token이 자동으로 호출된다.

run_manager.on_llm_new_token(token, chunk = chatGenerationChunk)

yield chatGenerationChunk

@property

def _llm_type(self) -> str:

"""Get the type of language model used by this chat model."""

return "echoing-chat-model-advanced"

@property

def _identifying_params(self) -> Dict[str, Any]:

"""Return a dictionary of identifying parameters.

This information is used by the LangChain callback system, which

is used for tracing purposes make it possible to monitor LLMs.

"""

return {

# 모델 이름을 사용하면 사용자가 사용자 정의 토큰 계산을 지정할 수 있다.

# LLM 모니터링 애플리케이션의 규칙(예 : LangSmith에서 사용자는 모델에 대한 토큰당 가격을 제공하고 지정된 LLM에 대한 비용을 모니터링할 수 있다)

"model_name": self.model_name,

}

customChatModel = CustomChatModel(n = 3, model_name = "custom_chat_model")

print("-" * 50)

responseAIMessage1 = customChatModel.invoke(

[

HumanMessage(content = "hello!"),

AIMessage(content = "Hi there human!"),

HumanMessage(content = "Meow!")

]

)

print(responseAIMessage1)

print("-" * 50)

responseAIMessage2 = customChatModel.invoke("hello")

print(responseAIMessage2)

print("-" * 50)

responseAIMessageList = customChatModel.batch(["hello", "goodbye"])

print(responseAIMessageList)

print("-" * 50)

for responseAIMessageChunk in customChatModel.stream("cat"):

print(responseAIMessageChunk.content, end = "|")

print()

print("-" * 50)

async def main():

async for eventDictionary in customChatModel.astream_events("cat", version = "v2"):

print(eventDictionary)

print("-" * 50)

asyncio.run(main())

"""

---------------------------------------------------

content='Meo' additional_kwargs={} response_metadata={'time_in_seconds': 3} id='run-dc5a9152-b42f-4767-b788-be2d845e5b6c-0'

---------------------------------------------------

content='hel' additional_kwargs={} response_metadata={'time_in_seconds': 3} id='run-27021883-ebfc-4419-b132-69a83f64bb17-0'

---------------------------------------------------

[AIMessage(content='hel', additional_kwargs={}, response_metadata={'time_in_seconds': 3}, id='run-e74822d3-b1b5-44f5-9bbf-190984467b03-0'), AIMessage(content='goo', additional_kwargs={}, response_metadata={'time_in_seconds': 3}, id='run-41dfe740-9519-49af-bb74-e3fab27539f0-0')]

---------------------------------------------------

c|a|t||

---------------------------------------------------

{'event': 'on_chat_model_start', 'data': {'input': 'cat'}, 'name': 'CustomChatModel', 'tags': [], 'run_id': '9c1a640c-8abc-4010-8aac-b950cc4cf3a8', 'metadata': {'ls_provider': 'customchatmodel', 'ls_model_type': 'chat', 'ls_model_name': 'custom_chat_model'}, 'parent_ids': []}

{'event': 'on_chat_model_stream', 'run_id': '9c1a640c-8abc-4010-8aac-b950cc4cf3a8', 'name': 'CustomChatModel', 'tags': [], 'metadata': {'ls_provider': 'customchatmodel', 'ls_model_type': 'chat', 'ls_model_name': 'custom_chat_model'}, 'data': {'chunk': AIMessageChunk(content='c', additional_kwargs={}, response_metadata={}, id='run-9c1a640c-8abc-4010-8aac-b950cc4cf3a8')}, 'parent_ids': []}

{'event': 'on_chat_model_stream', 'run_id': '9c1a640c-8abc-4010-8aac-b950cc4cf3a8', 'name': 'CustomChatModel', 'tags': [], 'metadata': {'ls_provider': 'customchatmodel', 'ls_model_type': 'chat', 'ls_model_name': 'custom_chat_model'}, 'data': {'chunk': AIMessageChunk(content='a', additional_kwargs={}, response_metadata={}, id='run-9c1a640c-8abc-4010-8aac-b950cc4cf3a8')}, 'parent_ids': []}

{'event': 'on_chat_model_stream', 'run_id': '9c1a640c-8abc-4010-8aac-b950cc4cf3a8', 'name': 'CustomChatModel', 'tags': [], 'metadata': {'ls_provider': 'customchatmodel', 'ls_model_type': 'chat', 'ls_model_name': 'custom_chat_model'}, 'data': {'chunk': AIMessageChunk(content='t', additional_kwargs={}, response_metadata={}, id='run-9c1a640c-8abc-4010-8aac-b950cc4cf3a8')}, 'parent_ids': []}

{'event': 'on_chat_model_stream', 'run_id': '9c1a640c-8abc-4010-8aac-b950cc4cf3a8', 'name': 'CustomChatModel', 'tags': [], 'metadata': {'ls_provider': 'customchatmodel', 'ls_model_type': 'chat', 'ls_model_name': 'custom_chat_model'}, 'data': {'chunk': AIMessageChunk(content='', additional_kwargs={}, response_metadata={'time_in_sec': 3}, id='run-9c1a640c-8abc-4010-8aac-b950cc4cf3a8')}, 'parent_ids': []}

{'event': 'on_chat_model_end', 'data': {'output': AIMessageChunk(content='cat', additional_kwargs={}, response_metadata={'time_in_sec': 3}, id='run-9c1a640c-8abc-4010-8aac-b950cc4cf3a8')}, 'run_id': '9c1a640c-8abc-4010-8aac-b950cc4cf3a8', 'name': 'CustomChatModel', 'tags': [], 'metadata': {'ls_provider': 'customchatmodel', 'ls_model_type': 'chat', 'ls_model_name': 'custom_chat_model'}, 'parent_ids': []}

---------------------------------------------------

"""

▶ requirements.txt


aiohappyeyeballs==2.4.3
aiohttp==3.11.7
aiosignal==1.3.1
annotated-types==0.7.0
anyio==4.6.2.post1
attrs==24.2.0
certifi==2024.8.30
charset-normalizer==3.4.0
frozenlist==1.5.0
greenlet==3.1.1
h11==0.14.0
httpcore==1.0.7
httpx==0.27.2
idna==3.10
jsonpatch==1.33
jsonpointer==3.0.0
langchain==0.3.7
langchain-core==0.3.19
langchain-text-splitters==0.3.2
langsmith==0.1.144
multidict==6.1.0
numpy==1.26.4
orjson==3.10.11
packaging==24.2
propcache==0.2.0
pydantic==2.10.1
pydantic_core==2.27.1
PyYAML==6.0.2
requests==2.32.3
requests-toolbelt==1.0.0
sniffio==1.3.1
SQLAlchemy==2.0.36
tenacity==9.0.0
typing_extensions==4.12.2
urllib3==2.2.3
yarl==1.18.0

aiohappyeyeballs==2.4.3

aiohttp==3.11.7

aiosignal==1.3.1

annotated-types==0.7.0

anyio==4.6.2.post1

attrs==24.2.0

certifi==2024.8.30

charset-normalizer==3.4.0

frozenlist==1.5.0

greenlet==3.1.1

h11==0.14.0

httpcore==1.0.7

httpx==0.27.2

idna==3.10

jsonpatch==1.33

jsonpointer==3.0.0

langchain==0.3.7

langchain-core==0.3.19

langchain-text-splitters==0.3.2

langsmith==0.1.144

multidict==6.1.0

numpy==1.26.4

orjson==3.10.11

packaging==24.2

propcache==0.2.0

pydantic==2.10.1

pydantic_core==2.27.1

PyYAML==6.0.2

requests==2.32.3

requests-toolbelt==1.0.0

sniffio==1.3.1

SQLAlchemy==2.0.36

tenacity==9.0.0

typing_extensions==4.12.2

urllib3==2.2.3

yarl==1.18.0

※ pip install langchain 명령을 실행했다.'