mirror of
https://github.com/hwchase17/langchain.git
synced 2025-06-16 11:58:51 +00:00
fake messages list model (#10152)
create a fake chat model that you can configure with list of messages
This commit is contained in:
parent
9e6cc7b236
commit
f5af756397
@ -7,9 +7,35 @@ from langchain.callbacks.manager import (
|
|||||||
AsyncCallbackManagerForLLMRun,
|
AsyncCallbackManagerForLLMRun,
|
||||||
CallbackManagerForLLMRun,
|
CallbackManagerForLLMRun,
|
||||||
)
|
)
|
||||||
from langchain.chat_models.base import SimpleChatModel
|
from langchain.chat_models.base import BaseChatModel, SimpleChatModel
|
||||||
|
from langchain.schema import ChatResult
|
||||||
from langchain.schema.messages import AIMessageChunk, BaseMessage
|
from langchain.schema.messages import AIMessageChunk, BaseMessage
|
||||||
from langchain.schema.output import ChatGenerationChunk
|
from langchain.schema.output import ChatGeneration, ChatGenerationChunk
|
||||||
|
|
||||||
|
|
||||||
|
class FakeMessagesListChatModel(BaseChatModel):
|
||||||
|
responses: List[BaseMessage]
|
||||||
|
sleep: Optional[float] = None
|
||||||
|
i: int = 0
|
||||||
|
|
||||||
|
def _generate(
|
||||||
|
self,
|
||||||
|
messages: List[BaseMessage],
|
||||||
|
stop: Optional[List[str]] = None,
|
||||||
|
run_manager: Optional[CallbackManagerForLLMRun] = None,
|
||||||
|
**kwargs: Any,
|
||||||
|
) -> ChatResult:
|
||||||
|
response = self.responses[self.i]
|
||||||
|
if self.i < len(self.responses) - 1:
|
||||||
|
self.i += 1
|
||||||
|
else:
|
||||||
|
self.i = 0
|
||||||
|
generation = ChatGeneration(message=response)
|
||||||
|
return ChatResult(generations=[generation])
|
||||||
|
|
||||||
|
@property
|
||||||
|
def _llm_type(self) -> str:
|
||||||
|
return "fake-messages-list-chat-model"
|
||||||
|
|
||||||
|
|
||||||
class FakeListChatModel(SimpleChatModel):
|
class FakeListChatModel(SimpleChatModel):
|
||||||
|
Loading…
Reference in New Issue
Block a user